SERVER-28346 Add more unit tests for the logical session cache

10 files changed, 327 insertions, 67 deletions

diff --git a/src/mongo/db/logical_session_cache.cpp b/src/mongo/db/logical_session_cache.cpp
index 9ab51495f52..1c94e019ac3 100644
--- a/src/mongo/db/logical_session_cache.cpp
+++ b/src/mongo/db/logical_session_cache.cpp
@@ -145,8 +145,8 @@ Status LogicalSessionCache::startSession(LogicalSessionRecord authoritativeRecor
 }
 
 void LogicalSessionCache::_refresh() {
-    SessionList activeSessions;
-    SessionList deadSessions;
+    LogicalSessionIdSet activeSessions;
+    LogicalSessionIdSet deadSessions;
 
     auto now = _service->now();
 
@@ -164,9 +164,9 @@ void LogicalSessionCache::_refresh() {
     for (auto& it : cacheCopy) {
         auto record = it.second;
         if (!_isDead(record, now)) {
-            activeSessions.push_back(record.getLsid());
+            activeSessions.insert(record.getLsid());
         } else {
-            deadSessions.push_back(record.getLsid());
+            deadSessions.insert(record.getLsid());
         }
     }
 
@@ -190,25 +190,28 @@ void LogicalSessionCache::_refresh() {
                 // by another thread.
                 it->second.setLastUse(now);
             }
+
+            activeSessions.insert(lsid);
         }
     }
 
-    activeSessions.splice(activeSessions.begin(), serviceSessions);
-
     // Query into the sessions collection to do the refresh. If any sessions have
     // failed to refresh, it means their authoritative records were removed, and
     // we should remove such records from our cache as well.
     auto failedToRefresh = _sessionsColl->refreshSessions(std::move(activeSessions));
-    deadSessions.splice(deadSessions.begin(), failedToRefresh);
 
     // Prune any dead records out of the cache. Dead records are ones that failed to
     // refresh, or ones that have expired locally. We don't make an effort to check
     // if the locally-expired records still have live authoritative records in the
     // sessions collection. We also don't attempt to resurrect our expired records.
+    // However, we *do* keep records alive if they are active on the service.
     {
         stdx::unique_lock<stdx::mutex> lk(_cacheMutex);
-        for (auto deadId : deadSessions) {
-            _cache.erase(deadId);
+        for (auto deadId : failedToRefresh) {
+            auto it = serviceSessions.find(deadId);
+            if (it == serviceSessions.end()) {
+                _cache.erase(deadId);
+            }
         }
     }
 }
