SERVER-41481 1 started new branch, added functionality for exhaustCursorId once again. Trying to kill the purple

3 files changed, 32 insertions, 5 deletions

diff --git a/src/mongo/db/dbmessage.h b/src/mongo/db/dbmessage.h
index f6281c3614c..c4c8bd3b243 100644
--- a/src/mongo/db/dbmessage.h
+++ b/src/mongo/db/dbmessage.h
@@ -445,6 +445,12 @@ Message makeGetMoreMessage(StringData ns, long long cursorId, int nToReturn, int
 struct DbResponse {
     Message response;       // If empty, nothing will be returned to the client.
     std::string exhaustNS;  // Namespace of cursor if exhaust mode, else "".
+
+    /**
+     * Cursor ID when running on exhaust mode. Defaults to '0', indicating
+     * that the cursor is exhausted.
+     */
+    long long exhaustCursorId = 0;
 };
 
 /**
diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index 89076c67f13..c0d131852bc 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -1006,11 +1006,21 @@ DbResponse receivedCommands(OperationContext* opCtx,
         return {};  // Don't reply.
     }
 
-    auto response = replyBuilder->done();
-    CurOp::get(opCtx)->debug().responseLength = response.header().dataLen();
+    DbResponse dbResponse;
+
+    if (OpMsg::isFlagSet(message, OpMsg::kExhaustSupported)) {
+        auto responseObj = replyBuilder->getBodyBuilder().asTempObj();
+        auto cursorObj = responseObj.getObjectField("cursor");
+        if (responseObj.getField("ok").trueValue() && !cursorObj.isEmpty()) {
+            dbResponse.exhaustNS = cursorObj.getField("ns").String();
+            dbResponse.exhaustCursorId = cursorObj.getField("id").numberLong();
+        }
+    }
+
+    dbResponse.response = replyBuilder->done();
+    CurOp::get(opCtx)->debug().responseLength = dbResponse.response.header().dataLen();
 
-    // TODO exhaust
-    return DbResponse{std::move(response)};
+    return dbResponse;
 }
 
 DbResponse receivedQuery(OperationContext* opCtx,
diff --git a/src/mongo/s/commands/strategy.cpp b/src/mongo/s/commands/strategy.cpp
index 1587c95c491..fcdf8a8299d 100644
--- a/src/mongo/s/commands/strategy.cpp
+++ b/src/mongo/s/commands/strategy.cpp
@@ -754,7 +754,18 @@ DbResponse Strategy::clientCommand(OperationContext* opCtx, const Message& m) {
         return {};  // Don't reply.
     }
 
-    return DbResponse{reply->done()};
+    DbResponse dbResponse;
+    if (OpMsg::isFlagSet(m, OpMsg::kExhaustSupported)) {
+        auto responseObj = reply->getBodyBuilder().asTempObj();
+        auto cursorObj = responseObj.getObjectField("cursor");
+        if (responseObj.getField("ok").trueValue() && !cursorObj.isEmpty()) {
+            dbResponse.exhaustNS = cursorObj.getField("ns").String();
+            dbResponse.exhaustCursorId = cursorObj.getField("id").numberLong();
+        }
+    }
+    dbResponse.response = reply->done();
+
+    return dbResponse;
 }
 
 void Strategy::commandOp(OperationContext* opCtx,