SERVER-40183 Create kill_session version of multi_statement_transaction_simple.js

4 files changed, 89 insertions, 52 deletions

diff --git a/jstests/concurrency/fsm_workload_helpers/kill_session.js b/jstests/concurrency/fsm_workload_helpers/kill_session.js
new file mode 100644
index 00000000000..38052adadde
--- /dev/null
+++ b/jstests/concurrency/fsm_workload_helpers/kill_session.js
@@ -0,0 +1,51 @@
+/**
+ * kill_session.js
+ *
+ * State function that kills a random session from config.system.sessions.
+ */
+function killSession(db, collName) {
+    let ourSessionWasKilled;
+    do {
+        ourSessionWasKilled = false;
+
+        try {
+            let res = db.adminCommand({refreshLogicalSessionCacheNow: 1});
+            if (res.ok === 1) {
+                assertAlways.commandWorked(res);
+            } else if (res.code === 18630 || res.code === 18631) {
+                // Refreshing the logical session cache may trigger sharding the sessions
+                // collection, which can fail with 18630 or 18631 if its session is killed while
+                // running DBClientBase::getCollectionInfos() or DBClientBase::getIndexSpecs(),
+                // respectively. This means the collection is not set up, so retry.
+                ourSessionWasKilled = true;
+                continue;
+            } else {
+                assertAlways.commandFailedWithCode(res, ErrorCodes.DuplicateKey);
+            }
+
+            const sessionToKill = db.getSiblingDB("config").system.sessions.aggregate([
+                {$listSessions: {}},
+                {$match: {"_id.id": {$ne: db.getSession().getSessionId().id}}},
+                {$sample: {size: 1}},
+            ]);
+
+            if (sessionToKill.toArray().length === 0) {
+                break;
+            }
+
+            const sessionUUID = sessionToKill.toArray()[0]._id.id;
+            res = db.runCommand({killSessions: [{id: sessionUUID}]});
+            assertAlways.commandWorked(res);
+        } catch (e) {
+            if (e.code == ErrorCodes.Interrupted || e.code == ErrorCodes.CursorKilled ||
+                e.code == ErrorCodes.CursorNotFound) {
+                // This session was killed when running either listSessions or killSesssions.
+                // We should retry.
+                ourSessionWasKilled = true;
+                continue;
+            }
+
+            throw e;
+        }
+    } while (ourSessionWasKilled);
+}
diff --git a/jstests/concurrency/fsm_workloads/multi_statement_transaction_kill_sessions_atomicity_isolation.js b/jstests/concurrency/fsm_workloads/multi_statement_transaction_kill_sessions_atomicity_isolation.js
index e5b490ab33e..f655afe8110 100644
--- a/jstests/concurrency/fsm_workloads/multi_statement_transaction_kill_sessions_atomicity_isolation.js
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_kill_sessions_atomicity_isolation.js
@@ -6,58 +6,14 @@
  * @tags: [uses_transactions, assumes_snapshot_transactions]
  */
 
