SERVER-35817 Allow shutdowns while prepared transactions are in-flight

6 files changed, 76 insertions, 4 deletions

diff --git a/jstests/replsets/shutdown_with_prepared_transaction.js b/jstests/replsets/shutdown_with_prepared_transaction.js
new file mode 100644
index 00000000000..387d7474ef3
--- /dev/null
+++ b/jstests/replsets/shutdown_with_prepared_transaction.js
@@ -0,0 +1,38 @@
+/**
+ * Tests that a server can still be shut down while it has prepared transactions pending.
+ *
+ * @tags: [uses_transactions]
+ */
+(function() {
+    "use strict";
+    load("jstests/core/txns/libs/prepare_helpers.js");
+
+    const replTest = new ReplSetTest({nodes: 1});
+    replTest.startSet();
+    replTest.initiate();
+
+    const conn = replTest.getPrimary();
+
+    const dbName = "test";
+    const collName = "shutdown_with_prepared_txn";
+    const testDB = conn.getDB(dbName);
+    const testColl = testDB.getCollection(collName);
+
+    assert.commandWorked(testDB.runCommand({create: collName, writeConcern: {w: "majority"}}));
+
+    const session = conn.startSession({causalConsistency: false});
+    const sessionDB = session.getDatabase(dbName);
+    const sessionColl = sessionDB.getCollection(collName);
+
+    jsTestLog("Starting a simple transaction and putting it into prepare");
+
+    session.startTransaction();
+    assert.commandWorked(sessionColl.insert({_id: 1}));
+
+    PrepareHelpers.prepareTransaction(session);
+
+    jsTestLog("Shutting down the set with the transaction still in prepare state");
+    // Skip validation during ReplSetTest cleanup since validate() will block behind the prepared
+    // transaction's locks when trying to take a collection X lock.
+    replTest.stopSet(null /*signal*/, false /*forRestart*/, {skipValidation: true});
+}());
diff --git a/src/mongo/db/db.cpp b/src/mongo/db/db.cpp
index cc570bcdae5..abe28452d2e 100644
--- a/src/mongo/db/db.cpp
+++ b/src/mongo/db/db.cpp
@@ -892,9 +892,7 @@ void shutdownTask() {
         ShardingInitializationMongoD::get(serviceContext)->shutDown(opCtx);
 
         // Destroy all stashed transaction resources, in order to release locks.
-        SessionKiller::Matcher matcherAllSessions(
-            KillAllSessionsByPatternSet{makeKillAllSessionsByPattern(opCtx)});
-        killSessionsLocalKillTransactions(opCtx, matcherAllSessions);
+        killSessionsLocalShutdownAllTransactions(opCtx);
     }
 
     serviceContext->setKillAllOperations();
diff --git a/src/mongo/db/kill_sessions_local.cpp b/src/mongo/db/kill_sessions_local.cpp
index 0c2b0cf9641..0406a818284 100644
--- a/src/mongo/db/kill_sessions_local.cpp
+++ b/src/mongo/db/kill_sessions_local.cpp
@@ -88,4 +88,13 @@ void killAllExpiredTransactions(OperationContext* opCtx) {
         });
 }
 
+void killSessionsLocalShutdownAllTransactions(OperationContext* opCtx) {
+    SessionKiller::Matcher matcherAllSessions(
+        KillAllSessionsByPatternSet{makeKillAllSessionsByPattern(opCtx)});
+    SessionCatalog::get(opCtx)->scanSessions(
+        opCtx, matcherAllSessions, [](OperationContext* opCtx, Session* session) {
+            TransactionParticipant::getFromNonCheckedOutSession(session)->shutdown();
+        });
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/kill_sessions_local.h b/src/mongo/db/kill_sessions_local.h
index a21bcec4ae5..ea8f442b325 100644
--- a/src/mongo/db/kill_sessions_local.h
+++ b/src/mongo/db/kill_sessions_local.h
@@ -53,4 +53,9 @@ void killSessionsLocalKillTransactions(OperationContext* opCtx,
  */
 void killAllExpiredTransactions(OperationContext* opCtx);
 
+/**
+ * Run during shutdown to kill all in-progress transactions, including those in prepare.
+ */
+void killSessionsLocalShutdownAllTransactions(OperationContext* opCtx);
+
 }  // namespace mongo
diff --git a/src/mongo/db/transaction_participant.cpp b/src/mongo/db/transaction_participant.cpp
index 33eb0d866ae..1d9fd5e4321 100644
--- a/src/mongo/db/transaction_participant.cpp
+++ b/src/mongo/db/transaction_participant.cpp
@@ -398,8 +398,11 @@ TransactionParticipant::SideTransactionBlock::~SideTransactionBlock() {
     }
 }
 void TransactionParticipant::_stashActiveTransaction(WithLock, OperationContext* opCtx) {
-    invariant(_activeTxnNumber == opCtx->getTxnNumber());
+    if (_inShutdown) {
+        return;
+    }
 
+    invariant(_activeTxnNumber == opCtx->getTxnNumber());
     {
         stdx::lock_guard<stdx::mutex> lm(_metricsMutex);
         _transactionMetricsObserver.onStash(ServerTransactionsMetrics::get(opCtx),
@@ -746,6 +749,13 @@ void TransactionParticipant::_commitTransaction(stdx::unique_lock<stdx::mutex> l
     _cleanUpTxnResourceOnOpCtx(lk, opCtx, TransactionState::kCommitted);
 }
 
+void TransactionParticipant::shutdown() {
+    stdx::lock_guard<stdx::mutex> lock(_mutex);
+
+    _inShutdown = true;
+    _txnResourceStash = boost::none;
+}
+
 void TransactionParticipant::abortArbitraryTransaction() {
     stdx::lock_guard<stdx::mutex> lock(_mutex);
 
diff --git a/src/mongo/db/transaction_participant.h b/src/mongo/db/transaction_participant.h
index 6946c1a24c0..a32b8d095d0 100644
--- a/src/mongo/db/transaction_participant.h
+++ b/src/mongo/db/transaction_participant.h
@@ -142,6 +142,16 @@ public:
     static TransactionParticipant* getFromNonCheckedOutSession(Session* session);
 
     /**
+     * Kills the transaction if it is running, ensuring that it releases all resources, even if the
+     * transaction is in prepare().  Avoids writing any oplog entries or making any changes to the
+     * transaction table.  State for prepared transactions will be re-constituted at startup.
+     * Note that we don't take any active steps to prevent continued use of this
+     * TransactionParticipant after shutdown() is called, but we rely on callers to not
+     * continue using the TransactionParticipant once we are in shutdown.
+     */
+    void shutdown();
+
+    /**
      * Called for speculative transactions to fix the optime of the snapshot to read from.
      */
     void setSpeculativeTransactionOpTimeToLastApplied(OperationContext* opCtx);
@@ -552,6 +562,8 @@ private:
     // Protects the member variables below.
     mutable stdx::mutex _mutex;
 
+    bool _inShutdown{false};
+
     // Holds transaction resources between network operations.
     boost::optional<TxnResources> _txnResourceStash;