SERVER-68215 add MongoDSessionCatalogTransactionInterface::invalidateSessionToKill()

4 files changed, 33 insertions, 9 deletions

diff --git a/src/mongo/db/session/session_catalog_mongod.cpp b/src/mongo/db/session/session_catalog_mongod.cpp
index 69b52d9f533..94b5584a146 100644
--- a/src/mongo/db/session/session_catalog_mongod.cpp
+++ b/src/mongo/db/session/session_catalog_mongod.cpp
@@ -99,12 +99,14 @@ auto getThreadPool(OperationContext* opCtx) {
  * set of kill tokens.
  */
 void killSessionTokens(OperationContext* opCtx,
+                       MongoDSessionCatalogTransactionInterface* ti,
                        std::vector<SessionCatalog::KillToken> sessionKillTokens) {
     if (sessionKillTokens.empty())
         return;
 
     getThreadPool(opCtx)->schedule(
         [service = opCtx->getServiceContext(),
+         ti,
          sessionKillTokens = std::move(sessionKillTokens)](auto status) mutable {
             invariant(status);
 
@@ -115,8 +117,7 @@ void killSessionTokens(OperationContext* opCtx,
 
             for (auto& sessionKillToken : sessionKillTokens) {
                 auto session = catalog->checkOutSessionForKill(opCtx, std::move(sessionKillToken));
-                auto participant = TransactionParticipant::get(session);
-                participant.invalidate(opCtx);
+                ti->invalidateSessionToKill(opCtx, session);
             }
         });
 }
@@ -569,7 +570,7 @@ void MongoDSessionCatalog::onStepUp(OperationContext* opCtx) {
             sessionsToReacquireLocks.emplace_back(sessionInfo);
         }
     });
-    killSessionTokens(opCtx, std::move(sessionKillTokens));
+    killSessionTokens(opCtx, _ti.get(), std::move(sessionKillTokens));
 
     {
         // Create a new opCtx because we need an empty locker to refresh the locks.
@@ -640,8 +641,9 @@ void MongoDSessionCatalog::observeDirectWriteToConfigTransactions(OperationConte
     class KillSessionTokenOnCommit : public RecoveryUnit::Change {
     public:
         KillSessionTokenOnCommit(OperationContext* opCtx,
+                                 MongoDSessionCatalogTransactionInterface* ti,
                                  SessionCatalog::KillToken sessionKillToken)
-            : _opCtx(opCtx), _sessionKillToken(std::move(sessionKillToken)) {}
+            : _opCtx(opCtx), _ti(ti), _sessionKillToken(std::move(sessionKillToken)) {}
 
         void commit(boost::optional<Timestamp>) override {
             rollback();
@@ -650,11 +652,12 @@ void MongoDSessionCatalog::observeDirectWriteToConfigTransactions(OperationConte
         void rollback() override {
             std::vector<SessionCatalog::KillToken> sessionKillTokenVec;
             sessionKillTokenVec.emplace_back(std::move(_sessionKillToken));
-            killSessionTokens(_opCtx, std::move(sessionKillTokenVec));
+            killSessionTokens(_opCtx, _ti, std::move(sessionKillTokenVec));
         }
 
     private:
         OperationContext* _opCtx;
+        MongoDSessionCatalogTransactionInterface* _ti;
         SessionCatalog::KillToken _sessionKillToken;
     };
 
@@ -662,7 +665,7 @@ void MongoDSessionCatalog::observeDirectWriteToConfigTransactions(OperationConte
 
     const auto lsid =
         LogicalSessionId::parse(IDLParserContext("lsid"), singleSessionDoc["_id"].Obj());
-    catalog->scanSession(lsid, [&](const ObservableSession& session) {
+    catalog->scanSession(lsid, [&, ti = _ti.get()](const ObservableSession& session) {
         const auto participant = TransactionParticipant::get(session);
         uassert(ErrorCodes::PreparedTransactionInProgress,
                 str::stream() << "Cannot modify the entry for session "
@@ -671,7 +674,7 @@ void MongoDSessionCatalog::observeDirectWriteToConfigTransactions(OperationConte
                 !participant.transactionIsPrepared());
 
         opCtx->recoveryUnit()->registerChange(
-            std::make_unique<KillSessionTokenOnCommit>(opCtx, session.kill()));
+            std::make_unique<KillSessionTokenOnCommit>(opCtx, ti, session.kill()));
     });
 }
 
@@ -688,7 +691,7 @@ void MongoDSessionCatalog::invalidateAllSessions(OperationContext* opCtx) {
         sessionKillTokens.emplace_back(session.kill());
     });
 
-    killSessionTokens(opCtx, std::move(sessionKillTokens));
+    killSessionTokens(opCtx, _ti.get(), std::move(sessionKillTokens));
 }
 
 int MongoDSessionCatalog::reapSessionsOlderThan(OperationContext* opCtx,
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 4a6af53d049..2cde4d52750 100644
--- a/src/mongo/db/session/session_catalog_mongod_transaction_interface.h
+++ b/src/mongo/db/session/session_catalog_mongod_transaction_interface.h
@@ -29,6 +29,9 @@
 
 #pragma once
 
+#include "mongo/db/operation_context.h"
+#include "mongo/db/transaction/transaction_participant.h"  // for SessionToKill
+
 namespace mongo {
 
 /**
@@ -39,6 +42,12 @@ namespace mongo {
 class MongoDSessionCatalogTransactionInterface {
 public:
     virtual ~MongoDSessionCatalogTransactionInterface() = default;
+
+    /**
+     * 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.
+     */
+    virtual void invalidateSessionToKill(OperationContext* opCtx, const SessionToKill& session) = 0;
 };
 
 }  // namespace mongo
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 a1f9ef06b89..3c570f34a48 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
@@ -29,4 +29,14 @@
 
 #include "mongo/db/transaction/session_catalog_mongod_transaction_interface_impl.h"
 
-namespace mongo {}  // namespace mongo
+#include "mongo/db/transaction/transaction_participant.h"
+
+namespace mongo {
+
+void MongoDSessionCatalogTransactionInterfaceImpl::invalidateSessionToKill(
+    OperationContext* opCtx, const SessionToKill& session) {
+    auto participant = TransactionParticipant::get(session);
+    participant.invalidate(opCtx);
+}
+
+}  // namespace mongo
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 b29427ba5c5..6d7751cb885 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
@@ -46,6 +46,8 @@ class MongoDSessionCatalogTransactionInterfaceImpl
 public:
     MongoDSessionCatalogTransactionInterfaceImpl() = default;
     virtual ~MongoDSessionCatalogTransactionInterfaceImpl() = default;
+
+    void invalidateSessionToKill(OperationContext* opCtx, const SessionToKill& session) override;
 };
 
 }  // namespace mongo