SERVER-66992 Optimize how often SessionCatalog mutex is held

(cherry picked from commit d457fc3a9723edc094bd7dc45bcf00112cad57c6)

12 files changed, 53 insertions, 41 deletions

diff --git a/src/mongo/db/initialize_operation_session_info.cpp b/src/mongo/db/initialize_operation_session_info.cpp
index 32cda8de3e1..a3b75ec443d 100644
--- a/src/mongo/db/initialize_operation_session_info.cpp
+++ b/src/mongo/db/initialize_operation_session_info.cpp
@@ -99,7 +99,7 @@ OperationSessionInfoFromClient initializeOperationSessionInfo(OperationContext*
             return {};
         }
 
-        if (getParentSessionId(lsid)) {
+        if (isChildSession(lsid)) {
             uassert(ErrorCodes::InvalidOptions,
                     "Internal sessions are only allowed for internal clients",
                     isAuthorizedForInternalClusterAction);
diff --git a/src/mongo/db/kill_sessions_util.h b/src/mongo/db/kill_sessions_util.h
index eef51275dbd..d7a9875e0a2 100644
--- a/src/mongo/db/kill_sessions_util.h
+++ b/src/mongo/db/kill_sessions_util.h
@@ -40,7 +40,7 @@ namespace kill_sessions {
  * Validation callback to check if a given lsid has an associated parent session.
  */
 inline Status validateLsid(const LogicalSessionId& lsid) {
-    if (getParentSessionId(lsid)) {
+    if (isChildSession(lsid)) {
         return {ErrorCodes::InvalidOptions, "Cannot kill a child session"};
     }
     return Status::OK();
diff --git a/src/mongo/db/logical_session_cache_impl.cpp b/src/mongo/db/logical_session_cache_impl.cpp
index 2f42b90960f..8281dad2b6f 100644
--- a/src/mongo/db/logical_session_cache_impl.cpp
+++ b/src/mongo/db/logical_session_cache_impl.cpp
@@ -380,8 +380,8 @@ void LogicalSessionCacheImpl::_refresh(Client* client) {
 void LogicalSessionCacheImpl::endSessions(const LogicalSessionIdSet& sessions) {
     for (const auto& lsid : sessions) {
         uassert(ErrorCodes::InvalidOptions,
-                str::stream() << "Cannot specify a session id with parent session id " << lsid,
-                !getParentSessionId(lsid));
+                str::stream() << "Cannot specify a child session id " << lsid,
+                isParentSessionId(lsid));
     }
     stdx::lock_guard<Latch> lk(_mutex);
     _endingSessions.insert(begin(sessions), end(sessions));
@@ -446,8 +446,8 @@ std::vector<LogicalSessionId> LogicalSessionCacheImpl::listIds(
 boost::optional<LogicalSessionRecord> LogicalSessionCacheImpl::peekCached(
     const LogicalSessionId& lsid) const {
     uassert(ErrorCodes::InvalidOptions,
-            str::stream() << "Cannot specify a session id with parent session id " << lsid,
-            !getParentSessionId(lsid));
+            str::stream() << "Cannot specify a child session id " << lsid,
+            isParentSessionId(lsid));
 
     stdx::lock_guard<Latch> lk(_mutex);
     const auto it = _activeSessions.find(lsid);
diff --git a/src/mongo/db/logical_session_id_helpers.cpp b/src/mongo/db/logical_session_id_helpers.cpp
index 14576c6bd96..9d1ac707eb1 100644
--- a/src/mongo/db/logical_session_id_helpers.cpp
+++ b/src/mongo/db/logical_session_id_helpers.cpp
@@ -80,6 +80,11 @@ bool isParentSessionId(const LogicalSessionId& sessionId) {
     return !sessionId.getTxnUUID();
 }
 
+bool isChildSession(const LogicalSessionId& sessionId) {
+    // All child sessions must have a txnUUID.
+    return bool(sessionId.getTxnUUID());
+}
+
 boost::optional<LogicalSessionId> getParentSessionId(const LogicalSessionId& sessionId) {
     if (sessionId.getTxnUUID()) {
         return LogicalSessionId{sessionId.getId(), sessionId.getUid()};
diff --git a/src/mongo/db/logical_session_id_helpers.h b/src/mongo/db/logical_session_id_helpers.h
index 3d692c701e5..6d1a398c41a 100644
--- a/src/mongo/db/logical_session_id_helpers.h
+++ b/src/mongo/db/logical_session_id_helpers.h
@@ -56,6 +56,12 @@ SHA256Block getLogicalSessionUserDigestFor(StringData user, StringData db);
 bool isParentSessionId(const LogicalSessionId& sessionId);
 
 /**
+ * Returns if the given session is a child session, ie it was created on behalf of an operation that
+ * already had a session.
+ */
+bool isChildSession(const LogicalSessionId& sessionId);
+
+/**
  * Returns the parent session id for the given session id if there is one.
  */
 boost::optional<LogicalSessionId> getParentSessionId(const LogicalSessionId& sessionId);
diff --git a/src/mongo/db/s/session_catalog_migration_source.cpp b/src/mongo/db/s/session_catalog_migration_source.cpp
index fd7d01784c3..5706e76ba44 100644
--- a/src/mongo/db/s/session_catalog_migration_source.cpp
+++ b/src/mongo/db/s/session_catalog_migration_source.cpp
@@ -718,7 +718,7 @@ SessionCatalogMigrationSource::SessionOplogIterator::SessionOplogIterator(
           }
           // The SessionCatalogMigrationSource should not try to create a SessionOplogIterator for
           // internal sessions for non-retryable writes.
-          invariant(!getParentSessionId(txnRecord.getSessionId()));
+          invariant(isParentSessionId(txnRecord.getSessionId()));
           return _record.getState() ? EntryType::kNonRetryableTransaction
                                     : EntryType::kRetryableWrite;
       }()) {
diff --git a/src/mongo/db/s/transaction_coordinator_service.cpp b/src/mongo/db/s/transaction_coordinator_service.cpp
index f936323dcee..8dc5cf0a3cd 100644
--- a/src/mongo/db/s/transaction_coordinator_service.cpp
+++ b/src/mongo/db/s/transaction_coordinator_service.cpp
@@ -381,7 +381,7 @@ TransactionCoordinatorService::getAllRemovalFuturesForCoordinatorsForInternalTra
                         const TxnNumberAndRetryCounter txnNumberAndRetryCounter,
                         const std::shared_ptr<TransactionCoordinator> transactionCoordinator) {
         TransactionCoordinator::Step step = transactionCoordinator->getStep();
-        if (step > TransactionCoordinator::Step::kInactive && getParentSessionId(lsid)) {
+        if (step > TransactionCoordinator::Step::kInactive && isChildSession(lsid)) {
             return true;
         }
         return false;
diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index 39fe678aa28..d3f772851b5 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -77,9 +77,9 @@ SessionCatalog* SessionCatalog::get(ServiceContext* service) {
 SessionCatalog::ScopedCheckedOutSession SessionCatalog::_checkOutSessionInner(
     OperationContext* opCtx, const LogicalSessionId& lsid, boost::optional<KillToken> killToken) {
     if (killToken) {
-        invariant(killToken->lsidToKill == lsid);
+        dassert(killToken->lsidToKill == lsid);
     } else {
-        invariant(opCtx->getLogicalSessionId() == lsid);
+        dassert(opCtx->getLogicalSessionId() == lsid);
     }
 
     stdx::unique_lock<Latch> ul(_mutex);
@@ -134,10 +134,15 @@ SessionCatalog::SessionToKill SessionCatalog::checkOutSessionForKill(OperationCo
 }
 
 void SessionCatalog::scanSession(const LogicalSessionId& lsid,
-                                 const ScanSessionsCallbackFn& workerFn) {
+                                 const ScanSessionsCallbackFn& workerFn,
+                                 ScanSessionCreateSession createSession) {
     stdx::lock_guard<Latch> lg(_mutex);
 
-    if (auto sri = _getSessionRuntimeInfo(lg, lsid)) {
+    auto sri = (createSession == ScanSessionCreateSession::kYes)
+        ? _getOrCreateSessionRuntimeInfo(lg, lsid)
+        : _getSessionRuntimeInfo(lg, lsid);
+
+    if (sri) {
         auto session = sri->getSession(lg, lsid);
         invariant(session);
 
@@ -194,7 +199,7 @@ LogicalSessionIdSet SessionCatalog::scanSessionsForReap(
     const LogicalSessionId& parentLsid,
     const ScanSessionsCallbackFn& parentSessionWorkerFn,
     const ScanSessionsCallbackFn& childSessionWorkerFn) {
-    invariant(!getParentSessionId(parentLsid));
+    invariant(isParentSessionId(parentLsid));
 
     std::unique_ptr<SessionRuntimeInfo> sriToReap;
     {
@@ -259,14 +264,9 @@ size_t SessionCatalog::size() const {
     return _sessions.size();
 }
 
-void SessionCatalog::createSessionIfDoesNotExist(const LogicalSessionId& lsid) {
-    stdx::lock_guard<Latch> lg(_mutex);
-    _getOrCreateSessionRuntimeInfo(lg, lsid);
-}
-
 SessionCatalog::SessionRuntimeInfo* SessionCatalog::_getSessionRuntimeInfo(
     WithLock wl, const LogicalSessionId& lsid) {
-    auto parentLsid = castToParentSessionId(lsid);
+    const auto& parentLsid = isParentSessionId(lsid) ? lsid : *getParentSessionId(lsid);
     auto sriIt = _sessions.find(parentLsid);
 
     if (sriIt == _sessions.end()) {
@@ -289,12 +289,12 @@ SessionCatalog::SessionRuntimeInfo* SessionCatalog::_getOrCreateSessionRuntimeIn
         return sri;
     }
 
-    auto parentLsid = castToParentSessionId(lsid);
+    const auto& parentLsid = isParentSessionId(lsid) ? lsid : *getParentSessionId(lsid);
     auto sriIt =
         _sessions.emplace(parentLsid, std::make_unique<SessionRuntimeInfo>(parentLsid)).first;
     auto sri = sriIt->second.get();
 
-    if (getParentSessionId(lsid)) {
+    if (isChildSession(lsid)) {
         auto [childSessionIt, inserted] = sri->childSessions.try_emplace(lsid, lsid);
         // Insert should always succeed since the session did not exist prior to this.
         invariant(inserted);
@@ -321,7 +321,7 @@ void SessionCatalog::_releaseSession(SessionRuntimeInfo* sri,
     invariant(_sessions[sri->parentSession.getSessionId()].get() == sri);
     invariant(sri->checkoutOpCtx);
     if (killToken) {
-        invariant(killToken->lsidToKill == session->getSessionId());
+        dassert(killToken->lsidToKill == session->getSessionId());
     }
 
     sri->checkoutOpCtx = nullptr;
@@ -356,11 +356,13 @@ void SessionCatalog::_releaseSession(SessionRuntimeInfo* sri,
 }
 
 Session* SessionCatalog::SessionRuntimeInfo::getSession(WithLock, const LogicalSessionId& lsid) {
-    if (lsid == parentSession._sessionId) {
+    if (isParentSessionId(lsid)) {
+        // We should have already compared the parent lsid when we found this SRI.
+        dassert(lsid == parentSession._sessionId);
         return &parentSession;
     }
 
-    invariant(getParentSessionId(lsid) == parentSession._sessionId);
+    dassert(getParentSessionId(lsid) == parentSession._sessionId);
     auto it = childSessions.find(lsid);
     if (it == childSessions.end()) {
         return nullptr;
@@ -378,8 +380,8 @@ SessionCatalog::KillToken ObservableSession::kill(ErrorCodes::Error reason) cons
         invariant(_clientLock.owns_lock());
         const auto checkedOutLsid = _sri->checkoutOpCtx->getLogicalSessionId();
         const auto lsidToKill = getSessionId();
-        const bool isKillingParentSession = !getParentSessionId(lsidToKill);
-        if ((checkedOutLsid == lsidToKill) || isKillingParentSession) {
+        const bool isKillingParentSession = isParentSessionId(lsidToKill);
+        if (isKillingParentSession || (checkedOutLsid == lsidToKill)) {
             const auto serviceContext = _sri->checkoutOpCtx->getServiceContext();
             serviceContext->killOperation(_clientLock, _sri->checkoutOpCtx, reason);
         }
@@ -389,7 +391,7 @@ SessionCatalog::KillToken ObservableSession::kill(ErrorCodes::Error reason) cons
 }
 
 void ObservableSession::markForReap(ReapMode reapMode) {
-    if (!getParentSessionId(getSessionId())) {
+    if (isParentSessionId(getSessionId())) {
         // By design, parent sessions are only safe to be reaped if all of their child sessions are.
         invariant(reapMode == ReapMode::kNonExclusive);
     }
diff --git a/src/mongo/db/session_catalog.h b/src/mongo/db/session_catalog.h
index 8186e2fcc60..b1a600989b0 100644
--- a/src/mongo/db/session_catalog.h
+++ b/src/mongo/db/session_catalog.h
@@ -103,7 +103,10 @@ public:
      * Iterates through the SessionCatalog and applies 'workerFn' to each Session. This locks the
      * SessionCatalog.
      */
-    void scanSession(const LogicalSessionId& lsid, const ScanSessionsCallbackFn& workerFn);
+    enum class ScanSessionCreateSession { kYes, kNo };
+    void scanSession(const LogicalSessionId& lsid,
+                     const ScanSessionsCallbackFn& workerFn,
+                     ScanSessionCreateSession createSession = ScanSessionCreateSession::kNo);
     void scanSessions(const SessionKiller::Matcher& matcher,
                       const ScanSessionsCallbackFn& workerFn);
 
@@ -135,11 +138,6 @@ public:
      */
     size_t size() const;
 
-    /**
-     * Creates the session runtime info for 'lsid' if it doesn't exist.
-     */
-    void createSessionIfDoesNotExist(const LogicalSessionId& lsid);
-
 private:
     /**
      * Tracks the runtime info for transaction sessions that corresponds to the same logical
@@ -149,7 +147,7 @@ private:
     struct SessionRuntimeInfo {
         SessionRuntimeInfo(LogicalSessionId lsid) : parentSession(std::move(lsid)) {
             // Can only create a SessionRuntimeInfo with a parent transaction session id.
-            invariant(!getParentSessionId(lsid));
+            invariant(isParentSessionId(lsid));
         }
 
         Session* getSession(WithLock, const LogicalSessionId& lsid);
diff --git a/src/mongo/db/session_catalog_mongod.cpp b/src/mongo/db/session_catalog_mongod.cpp
index 7170228af2a..eeb1ffb175a 100644
--- a/src/mongo/db/session_catalog_mongod.cpp
+++ b/src/mongo/db/session_catalog_mongod.cpp
@@ -139,7 +139,7 @@ LogicalSessionIdSet removeExpiredTransactionSessionsNotInUseFromMemory(
     // parent transaction session so they have the same last check-out time as the the parent's.
     catalog->scanParentSessions([&](const ObservableSession& session) {
         const auto sessionId = session.getSessionId();
-        invariant(!getParentSessionId(sessionId));
+        invariant(isParentSessionId(sessionId));
         if (session.getLastCheckout() < possiblyExpired) {
             possiblyExpiredLogicalSessionIds.insert(sessionId);
         }
@@ -153,7 +153,7 @@ LogicalSessionIdSet removeExpiredTransactionSessionsNotInUseFromMemory(
     // longer in use from the in-memory catalog.
     LogicalSessionIdSet expiredTransactionSessionIdsStillInUse;
     for (const auto& expiredLogicalSessionId : expiredLogicalSessionIds) {
-        invariant(!getParentSessionId(expiredLogicalSessionId));
+        invariant(isParentSessionId(expiredLogicalSessionId));
 
         // Scan all the transaction sessions for this logical session at once so reaping can be done
         // atomically.
diff --git a/src/mongo/db/transaction_participant.cpp b/src/mongo/db/transaction_participant.cpp
index ba117c17afc..a4bd9eb0a03 100644
--- a/src/mongo/db/transaction_participant.cpp
+++ b/src/mongo/db/transaction_participant.cpp
@@ -618,7 +618,7 @@ const LogicalSessionId& TransactionParticipant::Observer::_sessionId() const {
 }
 
 bool TransactionParticipant::Observer::_isInternalSession() const {
-    return getParentSessionId(_sessionId()).has_value();
+    return isChildSession(_sessionId());
 }
 
 bool TransactionParticipant::Observer::_isInternalSessionForRetryableWrite() const {
@@ -2930,13 +2930,14 @@ void TransactionParticipant::Participant::_refreshActiveTransactionParticipantsF
                                 << "transaction number " << *childLsid.getTxnNumber(),
                             *childLsid.getTxnNumber() == *activeRetryableWriteTxnNumber);
                         auto sessionCatalog = SessionCatalog::get(opCtx);
-                        sessionCatalog->createSessionIfDoesNotExist(childLsid);
                         sessionCatalog->scanSession(
-                            childLsid, [&](const ObservableSession& osession) {
+                            childLsid,
+                            [&](const ObservableSession& osession) {
                                 auto childTxnParticipant =
                                     TransactionParticipant::get(opCtx, osession.get());
                                 childTxnParticipants.push_back(childTxnParticipant);
-                            });
+                            },
+                            SessionCatalog::ScanSessionCreateSession::kYes);
                     }
                 });
             } catch (const ExceptionFor<ErrorCodes::BadValue>& ex) {
diff --git a/src/mongo/s/session_catalog_router.cpp b/src/mongo/s/session_catalog_router.cpp
index 6bbe4e8af65..414b22e78cb 100644
--- a/src/mongo/s/session_catalog_router.cpp
+++ b/src/mongo/s/session_catalog_router.cpp
@@ -49,7 +49,7 @@ int RouterSessionCatalog::reapSessionsOlderThan(OperationContext* opCtx,
     // parent transaction session so they have the same last check-out time as the parent's.
     catalog->scanParentSessions([&](const ObservableSession& session) {
         const auto sessionId = session.getSessionId();
-        invariant(!getParentSessionId(sessionId));
+        invariant(isParentSessionId(sessionId));
         if (session.getLastCheckout() < possiblyExpired) {
             possiblyExpiredLogicalSessionIds.insert(session.getSessionId());
         }