SERVER-36305 Add KillSessions stage to transactions FSM workloads

5 files changed, 79 insertions, 6 deletions

diff --git a/buildscripts/resmokeconfig/suites/concurrency_simultaneous_replication.yml b/buildscripts/resmokeconfig/suites/concurrency_simultaneous_replication.yml
index f16db546d66..812b9dece09 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_simultaneous_replication.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_simultaneous_replication.yml
@@ -18,6 +18,10 @@ selector:
   # This workload kills random cursors which takes a collection lock.
   # TODO: SERVER-35567.
   - jstests/concurrency/fsm_workloads/snapshot_read_kill_operations.js
+
+  # This workload kills random sessions and a different FSM workload wouldn't be able to handle 
+  # the error response from the op being killed.
+  - jstests/concurrency/fsm_workloads/multi_statement_transaction_kill_sessions_atomicity_isolation.js
   exclude_with_any_tags:
   - requires_sharding
 
diff --git a/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js b/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
index 6aa821966a3..402f59cc124 100644
--- a/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
+++ b/jstests/concurrency/fsm_workload_helpers/auto_retry_transaction.js
@@ -29,8 +29,10 @@ var {withTxnAndAutoRetry} = (function() {
      * transaction started by the withTxnAndAutoRetry() function is only known to have committed
      * after the withTxnAndAutoRetry() function returns.
      */
-    function withTxnAndAutoRetry(
-        session, func, {txnOptions: txnOptions = {readConcern: {level: 'snapshot'}}} = {}) {
+    function withTxnAndAutoRetry(session, func, {
+        txnOptions: txnOptions = {readConcern: {level: 'snapshot'}},
+        retryOnKilledSession: retryOnKilledSession = false
+    } = {}) {
         let hasTransientError;
 
         do {
@@ -49,6 +51,7 @@ var {withTxnAndAutoRetry} = (function() {
                     hasCommitTxnError = true;
                     throw e;
                 }
+
             } catch (e) {
                 if (!hasCommitTxnError) {
                     // Use the version of abortTransaction() that ignores errors. We ignore the
@@ -61,8 +64,10 @@ var {withTxnAndAutoRetry} = (function() {
                     session.abortTransaction();
                 }
 
-                if (e.hasOwnProperty('errorLabels') &&
-                    e.errorLabels.includes('TransientTransactionError')) {
+                if ((e.hasOwnProperty('errorLabels') &&
+                     e.errorLabels.includes('TransientTransactionError')) ||
+                    (retryOnKilledSession &&
+                     (e.code === ErrorCodes.Interrupted || e.code === ErrorCodes.CursorKilled))) {
                     hasTransientError = true;
                     continue;
                 }
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 bcbda5a5107..e044a5b2760 100644
--- a/jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation.js
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation.js
@@ -118,7 +118,7 @@ var $config = (function() {
                     const batchSize = Math.max(2, Math.floor(numDocs / 5));
                     allDocuments.push(...txnCollection.find().batchSize(batchSize).toArray());
                 }
-            });
+            }, {retryOnKilledSession: this.retryOnKilledSession});
         } else {
             for (let collection of this.collections) {
                 allDocuments.push(...collection.find().toArray());
@@ -176,16 +176,18 @@ var $config = (function() {
                                 }
                             }
                         };
+
                         const txnCollection =
                             this.session.getDatabase(txnDbName).getCollection(txnCollName);
                         const res = txnCollection.runCommand('update', {
                             updates: [{q: {_id: docId}, u: updateMods}],
                         });
+
                         assertAlways.commandWorked(res);
                         assertWhenOwnColl.eq(res.n, 1, () => tojson(res));
                         assertWhenOwnColl.eq(res.nModified, 1, () => tojson(res));
                     }
-                });
+                }, {retryOnKilledSession: this.retryOnKilledSession});
 
                 ++this.iteration;
             },
@@ -236,6 +238,7 @@ var $config = (function() {
             getAllDocuments,
             getDocIdsToUpdate,
             numDocs: 10,
+            retryOnKilledSession: false,
         },
         setup: setup,
         teardown: teardown,
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
new file mode 100644
index 00000000000..3747f1d31bb
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/multi_statement_transaction_kill_sessions_atomicity_isolation.js
@@ -0,0 +1,60 @@
+'use strict';
+
+/**
+ * Tests periodically killing sessions that are running transactions.
+ *
+ * @tags: [uses_transactions]
+ */
+
+load('jstests/concurrency/fsm_libs/extend_workload.js');  // for extendWorkload
+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 {
+                    assertAlways.commandFailedWithCode(res, ErrorCodes.DuplicateKey);
+                }
+
+                const sessionToKill = db.getSiblingDB("config").system.sessions.aggregate([
+                    {$listSessions: {}},
+                    {$match: {id: {$ne: this.session.getSessionId()}}},
+                    {$sample: {size: 1}},
+                ]);
+
+                assertAlways.eq(sessionToKill.toArray().length, 1);
+                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) {
+                    // This session was killed when running either listSessions or killSesssions.
+                    // We should retry.
+                    ourSessionWasKilled = true;
+                    continue;
+                }
+
+                throw e;
+            }
+        } while (ourSessionWasKilled);
+    };
+
+    $config.transitions = {
+        init: {update: 0.9, checkConsistency: 0.1},
+        update: {update: 0.8, checkConsistency: 0.1, killSession: 0.1},
+        checkConsistency: {update: 0.9, killSession: 0.1},
+        killSession: {update: 0.9, checkConsistency: 0.1}
+    };
+
+    return $config;
+});
diff --git a/jstests/noPassthrough/libs/backup_restore.js b/jstests/noPassthrough/libs/backup_restore.js
index 27e34fc158f..c4797fcf73b 100644
--- a/jstests/noPassthrough/libs/backup_restore.js
+++ b/jstests/noPassthrough/libs/backup_restore.js
@@ -145,6 +145,7 @@ var BackupRestoreTest = function(options) {
                 'multi_statement_transaction_atomicity_isolation.js',
                 'multi_statement_transaction_atomicity_isolation_multi_db.js',
                 'multi_statement_transaction_atomicity_isolation_repeated_reads.js',
+                'multi_statement_transaction_kill_sessions_atomicity_isolation.js',
                 'multi_statement_transaction_simple.js',
                 'multi_statement_transaction_simple_repeated_reads.js',
                 'reindex_background.js',