SERVER-68215 add MongoDSessionCatalogTransactionInterface::refreshLocksForPreparedTransaction()

4 files changed, 26 insertions, 11 deletions

diff --git a/src/mongo/db/session/session_catalog_mongod.cpp b/src/mongo/db/session/session_catalog_mongod.cpp
index b18b21190be..d02ed6d853d 100644
--- a/src/mongo/db/session/session_catalog_mongod.cpp
+++ b/src/mongo/db/session/session_catalog_mongod.cpp
@@ -551,16 +551,7 @@ void MongoDSessionCatalog::onStepUp(OperationContext* opCtx) {
             //   lock in the IS mode prior to reading the config.transactions collection but it
             //   cannot do that while the RSTL lock is being held by 'opCtx'.
             auto ocs = checkOutSessionWithoutRefresh(newOpCtx.get());
-            auto txnParticipant = TransactionParticipant::get(newOpCtx.get());
-            LOGV2_DEBUG(21979,
-                        3,
-                        "Restoring locks of prepared transaction. SessionId: {sessionId} "
-                        "TxnNumberAndRetryCounter: {txnNumberAndRetryCounter}",
-                        "Restoring locks of prepared transaction",
-                        "sessionId"_attr = sessionInfo.getSessionId()->getId(),
-                        "txnNumberAndRetryCounter"_attr =
-                            txnParticipant.getActiveTxnNumberAndRetryCounter());
-            txnParticipant.refreshLocksForPreparedTransaction(newOpCtx.get(), false);
+            _ti->refreshLocksForPreparedTransaction(newOpCtx.get(), sessionInfo);
         }
     }
 
diff --git a/src/mongo/db/session/session_catalog_mongod_transaction_interface.h b/src/mongo/db/session/session_catalog_mongod_transaction_interface.h
index 95b70e4bec6..aff43d5b4c7 100644
--- a/src/mongo/db/session/session_catalog_mongod_transaction_interface.h
+++ b/src/mongo/db/session/session_catalog_mongod_transaction_interface.h
@@ -30,7 +30,8 @@
 #pragma once
 
 #include "mongo/db/operation_context.h"
-#include "mongo/db/session/logical_session_id.h"  // for TxnNumberAndRetryCounter
+#include "mongo/db/session/logical_session_id.h"      // for TxnNumberAndRetryCounter
+#include "mongo/db/session/logical_session_id_gen.h"  // for OperationSessionInfo
 #include "mongo/db/session/session_catalog.h"  // for ObservableSession and ScanSessionsCallbackFn
 #include "mongo/db/session/session_txn_record_gen.h"       // for SessionTxnRecord
 #include "mongo/db/transaction/transaction_participant.h"  // for SessionToKill
@@ -89,6 +90,12 @@ public:
     virtual void abortTransaction(OperationContext* opCtx, const SessionTxnRecord& txnRecord) = 0;
 
     /**
+     * Yield or reacquire locks for prepared transactions, used on replication state transition.
+     */
+    virtual void refreshLocksForPreparedTransaction(OperationContext* opCtx,
+                                                    const OperationSessionInfo& sessionInfo) = 0;
+
+    /**
      * Marks the session as requiring refresh. Used when the session state has been modified
      * externally, such as through a direct write to the transactions table.
      */
diff --git a/src/mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.cpp b/src/mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.cpp
index 5fc2e6f2ce6..1706c932098 100644
--- a/src/mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.cpp
+++ b/src/mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.cpp
@@ -81,6 +81,20 @@ void MongoDSessionCatalogTransactionInterfaceImpl::abortTransaction(
     opCtx->resetMultiDocumentTransactionState();
 }
 
+void MongoDSessionCatalogTransactionInterfaceImpl::refreshLocksForPreparedTransaction(
+    OperationContext* opCtx, const OperationSessionInfo& sessionInfo) {
+    auto txnParticipant = TransactionParticipant::get(opCtx);
+    LOGV2_DEBUG(21979,
+                3,
+                "Restoring locks of prepared transaction. SessionId: {sessionId} "
+                "TxnNumberAndRetryCounter: {txnNumberAndRetryCounter}",
+                "Restoring locks of prepared transaction",
+                "sessionId"_attr = sessionInfo.getSessionId()->getId(),
+                "txnNumberAndRetryCounter"_attr =
+                    txnParticipant.getActiveTxnNumberAndRetryCounter());
+    txnParticipant.refreshLocksForPreparedTransaction(opCtx, /*yieldLocks=*/false);
+}
+
 void MongoDSessionCatalogTransactionInterfaceImpl::invalidateSessionToKill(
     OperationContext* opCtx, const SessionToKill& session) {
     auto participant = TransactionParticipant::get(session);
diff --git a/src/mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.h b/src/mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.h
index c5d43d1ffd1..4b226c51703 100644
--- a/src/mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.h
+++ b/src/mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.h
@@ -60,6 +60,9 @@ public:
 
     void abortTransaction(OperationContext* opCtx, const SessionTxnRecord& txnRecord) override;
 
+    void refreshLocksForPreparedTransaction(OperationContext* opCtx,
+                                            const OperationSessionInfo& sessionInfo) override;
+
     void invalidateSessionToKill(OperationContext* opCtx, const SessionToKill& session) override;
 
     ScanSessionsCallbackFn makeParentSessionWorkerFnForReap(