SERVER-30953 Add auto-retry logic to the mongo shell for testing during stepdown suites

4 files changed, 203 insertions, 5 deletions

diff --git a/jstests/libs/override_methods/auto_retry_on_network_error.js b/jstests/libs/override_methods/auto_retry_on_network_error.js
new file mode 100644
index 00000000000..9509eb372da
--- /dev/null
+++ b/jstests/libs/override_methods/auto_retry_on_network_error.js
@@ -0,0 +1,101 @@
+/**
+ * When a network connection to the mongo shell is closed, attempting to call
+ * Mongo.prototype.runCommand() and Mongo.prototype.runCommandWithMetadata() throws a JavaScript
+ * exception. This override catches these exceptions (i.e. ones where isNetworkError() returns true)
+ * and automatically re-sends the command request to the server, or propagates the error if the
+ * command should already be using the shell's existing retryability logic. The goal of this
+ * override is to implement retry logic such that the assertions within our existing JavaScript
+ * tests still pass despite stepdowns of the CSRS and replica set shards happening in the
+ * background.
+ */
+(function() {
+    "use strict";
+
+    const retryableWriteCommands =
+        new Set(["delete", "findandmodify", "findAndModify", "insert", "update"]);
+
+    // Store a session to access ServerSession#canRetryWrites.
+    let _serverSession;
+
+    const mongoRunCommandOriginal = Mongo.prototype.runCommand;
+    const mongoRunCommandWithMetadataOriginal = Mongo.prototype.runCommandWithMetadata;
+
+    Mongo.prototype.runCommand = function runCommand(dbName, cmdObj, options) {
+        if (typeof _serverSession === "undefined") {
+            _serverSession = this.startSession()._serverSession;
+        }
+
+        return runWithRetriesOnNetworkErrors(this, cmdObj, mongoRunCommandOriginal, arguments);
+    };
+
+    Mongo.prototype.runCommandWithMetadata = function runCommandWithMetadata(
+        dbName, metadata, cmdObj) {
+        if (typeof _serverSession === "undefined") {
+            _serverSession = this.startSession()._serverSession;
+        }
+
+        return runWithRetriesOnNetworkErrors(
+            this, cmdObj, mongoRunCommandWithMetadataOriginal, arguments);
+    };
+
+    function runWithRetriesOnNetworkErrors(mongo, cmdObj, clientFunction, clientFunctionArguments) {
+        let cmdName = Object.keys(cmdObj)[0];
+
+        // If the command is in a wrapped form, then we look for the actual command object
+        // inside the query/$query object.
+        if (cmdName === "query" || cmdName === "$query") {
+            cmdObj = cmdObj[cmdName];
+            cmdName = Object.keys(cmdObj)[0];
+        }
+
+        const isRetryableWriteCmd = retryableWriteCommands.has(cmdName);
+        const canRetryWrites = _serverSession.canRetryWrites(cmdObj);
+
+        let numRetries = !jsTest.options().skipRetryOnNetworkError ? 1 : 0;
+
+        do {
+            try {
+                return clientFunction.apply(mongo, clientFunctionArguments);
+            } catch (e) {
+                if (!isNetworkError(e) || numRetries === 0) {
+                    throw e;
+                } else if (isRetryableWriteCmd) {
+                    if (canRetryWrites) {
+                        // If the command is retryable, assume the command has already gone through
+                        // or will go through the retry logic in SessionAwareClient, so propagate
+                        // the error.
+                        throw e;
+                    } else {
+                        throw new Error(
+                            "Cowardly refusing to run a test that issues non-retryable write" +
+                            " operations since the test likely makes assertions on the write" +
+                            " results and can lead to spurious failures if a network error" +
+                            " occurs.");
+                    }
+                } else if (cmdName === "getMore") {
+                    throw new Error(
+                        "Cowardly refusing to run a test that issues a getMore command since if" +
+                        " a network error occurs during it then we won't know whether the cursor" +
+                        " was advanced or not.");
+                }
+
+                --numRetries;
+            }
+        } while (numRetries >= 0);
+    }
+
+    const startParallelShellOriginal = startParallelShell;
+
+    startParallelShell = function(jsCode, port, noConnect) {
+        let newCode;
+        const overridesFile = "jstests/libs/override_methods/auto_retry_on_network_error.js";
+        if (typeof(jsCode) === "function") {
+            // Load the override file and immediately invoke the supplied function.
+            newCode = `load("${overridesFile}"); (${jsCode})();`;
+        } else {
+            newCode = `load("${overridesFile}"); ${jsCode};`;
+        }
+
+        return startParallelShellOriginal(newCode, port, noConnect);
+    };
+})();
diff --git a/jstests/libs/retry_on_network_error.js b/jstests/libs/retry_on_network_error.js
index f47921f66cd..ab341818167 100644
--- a/jstests/libs/retry_on_network_error.js
+++ b/jstests/libs/retry_on_network_error.js
@@ -10,8 +10,7 @@ function retryOnNetworkError(func, numRetries = 1) {
         try {
             return func();
         } catch (e) {
-            if ((isNetworkError(e) || e.toString().indexOf("network error") > -1) &&
-                numRetries > 0) {
+            if (isNetworkError(e) && numRetries > 0) {
                 print("Network error occurred and the call will be retried: " +
                       tojson({error: e.toString(), stack: e.stack}));
                 numRetries--;
diff --git a/jstests/noPassthrough/auto_retry_on_network_error.js b/jstests/noPassthrough/auto_retry_on_network_error.js
new file mode 100644
index 00000000000..1d8e0c59601
--- /dev/null
+++ b/jstests/noPassthrough/auto_retry_on_network_error.js
@@ -0,0 +1,95 @@
+/**
+ * Tests that the auto_retry_on_network_error.js override automatically retries commands on network
+ * errors for commands run under a session.
+ */
+(function() {
+    "use strict";
+
+    load("jstests/libs/override_methods/auto_retry_on_network_error.js");
+
+    function stepDownPrimary(rst) {
+        // Since we expect the mongo shell's connection to get severed as a result of running the
+        // "replSetStepDown" command, we temporarily disable the retry on network error behavior.
+        TestData.skipRetryOnNetworkError = true;
+        try {
+            const primary = rst.getPrimary();
+            const error = assert.throws(function() {
+                const res = primary.adminCommand({replSetStepDown: 1, force: true});
+                print("replSetStepDown did not throw exception but returned: " + tojson(res));
+            });
+            assert(isNetworkError(error),
+                   "replSetStepDown did not disconnect client; failed with " + tojson(error));
+        } finally {
+            TestData.skipRetryOnNetworkError = false;
+        }
+    }
+
+    const rst = new ReplSetTest({nodes: 2});
+    rst.startSet();
+    rst.initiate();
+
+    const dbName = "test";
+    const collName = "auto_retry";
+
+    // The override requires the connection to be run under a session. Use the replica set URL to
+    // allow automatic re-targeting of the primary on NotMaster errors.
+    const db = new Mongo(rst.getURL()).startSession({retryWrites: true}).getDatabase(dbName);
+
+    // Commands with no stepdowns should work as normal.
+    assert.commandWorked(db.runCommand({ping: 1}));
+    assert.commandWorked(db.runCommandWithMetadata({ping: 1}, {}).commandReply);
+
+    // Read commands are automatically retried on network errors.
+    stepDownPrimary(rst);
+    assert.commandWorked(db.runCommand({find: collName}));
+
+    stepDownPrimary(rst);
+    assert.commandWorked(db.runCommandWithMetadata({find: collName}, {}).commandReply);
+
+    // Retryable write commands that can be retried succeed.
+    stepDownPrimary(rst);
+    assert.writeOK(db[collName].insert({x: 1}));
+
+    stepDownPrimary(rst);
+    assert.commandWorked(db.runCommandWithMetadata({
+                               insert: collName,
+                               documents: [{x: 2}, {x: 3}],
+                               txnNumber: NumberLong(10),
+                               lsid: {id: UUID()}
+                           },
+                                                   {})
+                             .commandReply);
+
+    // Retryable write commands that cannot be retried (i.e. no transaction number, no session id,
+    // or are unordered) throw.
+    stepDownPrimary(rst);
+    assert.throws(function() {
+        db.runCommand({insert: collName, documents: [{x: 1}, {x: 2}], ordered: false});
+    });
+
+    // The previous command shouldn't have been retried, so run a command to successfully re-target
+    // the primary, so the connection to it can be closed.
+    assert.commandWorked(db.runCommandWithMetadata({ping: 1}, {}).commandReply);
+
+    stepDownPrimary(rst);
+    assert.throws(function() {
+        db.runCommandWithMetadata({insert: collName, documents: [{x: 1}, {x: 2}], ordered: false},
+                                  {});
+    });
+
+    // getMore commands can't be retried because we won't know whether the cursor was advanced or
+    // not.
+    let cursorId = assert.commandWorked(db.runCommand({find: collName, batchSize: 0})).cursor.id;
+    stepDownPrimary(rst);
+    assert.throws(function() {
+        db.runCommand({getMore: cursorId, collection: collName});
+    });
+
+    cursorId = assert.commandWorked(db.runCommand({find: collName, batchSize: 0})).cursor.id;
+    stepDownPrimary(rst);
+    assert.throws(function() {
+        db.runCommandWithMetadata({getMore: cursorId, collection: collName}, {});
+    });
+
+    rst.stopSet();
+})();
diff --git a/jstests/sharding/retryable_writes.js b/jstests/sharding/retryable_writes.js
index cd1914e83bb..f83c1b9cdc7 100644
--- a/jstests/sharding/retryable_writes.js
+++ b/jstests/sharding/retryable_writes.js
@@ -238,7 +238,9 @@
 
         try {
             // If ran against mongos, the command will actually succeed, but only one of the writes
-            // would be executed
+            // would be executed. Set skipRetryOnNetworkError so the shell doesn't automatically
+            // retry, since the command has a txnNumber.
+            TestData.skipRetryOnNetworkError = true;
             var res = assert.commandWorked(testDb.runCommand({
                 insert: 'user',
                 documents: [{x: 0}, {x: 1}],
@@ -250,8 +252,9 @@
             assert.eq(1, res.writeErrors.length);
         } catch (e) {
             var exceptionMsg = e.toString();
-            assert(exceptionMsg.indexOf("network error") > -1,
-                   'Incorrect exception thrown: ' + exceptionMsg);
+            assert(isNetworkError(e), 'Incorrect exception thrown: ' + exceptionMsg);
+        } finally {
+            TestData.skipRetryOnNetworkError = false;
         }
 
         assert.eq(2, testDb.user.find({}).itcount());