SERVER-39991 Enable transaction workloads in concurrency sharded stepdown suites

2 files changed, 46 insertions, 11 deletions

diff --git a/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js b/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
index 02d6cfd5fc7..d60d53e796b 100644
--- a/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
+++ b/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
@@ -19,6 +19,32 @@ var {withTxnAndAutoRetry} = (function() {
         }
     }
 
+    // Returns true if the transaction can be retried with a higher transaction number after the
+    // given error.
+    function shouldRetryEntireTxnOnError(e, hasCommitTxnError, retryOnKilledSession) {
+        if ((e.hasOwnProperty('errorLabels') &&
+             e.errorLabels.includes('TransientTransactionError'))) {
+            return true;
+        }
+
+        if (isNetworkError(e) && !hasCommitTxnError) {
+            // A network error before commit is considered a transient txn error. Network
+            // errors during commit should be handled at the same level as retries of
+            // retryable writes.
+            return true;
+        }
+
+        // TODO SERVER-40186: Only retry the entire transaction on killed session errors from before
+        // commit.
+        if ((retryOnKilledSession &&
+             (e.code === ErrorCodes.Interrupted || e.code === ErrorCodes.CursorKilled ||
+              e.code == ErrorCodes.CursorNotFound))) {
+            return true;
+        }
+
+        return false;
+    }
+
     // Use a "signature" value that won't typically match a value assigned in normal use. This way
     // the wtimeout set by this override is distinguishable in the server logs.
     const kDefaultWtimeout = 5 * 60 * 1000 + 789;
@@ -26,8 +52,9 @@ var {withTxnAndAutoRetry} = (function() {
     /**
      * Runs 'func' inside of a transaction started with 'txnOptions', and automatically retries
      * until it either succeeds or the server returns a non-TransientTransactionError error
-     * response. There is a probability of 'prepareProbability' that the transaction is prepared
-     * before committing.
+     * response. If retryOnKilledSession is true, the transaction will be automatically retried on
+     * error codes that may come from a killed session as well. There is a probability of
+     * 'prepareProbability' that the transaction is prepared before committing.
      *
      * The caller should take care to ensure 'func' doesn't modify any captured variables in a
      * speculative fashion where calling it multiple times would lead to unintended behavior. The
@@ -81,11 +108,7 @@ var {withTxnAndAutoRetry} = (function() {
                     session.abortTransaction();
                 }
 
-                if ((e.hasOwnProperty('errorLabels') &&
-                     e.errorLabels.includes('TransientTransactionError')) ||
-                    (retryOnKilledSession &&
-                     (e.code === ErrorCodes.Interrupted || e.code === ErrorCodes.CursorKilled ||
-                      e.code == ErrorCodes.CursorNotFound))) {
+                if (shouldRetryEntireTxnOnError(e, hasCommitTxnError, retryOnKilledSession)) {
                     hasTransientError = true;
                     continue;
                 }
diff --git a/jstests/concurrency/fsm_workload_helpers/snapshot_read_utils.js b/jstests/concurrency/fsm_workload_helpers/snapshot_read_utils.js
index 8662e4f945e..18bbfc50455 100644
--- a/jstests/concurrency/fsm_workload_helpers/snapshot_read_utils.js
+++ b/jstests/concurrency/fsm_workload_helpers/snapshot_read_utils.js
@@ -3,6 +3,8 @@
  * spans a getmore.
  */
 load('jstests/concurrency/fsm_workload_helpers/cleanup_txns.js');
+load('jstests/libs/transactions_util.js');
+
 /**
  * Parses a cursor from cmdResult, if possible.
  */
@@ -41,9 +43,14 @@ function doSnapshotFind(sortByAscending, collName, data, findErrorCodes) {
         autocommit: false
     };
 
-    // Establish a snapshot batchSize:0 cursor.
     let res = data.sessionDb.runCommand(findCmd);
-    assert.commandWorkedOrFailedWithCode(res, findErrorCodes, () => `cmd: ${tojson(findCmd)}`);
+
+    // A transaction request can always fail with a transient transaction error, so only check the
+    // specific error code if it is not labeled as transient.
+    if (!TransactionsUtil.isTransientTransactionError(res)) {
+        assert.commandWorkedOrFailedWithCode(res, findErrorCodes, () => `cmd: ${tojson(findCmd)}`);
+    }
+
     const cursor = parseCursor(res);
 
     if (!cursor) {
@@ -77,8 +84,13 @@ function doSnapshotGetMore(collName, data, getMoreErrorCodes, commitTransactionE
         autocommit: false
     };
     let res = data.sessionDb.runCommand(getMoreCmd);
-    assert.commandWorkedOrFailedWithCode(
-        res, getMoreErrorCodes, () => `cmd: ${tojson(getMoreCmd)}`);
+
+    // A transaction request can always fail with a transient transaction error, so only check the
+    // specific error code if it is not labeled as transient.
+    if (!TransactionsUtil.isTransientTransactionError(res)) {
+        assert.commandWorkedOrFailedWithCode(
+            res, getMoreErrorCodes, () => `cmd: ${tojson(getMoreCmd)}`);
+    }
 
     const commitCmd = {
         commitTransaction: 1,