SERVER-35870 Allow more than one thread to block Session checkout at a time

3 files changed, 49 insertions, 10 deletions

diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index bbc201c12c3..d4d4634f36e 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -128,7 +128,7 @@ ScopedCheckedOutSession SessionCatalog::checkOutSession(OperationContext* opCtx)
 
     stdx::unique_lock<stdx::mutex> ul(_mutex);
 
-    while (!_allowCheckingOutSessions) {
+    while (!_isSessionCheckoutAllowed()) {
         opCtx->waitForConditionOrInterrupt(_checkingOutSessionsAllowedCond, ul);
     }
 
@@ -220,7 +220,7 @@ void SessionCatalog::scanSessions(OperationContext* opCtx,
 std::shared_ptr<SessionCatalog::SessionRuntimeInfo> SessionCatalog::_getOrCreateSessionRuntimeInfo(
     WithLock, OperationContext* opCtx, const LogicalSessionId& lsid) {
     invariant(!opCtx->lockState()->inAWriteUnitOfWork());
-    invariant(_allowCheckingOutSessions);
+    invariant(_isSessionCheckoutAllowed());
 
     auto it = _sessions.find(lsid);
     if (it == _sessions.end()) {
@@ -253,23 +253,24 @@ SessionCatalog::PreventCheckingOutSessionsBlock::PreventCheckingOutSessionsBlock
     invariant(sessionCatalog);
 
     stdx::lock_guard<stdx::mutex> lg(sessionCatalog->_mutex);
-    invariant(sessionCatalog->_allowCheckingOutSessions);
-    sessionCatalog->_allowCheckingOutSessions = false;
+    ++sessionCatalog->_preventSessionCheckoutRequests;
 }
 
 SessionCatalog::PreventCheckingOutSessionsBlock::~PreventCheckingOutSessionsBlock() {
     stdx::lock_guard<stdx::mutex> lg(_sessionCatalog->_mutex);
 
-    invariant(!_sessionCatalog->_allowCheckingOutSessions);
-    _sessionCatalog->_allowCheckingOutSessions = true;
-    _sessionCatalog->_checkingOutSessionsAllowedCond.notify_all();
+    invariant(_sessionCatalog->_preventSessionCheckoutRequests > 0);
+    --_sessionCatalog->_preventSessionCheckoutRequests;
+    if (_sessionCatalog->_preventSessionCheckoutRequests == 0) {
+        _sessionCatalog->_checkingOutSessionsAllowedCond.notify_all();
+    }
 }
 
 void SessionCatalog::PreventCheckingOutSessionsBlock::waitForAllSessionsToBeCheckedIn(
     OperationContext* opCtx) {
     stdx::unique_lock<stdx::mutex> ul(_sessionCatalog->_mutex);
 
-    invariant(!_sessionCatalog->_allowCheckingOutSessions);
+    invariant(!_sessionCatalog->_isSessionCheckoutAllowed());
     while (_sessionCatalog->_numCheckedOutSessions > 0) {
         opCtx->waitForConditionOrInterrupt(_sessionCatalog->_allSessionsCheckedInCond, ul);
     }
diff --git a/src/mongo/db/session_catalog.h b/src/mongo/db/session_catalog.h
index e4ca4dda0d6..97bad1e279e 100644
--- a/src/mongo/db/session_catalog.h
+++ b/src/mongo/db/session_catalog.h
@@ -163,6 +163,10 @@ private:
      */
     void _releaseSession(const LogicalSessionId& lsid);
 
+    bool _isSessionCheckoutAllowed() const {
+        return _preventSessionCheckoutRequests == 0;
+    };
+
     // Protects members below.
     stdx::mutex _mutex;
 
@@ -172,8 +176,8 @@ private:
     // Count of the number of Sessions that are currently checked out.
     uint32_t _numCheckedOutSessions{0};
 
-    // Set to false to cause all Session checkout or creation requests to block.
-    bool _allowCheckingOutSessions{true};
+    // When >0 all Session checkout or creation requests will block.
+    uint32_t _preventSessionCheckoutRequests{0};
 
     // Condition that is signaled when the number of checked out sessions goes to 0.
     stdx::condition_variable _allSessionsCheckedInCond;
diff --git a/src/mongo/db/session_catalog_test.cpp b/src/mongo/db/session_catalog_test.cpp
index e0343f8672a..3efb8a0dcb3 100644
--- a/src/mongo/db/session_catalog_test.cpp
+++ b/src/mongo/db/session_catalog_test.cpp
@@ -276,5 +276,39 @@ TEST_F(SessionCatalogTest, WaitForAllSessions) {
     future.get();
 }
 
+TEST_F(SessionCatalogTest, MultiplePreventCheckingOutSessionsBlocks) {
+    const auto lsid1 = makeLogicalSessionIdForTest();
+    opCtx()->setLogicalSessionId(lsid1);
+    opCtx()->setDeadlineAfterNowBy(Milliseconds(10), ErrorCodes::MaxTimeMSExpired);
+
+    boost::optional<OperationContextSession> ocs;
+
+    // Prevent new Sessions from being checked out.
+    boost::optional<SessionCatalog::PreventCheckingOutSessionsBlock> preventCheckoutBlock1,
+        preventCheckoutBlock2;
+    preventCheckoutBlock1.emplace(catalog());
+
+    // Ensure that checking out a Session fails
+    ASSERT_THROWS_CODE(
+        ocs.emplace(opCtx(), true), AssertionException, ErrorCodes::MaxTimeMSExpired);
+
+    // A second request to prevent checking out Sessions is legal.
+    preventCheckoutBlock2.emplace(catalog());
+    ASSERT_THROWS_CODE(
+        ocs.emplace(opCtx(), true), AssertionException, ErrorCodes::MaxTimeMSExpired);
+
+    // The first request completing before the second is valid and doesn't start allowing checkouts.
+    preventCheckoutBlock1.reset();
+    ASSERT_THROWS_CODE(
+        ocs.emplace(opCtx(), true), AssertionException, ErrorCodes::MaxTimeMSExpired);
+
+    // Releasing the last PreventCheckingOutSessionsBlock allows Session checkout to proceed.
+    preventCheckoutBlock2.reset();
+
+    ASSERT_TRUE(ocs == boost::none);
+    ocs.emplace(opCtx(), true);
+    ASSERT_EQ(lsid1, ocs->get(opCtx())->getSessionId());
+}
+
 }  // namespace
 }  // namespace mongo