-load('jstests/concurrency/fsm_libs/extend_workload.js');  // for extendWorkload
+load('jstests/concurrency/fsm_libs/extend_workload.js');           // for extendWorkload
+load('jstests/concurrency/fsm_workload_helpers/kill_session.js');  // for killSession
 load('jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation.js');
 
 var $config = extendWorkload($config, ($config, $super) => {
     $config.data.retryOnKilledSession = true;
 
-    $config.states.killSession = function killSession(db, collName) {
-        let ourSessionWasKilled;
-        do {
-            ourSessionWasKilled = false;
-
-            try {
-                let res = db.adminCommand({refreshLogicalSessionCacheNow: 1});
-                if (res.ok === 1) {
-                    assertAlways.commandWorked(res);
-                } else if (res.code === 18630 || res.code === 18631) {
-                    // Refreshing the logical session cache may trigger sharding the sessions
-                    // collection, which can fail with 18630 or 18631 if its session is killed while
-                    // running DBClientBase::getCollectionInfos() or DBClientBase::getIndexSpecs(),
-                    // respectively. This means the collection is not set up, so retry.
-                    ourSessionWasKilled = true;
-                    continue;
-                } else {
-                    assertAlways.commandFailedWithCode(res, ErrorCodes.DuplicateKey);
-                }
-
-                const sessionToKill = db.getSiblingDB("config").system.sessions.aggregate([
-                    {$listSessions: {}},
-                    {$match: {"_id.id": {$ne: db.getSession().getSessionId().id}}},
-                    {$sample: {size: 1}},
-                ]);
-
-                if (sessionToKill.toArray().length === 0) {
-                    break;
-                }
-
-                const sessionUUID = sessionToKill.toArray()[0]._id.id;
-                res = db.runCommand({killSessions: [{id: sessionUUID}]});
-                assertAlways.commandWorked(res);
-            } catch (e) {
-                if (e.code == ErrorCodes.Interrupted || e.code == ErrorCodes.CursorKilled ||
-                    e.code == ErrorCodes.CursorNotFound) {
-                    // This session was killed when running either listSessions or killSesssions.
-                    // We should retry.
-                    ourSessionWasKilled = true;
-                    continue;
-                }
-
-                throw e;
-            }
-        } while (ourSessionWasKilled);
-    };
+    $config.states.killSession = killSession;
 
     $config.transitions = {
         init: {update: 0.9, checkConsistency: 0.1},
diff --git a/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js b/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js
index 17acf7c4824..27c3cef5181 100644
--- a/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js
@@ -19,13 +19,13 @@ var $config = (function() {
 
     var states = (function() {
 
-        function getAllDocuments(session, collection, numDocs) {
+        function getAllDocuments(session, collection, numDocs, txnHelperOptions) {
             let documents;
             withTxnAndAutoRetry(session, () => {
                 documents = collection.find().toArray();
 
                 assertWhenOwnColl.eq(numDocs, documents.length, () => tojson(documents));
-            });
+            }, txnHelperOptions);
             return documents;
         }
 
@@ -35,7 +35,9 @@ var $config = (function() {
 
         function checkMoneyBalance(db, collName) {
             const collection = this.session.getDatabase(db.getName()).getCollection(collName);
-            const documents = getAllDocuments(this.session, collection, this.numAccounts);
+            const documents = getAllDocuments(this.session, collection, this.numAccounts, {
+                retryOnKilledSession: this.retryOnKilledSession
+            });
             assertWhenOwnColl.eq(this.numAccounts * this.initialValue,
                                  computeTotalOfAllBalances(documents),
                                  () => tojson(documents));
@@ -69,7 +71,7 @@ var $config = (function() {
                 assertAlways.commandWorked(res);
                 assertWhenOwnColl.eq(res.n, 1, () => tojson(res));
                 assertWhenOwnColl.eq(res.nModified, 1, () => tojson(res));
-            });
+            }, {retryOnKilledSession: this.retryOnKilledSession});
         }
 
         return {init: init, transferMoney: transferMoney, checkMoneyBalance: checkMoneyBalance};
@@ -113,7 +115,7 @@ var $config = (function() {
         startState: 'init',
         states: states,
         transitions: transitions,
-        data: {numAccounts: 20, initialValue: 2000},
+        data: {numAccounts: 20, initialValue: 2000, retryOnKilledSession: false},
         setup: setup,
         teardown: teardown
     };
diff --git a/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple_kill_sessions.js b/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple_kill_sessions.js
new file mode 100644
index 00000000000..359562be59b
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_simple_kill_sessions.js
@@ -0,0 +1,28 @@
+'use strict';
+
+/**
+ * Tests periodically killing sessions that are running transactions. The base workload runs
+ * transactions with two writes, which will require two phase commit in a sharded cluster if each
+ * write targets a different shard.
+ *
+ * @tags: [uses_transactions, assumes_snapshot_transactions]
+ */
+
+load('jstests/concurrency/fsm_libs/extend_workload.js');           // for extendWorkload
+load('jstests/concurrency/fsm_workload_helpers/kill_session.js');  // for killSession
+load('jstests/concurrency/fsm_workloads/multi_statement_transaction_simple.js');  // for $config
+
+var $config = extendWorkload($config, ($config, $super) => {
+    $config.data.retryOnKilledSession = true;
+
+    $config.states.killSession = killSession;
+
+    $config.transitions = {
+        init: {transferMoney: 1},
+        transferMoney: {transferMoney: 0.8, checkMoneyBalance: 0.1, killSession: 0.1},
+        checkMoneyBalance: {transferMoney: 0.9, killSession: 0.1},
+        killSession: {transferMoney: 0.8, checkMoneyBalance: 0.1, killSession: 0.1}
+    };
+
+    return $config;
+});