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

This reverts commit 4c3337abe512a37c8e4cb22e16eafcd04bc5bd3a.

10 files changed, 65 insertions, 330 deletions

diff --git a/src/mongo/db/logical_session_cache.cpp b/src/mongo/db/logical_session_cache.cpp
index 1c94e019ac3..9ab51495f52 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() {
-    LogicalSessionIdSet activeSessions;
-    LogicalSessionIdSet deadSessions;
+    SessionList activeSessions;
+    SessionList 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.insert(record.getLsid());
+            activeSessions.push_back(record.getLsid());
         } else {
-            deadSessions.insert(record.getLsid());
+            deadSessions.push_back(record.getLsid());
         }
     }
 
@@ -190,28 +190,25 @@ 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 : failedToRefresh) {
-            auto it = serviceSessions.find(deadId);
-            if (it == serviceSessions.end()) {
-                _cache.erase(deadId);
-            }
+        for (auto deadId : deadSessions) {
+            _cache.erase(deadId);
         }
     }
 }
diff --git a/src/mongo/db/logical_session_cache.h b/src/mongo/db/logical_session_cache.h
index 4e4bc4023af..f7aa671bf39 100644
--- a/src/mongo/db/logical_session_cache.h
+++ b/src/mongo/db/logical_session_cache.h
@@ -47,6 +47,8 @@ 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 043eebc85b8..93800d6c153 100644
--- a/src/mongo/db/logical_session_cache_test.cpp
+++ b/src/mongo/db/logical_session_cache_test.cpp
@@ -72,12 +72,6 @@ public:
         _service->join();
     }
 
-    void waitUntilRefreshScheduled() {
-        while (service()->jobs() < 2) {
-            sleepmillis(10);
-        }
-    }
-
     LogicalSessionRecord newRecord() {
         return LogicalSessionRecord::makeAuthoritativeRecord(
             LogicalSessionId::gen(), _user, _userId, _service->now());
@@ -242,7 +236,7 @@ TEST_F(LogicalSessionCacheTest, CacheRefreshesOwnRecords) {
 
     // Advance time to first refresh point, check that refresh happens, and
     // that it includes both our records
-    sessions()->setRefreshHook([&hitRefresh](LogicalSessionIdSet sessions) -> LogicalSessionIdSet {
+    sessions()->setRefreshHook([&hitRefresh](SessionList sessions) -> SessionList {
         hitRefresh.set_value(sessions.size());
         return {};
     });
@@ -264,300 +258,34 @@ 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](LogicalSessionIdSet sessions) -> LogicalSessionIdSet {
+    sessions()->setRefreshHook([&refresh2](SessionList sessions) -> SessionList {
         // We should only have one record here, the other should have expired
         ASSERT_EQ(sessions.size(), size_t(1));
-        refresh2.set_value(*(sessions.begin()));
+        refresh2.set_value(sessions.front());
         return {};
     });
 
     // Wait until the second job has been scheduled
-    waitUntilRefreshScheduled();
+    while (service()->jobs() < 2) {
+        sleepmillis(10);
+    }
 
     service()->fastForward(kSessionTimeout - kForceRefresh + Milliseconds(1));
     refresh2Future.wait();
     ASSERT_EQ(refresh2Future.get(), lsid);
 }
 
-// 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) -> LogicalSessionIdSet {
-            ASSERT_EQ(sessions.size(), size_t(2));
-            hitRefresh.set_value();
-            return {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) -> LogicalSessionIdSet {
-            ASSERT_EQ(sessions.size(), size_t(2));
-            hitRefresh.set_value();
-            return {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) -> LogicalSessionIdSet {
-            ASSERT_EQ(*(sessions.begin()), lsid);
-            {
-                stdx::unique_lock<stdx::mutex> lk(mutex);
-                count++;
-            }
-            cv.notify_all();
-
-            return {};
-        });
-
-    // 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) -> LogicalSessionIdSet {
-        ASSERT_EQ(sessions.size(), size_t(2));
-        hitRefresh.set_value();
-        return {};
-    });
-
-    // 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) -> LogicalSessionIdSet {
-            ASSERT_EQ(sessions.size(), size_t(count));
-            hitRefresh.set_value();
-            return {};
-        });
-
-    // 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) -> LogicalSessionIdSet {
-            ASSERT_EQ(sessions.size(), size_t(count));
-            hitRefresh.set_value();
-            return {};
-        });
-
-    // 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) -> LogicalSessionIdSet {
-        {
-            stdx::unique_lock<stdx::mutex> lk(mutex);
-            refreshes++;
-            nRefreshed = sessions.size();
-        }
-        cv.notify_all();
-
-        return {};
-    });
-
-    // 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) -> LogicalSessionIdSet {
-        {
-            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);
-}
+// 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
 
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/logical_session_id.h b/src/mongo/db/logical_session_id.h
index eca1d8a33f5..b67c935220e 100644
--- a/src/mongo/db/logical_session_id.h
+++ b/src/mongo/db/logical_session_id.h
@@ -114,9 +114,4 @@ 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 2389245f6a5..97535cea4e8 100644
--- a/src/mongo/db/service_liason.h
+++ b/src/mongo/db/service_liason.h
@@ -44,13 +44,15 @@ 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 LogicalSessionIdSet getActiveSessions() const = 0;
+    virtual SessionList 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 b7b03ba9ceb..6aed1c0306f 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();
 }
 
