SERVER-31194 Add a version of retryable_writes_jscore_passthrough.yml with stepdowns

2 files changed, 81 insertions, 4 deletions

diff --git a/src/mongo/shell/session.js b/src/mongo/shell/session.js
index 48e3eab554a..eb0cc9a6464 100644
--- a/src/mongo/shell/session.js
+++ b/src/mongo/shell/session.js
@@ -280,6 +280,47 @@ var {
             }
         }
 
+        /**
+         * Returns true if the error code is retryable, assuming the command is idempotent.
+         */
+        function isRetryableCode(code) {
+            return ErrorCodes.isNetworkError(code) || ErrorCodes.isNotMasterError(code) ||
+                // The driver's spec does not allow retrying on writeConcern errors, so only do so
+                // when testing retryable writes.
+                (jsTest.options().alwaysInjectTransactionNumber &&
+                 ErrorCodes.isWriteConcernError(code));
+        }
+
+        /**
+         * Returns the error code from a write response that should be used in the check for
+         * retryability.
+         */
+        function getEffectiveWriteErrorCode(res) {
+            let code;
+            if (res instanceof WriteResult) {
+                if (res.hasWriteError()) {
+                    code = res.getWriteError().code;
+                } else if (res.hasWriteConcernError()) {
+                    code = res.getWriteConcernError().code;
+                }
+            } else if (res instanceof BulkWriteResult) {
+                if (res.hasWriteErrors()) {
+                    code = res.getWriteErrorAt(0).code;
+                } else if (res.hasWriteConcernError()) {
+                    code = res.getWriteConcernError().code;
+                }
+            } else {
+                if (res.writeError) {
+                    code = res.writeError.code;
+                } else if (res.writeErrors) {
+                    code = res.writeErrors[0].code;
+                } else if (res.writeConcernError) {
+                    code = res.writeConcernError.code;
+                }
+            }
+            return code;
+        }
+
         function runClientFunctionWithRetries(
             driverSession, cmdObj, clientFunction, clientFunctionArguments) {
             let cmdName = Object.keys(cmdObj)[0];
@@ -296,13 +337,42 @@ var {
                 ? 1
                 : 0;
 
+            if (numRetries > 0 && jsTest.options().overrideRetryAttempts) {
+                numRetries = jsTest.options().overrideRetryAttempts;
+            }
+
             do {
                 try {
-                    const res = clientFunction.apply(client, clientFunctionArguments);
-                    if (res.ok === 1 || numRetries === 0 ||
-                        !ErrorCodes.isNotMasterError(res.code)) {
-                        return res;
+                    let res = clientFunction.apply(client, clientFunctionArguments);
+
+                    if (numRetries > 0) {
+                        if (!res.ok && isRetryableCode(res.code)) {
+                            // Don't decrement retries, because the command returned before the
+                            // connection was closed, so a subsequent attempt will receive a
+                            // network error (or NotMaster error) and need to retry.
+                            if (jsTest.options().logRetryAttempts) {
+                                print("=-=-=-= Retrying failed response with retryable code: " +
+                                      res.code + ", for command: " + cmdName +
+                                      ", retries remaining: " + numRetries);
+                            }
+                            continue;
+                        }
+
+                        let code = getEffectiveWriteErrorCode(res);
+                        if (isRetryableCode(code)) {
+                            // Don't decrement retries, because the command returned before the
+                            // connection was closed, so a subsequent attempt will receive a network
+                            // error (or NotMaster error) and need to retry.
+                            if (jsTest.options().logRetryAttempts) {
+                                print("=-=-=-= Retrying write with retryable write error code: " +
+                                      code + ", for command: " + cmdName + ", retries remaining: " +
+                                      numRetries);
+                            }
+                            continue;
+                        }
                     }
+
+                    return res;
                 } catch (e) {
                     if (!isNetworkError(e) || numRetries === 0) {
                         throw e;
@@ -329,6 +399,10 @@ var {
                 }
 
                 --numRetries;
+                if (jsTest.options().logRetryAttempts) {
+                    print("=-=-=-= Retrying on network error for command: " + cmdName +
+                          ", retries remaining: " + numRetries);
+                }
             } while (numRetries >= 0);
         }
 
diff --git a/src/mongo/shell/utils.js b/src/mongo/shell/utils.js
index 638580ebf3c..f7ec71ed8b6 100644
--- a/src/mongo/shell/utils.js
+++ b/src/mongo/shell/utils.js
@@ -261,6 +261,9 @@ jsTestOptions = function() {
             alwaysInjectTransactionNumber: TestData.alwaysInjectTransactionNumber,
             skipGossipingClusterTime: TestData.skipGossipingClusterTime || false,
             disableEnableSessions: TestData.disableEnableSessions,
+            overrideRetryAttempts: TestData.overrideRetryAttempts || 0,
+            logRetryAttempts: TestData.logRetryAttempts || false,
+            connectionString: TestData.connectionString || "",
         });
     }
     return _jsTestOptions;