SERVER-41838 Release RSTL if we choose not to commit a prepared transaction after reacquiring it

(cherry picked from commit 23bd9dc2cae50a9c1ae47fed7805db5da4f9583e)

3 files changed, 147 insertions, 14 deletions

diff --git a/src/mongo/db/transaction_participant.cpp b/src/mongo/db/transaction_participant.cpp
index 3898b3cf474..b6f36e7d7ab 100644
--- a/src/mongo/db/transaction_participant.cpp
+++ b/src/mongo/db/transaction_participant.cpp
@@ -1222,11 +1222,26 @@ void TransactionParticipant::Participant::commitPreparedTransaction(
     OperationContext* opCtx,
     Timestamp commitTimestamp,
     boost::optional<repl::OpTime> commitOplogEntryOpTime) {
+    // A correctly functioning coordinator could hit this uassert. This could happen if this
+    // participant shard failed over and the new primary majority committed prepare without this
+    // node in its majority. The coordinator could legally send commitTransaction with a
+    // commitTimestamp to this shard but target the old primary (this node) that has yet to prepare
+    // the transaction. We uassert since this node cannot commit the transaction.
+    if (!o().txnState.isPrepared()) {
+        uasserted(ErrorCodes::InvalidOptions,
+                  "commitTransaction cannot provide commitTimestamp to unprepared transaction.");
+    }
+
     // Re-acquire the RSTL to prevent state transitions while committing the transaction. When the
     // transaction was prepared, we dropped the RSTL. We do not need to reacquire the PBWM because
     // if we're not the primary we will uassert anyways.
     repl::ReplicationStateTransitionLockGuard rstl(opCtx, MODE_IX);
 
+    // Prepared transactions cannot hold the RSTL, or else they will deadlock with state
+    // transitions. If we do not commit the transaction we must unlock the RSTL explicitly so two
+    // phase locking doesn't hold onto it.
+    auto unlockGuard = makeGuard([&] { invariant(opCtx->lockState()->unlockRSTLforPrepare()); });
+
     const auto replCoord = repl::ReplicationCoordinator::get(opCtx);
 
     if (opCtx->writesAreReplicated()) {
@@ -1235,14 +1250,6 @@ void TransactionParticipant::Participant::commitPreparedTransaction(
                 replCoord->canAcceptWritesForDatabase(opCtx, "admin"));
     }
 
-    // A correctly functioning coordinator could hit this uassert. This could happen if this
-    // participant shard failed over and the new primary majority committed prepare without this
-    // node in its majority. The coordinator could legally send commitTransaction with a
-    // commitTimestamp to this shard but target the old primary (this node) that has yet to prepare
-    // the transaction. We uassert since this node cannot commit the transaction.
-    uassert(ErrorCodes::InvalidOptions,
-            "commitTransaction cannot provide commitTimestamp to unprepared transaction.",
-            o().txnState.isPrepared());
     uassert(
         ErrorCodes::InvalidOptions, "'commitTimestamp' cannot be null", !commitTimestamp.isNull());
 
@@ -1274,6 +1281,9 @@ void TransactionParticipant::Participant::commitPreparedTransaction(
     }
 
     try {
+        // We can no longer uassert without terminating.
+        unlockGuard.dismiss();
+
         // Once entering "committing with prepare" we cannot throw an exception.
         UninterruptibleLockGuard noInterrupt(opCtx->lockState());
         opCtx->recoveryUnit()->setCommitTimestamp(commitTimestamp);
@@ -1393,20 +1403,33 @@ bool TransactionParticipant::Observer::expiredAsOf(Date_t when) const {
 }
 
 void TransactionParticipant::Participant::abortActiveTransaction(OperationContext* opCtx) {
-    // Re-acquire the RSTL to prevent state transitions while aborting the transaction. If the
-    // transaction was prepared then we dropped it on preparing the transaction. We do not need to
+    if (o().txnState.isPrepared()) {
+        _abortActivePreparedTransaction(opCtx);
+    } else {
+        _abortActiveTransaction(opCtx, TransactionState::kInProgress, false);
+    }
+}
+
+void TransactionParticipant::Participant::_abortActivePreparedTransaction(OperationContext* opCtx) {
+    // Re-acquire the RSTL to prevent state transitions while aborting the transaction. Since the
+    // transaction was prepared, we dropped it on preparing the transaction. We do not need to
     // reacquire the PBWM because if we're not the primary we will uassert anyways.
     repl::ReplicationStateTransitionLockGuard rstl(opCtx, MODE_IX);
-    if (o().txnState.isPrepared() && opCtx->writesAreReplicated()) {
+
+    // Prepared transactions cannot hold the RSTL, or else they will deadlock with state
+    // transitions. If we do not abort the transaction we must unlock the RSTL explicitly so two
+    // phase locking doesn't hold onto it. Unlocking the RSTL may be a noop if it's already
+    // unlocked.
+    ON_BLOCK_EXIT([&] { opCtx->lockState()->unlockRSTLforPrepare(); });
+
+    if (opCtx->writesAreReplicated()) {
         auto replCoord = repl::ReplicationCoordinator::get(opCtx);
         uassert(ErrorCodes::NotMaster,
                 "Not primary so we cannot abort a prepared transaction",
                 replCoord->canAcceptWritesForDatabase(opCtx, "admin"));
     }
 
-    _abortActiveTransaction(opCtx,
-                            TransactionState::kInProgress | TransactionState::kPrepared,
-                            o().txnState.isPrepared());
+    _abortActiveTransaction(opCtx, TransactionState::kPrepared, true);
 }
 
 void TransactionParticipant::Participant::abortActiveUnpreparedOrStashPreparedTransaction(
diff --git a/src/mongo/db/transaction_participant.h b/src/mongo/db/transaction_participant.h
index dd2828911d7..2120c56c3cd 100644
--- a/src/mongo/db/transaction_participant.h
+++ b/src/mongo/db/transaction_participant.h
@@ -713,6 +713,9 @@ public:
                                      TransactionState::StateSet expectedStates,
                                      bool writeOplog);
 
+        // Aborts a prepared transaction.
+        void _abortActivePreparedTransaction(OperationContext* opCtx);
+
         // Releases stashed transaction resources to abort the transaction on the session.
         void _abortTransactionOnSession(OperationContext* opCtx);
 
diff --git a/src/mongo/db/transaction_participant_test.cpp b/src/mongo/db/transaction_participant_test.cpp
index 82f68661e2d..0841deb1318 100644
--- a/src/mongo/db/transaction_participant_test.cpp
+++ b/src/mongo/db/transaction_participant_test.cpp
@@ -32,6 +32,7 @@
 #include <boost/optional/optional_io.hpp>
 
 #include "mongo/db/client.h"
+#include "mongo/db/concurrency/replication_state_transition_lock_guard.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/dbdirectclient.h"
 #include "mongo/db/op_observer_noop.h"
@@ -838,6 +839,112 @@ TEST_F(TxnParticipantTest, StepDownDuringPreparedCommitFails) {
                        ErrorCodes::NotMaster);
 }
 
+TEST_F(TxnParticipantTest, StepDownDuringPreparedAbortReleasesRSTL) {
+    auto sessionCheckout = checkOutSession();
+    auto txnParticipant = TransactionParticipant::get(opCtx());
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+
+    txnParticipant.unstashTransactionResources(opCtx(), "insert");
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock), MODE_IX);
+
+    // Simulate the locking of an insert.
+    {
+        Lock::DBLock dbLock(opCtx(), "test", MODE_IX);
+        Lock::CollectionLock collLock(opCtx(), NamespaceString("test.foo"), MODE_IX);
+    }
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock), MODE_IX);
+    txnParticipant.stashTransactionResources(opCtx());
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+
+    txnParticipant.unstashTransactionResources(opCtx(), "prepareTransaction");
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock), MODE_IX);
+    txnParticipant.prepareTransaction(opCtx(), {});
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+    txnParticipant.stashTransactionResources(opCtx());
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+
+    txnParticipant.unstashTransactionResources(opCtx(), "abortTransaction");
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+    ASSERT_OK(repl::ReplicationCoordinator::get(opCtx())->setFollowerMode(
+        repl::MemberState::RS_SECONDARY));
+    ASSERT_THROWS_CODE(
+        txnParticipant.abortActiveTransaction(opCtx()), AssertionException, ErrorCodes::NotMaster);
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+
+    // Test that we can acquire the RSTL in mode X, and then immediately release it so the test can
+    // complete successfully.
+    auto func = [&](OperationContext* newOpCtx) {
+        newOpCtx->lockState()->lock(newOpCtx, resourceIdReplicationStateTransitionLock, MODE_X);
+        newOpCtx->lockState()->unlock(resourceIdReplicationStateTransitionLock);
+    };
+    runFunctionFromDifferentOpCtx(func);
+}
+
+TEST_F(TxnParticipantTest, StepDownDuringPreparedCommitReleasesRSTL) {
+    auto sessionCheckout = checkOutSession();
+    auto txnParticipant = TransactionParticipant::get(opCtx());
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+
+    txnParticipant.unstashTransactionResources(opCtx(), "insert");
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock), MODE_IX);
+
+    // Simulate the locking of an insert.
+    {
+        Lock::DBLock dbLock(opCtx(), "test", MODE_IX);
+        Lock::CollectionLock collLock(opCtx(), NamespaceString("test.foo"), MODE_IX);
+    }
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock), MODE_IX);
+    txnParticipant.stashTransactionResources(opCtx());
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+
+    txnParticipant.unstashTransactionResources(opCtx(), "prepareTransaction");
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock), MODE_IX);
+    auto prepareTimestamp = txnParticipant.prepareTransaction(opCtx(), {});
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+    txnParticipant.stashTransactionResources(opCtx());
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+
+    txnParticipant.unstashTransactionResources(opCtx(), "commitTransaction");
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+    ASSERT_OK(repl::ReplicationCoordinator::get(opCtx())->setFollowerMode(
+        repl::MemberState::RS_SECONDARY));
+    ASSERT_THROWS_CODE(
+        txnParticipant.commitPreparedTransaction(opCtx(), prepareTimestamp, boost::none),
+        AssertionException,
+        ErrorCodes::NotMaster);
+
+    ASSERT_EQ(opCtx()->lockState()->getLockMode(resourceIdReplicationStateTransitionLock),
+              MODE_NONE);
+
+    // Test that we can acquire the RSTL in mode X, and then immediately release it so the test can
+    // complete successfully.
+    auto func = [&](OperationContext* newOpCtx) {
+        newOpCtx->lockState()->lock(newOpCtx, resourceIdReplicationStateTransitionLock, MODE_X);
+        newOpCtx->lockState()->unlock(resourceIdReplicationStateTransitionLock);
+    };
+    runFunctionFromDifferentOpCtx(func);
+}
+
 TEST_F(TxnParticipantTest, ThrowDuringUnpreparedCommitLetsTheAbortAtEntryPointToCleanUp) {
     auto sessionCheckout = checkOutSession();
     auto txnParticipant = TransactionParticipant::get(opCtx());