SERVER-33697 Provide sanity check on the number of cached sessions

10 files changed, 52 insertions, 26 deletions

diff --git a/src/mongo/base/error_codes.err b/src/mongo/base/error_codes.err
index 44fd6c15a8c..00ee3fa697c 100644
--- a/src/mongo/base/error_codes.err
+++ b/src/mongo/base/error_codes.err
@@ -261,6 +261,7 @@ error_code("InvalidResumeToken", 260);
 error_code("TooManyFilesOpen", 261);
 error_code("FailPointSetFailed", 262)
 error_code("OperationNotSupportedInTransaction", 263)
+error_code("TooManyLogicalSessions", 264);
 
 # Error codes 4000-8999 are reserved.
 
diff --git a/src/mongo/db/commands/start_session_command.cpp b/src/mongo/db/commands/start_session_command.cpp
index 01ce7e0be6d..09652aacf42 100644
--- a/src/mongo/db/commands/start_session_command.cpp
+++ b/src/mongo/db/commands/start_session_command.cpp
@@ -79,7 +79,7 @@ public:
         auto lsCache = LogicalSessionCache::get(serviceContext);
         boost::optional<LogicalSessionRecord> record =
             makeLogicalSessionRecord(opCtx, lsCache->now());
-        lsCache->startSession(opCtx, record.get());
+        uassertStatusOK(lsCache->startSession(opCtx, record.get()));
 
         makeLogicalSessionToClient(record->getId()).serialize(&result);
 
