SERVER-21639 set ResultFlag_ShardConfigStale on OP_QUERY command replies if needed in LegacyReplyBuilder

4 files changed, 35 insertions, 4 deletions

diff --git a/src/mongo/rpc/legacy_reply_builder.cpp b/src/mongo/rpc/legacy_reply_builder.cpp
index 71ef4f9f5bf..393c3bb3519 100644
--- a/src/mongo/rpc/legacy_reply_builder.cpp
+++ b/src/mongo/rpc/legacy_reply_builder.cpp
@@ -51,6 +51,26 @@ LegacyReplyBuilder::LegacyReplyBuilder(Message&& message) : _message{std::move(m
 
 LegacyReplyBuilder::~LegacyReplyBuilder() {}
 
+LegacyReplyBuilder& LegacyReplyBuilder::setCommandReply(Status nonOKStatus,
+                                                        const BSONObj& extraErrorInfo) {
+    invariant(_state == State::kCommandReply);
+    if (nonOKStatus == ErrorCodes::SendStaleConfig) {
+        _staleConfigError = true;
+        // Need to use the special $err format for SendStaleConfig errors to be backwards
+        // compatible.
+        BSONObjBuilder err;
+        // $err must be the first field in object.
+        err.append("$err", nonOKStatus.reason());
+        err.append("code", nonOKStatus.code());
+        err.appendElements(extraErrorInfo);
+        setRawCommandReply(err.done());
+    } else {
+        // All other errors proceed through the normal path, which also handles state transitions.
+        ReplyBuilderInterface::setCommandReply(std::move(nonOKStatus), extraErrorInfo);
+    }
+    return *this;
+}
+
 LegacyReplyBuilder& LegacyReplyBuilder::setRawCommandReply(const BSONObj& commandReply) {
     invariant(_state == State::kCommandReply);
     commandReply.appendSelfToBufBuilder(_builder);
@@ -116,6 +136,7 @@ void LegacyReplyBuilder::reset() {
     _builder.skip(sizeof(QueryResult::Value));
     _message.reset();
     _state = State::kCommandReply;
+    _staleConfigError = false;
 }
 
 
@@ -124,7 +145,13 @@ Message LegacyReplyBuilder::done() {
 
     QueryResult::View qr = _builder.buf();
 
-    qr.setResultFlagsToOk();
+    if (_staleConfigError) {
+        // For compatibility with legacy mongos, we need to set this result flag on SendStaleConfig
+        qr.setResultFlags(ResultFlag_ErrSet | ResultFlag_ShardConfigStale);
+    } else {
+        qr.setResultFlagsToOk();
+    }
+
     qr.msgdata().setLen(_builder.len());
     qr.msgdata().setOperation(opReply);
     qr.setCursorId(0);
diff --git a/src/mongo/rpc/legacy_reply_builder.h b/src/mongo/rpc/legacy_reply_builder.h
index cf5f15c3d88..bb98661650a 100644
--- a/src/mongo/rpc/legacy_reply_builder.h
+++ b/src/mongo/rpc/legacy_reply_builder.h
@@ -49,7 +49,8 @@ public:
     LegacyReplyBuilder(Message&&);
     ~LegacyReplyBuilder() final;
 
-
+    // Override of setCommandReply specifically used to handle SendStaleConfigException.
+    LegacyReplyBuilder& setCommandReply(Status nonOKStatus, const BSONObj& extraErrorInfo) final;
     LegacyReplyBuilder& setRawCommandReply(const BSONObj& commandReply) final;
 
     BufBuilder& getInPlaceReplyBuilder(std::size_t) final;
@@ -71,6 +72,8 @@ private:
     BufBuilder _builder{};
     Message _message;
     State _state{State::kCommandReply};
+    // For stale config errors we need to set the correct ResultFlag.
+    bool _staleConfigError{false};
 };
 
 }  // namespace rpc
diff --git a/src/mongo/rpc/reply_builder_interface.h b/src/mongo/rpc/reply_builder_interface.h
index 57bf36b0fc0..1dc021faa49 100644
--- a/src/mongo/rpc/reply_builder_interface.h
+++ b/src/mongo/rpc/reply_builder_interface.h
@@ -92,7 +92,8 @@ public:
      * interfacing with legacy code that adds additional data to a failed command reply and
      * its use is discouraged in new code.
      */
-    ReplyBuilderInterface& setCommandReply(Status nonOKStatus, const BSONObj& extraErrorInfo);
+    virtual ReplyBuilderInterface& setCommandReply(Status nonOKStatus,
+                                                   const BSONObj& extraErrorInfo);
 
     /**
      * Add a range of output documents to the reply. This method can be called multiple times
diff --git a/src/mongo/rpc/reply_builder_test.cpp b/src/mongo/rpc/reply_builder_test.cpp
index fb7e8c1ed9e..ce37dc36141 100644
--- a/src/mongo/rpc/reply_builder_test.cpp
+++ b/src/mongo/rpc/reply_builder_test.cpp
@@ -104,7 +104,7 @@ TEST(LegacyReplyBuilder, MemAccess) {
     BSONObj metadata = buildMetadata();
     BSONObj commandReply = buildEmptyCommand();
     rpc::LegacyReplyBuilder replyBuilder;
-    replyBuilder.setCommandReply(commandReply);
+    replyBuilder.setRawCommandReply(commandReply);
     replyBuilder.setMetadata(metadata);
     auto msg = replyBuilder.done();