SERVER-35417 Concurrency workloads using transactions should use autoRetry for reads and writes

3 files changed, 31 insertions, 28 deletions

diff --git a/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js b/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
index 820d85d0162..8ca4c8aa8ea 100644
--- a/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
+++ b/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
@@ -1,6 +1,6 @@
 'use strict';
 
-var {withTxnAndAutoRetryOnWriteConflict} = (function() {
+var {withTxnAndAutoRetry} = (function() {
 
     /**
      * Calls 'func' with the print() function overridden to be a no-op.
@@ -21,20 +21,21 @@ var {withTxnAndAutoRetryOnWriteConflict} = (function() {
 
     /**
      * Runs 'func' inside of a transaction started with 'txnOptions', and automatically retries
-     * until it either succeeds or the server returns a non-WriteConflict error response.
+     * until it either succeeds or the server returns a non-TransientTransactionError error
+     * response.
      *
      * 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
-     * transaction started by the withTxnAndAutoRetryOnWriteConflict() function is only known to
-     * have committed after the withTxnAndAutoRetryOnWriteConflict() function returns.
+     * transaction started by the withTxnAndAutoRetry() function is only known to have committed
+     * after the withTxnAndAutoRetry() function returns.
      */
-    function withTxnAndAutoRetryOnWriteConflict(
+    function withTxnAndAutoRetry(
         session, func, {txnOptions: txnOptions = {readConcern: {level: 'snapshot'}}} = {}) {
-        let hasWriteConflict;
+        let hasTransientError;
 
         do {
             session.startTransaction(txnOptions);
-            hasWriteConflict = false;
+            hasTransientError = false;
 
             try {
                 func();
@@ -50,14 +51,15 @@ var {withTxnAndAutoRetryOnWriteConflict} = (function() {
                 // such that it agrees no transaction is currently in progress on this session.
                 session.abortTransaction();
 
-                if (e.code !== ErrorCodes.WriteConflict) {
+                if (!e.hasOwnProperty('errorLabels') ||
+                    !e.errorLabels.includes('TransientTransactionError')) {
                     throw e;
                 }
 
-                hasWriteConflict = true;
+                hasTransientError = true;
             }
-        } while (hasWriteConflict);
+        } while (hasTransientError);
     }
 
-    return {withTxnAndAutoRetryOnWriteConflict};
+    return {withTxnAndAutoRetry};
 })();
diff --git a/jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation.js b/jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation.js
index a6a2f41f5ac..dcafb5dd508 100644
--- a/jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation.js
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation.js
@@ -41,7 +41,7 @@
 
 load('jstests/libs/cycle_detection.js');  // for Graph
 
-// For withTxnAndAutoRetryOnWriteConflict.
+// For withTxnAndAutoRetry.
 load('jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js');
 
 var $config = (function() {
@@ -99,18 +99,18 @@ var $config = (function() {
 
     const states = (function() {
 
-        function getAllDocuments(collection, numDocs) {
+        function getAllDocuments(session, collection, numDocs) {
             // We intentionally use a smaller batch size when fetching all of the documents in the
             // collection in order to stress the behavior of reading from the same snapshot over the
             // course of multiple network roundtrips.
             const batchSize = Math.max(2, Math.floor(numDocs / 5));
 
-            collection.getDB().getSession().startTransaction();
-            const documents =
-                collection.find().batchSize(batchSize).readConcern('snapshot').toArray();
-            collection.getDB().getSession().commitTransaction();
+            let documents;
+            withTxnAndAutoRetry(session, () => {
+                documents = collection.find().batchSize(batchSize).toArray();
 
-            assertWhenOwnColl.eq(numDocs, documents.length, () => tojson(documents));
+                assertWhenOwnColl.eq(numDocs, documents.length, () => tojson(documents));
+            });
             return documents;
         }
 
@@ -148,7 +148,7 @@ var $config = (function() {
                     }
                 };
 
-                withTxnAndAutoRetryOnWriteConflict(this.session, () => {
+                withTxnAndAutoRetry(this.session, () => {
                     for (let [i, docId] of docIds.entries()) {
                         const res = collection.runCommand('update', {
                             updates: [{q: {_id: docId}, u: updateMods}],
@@ -164,7 +164,7 @@ var $config = (function() {
 
             checkConsistency: function checkConsistency(db, collName) {
                 const collection = this.session.getDatabase(db.getName()).getCollection(collName);
-                const documents = getAllDocuments(collection, this.numDocs);
+                const documents = getAllDocuments(this.session, collection, this.numDocs);
                 checkTransactionCommitOrder(documents);
                 checkNumUpdatedByEachTransaction(documents);
             }
diff --git a/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js b/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js
index 3c7362cb205..42646e1b9db 100644
--- a/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js
@@ -8,7 +8,7 @@
  * @tags: [uses_transactions]
  */
 
-// For withTxnAndAutoRetryOnWriteConflict.
+// For withTxnAndAutoRetry.
 load('jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js');
 
 var $config = (function() {
@@ -19,12 +19,13 @@ var $config = (function() {
 
     var states = (function() {
 
-        function getAllDocuments(collection, numDocs) {
-            collection.getDB().getSession().startTransaction();
-            const documents = collection.find().readConcern('snapshot').toArray();
-            collection.getDB().getSession().commitTransaction();
+        function getAllDocuments(session, collection, numDocs) {
+            let documents;
+            withTxnAndAutoRetry(session, () => {
+                documents = collection.find().toArray();
 
-            assertWhenOwnColl.eq(numDocs, documents.length, () => tojson(documents));
+                assertWhenOwnColl.eq(numDocs, documents.length, () => tojson(documents));
+            });
             return documents;
         }
 
@@ -34,7 +35,7 @@ var $config = (function() {
 
         function checkMoneyBalance(db, collName) {
             const collection = this.session.getDatabase(db.getName()).getCollection(collName);
-            const documents = getAllDocuments(collection, this.numAccounts);
+            const documents = getAllDocuments(this.session, collection, this.numAccounts);
             assertWhenOwnColl.eq(this.numAccounts * this.initialValue,
                                  computeTotalOfAllBalances(documents),
                                  () => tojson(documents));
@@ -52,7 +53,7 @@ var $config = (function() {
             const transferAmount = Random.randInt(this.initialValue / 10) + 1;
 
             const collection = this.session.getDatabase(db.getName()).getCollection(collName);
-            withTxnAndAutoRetryOnWriteConflict(this.session, () => {
+            withTxnAndAutoRetry(this.session, () => {
                 let res = collection.runCommand('update', {
                     updates: [{q: {_id: transferFrom}, u: {$inc: {balance: -transferAmount}}}],
                 });