diff --git a/src/mongo/db/cursor_manager.cpp b/src/mongo/db/cursor_manager.cpp
index cac5d43de1d..6cfc442ed67 100644
--- a/src/mongo/db/cursor_manager.cpp
+++ b/src/mongo/db/cursor_manager.cpp
@@ -593,7 +593,11 @@ StatusWith<ClientCursorPin> CursorManager::pinCursor(OperationContext* opCtx,
     // We use pinning of a cursor as a proxy for active, user-initiated use of a cursor.  Therefor,
     // we pass down to the logical session cache and vivify the record (updating last use).
     if (cursor->getSessionId()) {
-        LogicalSessionCache::get(opCtx)->vivify(opCtx, cursor->getSessionId().get());
+        auto vivifyCursorStatus =
+            LogicalSessionCache::get(opCtx)->vivify(opCtx, cursor->getSessionId().get());
+        if (!vivifyCursorStatus.isOK()) {
+            return vivifyCursorStatus;
+        }
     }
 
     return ClientCursorPin(opCtx, cursor);
diff --git a/src/mongo/db/initialize_operation_session_info.cpp b/src/mongo/db/initialize_operation_session_info.cpp
index 137fba1196c..965dc0f1e53 100644
--- a/src/mongo/db/initialize_operation_session_info.cpp
+++ b/src/mongo/db/initialize_operation_session_info.cpp
@@ -73,7 +73,7 @@ boost::optional<OperationSessionInfoFromClient> initializeOperationSessionInfo(
         }
 
         opCtx->setLogicalSessionId(makeLogicalSessionId(osi.getSessionId().get(), opCtx));
-        lsc->vivify(opCtx, opCtx->getLogicalSessionId().get());
+        uassertStatusOK(lsc->vivify(opCtx, opCtx->getLogicalSessionId().get()));
     } else {
         uassert(ErrorCodes::InvalidOptions,
                 "Transaction number requires a session ID to also be specified",
diff --git a/src/mongo/db/logical_session_cache.h b/src/mongo/db/logical_session_cache.h
index ece611339ba..2827fd8cd98 100644
--- a/src/mongo/db/logical_session_cache.h
+++ b/src/mongo/db/logical_session_cache.h
@@ -69,7 +69,7 @@ public:
      * should only be used when starting new sessions and should not be used to
      * insert records for existing sessions.
      */
-    virtual void startSession(OperationContext* opCtx, LogicalSessionRecord record) = 0;
+    virtual Status startSession(OperationContext* opCtx, LogicalSessionRecord record) = 0;
 
     /**
      * Refresh the given sessions. Updates the timestamps of these records in
@@ -84,7 +84,7 @@ public:
      * Vivifies the session in the cache. I.e. creates it if it isn't there, updates last use if it
      * is.
      */
-    virtual void vivify(OperationContext* opCtx, const LogicalSessionId& lsid) = 0;
+    virtual Status vivify(OperationContext* opCtx, const LogicalSessionId& lsid) = 0;
 
     /**
      * enqueues LogicalSessionIds for removal during the next _refresh()
diff --git a/src/mongo/db/logical_session_cache_impl.cpp b/src/mongo/db/logical_session_cache_impl.cpp
index 8b46cbb4638..fa7e533cad5 100644
--- a/src/mongo/db/logical_session_cache_impl.cpp
+++ b/src/mongo/db/logical_session_cache_impl.cpp
@@ -52,6 +52,8 @@ MONGO_EXPORT_STARTUP_SERVER_PARAMETER(
 
 MONGO_EXPORT_STARTUP_SERVER_PARAMETER(disableLogicalSessionCacheRefresh, bool, false);
 
+MONGO_EXPORT_STARTUP_SERVER_PARAMETER(maxSessions, int, 1'000'000);
+
 constexpr Minutes LogicalSessionCacheImpl::kLogicalSessionDefaultRefresh;
 
 LogicalSessionCacheImpl::LogicalSessionCacheImpl(
@@ -99,11 +101,11 @@ Status LogicalSessionCacheImpl::promote(LogicalSessionId lsid) {
     return Status::OK();
 }
 
-void LogicalSessionCacheImpl::startSession(OperationContext* opCtx, LogicalSessionRecord record) {
+Status LogicalSessionCacheImpl::startSession(OperationContext* opCtx, LogicalSessionRecord record) {
     // Add the new record to our local cache. We will insert it into the sessions collection
     // the next time _refresh is called. If there is already a record in the cache for this
     // session, we'll just write over it with our newer, more recent one.
-    _addToCache(record);
+    return _addToCache(record);
 }
 
 Status LogicalSessionCacheImpl::refreshSessions(OperationContext* opCtx,
@@ -113,7 +115,10 @@ Status LogicalSessionCacheImpl::refreshSessions(OperationContext* opCtx,
     for (const auto& lsid : sessions) {
         if (!promote(lsid).isOK()) {
             // This is a new record, insert it.
-            _addToCache(makeLogicalSessionRecord(opCtx, lsid, now()));
+            auto addToCacheStatus = _addToCache(makeLogicalSessionRecord(opCtx, lsid, now()));
+            if (!addToCacheStatus.isOK()) {
+                return addToCacheStatus;
+            }
         }
     }
 
@@ -127,17 +132,21 @@ Status LogicalSessionCacheImpl::refreshSessions(OperationContext* opCtx,
     for (const auto& record : records) {
         if (!promote(record.getId()).isOK()) {
             // This is a new record, insert it.
-            _addToCache(record);
+            auto addToCacheStatus = _addToCache(record);
+            if (!addToCacheStatus.isOK()) {
+                return addToCacheStatus;
+            }
         }
     }
 
     return Status::OK();
 }
 
-void LogicalSessionCacheImpl::vivify(OperationContext* opCtx, const LogicalSessionId& lsid) {
+Status LogicalSessionCacheImpl::vivify(OperationContext* opCtx, const LogicalSessionId& lsid) {
     if (!promote(lsid).isOK()) {
-        startSession(opCtx, makeLogicalSessionRecord(opCtx, lsid, now()));
+        return startSession(opCtx, makeLogicalSessionRecord(opCtx, lsid, now()));
     }
+    return Status::OK();
 }
 
 Status LogicalSessionCacheImpl::refreshNow(Client* client) {
@@ -401,9 +410,13 @@ LogicalSessionCacheStats LogicalSessionCacheImpl::getStats() {
     return _stats;
 }
 
-void LogicalSessionCacheImpl::_addToCache(LogicalSessionRecord record) {
+Status LogicalSessionCacheImpl::_addToCache(LogicalSessionRecord record) {
     stdx::lock_guard<stdx::mutex> lk(_cacheMutex);
+    if (_activeSessions.size() >= static_cast<size_t>(maxSessions)) {
+        return {ErrorCodes::TooManyLogicalSessions, "cannot add session into the cache"};
+    }
     _activeSessions.insert(std::make_pair(record.getId(), record));
+    return Status::OK();
 }
 
 std::vector<LogicalSessionId> LogicalSessionCacheImpl::listIds() const {
diff --git a/src/mongo/db/logical_session_cache_impl.h b/src/mongo/db/logical_session_cache_impl.h
index 399753e4cba..1703ab05cad 100644
--- a/src/mongo/db/logical_session_cache_impl.h
+++ b/src/mongo/db/logical_session_cache_impl.h
@@ -99,14 +99,14 @@ public:
 
     Status promote(LogicalSessionId lsid) override;
 
-    void startSession(OperationContext* opCtx, LogicalSessionRecord record) override;
+    Status startSession(OperationContext* opCtx, LogicalSessionRecord record) override;
 
     Status refreshSessions(OperationContext* opCtx,
                            const RefreshSessionsCmdFromClient& cmd) override;
     Status refreshSessions(OperationContext* opCtx,
                            const RefreshSessionsCmdFromClusterMember& cmd) override;
 
-    void vivify(OperationContext* opCtx, const LogicalSessionId& lsid) override;
+    Status vivify(OperationContext* opCtx, const LogicalSessionId& lsid) override;
 
     Status refreshNow(Client* client) override;
 
@@ -146,7 +146,7 @@ private:
     /**
      * Takes the lock and inserts the given record into the cache.
      */
-    void _addToCache(LogicalSessionRecord record);
+    Status _addToCache(LogicalSessionRecord record);
 
     const Minutes _refreshInterval;
     const Minutes _sessionTimeout;
diff --git a/src/mongo/db/logical_session_cache_noop.h b/src/mongo/db/logical_session_cache_noop.h
index acb6531a388..7d0c5d1051e 100644
--- a/src/mongo/db/logical_session_cache_noop.h
+++ b/src/mongo/db/logical_session_cache_noop.h
@@ -45,7 +45,9 @@ public:
         return Status::OK();
     }
 
-    void startSession(OperationContext* opCtx, LogicalSessionRecord record) override {}
+    Status startSession(OperationContext* opCtx, LogicalSessionRecord record) override {
+        return Status::OK();
+    }
 
     Status refreshSessions(OperationContext* opCtx,
                            const RefreshSessionsCmdFromClient& cmd) override {
@@ -56,7 +58,9 @@ public:
         return Status::OK();
     }
 
-    void vivify(OperationContext* opCtx, const LogicalSessionId& lsid) override {}
+    Status vivify(OperationContext* opCtx, const LogicalSessionId& lsid) override {
+        return Status::OK();
+    }
 
     Status refreshNow(Client* client) override {
         return Status::OK();
diff --git a/src/mongo/db/logical_session_cache_test.cpp b/src/mongo/db/logical_session_cache_test.cpp
index 37256c26670..b8ec36827fe 100644
--- a/src/mongo/db/logical_session_cache_test.cpp
+++ b/src/mongo/db/logical_session_cache_test.cpp
@@ -142,7 +142,7 @@ TEST_F(LogicalSessionCacheTest, PromoteUpdatesLastUse) {
     auto start = service()->now();
 
     // Insert the record into the sessions collection with 'start'
-    cache()->startSession(opCtx(), makeLogicalSessionRecord(lsid, start));
+    ASSERT_OK(cache()->startSession(opCtx(), makeLogicalSessionRecord(lsid, start)));
 
     // Fast forward time and promote
     service()->fastForward(Milliseconds(500));
@@ -182,7 +182,7 @@ TEST_F(LogicalSessionCacheTest, StartSession) {
     auto lsid = record.getId();
 
     // Test starting a new session
-    cache()->startSession(opCtx(), record);
+    ASSERT_OK(cache()->startSession(opCtx(), record));
 
     // Record will not be in the collection yet; refresh must happen first.
     ASSERT(!sessions()->has(lsid));
@@ -193,27 +193,27 @@ TEST_F(LogicalSessionCacheTest, StartSession) {
     ASSERT(sessions()->has(lsid));
 
     // Try to start the same session again, should succeed.
-    cache()->startSession(opCtx(), record);
+    ASSERT_OK(cache()->startSession(opCtx(), record));
 
     // Try to start a session that is already in the sessions collection but
     // is not in our local cache, should succeed.
     auto record2 = makeLogicalSessionRecord(makeLogicalSessionIdForTest(), service()->now());
     sessions()->add(record2);
-    cache()->startSession(opCtx(), record2);
+    ASSERT_OK(cache()->startSession(opCtx(), record2));
 
     // Try to start a session that has expired from our cache, and is no
     // longer in the sessions collection, should succeed
     service()->fastForward(Milliseconds(kSessionTimeout.count() + 5));
     sessions()->remove(lsid);
     ASSERT(!sessions()->has(lsid));
-    cache()->startSession(opCtx(), record);
+    ASSERT_OK(cache()->startSession(opCtx(), record));
 }
 
 // Test that session cache properly expires lsids after 30 minutes of no use
 TEST_F(LogicalSessionCacheTest, BasicSessionExpiration) {
     // Insert a lsid
     auto record = makeLogicalSessionRecordForTest();
-    cache()->startSession(opCtx(), record);
+    ASSERT_OK(cache()->startSession(opCtx(), record));
     auto res = cache()->promote(record.getId());
     ASSERT(res.isOK());
 
@@ -232,7 +232,7 @@ TEST_F(LogicalSessionCacheTest, ManySignedLsidsInCacheRefresh) {
     int count = 10000;
     for (int i = 0; i < count; i++) {
         auto record = makeLogicalSessionRecordForTest();
-        cache()->startSession(opCtx(), record);
+        ASSERT_OK(cache()->startSession(opCtx(), record));
     }
 
     // Check that all signedLsids refresh
@@ -344,7 +344,7 @@ TEST_F(LogicalSessionCacheTest, RefreshMatrixSessionState) {
             service()->add(lsid);
         }
         if (active) {
-            cache()->startSession(opCtx(), lsRecord);
+            ASSERT_OK(cache()->startSession(opCtx(), lsRecord));
         }
         if (!expired) {
             sessions()->add(lsRecord);
diff --git a/src/mongo/s/query/cluster_cursor_manager.cpp b/src/mongo/s/query/cluster_cursor_manager.cpp
index f42465e845d..cbfebdb5d80 100644
--- a/src/mongo/s/query/cluster_cursor_manager.cpp
+++ b/src/mongo/s/query/cluster_cursor_manager.cpp
@@ -319,7 +319,11 @@ StatusWith<ClusterCursorManager::PinnedCursor> ClusterCursorManager::checkOutCur
     // We use pinning of a cursor as a proxy for active, user-initiated use of a cursor.  Therefore,
     // we pass down to the logical session cache and vivify the record (updating last use).
     if (cursor->getLsid()) {
-        LogicalSessionCache::get(opCtx)->vivify(opCtx, cursor->getLsid().get());
+        auto vivifyCursorStatus =
+            LogicalSessionCache::get(opCtx)->vivify(opCtx, cursor->getLsid().get());
+        if (!vivifyCursorStatus.isOK()) {
+            return vivifyCursorStatus;
+        }
     }
 
     cursor->reattachToOperationContext(opCtx);