SERVER-23219 DBCommandCursor should route getMore operations to original server

11 files changed, 145 insertions, 31 deletions

diff --git a/jstests/noPassthrough/replica_set_connection_getmore.js b/jstests/noPassthrough/replica_set_connection_getmore.js
index fee090e2aeb..aa63a05f87f 100644
--- a/jstests/noPassthrough/replica_set_connection_getmore.js
+++ b/jstests/noPassthrough/replica_set_connection_getmore.js
@@ -4,7 +4,6 @@
  */
 (function() {
     "use strict";
-
     var rst = new ReplSetTest({nodes: 2});
     rst.startSet();
     rst.initiate();
@@ -17,13 +16,12 @@
     var conn = new Mongo(rst.getURL());
 
     // We force a read mode of "compatibility" so that we can test Mongo.prototype.readMode()
-    // resolves to "legacy" independently of the --readMode passed to the mongo shell running this
+    // resolves to "commands" independently of the --readMode passed to the mongo shell running this
     // test.
     conn.forceReadMode("compatibility");
-    assert.eq("legacy",
+    assert.eq("commands",
               conn.readMode(),
-              "replica set connections created by the mongo shell should use 'legacy' read mode");
-
+              "replica set connections created by the mongo shell should use 'commands' read mode");
     var coll = conn.getDB(dbName)[collName];
     coll.drop();
 
@@ -36,7 +34,11 @@
     rst.awaitReplication();
 
     // Establish a cursor on the secondary and verify that the getMore operations are routed to it.
-    conn.forceReadMode("compatibility");
+    var cursor = coll.find().readPref("secondary").batchSize(2);
+    assert.eq(5, cursor.itcount(), "failed to read the documents from the secondary");
+
+    // Verify that queries work when the read mode is forced to "legacy" reads.
+    conn.forceReadMode("legacy");
     var cursor = coll.find().readPref("secondary").batchSize(2);
     assert.eq(5, cursor.itcount(), "failed to read the documents from the secondary");
 
diff --git a/jstests/noPassthrough/stepdown_query.js b/jstests/noPassthrough/stepdown_query.js
index 05d22f34a63..afeb14c4524 100644
--- a/jstests/noPassthrough/stepdown_query.js
+++ b/jstests/noPassthrough/stepdown_query.js
@@ -5,7 +5,7 @@
     var dbName = "test";
     var collName = jsTest.name();
 
-    function runTest(host, rst) {
+    function runTest(host, rst, waitForPrimary) {
         // We create a new connection to 'host' here instead of passing in the original connection.
         // This to work around the fact that connections created by ReplSetTest already have slaveOk
         // set on them, but we need a connection with slaveOk not set for this test.
@@ -19,10 +19,14 @@
         cursor.next();
         assert.eq(0, cursor.objsLeftInBatch());
         var primary = rst.getPrimary();
+        var secondary = rst.getSecondary();
         assert.throws(function() {
             primary.getDB("admin").runCommand({replSetStepDown: 60, force: true});
         });
         rst.waitForState(primary, ReplSetTest.State.SECONDARY, 60 * 1000);
+        if (waitForPrimary) {
+            rst.waitForState(secondary, ReplSetTest.State.PRIMARY, 60 * 1000);
+        }
         // When the primary steps down, it closes all client connections. Since 'conn' may be a
         // direct connection to the primary and the shell doesn't automatically retry operations on
         // network errors, we run a dummy operation here to force the shell to reconnect.
@@ -43,7 +47,13 @@
     var rst = new ReplSetTest({nodes: 1});
     rst.startSet();
     rst.initiate();
-    runTest(rst.getPrimary().host, rst);
+    runTest(rst.getPrimary().host, rst, false);
+    rst.stopSet();
+
+    rst = new ReplSetTest({nodes: 2});
+    rst.startSet();
+    rst.initiate();
+    runTest(rst.getURL(), rst, true);
     rst.stopSet();
 
     // Test querying a replica set primary through mongos.
diff --git a/src/mongo/client/dbclient.cpp b/src/mongo/client/dbclient.cpp
index a6735fea214..faace525444 100644
--- a/src/mongo/client/dbclient.cpp
+++ b/src/mongo/client/dbclient.cpp
@@ -236,10 +236,16 @@ rpc::UniqueReply DBClientWithCommands::runCommandWithMetadata(StringData databas
     return rpc::UniqueReply(std::move(replyMsg), std::move(commandReply));
 }
 
-bool DBClientWithCommands::runCommand(const string& dbname,
-                                      const BSONObj& cmd,
-                                      BSONObj& info,
-                                      int options) {
+std::tuple<rpc::UniqueReply, DBClientWithCommands*>
+DBClientWithCommands::runCommandWithMetadataAndTarget(StringData database,
+                                                      StringData command,
+                                                      const BSONObj& metadata,
+                                                      const BSONObj& commandArgs) {
+    return std::make_tuple(runCommandWithMetadata(database, command, metadata, commandArgs), this);
+}
+
+std::tuple<bool, DBClientWithCommands*> DBClientWithCommands::runCommandWithTarget(
+    const string& dbname, const BSONObj& cmd, BSONObj& info, int options) {
     BSONObj upconvertedCmd;
     BSONObj upconvertedMetadata;
 
@@ -251,13 +257,24 @@ bool DBClientWithCommands::runCommand(const string& dbname,
 
     auto commandName = upconvertedCmd.firstElementFieldName();
 
-    auto result = runCommandWithMetadata(dbname, commandName, upconvertedMetadata, upconvertedCmd);
+    auto resultTuple =
+        runCommandWithMetadataAndTarget(dbname, commandName, upconvertedMetadata, upconvertedCmd);
+    auto result = std::move(std::get<0>(resultTuple));
 
     info = result->getCommandReply().getOwned();
 
-    return isOk(info);
+    return std::make_tuple(isOk(info), std::get<1>(resultTuple));
 }
 
+bool DBClientWithCommands::runCommand(const string& dbname,
+                                      const BSONObj& cmd,
+                                      BSONObj& info,
+                                      int options) {
+    auto res = runCommandWithTarget(dbname, cmd, info, options);
+    return std::get<0>(res);
+}
+
+
 /* note - we build a bson obj here -- for something that is super common like getlasterror you
           should have that object prebuilt as that would be faster.
 */
diff --git a/src/mongo/client/dbclient_rs.cpp b/src/mongo/client/dbclient_rs.cpp
index ae18bd070d0..ffca94992e5 100644
--- a/src/mongo/client/dbclient_rs.cpp
+++ b/src/mongo/client/dbclient_rs.cpp
@@ -898,6 +898,15 @@ rpc::UniqueReply DBClientReplicaSet::runCommandWithMetadata(StringData database,
                                                             StringData command,
                                                             const BSONObj& metadata,
                                                             const BSONObj& commandArgs) {
+    auto ret = runCommandWithMetadataAndTarget(database, command, metadata, commandArgs);
+    return std::move(std::get<0>(ret));
+}
+
+std::tuple<rpc::UniqueReply, DBClientWithCommands*>
+DBClientReplicaSet::runCommandWithMetadataAndTarget(StringData database,
+                                                    StringData command,
+                                                    const BSONObj& metadata,
+                                                    const BSONObj& commandArgs) {
     // This overload exists so we can parse out the read preference and then use server
     // selection directly without having to re-parse the raw message.
 
@@ -920,8 +929,9 @@ rpc::UniqueReply DBClientReplicaSet::runCommandWithMetadata(StringData database,
         // If the command is not runnable on a secondary, we run it on the primary
         // regardless of the read preference.
         !_isSecondaryCommand(command, commandArgs)) {
-        return checkMaster()->runCommandWithMetadata(
-            std::move(database), std::move(command), metadata, commandArgs);
+        auto conn = checkMaster();
+        return std::make_tuple(
+            conn->runCommandWithMetadata(database, command, metadata, commandArgs), conn);
     }
 
     auto rpShared = std::make_shared<ReadPreferenceSetting>(std::move(readPref));
@@ -934,7 +944,8 @@ rpc::UniqueReply DBClientReplicaSet::runCommandWithMetadata(StringData database,
             }
             // We can't move database and command in case this throws
             // and we retry.
-            return conn->runCommandWithMetadata(database, command, metadata, commandArgs);
+            return std::make_tuple(
+                conn->runCommandWithMetadata(database, command, metadata, commandArgs), conn);
         } catch (const DBException& ex) {
             log() << exceptionToStatus();
             invalidateLastSlaveOkCache();
diff --git a/src/mongo/client/dbclient_rs.h b/src/mongo/client/dbclient_rs.h
index 60e615d54dc..1876ac20068 100644
--- a/src/mongo/client/dbclient_rs.h
+++ b/src/mongo/client/dbclient_rs.h
@@ -189,6 +189,12 @@ public:
                                             const BSONObj& metadata,
                                             const BSONObj& commandArgs) final;
 
+    std::tuple<rpc::UniqueReply, DBClientWithCommands*> runCommandWithMetadataAndTarget(
+        StringData database,
+        StringData command,
+        const BSONObj& metadata,
+        const BSONObj& commandArgs) final;
+
     void setRequestMetadataWriter(rpc::RequestMetadataWriter writer) final;
 
     void setReplyMetadataReader(rpc::ReplyMetadataReader reader) final;
diff --git a/src/mongo/client/dbclientinterface.h b/src/mongo/client/dbclientinterface.h
index 880f94aa4e4..62e09660524 100644
--- a/src/mongo/client/dbclientinterface.h
+++ b/src/mongo/client/dbclientinterface.h
@@ -419,6 +419,19 @@ public:
                                                     const BSONObj& metadata,
                                                     const BSONObj& commandArgs);
 
+    /*
+     * This wraps up the runCommandWithMetadata function above, but returns the DBClient that
+     * actually ran the command. When called against a replica set, this will return the specific
+     * replica set member the command ran against.
+     *
+     * This is used in the shell so that cursors can send getMore through the correct connection.
+     */
+    virtual std::tuple<rpc::UniqueReply, DBClientWithCommands*> runCommandWithMetadataAndTarget(
+        StringData database,
+        StringData command,
+        const BSONObj& metadata,
+        const BSONObj& commandArgs);
+
     /** Run a database command.  Database commands are represented as BSON objects.  Common database
         commands have prebuilt helper functions -- see below.  If a helper is not available you can
         directly call runCommand.
@@ -437,6 +450,18 @@ public:
                             BSONObj& info,
                             int options = 0);
 
+    /*
+     * This wraps up the runCommand function avove, but returns the DBClient that actually ran
+     * the command. When called against a replica set, this will return the specific
+     * replica set member the command ran against.
+     *
+     * This is used in the shell so that cursors can send getMore through the correct connection.
+     */
+    virtual std::tuple<bool, DBClientWithCommands*> runCommandWithTarget(const std::string& dbname,
+                                                                         const BSONObj& cmd,
+                                                                         BSONObj& info,
+                                                                         int options = 0);
+
     /**
     * Authenticates to another cluster member using appropriate authentication data.
     * Uses getInternalUserAuthParams() to retrive authentication parameters.
diff --git a/src/mongo/scripting/mozjs/mongo.cpp b/src/mongo/scripting/mozjs/mongo.cpp
index da62b092f34..8df2629c08f 100644
--- a/src/mongo/scripting/mozjs/mongo.cpp
+++ b/src/mongo/scripting/mozjs/mongo.cpp
@@ -112,6 +112,50 @@ void setCursorHandle(JS::HandleObject target, long long cursorId, JS::CallArgs&
     // Copy the client shared pointer to up the refcount.
     JS_SetPrivate(target, new CursorHandleInfo::CursorTracker(cursorId, *client));
 }
+
+void setHiddenMongo(JSContext* cx,
+                    DBClientWithCommands* resPtr,
+                    DBClientWithCommands* origConn,
+                    JS::CallArgs& args) {
+    ObjectWrapper o(cx, args.rval());
+    // If the connection that ran the command is the same as conn, then we set a hidden "_mongo"
+    // property on the returned object that is just "this" Mongo object.
+    if (resPtr == origConn) {
+        o.defineProperty(InternedString::_mongo, args.thisv(), JSPROP_READONLY | JSPROP_PERMANENT);
+    } else {
+        // Otherwise, we construct a new Mongo object that is a copy of "this", but has a different
+        // private value which is the specific DBClientBase that should be used for getMore calls.
+        auto& connSharedPtr = *(static_cast<std::shared_ptr<DBClientBase>*>(
+            JS_GetPrivate(args.thisv().toObjectOrNull())));
+
+        JS::RootedObject newMongo(cx);
+
+        auto scope = getScope(cx);
+        auto isLocalInfo = scope->getProto<MongoLocalInfo>().instanceOf(args.thisv());
+        if (isLocalInfo) {
+            scope->getProto<MongoLocalInfo>().newObject(&newMongo);
+        } else {
+            scope->getProto<MongoExternalInfo>().newObject(&newMongo);
+        }
+        JS_SetPrivate(
+            newMongo,
+            new std::shared_ptr<DBClientBase>(connSharedPtr, static_cast<DBClientBase*>(resPtr)));
+
+        ObjectWrapper from(cx, args.thisv());
+        ObjectWrapper to(cx, newMongo);
+        for (const auto& k :
+             {InternedString::slaveOk, InternedString::defaultDB, InternedString::host}) {
+            JS::RootedValue tmpValue(cx);
+            from.getValue(k, &tmpValue);
+            to.setValue(k, tmpValue);
+        }
+
+        JS::RootedValue value(cx);
+        value.setObjectOrNull(newMongo);
+
+        o.defineProperty(InternedString::_mongo, value, JSPROP_READONLY | JSPROP_PERMANENT);
+    }
+}
 }  // namespace
 
 void MongoBase::finalize(JSFreeOp* fop, JSObject* obj) {
@@ -143,12 +187,13 @@ void MongoBase::Functions::runCommand::call(JSContext* cx, JS::CallArgs args) {
 
     int queryOptions = ValueWriter(cx, args.get(2)).toInt32();
     BSONObj cmdRes;
-    conn->runCommand(database, cmdObj, cmdRes, queryOptions);
+    auto resTuple = conn->runCommandWithTarget(database, cmdObj, cmdRes, queryOptions);
 
     // the returned object is not read only as some of our tests depend on modifying it.
     //
     // Also, we make a copy here because we want a copy after we dump cmdRes
     ValueReader(cx, args.rval()).fromBSON(cmdRes.getOwned(), nullptr, false /* read only */);
+    setHiddenMongo(cx, std::get<1>(resTuple), conn, args);
 }
 
 void MongoBase::Functions::runCommandWithMetadata::call(JSContext* cx, JS::CallArgs args) {
@@ -177,14 +222,16 @@ void MongoBase::Functions::runCommandWithMetadata::call(JSContext* cx, JS::CallA
     BSONObj commandArgs = ValueWriter(cx, args.get(3)).toBSON();
 
     auto conn = getConnection(args);
-    auto res = conn->runCommandWithMetadata(database, commandName, metadata, commandArgs);
+    auto resTuple =
+        conn->runCommandWithMetadataAndTarget(database, commandName, metadata, commandArgs);
+    auto res = std::move(std::get<0>(resTuple));
 
     BSONObjBuilder mergedResultBob;
     mergedResultBob.append("commandReply", res->getCommandReply());
     mergedResultBob.append("metadata", res->getMetadata());
 
-    auto mergedResult = mergedResultBob.obj();
-    ValueReader(cx, args.rval()).fromBSON(mergedResult, nullptr, false);
+    ValueReader(cx, args.rval()).fromBSON(mergedResultBob.obj(), nullptr, false);
+    setHiddenMongo(cx, std::get<1>(resTuple), conn, args);
 }
 
 void MongoBase::Functions::find::call(JSContext* cx, JS::CallArgs args) {
diff --git a/src/mongo/shell/collection.js b/src/mongo/shell/collection.js
index 68bd6f7014e..6c70ee8da87 100644
--- a/src/mongo/shell/collection.js
+++ b/src/mongo/shell/collection.js
@@ -1033,7 +1033,7 @@ DBCollection.prototype._getIndexesCommand = function(filter) {
         throw _getErrorWithCode(res, "listIndexes failed: " + tojson(res));
     }
 
-    return new DBCommandCursor(this._mongo, res).toArray();
+    return new DBCommandCursor(res._mongo, res).toArray();
 };
 
 DBCollection.prototype.getIndexes = function(filter) {
@@ -1208,7 +1208,7 @@ DBCollection.prototype.convertToCapped = function(bytes) {
 DBCollection.prototype.exists = function() {
     var res = this._db.runCommand("listCollections", {filter: {name: this._shortName}});
     if (res.ok) {
-        var cursor = new DBCommandCursor(this._mongo, res);
+        var cursor = new DBCommandCursor(res._mongo, res);
         if (!cursor.hasNext())
             return null;
         return cursor.next();
@@ -1306,7 +1306,7 @@ DBCollection.prototype.aggregate = function(pipeline, aggregateOptions) {
     assert.commandWorked(res, "aggregate failed");
 
     if ("cursor" in res) {
-        return new DBCommandCursor(this._mongo, res);
+        return new DBCommandCursor(res._mongo, res);
     }
 
     return res;
diff --git a/src/mongo/shell/db.js b/src/mongo/shell/db.js
index f6eb72e10fb..53911769c49 100644
--- a/src/mongo/shell/db.js
+++ b/src/mongo/shell/db.js
@@ -804,7 +804,7 @@ var DB;
             throw _getErrorWithCode(res, "listCollections failed: " + tojson(res));
         }
 
-        return new DBCommandCursor(this._mongo, res).toArray().sort(compareOn("name"));
+        return new DBCommandCursor(res._mongo, res).toArray().sort(compareOn("name"));
     };
 
     /**
diff --git a/src/mongo/shell/mongo.js b/src/mongo/shell/mongo.js
index 2cadfcb77ed..2ba2d7ffcd9 100644
--- a/src/mongo/shell/mongo.js
+++ b/src/mongo/shell/mongo.js
@@ -341,11 +341,7 @@ Mongo.prototype.readMode = function() {
         // commands. If it does, use commands mode. If not, degrade to legacy mode.
         try {
             var hasReadCommands = (this.getMinWireVersion() <= 4 && 4 <= this.getMaxWireVersion());
-            // TODO SERVER-23219: DBCommandCursor doesn't route getMore and killCursors operations
-            // to the server that the cursor was originally established on. As a workaround, we make
-            // replica set connections use 'legacy' read mode because the underlying DBClientCursor
-            // will correctly route operations to the original server.
-            if (hasReadCommands && !this.isReplicaSetConnection()) {
+            if (hasReadCommands) {
                 this._readMode = "commands";
             } else {
                 print("Cannot use 'commands' readMode, degrading to 'legacy' mode");
diff --git a/src/mongo/shell/query.js b/src/mongo/shell/query.js
index 3a32b7a951a..ec010f4d6ad 100644
--- a/src/mongo/shell/query.js
+++ b/src/mongo/shell/query.js
@@ -114,7 +114,7 @@ DBQuery.prototype._exec = function() {
             var canAttachReadPref = true;
             var findCmd = this._convertToCommand(canAttachReadPref);
             var cmdRes = this._db.runReadCommand(findCmd, null, this._options);
-            this._cursor = new DBCommandCursor(this._mongo, cmdRes, this._batchSize);
+            this._cursor = new DBCommandCursor(cmdRes._mongo, cmdRes, this._batchSize);
         } else {
             if (this._special && this._query.readConcern) {
                 throw new Error("readConcern requires use of read commands");