SERVER-68215 check out mongod sessions through MongoDSessionCatalog

2 files changed, 157 insertions, 19 deletions

diff --git a/src/mongo/db/session/session_catalog_mongod.cpp b/src/mongo/db/session/session_catalog_mongod.cpp
index 82198fdf50c..2822274690b 100644
--- a/src/mongo/db/session/session_catalog_mongod.cpp
+++ b/src/mongo/db/session/session_catalog_mongod.cpp
@@ -448,7 +448,8 @@ void createRetryableFindAndModifyTable(OperationContext* opCtx) {
 }
 
 
-void abortInProgressTransactions(OperationContext* opCtx) {
+void abortInProgressTransactions(OperationContext* opCtx,
+                                 MongoDSessionCatalog* mongoDSessionCatalog) {
     DBDirectClient client(opCtx);
     FindCommandRequest findRequest{NamespaceString::kSessionTransactionsTableNamespace};
     findRequest.setFilter(BSON(SessionTxnRecord::kStateFieldName
@@ -464,7 +465,7 @@ void abortInProgressTransactions(OperationContext* opCtx) {
         opCtx->setLogicalSessionId(txnRecord.getSessionId());
         opCtx->setTxnNumber(txnRecord.getTxnNum());
         opCtx->setInMultiDocumentTransaction();
-        MongoDOperationContextSessionWithoutRefresh ocs(opCtx);
+        auto ocs = mongoDSessionCatalog->checkOutSessionWithoutRefresh(opCtx);
         auto txnParticipant = TransactionParticipant::get(opCtx);
         LOGV2_DEBUG(21978,
                     3,
@@ -476,6 +477,18 @@ void abortInProgressTransactions(OperationContext* opCtx) {
         opCtx->resetMultiDocumentTransactionState();
     }
 }
+
+void _checkInUnscopedSession(OperationContext* opCtx,
+                             OperationContextSession::CheckInReason reason) {
+    OperationContextSession::checkIn(opCtx, reason);
+}
+
+void _checkOutUnscopedSession(OperationContext* opCtx) {
+    OperationContextSession::checkOut(opCtx);
+    auto txnParticipant = TransactionParticipant::get(opCtx);
+    txnParticipant.refreshFromStorageIfNeeded(opCtx);
+}
+
 }  // namespace
 
 const std::string MongoDSessionCatalog::kConfigTxnsPartialIndexName = "parent_lsid";
@@ -566,7 +579,7 @@ void MongoDSessionCatalog::onStepUp(OperationContext* opCtx) {
             //   expected to cause a deadlock since this 'newOpCtx' will need to acquire the global
             //   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'.
-            MongoDOperationContextSessionWithoutRefresh ocs(newOpCtx.get());
+            auto ocs = checkOutSessionWithoutRefresh(newOpCtx.get());
             auto txnParticipant = TransactionParticipant::get(newOpCtx.get());
             LOGV2_DEBUG(21979,
                         3,
@@ -580,7 +593,7 @@ void MongoDSessionCatalog::onStepUp(OperationContext* opCtx) {
         }
     }
 
-    abortInProgressTransactions(opCtx);
+    abortInProgressTransactions(opCtx, this);
 
     createTransactionTable(opCtx);
     if (repl::feature_flags::gFeatureFlagRetryableFindAndModify.isEnabledAndIgnoreFCV()) {
@@ -691,7 +704,32 @@ int MongoDSessionCatalog::removeSessionsTransactionRecords(
     return numReaped;
 }
 
-MongoDOperationContextSession::MongoDOperationContextSession(OperationContext* opCtx)
+std::unique_ptr<MongoDSessionCatalog::Session> MongoDSessionCatalog::checkOutSession(
+    OperationContext* opCtx) {
+    return std::make_unique<MongoDOperationContextSession>(opCtx, CheckoutTag());
+}
+
+std::unique_ptr<MongoDSessionCatalog::Session> MongoDSessionCatalog::checkOutSessionWithoutRefresh(
+    OperationContext* opCtx) {
+    return std::make_unique<MongoDOperationContextSessionWithoutRefresh>(opCtx, CheckoutTag());
+}
+
+std::unique_ptr<MongoDSessionCatalog::Session>
+MongoDSessionCatalog::checkOutSessionWithoutOplogRead(OperationContext* opCtx) {
+    return std::make_unique<MongoDOperationContextSessionWithoutOplogRead>(opCtx, CheckoutTag());
+}
+
+void MongoDSessionCatalog::checkInUnscopedSession(OperationContext* opCtx,
+                                                  OperationContextSession::CheckInReason reason) {
+    _checkInUnscopedSession(opCtx, reason);
+}
+
+void MongoDSessionCatalog::checkOutUnscopedSession(OperationContext* opCtx) {
+    _checkOutUnscopedSession(opCtx);
+}
+
+MongoDOperationContextSession::MongoDOperationContextSession(OperationContext* opCtx,
+                                                             MongoDSessionCatalog::CheckoutTag tag)
     : _operationContextSession(opCtx) {
     invariant(!opCtx->getClient()->isInDirectClient());
 
@@ -703,17 +741,15 @@ MongoDOperationContextSession::~MongoDOperationContextSession() = default;
 
 void MongoDOperationContextSession::checkIn(OperationContext* opCtx,
                                             OperationContextSession::CheckInReason reason) {
-    OperationContextSession::checkIn(opCtx, reason);
+    _checkInUnscopedSession(opCtx, reason);
 }
 
 void MongoDOperationContextSession::checkOut(OperationContext* opCtx) {
-    OperationContextSession::checkOut(opCtx);
-    auto txnParticipant = TransactionParticipant::get(opCtx);
-    txnParticipant.refreshFromStorageIfNeeded(opCtx);
+    _checkOutUnscopedSession(opCtx);
 }
 
 MongoDOperationContextSessionWithoutRefresh::MongoDOperationContextSessionWithoutRefresh(
-    OperationContext* opCtx)
+    OperationContext* opCtx, MongoDSessionCatalog::CheckoutTag tag)
     : _operationContextSession(opCtx), _opCtx(opCtx) {
     invariant(!opCtx->getClient()->isInDirectClient());
     const auto clientTxnNumber = *opCtx->getTxnNumber();
@@ -732,7 +768,7 @@ MongoDOperationContextSessionWithoutRefresh::~MongoDOperationContextSessionWitho
 }
 
 MongoDOperationContextSessionWithoutOplogRead::MongoDOperationContextSessionWithoutOplogRead(
-    OperationContext* opCtx)
+    OperationContext* opCtx, MongoDSessionCatalog::CheckoutTag tag)
     : _operationContextSession(opCtx), _opCtx(opCtx) {
     invariant(!opCtx->getClient()->isInDirectClient());
 
diff --git a/src/mongo/db/session/session_catalog_mongod.h b/src/mongo/db/session/session_catalog_mongod.h
index 7b65df78193..ce5de294b51 100644
--- a/src/mongo/db/session/session_catalog_mongod.h
+++ b/src/mongo/db/session/session_catalog_mongod.h
@@ -40,6 +40,8 @@ class MongoDSessionCatalog {
     MongoDSessionCatalog& operator=(const MongoDSessionCatalog&) = delete;
 
 public:
+    class CheckoutTag {};
+
     /**
      * Retrieves the mongod session transaction table associated with the service or operation
      * context.
@@ -109,6 +111,75 @@ public:
      */
     int removeSessionsTransactionRecords(OperationContext* opCtx,
                                          const std::vector<LogicalSessionId>& lsidsToRemove);
+
+    /**
+     * Functions to check out a session. Returns a scoped object that checks in the session on
+     * destruction.
+     */
+    class Session {
+        Session(const Session&) = delete;
+        Session& operator=(const Session&) = delete;
+
+    public:
+        Session() = default;
+        virtual ~Session() = default;
+
+        /**
+         * This method allows a checked-out session to be temporarily or permanently checked
+         * back in, in order to allow other operations to use it.
+         *
+         * Applies to Session objects returned by checkOutSession() only.
+         *
+         * May only be called if the session has actually been checked out previously.
+         */
+        virtual void checkIn(OperationContext* opCtx,
+                             OperationContextSession::CheckInReason reason) = 0;
+
+        /**
+         * Applies to Session objects returned by checkOutSession() only.
+         *
+         * May only be called if the session is not checked out already.
+         */
+        virtual void checkOut(OperationContext* opCtx) = 0;
+    };
+
+    /**
+     * Checks out the session specified in the passed operation context and stores it
+     * for later access by the command. The session is installed when this method returns
+     * and is removed at when the returned Session object goes out of scope.
+     */
+    std::unique_ptr<Session> checkOutSession(OperationContext* opCtx);
+
+    /**
+     * Similar to checkOutSession(), but marks the TransactionParticipant as valid without
+     * refreshing from disk and starts a new transaction unconditionally.
+     *
+     * Returns a scoped Session object that does not support checkIn() or checkOut().
+     *
+     * NOTE: Only used by the replication oplog application logic on secondaries in order to replay
+     * prepared transactions.
+     */
+    std::unique_ptr<Session> checkOutSessionWithoutRefresh(OperationContext* opCtx);
+
+    /**
+     * Similar to checkOutSession(), but marks the TransactionParticipant as valid without
+     * loading the retryable write oplog history.  If the last operation was a multi-document
+     * transaction, is equivalent to MongoDOperationContextSession.
+     *
+     * Returns a scoped Session object that does not support checkIn() or checkOut().
+     *
+     * NOTE: Should only be used when reading the oplog history is not possible.
+     */
+    std::unique_ptr<Session> checkOutSessionWithoutOplogRead(OperationContext* opCtx);
+
+    /**
+     * These are lower-level functions for checking in or out sessions without a scoped Session
+     * object (see checkOutSession*() functions above).
+     * Used to implement checkIn()/checkOut() in MongoDOperationContextSession.
+     */
+    void checkInUnscopedSession(OperationContext* opCtx,
+                                OperationContextSession::CheckInReason reason);
+    void checkOutUnscopedSession(OperationContext* opCtx);
 };
 
 /**
@@ -116,9 +187,12 @@ public:
  * it for later access by the command. The session is installed at construction time and is removed
  * at destruction.
  */
-class MongoDOperationContextSession {
+class MongoDOperationContextSession : public MongoDSessionCatalog::Session {
+    MongoDOperationContextSession(const MongoDOperationContextSession&) = delete;
+    MongoDOperationContextSession& operator=(const MongoDOperationContextSession&) = delete;
+
 public:
-    MongoDOperationContextSession(OperationContext* opCtx);
+    MongoDOperationContextSession(OperationContext* opCtx, MongoDSessionCatalog::CheckoutTag tag);
     ~MongoDOperationContextSession();
 
     /**
@@ -127,12 +201,12 @@ public:
      *
      * May only be called if the session has actually been checked out previously.
      */
-    static void checkIn(OperationContext* opCtx, OperationContextSession::CheckInReason reason);
+    void checkIn(OperationContext* opCtx, OperationContextSession::CheckInReason reason) override;
 
     /**
      * May only be called if the session is not checked out already.
      */
-    static void checkOut(OperationContext* opCtx);
+    void checkOut(OperationContext* opCtx) override;
 
 private:
     OperationContextSession _operationContextSession;
@@ -145,11 +219,25 @@ private:
  * NOTE: Only used by the replication oplog application logic on secondaries in order to replay
  * prepared transactions.
  */
-class MongoDOperationContextSessionWithoutRefresh {
+class MongoDOperationContextSessionWithoutRefresh : public MongoDSessionCatalog::Session {
+    MongoDOperationContextSessionWithoutRefresh(
+        const MongoDOperationContextSessionWithoutRefresh&) = delete;
+    MongoDOperationContextSessionWithoutRefresh& operator=(
+        const MongoDOperationContextSessionWithoutRefresh&) = delete;
+
 public:
-    MongoDOperationContextSessionWithoutRefresh(OperationContext* opCtx);
+    MongoDOperationContextSessionWithoutRefresh(OperationContext* opCtx,
+                                                MongoDSessionCatalog::CheckoutTag tag);
     ~MongoDOperationContextSessionWithoutRefresh();
 
+    void checkIn(OperationContext* opCtx, OperationContextSession::CheckInReason reason) override {
+        MONGO_UNREACHABLE;
+    }
+
+    void checkOut(OperationContext* opCtx) override {
+        MONGO_UNREACHABLE;
+    }
+
 private:
     OperationContextSession _operationContextSession;
     OperationContext* const _opCtx;
@@ -162,11 +250,25 @@ private:
  *
  * NOTE: Should only be used when reading the oplog history is not possible.
  */
-class MongoDOperationContextSessionWithoutOplogRead {
+class MongoDOperationContextSessionWithoutOplogRead : public MongoDSessionCatalog::Session {
+    MongoDOperationContextSessionWithoutOplogRead(
+        const MongoDOperationContextSessionWithoutOplogRead&) = delete;
+    MongoDOperationContextSessionWithoutOplogRead& operator=(
+        const MongoDOperationContextSessionWithoutOplogRead&) = delete;
+
 public:
-    MongoDOperationContextSessionWithoutOplogRead(OperationContext* opCtx);
+    MongoDOperationContextSessionWithoutOplogRead(OperationContext* opCtx,
+                                                  MongoDSessionCatalog::CheckoutTag tag);
     ~MongoDOperationContextSessionWithoutOplogRead();
 
+    void checkIn(OperationContext* opCtx, OperationContextSession::CheckInReason reason) override {
+        MONGO_UNREACHABLE;
+    }
+
+    void checkOut(OperationContext* opCtx) override {
+        MONGO_UNREACHABLE;
+    }
+
 private:
     OperationContextSession _operationContextSession;
     OperationContext* const _opCtx;