diff --git a/src/mongo/db/logical_session_cache.h b/src/mongo/db/logical_session_cache.h
index f7aa671bf39..4e4bc4023af 100644
--- a/src/mongo/db/logical_session_cache.h
+++ b/src/mongo/db/logical_session_cache.h
@@ -47,8 +47,6 @@ namespace mongo {
  */
 class LogicalSessionCache {
 public:
-    using SessionList = std::list<LogicalSessionId>;
-
     static constexpr int kLogicalSessionCacheDefaultCapacity = 10000;
     static constexpr Minutes kLogicalSessionDefaultTimeout = Minutes(30);
     static constexpr Minutes kLogicalSessionDefaultRefresh = Minutes(5);
diff --git a/src/mongo/db/logical_session_cache_test.cpp b/src/mongo/db/logical_session_cache_test.cpp
index 93800d6c153..c0f86306ef4 100644
--- a/src/mongo/db/logical_session_cache_test.cpp
+++ b/src/mongo/db/logical_session_cache_test.cpp
@@ -72,6 +72,12 @@ public:
         _service->join();
     }
 
+    void waitUntilRefreshScheduled() {
+        while (service()->jobs() < 2) {
+            sleepmillis(10);
+        }
+    }
+
     LogicalSessionRecord newRecord() {
         return LogicalSessionRecord::makeAuthoritativeRecord(
             LogicalSessionId::gen(), _user, _userId, _service->now());
@@ -236,9 +242,9 @@ TEST_F(LogicalSessionCacheTest, CacheRefreshesOwnRecords) {
 
     // Advance time to first refresh point, check that refresh happens, and
     // that it includes both our records
-    sessions()->setRefreshHook([&hitRefresh](SessionList sessions) -> SessionList {
+    sessions()->setRefreshHook([&hitRefresh](LogicalSessionIdSet sessions) {
         hitRefresh.set_value(sessions.size());
-        return {};
+        return LogicalSessionIdSet{};
     });
 
     // Wait for the refresh to happen
@@ -258,34 +264,295 @@ TEST_F(LogicalSessionCacheTest, CacheRefreshesOwnRecords) {
 
     // Advance time so that one record expires
     // Ensure that first record was refreshed, and second was thrown away
-    sessions()->setRefreshHook([&refresh2](SessionList sessions) -> SessionList {
+    sessions()->setRefreshHook([&refresh2](LogicalSessionIdSet sessions) {
         // We should only have one record here, the other should have expired
         ASSERT_EQ(sessions.size(), size_t(1));
-        refresh2.set_value(sessions.front());
-        return {};
+        refresh2.set_value(*(sessions.begin()));
+        return LogicalSessionIdSet{};
     });
 
     // Wait until the second job has been scheduled
-    while (service()->jobs() < 2) {
-        sleepmillis(10);
-    }
+    waitUntilRefreshScheduled();
 
     service()->fastForward(kSessionTimeout - kForceRefresh + Milliseconds(1));
     refresh2Future.wait();
     ASSERT_EQ(refresh2Future.get(), lsid);
 }
 
-// Additional tests:
-// SERVER-28346
-// - Test that cache deletes records that fail to refresh
-// - Test that session cache properly expires records after 30 minutes of no use
-// - Test that we keep refreshing sessions that are active on the service
-// - Test that if we try to refresh a record and it is not in the sessions collection,
-//   we remove it from the cache (unless it is active on the service)
-// - Test small mixed set of cache/service active sessions
-// - Test larger sets of cache-only session records
-// - Test larger sets of service-only session records
-// - Test larger mixed sets of cache/service active sessions
+// Test that cache deletes records that fail to refresh
+TEST_F(LogicalSessionCacheTest, CacheDeletesRecordsThatFailToRefresh) {
+    // Put two sessions into the cache
+    auto record1 = newRecord();
+    auto record2 = newRecord();
+    cache()->startSession(record1);
+    cache()->startSession(record2);
+
+    stdx::promise<void> hitRefresh;
+    auto refreshFuture = hitRefresh.get_future();
+
+    // Record 1 fails to refresh
+    sessions()->setRefreshHook([&hitRefresh, &record1](LogicalSessionIdSet sessions) {
+        ASSERT_EQ(sessions.size(), size_t(2));
+        hitRefresh.set_value();
+        return LogicalSessionIdSet{record1.getLsid()};
+    });
+
+    // Force a refresh
+    service()->fastForward(kForceRefresh);
+    refreshFuture.wait();
+
+    // Ensure that one record is still there and the other is gone
+    auto res = cache()->getOwnerFromCache(record1.getLsid());
+    ASSERT(!res.isOK());
+    res = cache()->getOwnerFromCache(record2.getLsid());
+    ASSERT(res.isOK());
+}
+
+// Test that we don't remove records that fail to refresh if they are active on the service
+TEST_F(LogicalSessionCacheTest, KeepActiveSessionAliveEvenIfRefreshFails) {
+    // Put two sessions into the cache, one into the service
+    auto record1 = newRecord();
+    auto record2 = newRecord();
+    cache()->startSession(record1);
+    service()->add(record1.getLsid());
+    cache()->startSession(record2);
+
+    stdx::promise<void> hitRefresh;
+    auto refreshFuture = hitRefresh.get_future();
+
+    // Record 1 fails to refresh
+    sessions()->setRefreshHook([&hitRefresh, &record1](LogicalSessionIdSet sessions) {
+        ASSERT_EQ(sessions.size(), size_t(2));
+        hitRefresh.set_value();
+        return LogicalSessionIdSet{record1.getLsid()};
+    });
+
+    // Force a refresh
+    service()->fastForward(kForceRefresh);
+    refreshFuture.wait();
+
+    // Ensure that both records are still there
+    auto res = cache()->getOwnerFromCache(record1.getLsid());
+    ASSERT(res.isOK());
+    res = cache()->getOwnerFromCache(record2.getLsid());
+    ASSERT(res.isOK());
+}
+
+// Test that session cache properly expires records after 30 minutes of no use
+TEST_F(LogicalSessionCacheTest, BasicSessionExpiration) {
+    // Insert a record
+    auto record = newRecord();
+    cache()->startSession(record);
+    auto res = cache()->getOwnerFromCache(record.getLsid());
+    ASSERT(res.isOK());
+
+    // Force it to expire
+    service()->fastForward(Milliseconds(kSessionTimeout.count() + 5));
+
+    // Check that it is no longer in the cache
+    res = cache()->getOwnerFromCache(record.getLsid());
+    ASSERT(!res.isOK());
+}
+
+// Test that we keep refreshing sessions that are active on the service
+TEST_F(LogicalSessionCacheTest, LongRunningQueriesAreRefreshed) {
+    auto record = newRecord();
+    auto lsid = record.getLsid();
+
+    // Insert one active record on the service, none in the cache
+    service()->add(lsid);
+
+    stdx::mutex mutex;
+    stdx::condition_variable cv;
+    int count = 0;
+
+    sessions()->setRefreshHook([&cv, &mutex, &count, &lsid](LogicalSessionIdSet sessions) {
+        ASSERT_EQ(*(sessions.begin()), lsid);
+        {
+            stdx::unique_lock<stdx::mutex> lk(mutex);
+            count++;
+        }
+        cv.notify_all();
+
+        return LogicalSessionIdSet{};
+    });
+
+    // Force a refresh, it should refresh our active session
+    service()->fastForward(kForceRefresh);
+    {
+        stdx::unique_lock<stdx::mutex> lk(mutex);
+        cv.wait(lk, [&count] { return count == 1; });
+    }
+
+    // Wait until the next job has been scheduled
+    waitUntilRefreshScheduled();
+
+    // Force a session timeout, session is still on the service
+    service()->fastForward(kSessionTimeout);
+    {
+        stdx::unique_lock<stdx::mutex> lk(mutex);
+        cv.wait(lk, [&count] { return count == 2; });
+    }
+
+    // Wait until the next job has been scheduled
+    waitUntilRefreshScheduled();
+
+    // Force another refresh, check that it refreshes that active record again
+    service()->fastForward(kForceRefresh);
+    {
+        stdx::unique_lock<stdx::mutex> lk(mutex);
+        cv.wait(lk, [&count] { return count == 3; });
+    }
+}
+
+// Test that the set of records we refresh is a sum of cached + active records
+TEST_F(LogicalSessionCacheTest, RefreshCachedAndServiceRecordsTogether) {
+    // Put one session into the cache, one into the service
+    auto record1 = newRecord();
+    service()->add(record1.getLsid());
+    auto record2 = newRecord();
+    cache()->startSession(record2);
+
+    stdx::promise<void> hitRefresh;
+    auto refreshFuture = hitRefresh.get_future();
+
+    // Both records refresh
+    sessions()->setRefreshHook([&hitRefresh](LogicalSessionIdSet sessions) {
+        ASSERT_EQ(sessions.size(), size_t(2));
+        hitRefresh.set_value();
+        return LogicalSessionIdSet{};
+    });
+
+    // Force a refresh
+    service()->fastForward(kForceRefresh);
+    refreshFuture.wait();
+}
+
+// Test large sets of cache-only session records
+TEST_F(LogicalSessionCacheTest, ManyRecordsInCacheRefresh) {
+    int count = LogicalSessionCache::kLogicalSessionCacheDefaultCapacity;
+    for (int i = 0; i < count; i++) {
+        auto record = newRecord();
+        cache()->startSession(record);
+    }
+
+    stdx::promise<void> hitRefresh;
+    auto refreshFuture = hitRefresh.get_future();
+
+    // Check that all records refresh
+    sessions()->setRefreshHook([&hitRefresh, &count](LogicalSessionIdSet sessions) {
+        ASSERT_EQ(sessions.size(), size_t(count));
+        hitRefresh.set_value();
+        return LogicalSessionIdSet{};
+    });
+
+    // Force a refresh
+    service()->fastForward(kForceRefresh);
+    refreshFuture.wait();
+}
+
+// Test larger sets of service-only session records
+TEST_F(LogicalSessionCacheTest, ManyLongRunningSessionsRefresh) {
+    int count = LogicalSessionCache::kLogicalSessionCacheDefaultCapacity;
+    for (int i = 0; i < count; i++) {
+        auto record = newRecord();
+        service()->add(record.getLsid());
+    }
+
+    stdx::promise<void> hitRefresh;
+    auto refreshFuture = hitRefresh.get_future();
+
+    // Check that all records refresh
+    sessions()->setRefreshHook([&hitRefresh, &count](LogicalSessionIdSet sessions) {
+        ASSERT_EQ(sessions.size(), size_t(count));
+        hitRefresh.set_value();
+        return LogicalSessionIdSet{};
+    });
+
+    // Force a refresh
+    service()->fastForward(kForceRefresh);
+    refreshFuture.wait();
+}
+
+// Test larger mixed sets of cache/service active sessions
+TEST_F(LogicalSessionCacheTest, ManySessionsRefreshComboDeluxe) {
+    int count = LogicalSessionCache::kLogicalSessionCacheDefaultCapacity;
+    for (int i = 0; i < count; i++) {
+        auto record = newRecord();
+        service()->add(record.getLsid());
+
+        auto record2 = newRecord();
+        cache()->startSession(record2);
+    }
+
+    stdx::mutex mutex;
+    stdx::condition_variable cv;
+    int refreshes = 0;
+    int nRefreshed = 0;
+
+    // Check that all records refresh successfully
+    sessions()->setRefreshHook(
+        [&refreshes, &mutex, &cv, &nRefreshed](LogicalSessionIdSet sessions) {
+            {
+                stdx::unique_lock<stdx::mutex> lk(mutex);
+                refreshes++;
+                nRefreshed = sessions.size();
+            }
+            cv.notify_all();
+
+            return LogicalSessionIdSet{};
+        });
+
+    // Force a refresh
+    service()->fastForward(kForceRefresh);
+    {
+        stdx::unique_lock<stdx::mutex> lk(mutex);
+        cv.wait(lk, [&refreshes] { return refreshes == 1; });
+    }
+    ASSERT_EQ(nRefreshed, count * 2);
+
+    // Remove all of the service sessions, should just refresh the cache entries
+    // (and make all but one fail to refresh)
+    service()->clear();
+    sessions()->setRefreshHook(
+        [&refreshes, &mutex, &cv, &nRefreshed](LogicalSessionIdSet sessions) {
+            {
+                stdx::unique_lock<stdx::mutex> lk(mutex);
+                refreshes++;
+                nRefreshed = sessions.size();
+            }
+            cv.notify_all();
+
+            sessions.erase(sessions.begin());
+            return sessions;
+        });
+
+    // Wait for job to be scheduled
+    waitUntilRefreshScheduled();
+
+    // Force another refresh
+    service()->fastForward(kForceRefresh);
+    {
+        stdx::unique_lock<stdx::mutex> lk(mutex);
+        cv.wait(lk, [&refreshes] { return refreshes == 2; });
+    }
+
+    // We should not have refreshed any sessions from the service, only the cache
+    ASSERT_EQ(nRefreshed, count);
+
+    // Wait for job to be scheduled
+    waitUntilRefreshScheduled();
+
+    // Force a third refresh
+    service()->fastForward(kForceRefresh);
+    {
+        stdx::unique_lock<stdx::mutex> lk(mutex);
+        cv.wait(lk, [&refreshes] { return refreshes == 3; });
+    }
+
+    // Since all but one record failed to refresh, third set should just have one record
+    ASSERT_EQ(nRefreshed, 1);
+}
 
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/logical_session_id.h b/src/mongo/db/logical_session_id.h
index b67c935220e..eca1d8a33f5 100644
--- a/src/mongo/db/logical_session_id.h
+++ b/src/mongo/db/logical_session_id.h
@@ -114,4 +114,9 @@ inline StringBuilder& operator<<(StringBuilder& s, const LogicalSessionId& lsid)
     return (s << lsid.toString());
 }
 
+/**
+ * An alias for sets of session ids.
+ */
+using LogicalSessionIdSet = stdx::unordered_set<LogicalSessionId, LogicalSessionId::Hash>;
+
 }  // namespace mongo
diff --git a/src/mongo/db/service_liason.h b/src/mongo/db/service_liason.h
index 97535cea4e8..2389245f6a5 100644
--- a/src/mongo/db/service_liason.h
+++ b/src/mongo/db/service_liason.h
@@ -44,15 +44,13 @@ namespace mongo {
  */
 class ServiceLiason {
 public:
-    using SessionList = std::list<LogicalSessionId>;
-
     virtual ~ServiceLiason();
 
     /**
      * Return a list of sessions that are currently being used to run operations
      * on this service.
      */
-    virtual SessionList getActiveSessions() const = 0;
+    virtual LogicalSessionIdSet getActiveSessions() const = 0;
 
     /**
      * Schedule a job to be run at regular intervals until the server shuts down.
diff --git a/src/mongo/db/service_liason_mock.cpp b/src/mongo/db/service_liason_mock.cpp
index 6aed1c0306f..b7b03ba9ceb 100644
--- a/src/mongo/db/service_liason_mock.cpp
+++ b/src/mongo/db/service_liason_mock.cpp
@@ -42,7 +42,7 @@ MockServiceLiasonImpl::MockServiceLiasonImpl() {
     _runner->startup();
 }
 
-MockServiceLiasonImpl::SessionList MockServiceLiasonImpl::getActiveSessions() const {
+LogicalSessionIdSet MockServiceLiasonImpl::getActiveSessions() const {
     stdx::unique_lock<stdx::mutex> lk(_mutex);
     return _activeSessions;
 }
@@ -61,13 +61,12 @@ void MockServiceLiasonImpl::scheduleJob(PeriodicRunner::PeriodicJob job) {
 
 void MockServiceLiasonImpl::add(LogicalSessionId lsid) {
     stdx::unique_lock<stdx::mutex> lk(_mutex);
-    _activeSessions.push_back(std::move(lsid));
+    _activeSessions.insert(std::move(lsid));
 }
 
 void MockServiceLiasonImpl::remove(LogicalSessionId lsid) {
     stdx::unique_lock<stdx::mutex> lk(_mutex);
-    _activeSessions.erase(std::remove(_activeSessions.begin(), _activeSessions.end(), lsid),
-                          _activeSessions.end());
+    _activeSessions.erase(lsid);
 }
 
 void MockServiceLiasonImpl::clear() {
diff --git a/src/mongo/db/service_liason_mock.h b/src/mongo/db/service_liason_mock.h
index 577057708ad..bfd2d281ca1 100644
--- a/src/mongo/db/service_liason_mock.h
+++ b/src/mongo/db/service_liason_mock.h
@@ -54,12 +54,10 @@ namespace mongo {
  */
 class MockServiceLiasonImpl {
 public:
-    using SessionList = std::list<LogicalSessionId>;
-
     MockServiceLiasonImpl();
 
     // Forwarding methods from the MockServiceLiason
-    SessionList getActiveSessions() const;
+    LogicalSessionIdSet getActiveSessions() const;
     Date_t now() const;
     void scheduleJob(PeriodicRunner::PeriodicJob job);
     void join();
@@ -76,7 +74,7 @@ private:
     std::unique_ptr<PeriodicRunnerASIO> _runner;
 
     mutable stdx::mutex _mutex;
-    SessionList _activeSessions;
+    LogicalSessionIdSet _activeSessions;
 };
 
 /**
@@ -84,12 +82,10 @@ private:
  */
 class MockServiceLiason : public ServiceLiason {
 public:
-    using SessionList = std::list<LogicalSessionId>;
-
     explicit MockServiceLiason(std::shared_ptr<MockServiceLiasonImpl> impl)
         : _impl(std::move(impl)) {}
 
-    SessionList getActiveSessions() const override {
+    LogicalSessionIdSet getActiveSessions() const override {
         return _impl->getActiveSessions();
     }
 
diff --git a/src/mongo/db/sessions_collection.h b/src/mongo/db/sessions_collection.h
index 8e65828033d..ec59f28c597 100644
--- a/src/mongo/db/sessions_collection.h
+++ b/src/mongo/db/sessions_collection.h
@@ -28,6 +28,7 @@
 
 #pragma once
 
+#include "mongo/db/logical_session_id.h"
 #include "mongo/db/logical_session_record.h"
 
 namespace mongo {
@@ -40,8 +41,6 @@ namespace mongo {
  */
 class SessionsCollection {
 public:
-    using SessionList = std::list<LogicalSessionId>;
-
     virtual ~SessionsCollection();
 
     /**
@@ -66,7 +65,7 @@ public:
      * Returns a list of sessions for which no authoritative record was found,
      * and hence were not refreshed.
      */
-    virtual SessionList refreshSessions(SessionList sessions) = 0;
+    virtual LogicalSessionIdSet refreshSessions(LogicalSessionIdSet sessions) = 0;
 
     /**
      * Removes the authoritative records for the specified sessions.
@@ -74,7 +73,7 @@ public:
      * Implementations should perform authentication checks to ensure that
      * session records may only be removed if their owner is logged in.
      */
-    virtual void removeRecords(SessionList sessions) = 0;
+    virtual void removeRecords(LogicalSessionIdSet sessions) = 0;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/sessions_collection_mock.cpp b/src/mongo/db/sessions_collection_mock.cpp
index a56f6dec21a..4abc5a326d1 100644
--- a/src/mongo/db/sessions_collection_mock.cpp
+++ b/src/mongo/db/sessions_collection_mock.cpp
@@ -73,12 +73,11 @@ Status MockSessionsCollectionImpl::insertRecord(LogicalSessionRecord record) {
     return _insert(std::move(record));
 }
 
-MockSessionsCollectionImpl::SessionList MockSessionsCollectionImpl::refreshSessions(
-    SessionList sessions) {
+LogicalSessionIdSet MockSessionsCollectionImpl::refreshSessions(LogicalSessionIdSet sessions) {
     return _refresh(std::move(sessions));
 }
 
-void MockSessionsCollectionImpl::removeRecords(SessionList sessions) {
+void MockSessionsCollectionImpl::removeRecords(LogicalSessionIdSet sessions) {
     _remove(std::move(sessions));
 }
 
@@ -131,16 +130,15 @@ Status MockSessionsCollectionImpl::_insertRecord(LogicalSessionRecord record) {
     return Status::OK();
 }
 
-MockSessionsCollectionImpl::SessionList MockSessionsCollectionImpl::_refreshSessions(
-    SessionList sessions) {
-    SessionList notFound{};
+LogicalSessionIdSet MockSessionsCollectionImpl::_refreshSessions(LogicalSessionIdSet sessions) {
+    LogicalSessionIdSet notFound{};
 
     {
         stdx::unique_lock<stdx::mutex> lk(_mutex);
         for (auto& lsid : sessions) {
             auto it = _sessions.find(lsid);
             if (it == _sessions.end()) {
-                notFound.push_back(lsid);
+                notFound.insert(lsid);
             }
         }
     }
@@ -148,7 +146,7 @@ MockSessionsCollectionImpl::SessionList MockSessionsCollectionImpl::_refreshSess
     return notFound;
 }
 
-void MockSessionsCollectionImpl::_removeRecords(SessionList sessions) {
+void MockSessionsCollectionImpl::_removeRecords(LogicalSessionIdSet sessions) {
     stdx::unique_lock<stdx::mutex> lk(_mutex);
     for (auto& lsid : sessions) {
         _sessions.erase(lsid);
diff --git a/src/mongo/db/sessions_collection_mock.h b/src/mongo/db/sessions_collection_mock.h
index 05acb1f256f..64e2a76fc7b 100644
--- a/src/mongo/db/sessions_collection_mock.h
+++ b/src/mongo/db/sessions_collection_mock.h
@@ -54,7 +54,6 @@ namespace mongo {
  */
 class MockSessionsCollectionImpl {
 public:
-    using SessionList = std::list<LogicalSessionId>;
     using SessionMap =
         stdx::unordered_map<LogicalSessionId, LogicalSessionRecord, LogicalSessionId::Hash>;
 
@@ -62,8 +61,8 @@ public:
 
     using FetchHook = stdx::function<StatusWith<LogicalSessionRecord>(LogicalSessionId)>;
     using InsertHook = stdx::function<Status(LogicalSessionRecord)>;
-    using RefreshHook = stdx::function<SessionList(SessionList)>;
-    using RemoveHook = stdx::function<void(SessionList)>;
+    using RefreshHook = stdx::function<LogicalSessionIdSet(LogicalSessionIdSet)>;
+    using RemoveHook = stdx::function<void(LogicalSessionIdSet)>;
 
     // Set custom hooks to override default behavior
     void setFetchHook(FetchHook hook);
@@ -77,8 +76,8 @@ public:
     // Forwarding methods from the MockSessionsCollection
     StatusWith<LogicalSessionRecord> fetchRecord(LogicalSessionId lsid);
     Status insertRecord(LogicalSessionRecord record);
-    SessionList refreshSessions(SessionList sessions);
-    void removeRecords(SessionList sessions);
+    LogicalSessionIdSet refreshSessions(LogicalSessionIdSet sessions);
+    void removeRecords(LogicalSessionIdSet sessions);
 
     // Test-side methods that operate on the _sessions map
     void add(LogicalSessionRecord record);
@@ -91,8 +90,8 @@ private:
     // Default implementations, may be overridden with custom hooks.
     StatusWith<LogicalSessionRecord> _fetchRecord(LogicalSessionId lsid);
     Status _insertRecord(LogicalSessionRecord record);
-    SessionList _refreshSessions(SessionList sessions);
-    void _removeRecords(SessionList sessions);
+    LogicalSessionIdSet _refreshSessions(LogicalSessionIdSet sessions);
+    void _removeRecords(LogicalSessionIdSet sessions);
 
     stdx::mutex _mutex;
     SessionMap _sessions;
@@ -110,8 +109,6 @@ private:
  */
 class MockSessionsCollection : public SessionsCollection {
 public:
-    using SessionList = std::list<LogicalSessionId>;
-
     explicit MockSessionsCollection(std::shared_ptr<MockSessionsCollectionImpl> impl)
         : _impl(std::move(impl)) {}
 
@@ -123,11 +120,11 @@ public:
         return _impl->insertRecord(std::move(record));
     }
 
-    SessionList refreshSessions(SessionList sessions) override {
+    LogicalSessionIdSet refreshSessions(LogicalSessionIdSet sessions) override {
         return _impl->refreshSessions(std::move(sessions));
     }
 
-    void removeRecords(SessionList sessions) override {
+    void removeRecords(LogicalSessionIdSet sessions) override {
         return _impl->removeRecords(std::move(sessions));
     }