SERVER-36306 Add FSM workloads that run transactions workloads with all threads using same session

6 files changed, 331 insertions, 79 deletions

diff --git a/jstests/concurrency/fsm_workload_helpers/cleanup_txns.js b/jstests/concurrency/fsm_workload_helpers/cleanup_txns.js
new file mode 100644
index 00000000000..ee94c835d8a
--- /dev/null
+++ b/jstests/concurrency/fsm_workload_helpers/cleanup_txns.js
@@ -0,0 +1,48 @@
+/**
+ * Helpers for aborting transactions in concurrency workloads.
+ */
+
+/**
+ * Abort the transaction on the session and return result.
+ */
+function abortTransaction(db, txnNumber, errorCodes) {
+    const abortCmd = {abortTransaction: 1, txnNumber: NumberLong(txnNumber), autocommit: false};
+    const res = db.adminCommand(abortCmd);
+    return assert.commandWorkedOrFailedWithCode(res, errorCodes, () => `cmd: ${tojson(cmd)}`);
+}
+
+/**
+ * This function operates on the last iteration of each thread to abort any active transactions.
+ */
+var {cleanupOnLastIteration} = (function() {
+    function cleanupOnLastIteration(data, func) {
+        const abortErrorCodes = [
+            ErrorCodes.NoSuchTransaction,
+            ErrorCodes.TransactionCommitted,
+            ErrorCodes.TransactionTooOld
+        ];
+
+        let lastIteration = ++data.iteration >= data.iterations;
+        try {
+            func();
+        } catch (e) {
+            lastIteration = true;
+            throw e;
+        } finally {
+            if (lastIteration) {
+                // Abort the latest transactions for this session as some may have been skipped due
+                // to incrementing data.txnNumber. Go in increasing order, so as to avoid bumping
+                // the txnNumber on the server past that of an in-progress transaction. See
+                // SERVER-36847.
+                for (let i = 0; i <= data.txnNumber; i++) {
+                    let res = abortTransaction(data.sessionDb, i, abortErrorCodes);
+                    if (res.ok === 1) {
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
+    return {cleanupOnLastIteration};
+})();
\ No newline at end of file
diff --git a/jstests/concurrency/fsm_workload_helpers/snapshot_read_utils.js b/jstests/concurrency/fsm_workload_helpers/snapshot_read_utils.js
index 2d42d19e2d1..901d7da7d3f 100644
--- a/jstests/concurrency/fsm_workload_helpers/snapshot_read_utils.js
+++ b/jstests/concurrency/fsm_workload_helpers/snapshot_read_utils.js
@@ -2,7 +2,7 @@
  * Helpers for doing a snapshot read in concurrency suites. Specifically, the read is a find that
  * spans a getmore.
  */
-
+load('jstests/concurrency/fsm_workload_helpers/cleanup_txns.js');
 /**
  * Parses a cursor from cmdResult, if possible.
  */
@@ -15,21 +15,6 @@ function parseCursor(cmdResult) {
 }
 
 /**
- * Asserts cmd has either failed with a code in a specified set of codes or has succeeded.
- */
-function assertWorkedOrFailed(cmd, cmdResult, errorCodeSet) {
-    if (!cmdResult.ok) {
-        assert.commandFailedWithCode(cmdResult,
-                                     errorCodeSet,
-                                     "expected command to fail with one of " + errorCodeSet +
-                                         ", cmd: " + tojson(cmd) + ", result: " +
-                                         tojson(cmdResult));
-    } else {
-        assert.commandWorked(cmdResult);
-    }
-}
-
-/**
  * Performs a snapshot find.
  */
 function doSnapshotFind(sortByAscending, collName, data, findErrorCodes) {
@@ -51,7 +36,7 @@ function doSnapshotFind(sortByAscending, collName, data, findErrorCodes) {
 
     // Establish a snapshot batchSize:0 cursor.
     let res = data.sessionDb.runCommand(findCmd);
-    assertWorkedOrFailed(findCmd, res, findErrorCodes);
+    assert.commandWorkedOrFailedWithCode(res, findErrorCodes, () => `cmd: ${tojson(findCmd)}`);
     const cursor = parseCursor(res);
 
     if (!cursor) {
@@ -84,7 +69,8 @@ function doSnapshotGetMore(collName, data, getMoreErrorCodes, commitTransactionE
         autocommit: false
     };
     let res = data.sessionDb.runCommand(getMoreCmd);
-    assertWorkedOrFailed(getMoreCmd, res, getMoreErrorCodes);
+    assert.commandWorkedOrFailedWithCode(
+        res, getMoreErrorCodes, () => `cmd: ${tojson(getMoreCmd)}`);
 
     const commitCmd = {
         commitTransaction: 1,
@@ -93,7 +79,8 @@ function doSnapshotGetMore(collName, data, getMoreErrorCodes, commitTransactionE
         autocommit: false
     };
     res = data.sessionDb.adminCommand(commitCmd);
-    assertWorkedOrFailed(commitCmd, res, commitTransactionErrorCodes);
+    assert.commandWorkedOrFailedWithCode(
+        res, commitTransactionErrorCodes, () => `cmd: ${tojson(commitCmd)}`);
 }
 
 /**
@@ -119,48 +106,3 @@ function killSessionsFromDocs(db, collName, tid) {
     let sessionIds = db[collName].find({"_id": docs}, {_id: 0, id: 1}).toArray();
     assert.commandWorked(db.runCommand({killSessions: sessionIds}));
 }
-
-/**
- * Abort the transaction on the session and return result.
- */
-function abortTransaction(db, txnNumber, errorCodes) {
-    abortCmd = {abortTransaction: 1, txnNumber: NumberLong(txnNumber), autocommit: false};
-    res = db.adminCommand(abortCmd);
-    assertWorkedOrFailed(abortCmd, res, errorCodes);
-    return res;
-}
-
-/**
- * This function operates on the last iteration of each thread to abort any active transactions.
- */
-var {cleanupOnLastIteration} = (function() {
-    function cleanupOnLastIteration(data, func) {
-        const abortErrorCodes = [
-            ErrorCodes.NoSuchTransaction,
-            ErrorCodes.TransactionCommitted,
-            ErrorCodes.TransactionTooOld
-        ];
-        let lastIteration = ++data.iteration >= data.iterations;
-        try {
-            func();
-        } catch (e) {
-            lastIteration = true;
-            throw e;
-        } finally {
-            if (lastIteration) {
-                // Abort the latest transactions for this session as some may have been skipped due
-                // to incrementing data.txnNumber. Go in increasing order, so as to avoid bumping
-                // the txnNumber on the server past that of an in-progress transaction. See
-                // SERVER-36847.
-                for (let i = 0; i <= data.txnNumber; i++) {
-                    let res = abortTransaction(data.sessionDb, i, abortErrorCodes);
-                    if (res.ok === 1) {
-                        break;
-                    }
-                }
-            }
-        }
-    }
-
-    return {cleanupOnLastIteration};
-})();
diff --git a/jstests/concurrency/fsm_workloads/multi_statement_transaction_all_commands.js b/jstests/concurrency/fsm_workloads/multi_statement_transaction_all_commands.js
new file mode 100644
index 00000000000..66c5d262cd7
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_all_commands.js
@@ -0,0 +1,220 @@
+'use strict';
+
+/**
+ * Runs findAndModify, update, delete, find, and getMore within a transaction.
+ *
+ * @tags: [uses_transactions]
+ */
+load('jstests/concurrency/fsm_workload_helpers/cleanup_txns.js');
+var $config = (function() {
+
+    function quietly(func) {
+        const printOriginal = print;
+        try {
+            print = Function.prototype;
+            func();
+        } finally {
+            print = printOriginal;
+        }
+    }
+
+    function autoRetryTxn(data, func) {
+        // conflictingOp is true when startTransaction fails with ConflictingOperationInProgress.
+        // This occurs when we attempt to start a transaction with a txnNumber that is already
+        // active on this session. In this case, we will re-run the command with this txnNumber
+        // without calling startTransaction, and essentially join the already running transaction.
+        let conflictingOp = false;
+
+        // startNewTxn is true if the transaction fails with TransactionTooOld or NoSuchTransaction.
+        // TransactionTooOld occurs when a transaction on this session with a higher txnNumber has
+        // started and NoSuchTransaction can occur if a transaction on this session with the same
+        // txnNumber was aborted. In this case, we will start a new transaction to bump the
+        // txnNumber and then re-run the command.
+        let startNewTxn = true;
+
+        do {
+            try {
+                if (startNewTxn) {
+                    // We pass `ignoreActiveTxn = true` to startTransaction so that we will not
+                    // throw `Transaction already in progress on this session` when trying to start
+                    // a new transaction on a session that already has a transaction running on it.
+                    // We instead will catch the error that the server later throws, and will re-run
+                    // the command with 'startTransaction = false' so that we join the already
+                    // running transaction.
+                    data.session.startTransaction_forTesting({readConcern: {level: 'snapshot'}},
+                                                             {ignoreActiveTxn: true});
+                    data.txnNumber++;
+                }
+                startNewTxn = false;
+                conflictingOp = false;
+
+                func();
+
+            } catch (e) {
+                if (e.code === ErrorCodes.TransactionTooOld ||
+                    e.code === ErrorCodes.NoSuchTransaction) {
+                    startNewTxn = true;
+                    continue;
+                }
+
+                if (e.code === ErrorCodes.ConflictingOperationInProgress) {
+                    conflictingOp = true;
+                    continue;
+                }
+
+                if (e.code === ErrorCodes.TransactionCommitted) {
+                    // If running in the same_session workload, it is possible another worker thread
+                    // has already committed this transaction, but a new one has not yet been
+                    // started.
+                    break;
+                }
+
+                throw e;
+            }
+        } while (startNewTxn || conflictingOp);
+    }
+
+    const states = {
+
+        init: function init(db, collName) {
+            this.session = db.getMongo().startSession({causalConsistency: false});
+            this.txnNumber = -1;
+            this.sessionDb = this.session.getDatabase(db.getName());
+            this.iteration = 1;
+        },
+
+        runFindAndModify: function runFindAndModify(db, collName) {
+            autoRetryTxn(this, () => {
+                const collection = this.session.getDatabase(db.getName()).getCollection(collName);
+                assertAlways.commandWorked(collection.runCommand(
+                    'findAndModify', {query: {_id: this.tid}, update: {$inc: {x: 1}}, new: true}));
+            });
+        },
+
+        runUpdate: function runUpdate(db, collName) {
+            autoRetryTxn(this, () => {
+                const collection = this.session.getDatabase(db.getName()).getCollection(collName);
+                assertAlways.commandWorked(collection.runCommand('update', {
+                    updates: [{q: {_id: this.tid}, u: {$inc: {x: 1}}}],
+                }));
+            });
+        },
+
+        runDelete: function runDelete(db, collName) {
+            autoRetryTxn(this, () => {
+                const collection = this.session.getDatabase(db.getName()).getCollection(collName);
+                assertAlways.commandWorked(collection.runCommand('delete', {
+                    deletes: [{q: {_id: this.tid}, limit: 1}],
+                }));
+            });
+        },
+
+        runFindAndGetMore: function runFindAndGetMore(db, collName) {
+            autoRetryTxn(this, () => {
+                const collection = this.session.getDatabase(db.getName()).getCollection(collName);
+                const documents = collection.find().batchSize(2).toArray();
+            });
+        },
+
+        commitTxn: function commitTxn(db, collName) {
+            // shouldJoin is true when commitTransaction fails with ConflictingOperationInProgress.
+            // This occurs when there's a transaction with the same txnNumber running on this
+            // session. In this case we "join" this other transaction and retry the commit, meaning
+            // all operations that were run on this thread will be committed in the same transaction
+            // as the transaction we join.
+            let shouldJoin;
+            do {
+                try {
+                    shouldJoin = false;
+                    quietly(() => this.session.commitTransaction());
+                } catch (e) {
+                    if (e.code === ErrorCodes.TransactionTooOld ||
+                        e.code === ErrorCodes.TransactionCommitted ||
+                        e.code === ErrorCodes.NoSuchTransaction) {
+                        // If we get TransactionTooOld, TransactionCommitted, or NoSuchTransaction
+                        // we do not try to commit this transaction.
+                        break;
+                    }
+
+                    if (e.code === ErrorCodes.ConflictingOperationInProgress) {
+                        shouldJoin = true;
+                        continue;
+                    }
+
+                    throw e;
+                }
+            } while (shouldJoin);
+        },
+    };
+
+    // Wrap each state in a cleanupOnLastIteration() invocation.
+    for (let stateName of Object.keys(states)) {
+        const stateFn = states[stateName];
+        states[stateName] = function(db, collName) {
+            cleanupOnLastIteration(this, () => stateFn.apply(this, arguments));
+        };
+    }
+
+    function setup(db, collName) {
+        assertWhenOwnColl.commandWorked(db.runCommand({create: collName}));
+        const bulk = db[collName].initializeUnorderedBulkOp();
+
+        for (let i = 0; i < this.numDocs; ++i) {
+            bulk.insert({_id: i, x: i});
+        }
+
+        const res = bulk.execute({w: 'majority'});
+        assertWhenOwnColl.commandWorked(res);
+        assertWhenOwnColl.eq(this.numDocs, res.nInserted);
+    }
+
+    function teardown(db, collName, cluster) {
+    }
+
+    const transitions = {
+        init: {runFindAndModify: .25, runUpdate: .25, runDelete: .25, runFindAndGetMore: .25},
+        runFindAndModify: {
+            runFindAndModify: .2,
+            runUpdate: .2,
+            runDelete: .2,
+            runFindAndGetMore: .2,
+            commitTxn: .2
+        },
+        runUpdate: {
+            runFindAndModify: .2,
+            runUpdate: .2,
+            runDelete: .2,
+            runFindAndGetMore: .2,
+            commitTxn: .2
+        },
+        runDelete: {
+            runFindAndModify: .2,
+            runUpdate: .2,
+            runDelete: .2,
+            runFindAndGetMore: .2,
+            commitTxn: .2
+        },
+        runFindAndGetMore: {
+            runFindAndModify: .2,
+            runUpdate: .2,
+            runDelete: .2,
+            runFindAndGetMore: .2,
+            commitTxn: .2
+        },
+        commitTxn: {runFindAndModify: .25, runUpdate: .25, runDelete: .25, runFindAndGetMore: .25},
+    };
+
+    return {
+        threadCount: 5,
+        iterations: 10,
+        states: states,
+        transitions: transitions,
+        data: {
+            numDocs: 20,
+            ignoreActiveTxn: false,
+        },
+        setup: setup,
+        teardown: teardown
+    };
+
+})();
diff --git a/jstests/concurrency/fsm_workloads/multi_statement_transaction_all_commands_same_session.js b/jstests/concurrency/fsm_workloads/multi_statement_transaction_all_commands_same_session.js
new file mode 100644
index 00000000000..c82058e8981
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_all_commands_same_session.js
@@ -0,0 +1,34 @@
+'use strict';
+
+/**
+ * Runs update, findAndModify, delete, find, and getMore in a transaction with all threads using the
+ * same session.
+ *
+ * @tags: [uses_transactions]
+ */
+
+load('jstests/concurrency/fsm_libs/extend_workload.js');  // for extendWorkload
+load('jstests/concurrency/fsm_workloads/multi_statement_transaction_all_commands.js');  // for
+                                                                                        // $config
+
+var $config = extendWorkload($config, function($config, $super) {
+
+    $config.setup = function(db, collName, cluster) {
+        $super.setup.apply(this, arguments);
+        this.lsid = tojson({id: UUID()});
+    };
+
+    $config.states.init = function init(db, collName) {
+        const lsid = eval(`(${this.lsid})`);
+        this.session = db.getMongo().startSession({causalConsistency: false});
+        // Force the session to use `lsid` for its session id. This way all threads will use
+        // the same session.
+        this.session._serverSession.handle.getId = () => lsid;
+
+        this.txnNumber = -1;
+        this.sessionDb = this.session.getDatabase(db.getName());
+        this.iteration = 1;
+    };
+
+    return $config;
+});
diff --git a/jstests/concurrency/fsm_workloads/snapshot_read_kill_operations.js b/jstests/concurrency/fsm_workloads/snapshot_read_kill_operations.js
index 13d0d75325c..7458be7ec44 100644
--- a/jstests/concurrency/fsm_workloads/snapshot_read_kill_operations.js
+++ b/jstests/concurrency/fsm_workloads/snapshot_read_kill_operations.js
@@ -69,7 +69,8 @@ var $config = (function() {
         killCursors: function killCursors(db, collName) {
             const killCursorCmd = {killCursors: collName, cursors: [this.cursorId]};
             const res = this.sessionDb.runCommand(killCursorCmd);
-            assertWorkedOrFailed(killCursorCmd, res, [ErrorCodes.CursorNotFound]);
+            assert.commandWorkedOrFailedWithCode(
+                res, [ErrorCodes.CursorNotFound], () => `cmd: ${tojson(killCursorCmd)}`);
         },
 
     };
diff --git a/jstests/concurrency/fsm_workloads/view_catalog_cycle_with_drop.js b/jstests/concurrency/fsm_workloads/view_catalog_cycle_with_drop.js
index 247361ef525..4673cdbeed5 100644
--- a/jstests/concurrency/fsm_workloads/view_catalog_cycle_with_drop.js
+++ b/jstests/concurrency/fsm_workloads/view_catalog_cycle_with_drop.js
@@ -15,9 +15,6 @@ var $config = (function() {
 
     var data = {
         viewList: ['viewA', 'viewB', 'viewC'].map(viewName => prefix + viewName),
-        assertCommandWorkedOrFailedWithCode: function(result, codeArr) {
-            assertAlways(result.ok === 1 || codeArr.indexOf(result.code) > -1, tojson(result));
-        },
         getRandomView: function(viewList) {
             return viewList[Random.randInt(viewList.length)];
         },
@@ -33,9 +30,11 @@ var $config = (function() {
         function remapViewToView(db, collName) {
             const fromName = this.getRandomView(this.viewList);
             const toName = this.getRandomView(this.viewList);
-            const res = db.runCommand({collMod: fromName, viewOn: toName, pipeline: []});
-            this.assertCommandWorkedOrFailedWithCode(
-                res, [ErrorCodes.GraphContainsCycle, ErrorCodes.NamespaceNotFound]);
+            const cmd = {collMod: fromName, viewOn: toName, pipeline: []};
+            const res = db.runCommand(cmd);
+            const errorCodes = [ErrorCodes.GraphContainsCycle, ErrorCodes.NamespaceNotFound];
+            assertAlways.commandWorkedOrFailedWithCode(
+                res, errorCodes, () => `cmd: ${tojson(cmd)}`);
         }
 
         /**
@@ -45,11 +44,16 @@ var $config = (function() {
          */
         function recreateViewOnCollection(db, collName) {
             const viewName = this.getRandomView(this.viewList);
-            this.assertCommandWorkedOrFailedWithCode(db.runCommand({drop: viewName}),
-                                                     [ErrorCodes.NamespaceNotFound]);
-            this.assertCommandWorkedOrFailedWithCode(
-                db.createView(viewName, collName, []),
-                [ErrorCodes.NamespaceExists, ErrorCodes.NamespaceNotFound]);
+            const dropCmd = {drop: viewName};
+            let res = db.runCommand(dropCmd);
+            let errorCodes = [ErrorCodes.NamespaceNotFound];
+            assertAlways.commandWorkedOrFailedWithCode(
+                db.runCommand(dropCmd), errorCodes, () => `cmd: ${tojson(cmd)}`);
+
+            res = db.createView(viewName, collName, []);
+            errorCodes = [ErrorCodes.NamespaceExists, ErrorCodes.NamespaceNotFound];
+            assertAlways.commandWorkedOrFailedWithCode(
+                res, errorCodes, () => `cmd: ${tojson(cmd)}`);
         }
 
         /**
@@ -60,9 +64,12 @@ var $config = (function() {
          */
         function readFromView(db, collName) {
             const viewName = this.getRandomView(this.viewList);
-            const res = db.runCommand({find: viewName});
+            const cmd = {find: viewName};
+            const res = db.runCommand(cmd);
+            const errorCodes = [ErrorCodes.CommandNotSupportedOnView];
             // TODO SERVER-26037: Replace with the appropriate error code. See ticket for details.
-            this.assertCommandWorkedOrFailedWithCode(res, [ErrorCodes.CommandNotSupportedOnView]);
+            assertAlways.commandWorkedOrFailedWithCode(
+                res, errorCodes, () => `cmd: ${tojson(cmd)}`);
         }
 
         return {
@@ -86,7 +93,7 @@ var $config = (function() {
         assertAlways.writeOK(coll.insert({x: 1}));
 
         for (let viewName of this.viewList) {
-            assert.commandWorked(db.createView(viewName, collName, []));
+            assertAlways.commandWorked(db.createView(viewName, collName, []));
         }
     }