-LogicalSessionIdSet MockServiceLiasonImpl::getActiveSessions() const {
+MockServiceLiasonImpl::SessionList MockServiceLiasonImpl::getActiveSessions() const {
     stdx::unique_lock<stdx::mutex> lk(_mutex);
     return _activeSessions;
 }
@@ -61,12 +61,13 @@ void MockServiceLiasonImpl::scheduleJob(PeriodicRunner::PeriodicJob job) {
 
 void MockServiceLiasonImpl::add(LogicalSessionId lsid) {
     stdx::unique_lock<stdx::mutex> lk(_mutex);
-    _activeSessions.insert(std::move(lsid));
+    _activeSessions.push_back(std::move(lsid));
 }
 
 void MockServiceLiasonImpl::remove(LogicalSessionId lsid) {
     stdx::unique_lock<stdx::mutex> lk(_mutex);
-    _activeSessions.erase(lsid);
+    _activeSessions.erase(std::remove(_activeSessions.begin(), _activeSessions.end(), lsid),
+                          _activeSessions.end());
 }
 
 void MockServiceLiasonImpl::clear() {
diff --git a/src/mongo/db/service_liason_mock.h b/src/mongo/db/service_liason_mock.h
index bfd2d281ca1..577057708ad 100644
--- a/src/mongo/db/service_liason_mock.h
+++ b/src/mongo/db/service_liason_mock.h
@@ -54,10 +54,12 @@ namespace mongo {
  */
 class MockServiceLiasonImpl {
 public:
+    using SessionList = std::list<LogicalSessionId>;
+
     MockServiceLiasonImpl();
 
     // Forwarding methods from the MockServiceLiason
-    LogicalSessionIdSet getActiveSessions() const;
+    SessionList getActiveSessions() const;
     Date_t now() const;
     void scheduleJob(PeriodicRunner::PeriodicJob job);
     void join();
@@ -74,7 +76,7 @@ private:
     std::unique_ptr<PeriodicRunnerASIO> _runner;
 
     mutable stdx::mutex _mutex;
-    LogicalSessionIdSet _activeSessions;
+    SessionList _activeSessions;
 };
 
 /**
@@ -82,10 +84,12 @@ private:
  */
 class MockServiceLiason : public ServiceLiason {
 public:
+    using SessionList = std::list<LogicalSessionId>;
+
     explicit MockServiceLiason(std::shared_ptr<MockServiceLiasonImpl> impl)
         : _impl(std::move(impl)) {}
 
-    LogicalSessionIdSet getActiveSessions() const override {
+    SessionList getActiveSessions() const override {
         return _impl->getActiveSessions();
     }
 
diff --git a/src/mongo/db/sessions_collection.h b/src/mongo/db/sessions_collection.h
index ec59f28c597..8e65828033d 100644
--- a/src/mongo/db/sessions_collection.h
+++ b/src/mongo/db/sessions_collection.h
@@ -28,7 +28,6 @@
 
 #pragma once
 
-#include "mongo/db/logical_session_id.h"
 #include "mongo/db/logical_session_record.h"
 
 namespace mongo {
@@ -41,6 +40,8 @@ namespace mongo {
  */
 class SessionsCollection {
 public:
+    using SessionList = std::list<LogicalSessionId>;
+
     virtual ~SessionsCollection();
 
     /**
@@ -65,7 +66,7 @@ public:
      * Returns a list of sessions for which no authoritative record was found,
      * and hence were not refreshed.
      */
-    virtual LogicalSessionIdSet refreshSessions(LogicalSessionIdSet sessions) = 0;
+    virtual SessionList refreshSessions(SessionList sessions) = 0;
 
     /**
      * Removes the authoritative records for the specified sessions.
@@ -73,7 +74,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(LogicalSessionIdSet sessions) = 0;
+    virtual void removeRecords(SessionList sessions) = 0;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/sessions_collection_mock.cpp b/src/mongo/db/sessions_collection_mock.cpp
index 4abc5a326d1..a56f6dec21a 100644
--- a/src/mongo/db/sessions_collection_mock.cpp
+++ b/src/mongo/db/sessions_collection_mock.cpp
@@ -73,11 +73,12 @@ Status MockSessionsCollectionImpl::insertRecord(LogicalSessionRecord record) {
     return _insert(std::move(record));
 }
 
-LogicalSessionIdSet MockSessionsCollectionImpl::refreshSessions(LogicalSessionIdSet sessions) {
+MockSessionsCollectionImpl::SessionList MockSessionsCollectionImpl::refreshSessions(
+    SessionList sessions) {
     return _refresh(std::move(sessions));
 }
 
-void MockSessionsCollectionImpl::removeRecords(LogicalSessionIdSet sessions) {
+void MockSessionsCollectionImpl::removeRecords(SessionList sessions) {
     _remove(std::move(sessions));
 }
 
@@ -130,15 +131,16 @@ Status MockSessionsCollectionImpl::_insertRecord(LogicalSessionRecord record) {
     return Status::OK();
 }
 
-LogicalSessionIdSet MockSessionsCollectionImpl::_refreshSessions(LogicalSessionIdSet sessions) {
-    LogicalSessionIdSet notFound{};
+MockSessionsCollectionImpl::SessionList MockSessionsCollectionImpl::_refreshSessions(
+    SessionList sessions) {
+    SessionList notFound{};
 
     {
         stdx::unique_lock<stdx::mutex> lk(_mutex);
         for (auto& lsid : sessions) {
             auto it = _sessions.find(lsid);
             if (it == _sessions.end()) {
-                notFound.insert(lsid);
+                notFound.push_back(lsid);
             }
         }
     }
@@ -146,7 +148,7 @@ LogicalSessionIdSet MockSessionsCollectionImpl::_refreshSessions(LogicalSessionI
     return notFound;
 }
 
-void MockSessionsCollectionImpl::_removeRecords(LogicalSessionIdSet sessions) {
+void MockSessionsCollectionImpl::_removeRecords(SessionList 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 64e2a76fc7b..05acb1f256f 100644
--- a/src/mongo/db/sessions_collection_mock.h
+++ b/src/mongo/db/sessions_collection_mock.h
@@ -54,6 +54,7 @@ namespace mongo {
  */
 class MockSessionsCollectionImpl {
 public:
+    using SessionList = std::list<LogicalSessionId>;
     using SessionMap =
         stdx::unordered_map<LogicalSessionId, LogicalSessionRecord, LogicalSessionId::Hash>;
 
@@ -61,8 +62,8 @@ public:
 
     using FetchHook = stdx::function<StatusWith<LogicalSessionRecord>(LogicalSessionId)>;
     using InsertHook = stdx::function<Status(LogicalSessionRecord)>;
-    using RefreshHook = stdx::function<LogicalSessionIdSet(LogicalSessionIdSet)>;
-    using RemoveHook = stdx::function<void(LogicalSessionIdSet)>;
+    using RefreshHook = stdx::function<SessionList(SessionList)>;
+    using RemoveHook = stdx::function<void(SessionList)>;
 
     // Set custom hooks to override default behavior
     void setFetchHook(FetchHook hook);
@@ -76,8 +77,8 @@ public:
     // Forwarding methods from the MockSessionsCollection
     StatusWith<LogicalSessionRecord> fetchRecord(LogicalSessionId lsid);
     Status insertRecord(LogicalSessionRecord record);
-    LogicalSessionIdSet refreshSessions(LogicalSessionIdSet sessions);
-    void removeRecords(LogicalSessionIdSet sessions);
+    SessionList refreshSessions(SessionList sessions);
+    void removeRecords(SessionList sessions);
 
     // Test-side methods that operate on the _sessions map
     void add(LogicalSessionRecord record);
@@ -90,8 +91,8 @@ private:
     // Default implementations, may be overridden with custom hooks.
     StatusWith<LogicalSessionRecord> _fetchRecord(LogicalSessionId lsid);
     Status _insertRecord(LogicalSessionRecord record);
-    LogicalSessionIdSet _refreshSessions(LogicalSessionIdSet sessions);
-    void _removeRecords(LogicalSessionIdSet sessions);
+    SessionList _refreshSessions(SessionList sessions);
+    void _removeRecords(SessionList sessions);
 
     stdx::mutex _mutex;
     SessionMap _sessions;
@@ -109,6 +110,8 @@ private:
  */
 class MockSessionsCollection : public SessionsCollection {
 public:
+    using SessionList = std::list<LogicalSessionId>;
+
     explicit MockSessionsCollection(std::shared_ptr<MockSessionsCollectionImpl> impl)
         : _impl(std::move(impl)) {}
 
@@ -120,11 +123,11 @@ public:
         return _impl->insertRecord(std::move(record));
     }
 
-    LogicalSessionIdSet refreshSessions(LogicalSessionIdSet sessions) override {
+    SessionList refreshSessions(SessionList sessions) override {
         return _impl->refreshSessions(std::move(sessions));
     }
 
-    void removeRecords(LogicalSessionIdSet sessions) override {
+    void removeRecords(SessionList sessions) override {
         return _impl->removeRecords(std::move(sessions));
     }