SERVER-37665 Add interface to explicitly checkin/checkout the OpContext session

1 files changed, 28 insertions, 14 deletions

diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index bc6d09710a5..94ee9fa940b 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -195,13 +195,7 @@ OperationContextSession::OperationContextSession(OperationContext* opCtx) : _opC
         return;
     }
 
-    const auto catalog = SessionCatalog::get(opCtx);
-    auto scopedCheckedOutSession = catalog->checkOutSession(opCtx);
-
-    // We acquire a Client lock here to guard the construction of this session so that references to
-    // this session are safe to use while the lock is held
-    stdx::lock_guard<Client> lk(*opCtx->getClient());
-    checkedOutSession.emplace(std::move(scopedCheckedOutSession));
+    checkOut(opCtx);
 }
 
 OperationContextSession::~OperationContextSession() {
@@ -215,24 +209,44 @@ OperationContextSession::~OperationContextSession() {
     if (!checkedOutSession)
         return;
 
+    checkIn(_opCtx);
+}
+
+Session* OperationContextSession::get(OperationContext* opCtx) {
+    auto& checkedOutSession = operationSessionDecoration(opCtx);
+    if (checkedOutSession) {
+        return checkedOutSession->get();
+    }
+
+    return nullptr;
+}
+
+void OperationContextSession::checkIn(OperationContext* opCtx) {
+    auto& checkedOutSession = operationSessionDecoration(opCtx);
+    invariant(checkedOutSession);
+
     // Removing the checkedOutSession from the OperationContext must be done under the Client lock,
     // but destruction of the checkedOutSession must not be, as it takes the SessionCatalog mutex,
     // and other code may take the Client lock while holding that mutex.
-    stdx::unique_lock<Client> lk(*_opCtx->getClient());
-    ScopedCheckedOutSession sessionToDelete(std::move(checkedOutSession.get()));
+    stdx::unique_lock<Client> lk(*opCtx->getClient());
+    ScopedCheckedOutSession sessionToReleaseOutOfLock(std::move(*checkedOutSession));
 
     // This destroys the moved-from ScopedCheckedOutSession, and must be done within the client lock
     checkedOutSession = boost::none;
     lk.unlock();
 }
 
-Session* OperationContextSession::get(OperationContext* opCtx) {
+void OperationContextSession::checkOut(OperationContext* opCtx) {
     auto& checkedOutSession = operationSessionDecoration(opCtx);
-    if (checkedOutSession) {
-        return checkedOutSession->get();
-    }
+    invariant(!checkedOutSession);
 
-    return nullptr;
+    const auto catalog = SessionCatalog::get(opCtx);
+    auto scopedCheckedOutSession = catalog->checkOutSession(opCtx);
+
+    // We acquire a Client lock here to guard the construction of this session so that references to
+    // this session are safe to use while the lock is held
+    stdx::lock_guard<Client> lk(*opCtx->getClient());
+    checkedOutSession.emplace(std::move(scopedCheckedOutSession));
 }
 
 }  // namespace mongo