SERVER-28300 Implement mock libraries to test logical session cache

8 files changed, 849 insertions, 0 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 79aaf13f2c2..6c7c43215ff 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -928,6 +928,21 @@ env.Library(
     ],
 )
 
+envWithAsio = env.Clone()
+envWithAsio.InjectThirdPartyIncludePaths(libraries=['asio'])
+
+envWithAsio.Library(
+    target='service_liason_mock',
+    source=[
+        'service_liason_mock.cpp',
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/executor/async_timer_mock',
+        '$BUILD_DIR/mongo/util/periodic_runner_asio',
+        'service_liason',
+    ],
+)
+
 env.Library(
     target='sessions_collection',
     source=[
@@ -938,6 +953,18 @@ env.Library(
 )
 
 env.Library(
+    target='sessions_collection_mock',
+    source=[
+        'sessions_collection_mock.cpp',
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/base',
+        '$BUILD_DIR/mongo/db/logical_session_id',
+        'sessions_collection',
+    ],
+)
+
+env.Library(
     target='logical_session_cache',
     source=[
         'logical_session_cache.cpp',
@@ -949,6 +976,19 @@ env.Library(
     ],
 )
 
+envWithAsio.CppUnitTest(
+    target='logical_session_cache_test',
+    source=[
+        'logical_session_cache_test.cpp',
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/executor/async_timer_mock',
+        'logical_session_cache',
+        'service_liason_mock',
+        'sessions_collection_mock',
+    ],
+)
+
 env.Library(
     target='logical_time',
     source=[
diff --git a/src/mongo/db/logical_session_cache_test.cpp b/src/mongo/db/logical_session_cache_test.cpp
new file mode 100644
index 00000000000..32e0771e965
--- /dev/null
+++ b/src/mongo/db/logical_session_cache_test.cpp
@@ -0,0 +1,290 @@
+/**
+ *    Copyright (C) 2017 MongoDB Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/bson/oid.h"
+#include "mongo/db/auth/user_name.h"
+#include "mongo/db/logical_session_cache.h"
+#include "mongo/db/logical_session_id.h"
+#include "mongo/db/logical_session_record.h"
+#include "mongo/db/service_liason_mock.h"
+#include "mongo/db/sessions_collection_mock.h"
+#include "mongo/stdx/future.h"
+#include "mongo/unittest/unittest.h"
+
+namespace mongo {
+namespace {
+
+const Milliseconds kSessionTimeout =
+    duration_cast<Milliseconds>(LogicalSessionCache::kLogicalSessionDefaultTimeout);
+const Milliseconds kForceRefresh =
+    duration_cast<Milliseconds>(LogicalSessionCache::kLogicalSessionDefaultRefresh);
+
+using SessionList = std::list<LogicalSessionId>;
+
+/**
+ * Test fixture that sets up a session cache attached to a mock service liason
+ * and mock sessions collection implementation.
+ */
+class LogicalSessionCacheTest : public unittest::Test {
+public:
+    LogicalSessionCacheTest()
+        : _service(std::make_shared<MockServiceLiasonImpl>()),
+          _sessions(std::make_shared<MockSessionsCollectionImpl>()),
+          _user("sam", "test"),
+          _userId(OID::gen()) {}
+
+    void setUp() override {
+        auto mockService = std::make_unique<MockServiceLiason>(_service);
+        auto mockSessions = std::make_unique<MockSessionsCollection>(_sessions);
+        _cache =
+            std::make_unique<LogicalSessionCache>(std::move(mockService), std::move(mockSessions));
+    }
+
+    void tearDown() override {
+        _service->join();
+    }
+
+    LogicalSessionRecord newRecord() {
+        return LogicalSessionRecord::makeAuthoritativeRecord(
+            LogicalSessionId::gen(), _user, _userId, _service->now());
+    }
+
+    std::unique_ptr<LogicalSessionCache>& cache() {
+        return _cache;
+    }
+
+    std::shared_ptr<MockServiceLiasonImpl> service() {
+        return _service;
+    }
+
+    std::shared_ptr<MockSessionsCollectionImpl> sessions() {
+        return _sessions;
+    }
+
+private:
+    std::shared_ptr<MockServiceLiasonImpl> _service;
+    std::shared_ptr<MockSessionsCollectionImpl> _sessions;
+
+    std::unique_ptr<LogicalSessionCache> _cache;
+
+    UserName _user;
+    boost::optional<OID> _userId;
+};
+
+// Test that session cache fetches new records from the sessions collection
+TEST_F(LogicalSessionCacheTest, CacheFetchesNewRecords) {
+    auto record = newRecord();
+    auto lsid = record.getLsid();
+
+    // When the record is not present (and not in the sessions collection) returns an error
+    auto res = cache()->getOwner(lsid);
+    ASSERT(!res.isOK());
+
+    // When the record is not present (but is in the sessions collection) returns it
+    sessions()->add(record);
+    res = cache()->getOwner(lsid);
+    ASSERT(res.isOK());
+    ASSERT(res.getValue() == record.getSessionOwner());
+
+    // When the record is present in the cache, returns it
+    sessions()->setFetchHook([](LogicalSessionId id) -> StatusWith<LogicalSessionRecord> {
+        // We should not be querying the sessions collection on the next call
+        ASSERT(false);
+        return {ErrorCodes::NoSuchSession, "no such session"};
+    });
+
+    res = cache()->getOwner(lsid);
+    ASSERT(res.isOK());
+    ASSERT(res.getValue() == record.getSessionOwner());
+}
+
+// Test that the getFromCache method does not make calls to the sessions collection
+TEST_F(LogicalSessionCacheTest, TestCacheHitsOnly) {
+    auto record = newRecord();
+    auto lsid = record.getLsid();
+
+    // When the record is not present (and not in the sessions collection), returns an error
+    auto res = cache()->getOwnerFromCache(lsid);
+    ASSERT(!res.isOK());
+
+    // When the record is not present (but is in the sessions collection), returns an error
+    sessions()->add(record);
+    res = cache()->getOwnerFromCache(lsid);
+    ASSERT(!res.isOK());
+
+    // When the record is present, returns the owner
+    cache()->getOwner(lsid);
+    res = cache()->getOwnerFromCache(lsid);
+    ASSERT(res.isOK());
+    auto fetched = res.getValue();
+    ASSERT(res.getValue() == record.getSessionOwner());
+}
+
+// Test that fetching from the cache updates the lastUse date of records
+TEST_F(LogicalSessionCacheTest, FetchUpdatesLastUse) {
+    auto record = newRecord();
+    auto lsid = record.getLsid();
+
+    auto start = service()->now();
+
+    // Insert the record into the sessions collection with 'start'
+    record.setLastUse(start);
+    sessions()->add(record);
+
+    // Fast forward time and fetch
+    service()->fastForward(Milliseconds(500));
+    ASSERT(start != service()->now());
+    auto res = cache()->getOwner(lsid);
+    ASSERT(res.isOK());
+
+    // Now that we fetched, lifetime of session should be extended
+    service()->fastForward(kSessionTimeout - Milliseconds(500));
+    res = cache()->getOwner(lsid);
+    ASSERT(res.isOK());
+
+    // We fetched again, so lifetime extended again
+    service()->fastForward(kSessionTimeout - Milliseconds(10));
+    res = cache()->getOwner(lsid);
+    ASSERT(res.isOK());
+
+    // Fast forward and hit-only fetch
+    service()->fastForward(kSessionTimeout - Milliseconds(10));
+    res = cache()->getOwnerFromCache(lsid);
+    ASSERT(res.isOK());
+
+    // Lifetime extended again
+    service()->fastForward(Milliseconds(11));
+    res = cache()->getOwnerFromCache(lsid);
+    ASSERT(res.isOK());
+
+    // Let record expire, we should not be able to get it from the cache
+    service()->fastForward(kSessionTimeout + Milliseconds(1));
+    res = cache()->getOwnerFromCache(lsid);
+    ASSERT(!res.isOK());
+}
+
+// Test the startSession method
+TEST_F(LogicalSessionCacheTest, StartSession) {
+    auto record = newRecord();
+    auto lsid = record.getLsid();
+
+    // Test starting a new session
+    auto res = cache()->startSession(record);
+    ASSERT(res.isOK());
+    ASSERT(sessions()->has(lsid));
+
+    // Try to start a session that is already in the sessions collection and our
+    // local cache, should fail
+    res = cache()->startSession(record);
+    ASSERT(!res.isOK());
+
+    // Try to start a session that is already in the sessions collection but
+    // is not in our local cache, should fail
+    auto record2 = newRecord();
+    sessions()->add(record2);
+    res = cache()->startSession(record2);
+    ASSERT(!res.isOK());
+
+    // 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));
+    res = cache()->startSession(record);
+    ASSERT(res.isOK());
+    ASSERT(sessions()->has(lsid));
+}
+
+// Test that records in the cache are properly refreshed until they expire
+TEST_F(LogicalSessionCacheTest, CacheRefreshesOwnRecords) {
+    // Insert two records into the cache
+    auto record1 = newRecord();
+    auto record2 = newRecord();
+    cache()->startSession(record1);
+    cache()->startSession(record2);
+
+    stdx::promise<int> hitRefresh;
+    auto refreshFuture = hitRefresh.get_future();
+
+    // Advance time to first refresh point, check that refresh happens, and
+    // that it includes both our records
+    sessions()->setRefreshHook([&hitRefresh](SessionList sessions) -> SessionList {
+        hitRefresh.set_value(sessions.size());
+        return {};
+    });
+
+    // Wait for the refresh to happen
+    service()->fastForward(kForceRefresh);
+    refreshFuture.wait();
+    ASSERT_EQ(refreshFuture.get(), 2);
+
+    sessions()->clearHooks();
+
+    stdx::promise<LogicalSessionId> refresh2;
+    auto refresh2Future = refresh2.get_future();
+
+    // Use one of the records
+    auto lsid = record1.getLsid();
+    auto res = cache()->getOwner(lsid);
+    ASSERT(res.isOK());
+
+    // Advance time so that one record expires
+    // Ensure that first record was refreshed, and second was thrown away
+    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.front());
+        return {};
+    });
+
+    // Wait until the second job has been scheduled
+    while (service()->jobs() < 2) {
+        sleepmillis(10);
+    }
+
+    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
+
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/db/service_liason_mock.cpp b/src/mongo/db/service_liason_mock.cpp
new file mode 100644
index 00000000000..99c63101d15
--- /dev/null
+++ b/src/mongo/db/service_liason_mock.cpp
@@ -0,0 +1,85 @@
+/**
+ *    Copyright (C) 2017 MongoDB Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include <algorithm>
+
+#include "mongo/db/service_liason_mock.h"
+
+namespace mongo {
+
+MockServiceLiasonImpl::MockServiceLiasonImpl() {
+    auto timerFactory = std::make_unique<executor::AsyncTimerFactoryMock>();
+    _timerFactory = timerFactory.get();
+    _runner = std::make_unique<PeriodicRunnerASIO>(std::move(timerFactory));
+    _runner->startup();
+}
+
+MockServiceLiasonImpl::SessionList MockServiceLiasonImpl::getActiveSessions() const {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    return _activeSessions;
+}
+
+void MockServiceLiasonImpl::join() {
+    _runner->shutdown();
+}
+
+Date_t MockServiceLiasonImpl::now() const {
+    return _timerFactory->now();
+}
+
+void MockServiceLiasonImpl::scheduleJob(PeriodicRunner::PeriodicJob job) {
+    _runner->scheduleJob(std::move(job));
+}
+
+void MockServiceLiasonImpl::add(LogicalSessionId lsid) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    _activeSessions.push_back(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());
+}
+
+void MockServiceLiasonImpl::clear() {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    _activeSessions.clear();
+}
+
+void MockServiceLiasonImpl::fastForward(Milliseconds time) {
+    _timerFactory->fastForward(time);
+}
+
+int MockServiceLiasonImpl::jobs() {
+    return _timerFactory->jobs();
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/service_liason_mock.h b/src/mongo/db/service_liason_mock.h
new file mode 100644
index 00000000000..577057708ad
--- /dev/null
+++ b/src/mongo/db/service_liason_mock.h
@@ -0,0 +1,112 @@
+/**
+ *    Copyright (C) 2017 MongoDB Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/service_liason.h"
+#include "mongo/executor/async_timer_mock.h"
+#include "mongo/platform/atomic_word.h"
+#include "mongo/stdx/condition_variable.h"
+#include "mongo/stdx/mutex.h"
+#include "mongo/util/periodic_runner_asio.h"
+#include "mongo/util/time_support.h"
+
+namespace mongo {
+
+/**
+ * To allow us to move a MockServiceLiason into the session cache while
+ * maintaining a hold on it from within our unit tests, the MockServiceLiason
+ * will have an internal pointer to a MockServiceLiasonImpl object that the
+ * test creates and controls.
+ *
+ * This class maintains an internal _activeSessions list that may be modified
+ * by the test caller. It also maintains an internal mocked representation of
+ * time, which the caller can fastForward(). The Date_t returned by now() will
+ * be the epoch + the amount of minutes this object has been fast-forwarded over
+ * course of its life.
+ *
+ * This service liason starts up its internal periodic runner on construction.
+ */
+class MockServiceLiasonImpl {
+public:
+    using SessionList = std::list<LogicalSessionId>;
+
+    MockServiceLiasonImpl();
+
+    // Forwarding methods from the MockServiceLiason
+    SessionList getActiveSessions() const;
+    Date_t now() const;
+    void scheduleJob(PeriodicRunner::PeriodicJob job);
+    void join();
+
+    // Test-side methods that operate on the _activeSessions list.
+    void add(LogicalSessionId lsid);
+    void remove(LogicalSessionId lsid);
+    void clear();
+    void fastForward(Milliseconds time);
+    int jobs();
+
+private:
+    executor::AsyncTimerFactoryMock* _timerFactory;
+    std::unique_ptr<PeriodicRunnerASIO> _runner;
+
+    mutable stdx::mutex _mutex;
+    SessionList _activeSessions;
+};
+
+/**
+ * A mock service liason for testing the logical session cache.
+ */
+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 {
+        return _impl->getActiveSessions();
+    }
+
+    Date_t now() const override {
+        return _impl->now();
+    }
+
+    void scheduleJob(PeriodicRunner::PeriodicJob job) override {
+        _impl->scheduleJob(std::move(job));
+    }
+
+    void join() override {
+        return _impl->join();
+    }
+
+private:
+    std::shared_ptr<MockServiceLiasonImpl> _impl;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/sessions_collection_mock.cpp b/src/mongo/db/sessions_collection_mock.cpp
new file mode 100644
index 00000000000..a56f6dec21a
--- /dev/null
+++ b/src/mongo/db/sessions_collection_mock.cpp
@@ -0,0 +1,158 @@
+/**
+ *    Copyright (C) 2017 MongoDB Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/db/sessions_collection_mock.h"
+#include "mongo/platform/basic.h"
+#include "mongo/stdx/functional.h"
+
+namespace mongo {
+
+MockSessionsCollectionImpl::MockSessionsCollectionImpl()
+    : _sessions(),
+      _fetch(stdx::bind(&MockSessionsCollectionImpl::_fetchRecord, this, stdx::placeholders::_1)),
+      _insert(stdx::bind(&MockSessionsCollectionImpl::_insertRecord, this, stdx::placeholders::_1)),
+      _refresh(
+          stdx::bind(&MockSessionsCollectionImpl::_refreshSessions, this, stdx::placeholders::_1)),
+      _remove(
+          stdx::bind(&MockSessionsCollectionImpl::_removeRecords, this, stdx::placeholders::_1)) {}
+
+void MockSessionsCollectionImpl::setFetchHook(FetchHook hook) {
+    _fetch = std::move(hook);
+}
+
+void MockSessionsCollectionImpl::setInsertHook(InsertHook hook) {
+    _insert = std::move(hook);
+}
+
+void MockSessionsCollectionImpl::setRefreshHook(RefreshHook hook) {
+    _refresh = std::move(hook);
+}
+
+void MockSessionsCollectionImpl::setRemoveHook(RemoveHook hook) {
+    _remove = std::move(hook);
+}
+
+void MockSessionsCollectionImpl::clearHooks() {
+    _fetch = stdx::bind(&MockSessionsCollectionImpl::_fetchRecord, this, stdx::placeholders::_1);
+    _insert = stdx::bind(&MockSessionsCollectionImpl::_insertRecord, this, stdx::placeholders::_1);
+    _refresh =
+        stdx::bind(&MockSessionsCollectionImpl::_refreshSessions, this, stdx::placeholders::_1);
+    _remove = stdx::bind(&MockSessionsCollectionImpl::_removeRecords, this, stdx::placeholders::_1);
+}
+
+StatusWith<LogicalSessionRecord> MockSessionsCollectionImpl::fetchRecord(LogicalSessionId lsid) {
+    return _fetch(std::move(lsid));
+}
+
+Status MockSessionsCollectionImpl::insertRecord(LogicalSessionRecord record) {
+    return _insert(std::move(record));
+}
+
+MockSessionsCollectionImpl::SessionList MockSessionsCollectionImpl::refreshSessions(
+    SessionList sessions) {
+    return _refresh(std::move(sessions));
+}
+
+void MockSessionsCollectionImpl::removeRecords(SessionList sessions) {
+    _remove(std::move(sessions));
+}
+
+void MockSessionsCollectionImpl::add(LogicalSessionRecord record) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    _sessions.insert({record.getLsid(), std::move(record)});
+}
+
+void MockSessionsCollectionImpl::remove(LogicalSessionId lsid) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    _sessions.erase(lsid);
+}
+
+bool MockSessionsCollectionImpl::has(LogicalSessionId lsid) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    return _sessions.find(lsid) != _sessions.end();
+}
+
+void MockSessionsCollectionImpl::clearSessions() {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    _sessions.clear();
+}
+
+const MockSessionsCollectionImpl::SessionMap& MockSessionsCollectionImpl::sessions() const {
+    return _sessions;
+}
+
+StatusWith<LogicalSessionRecord> MockSessionsCollectionImpl::_fetchRecord(LogicalSessionId lsid) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+
+    // If we do not have this record, return an error
+    auto it = _sessions.find(lsid);
+    if (it == _sessions.end()) {
+        return {ErrorCodes::NoSuchSession, "No matching record in the sessions collection"};
+    }
+
+    return it->second;
+}
+
+Status MockSessionsCollectionImpl::_insertRecord(LogicalSessionRecord record) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    auto res = _sessions.insert({record.getLsid(), std::move(record)});
+
+    // We should never try to insert the same record twice. In theory this could
+    // happen because of a UUID conflict.
+    if (!res.second) {
+        return {ErrorCodes::DuplicateSession, "Session already exists in the sessions collection"};
+    }
+
+    return Status::OK();
+}
+
+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.push_back(lsid);
+            }
+        }
+    }
+
+    return notFound;
+}
+
+void MockSessionsCollectionImpl::_removeRecords(SessionList sessions) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    for (auto& lsid : sessions) {
+        _sessions.erase(lsid);
+    }
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/sessions_collection_mock.h b/src/mongo/db/sessions_collection_mock.h
new file mode 100644
index 00000000000..05acb1f256f
--- /dev/null
+++ b/src/mongo/db/sessions_collection_mock.h
@@ -0,0 +1,138 @@
+/**
+ *    Copyright (C) 2017 MongoDB Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/logical_session_id.h"
+#include "mongo/db/logical_session_record.h"
+#include "mongo/db/sessions_collection.h"
+#include "mongo/stdx/functional.h"
+#include "mongo/stdx/mutex.h"
+#include "mongo/stdx/unordered_map.h"
+
+namespace mongo {
+
+/**
+ * To allow us to move a MockSessionCollection into the session cache while
+ * maintaining a hold on it from within our unit tests, the MockSessionCollection
+ * will have an internal pointer to a MockSessionsCollectionImpl object that the
+ * test creates and controls.
+ *
+ * This class can operate in two modes:
+ *
+ * - if no custom hooks are set, then the test caller may add() and remove() items
+ *   from the _sessions map, and this class will simply perform the required
+ *   operations against that set.
+ *
+ * - if custom hooks are set, then the hooks will be run instead of the provided
+ *   defaults, and the internal _sessions map will NOT be updated.
+ */
+class MockSessionsCollectionImpl {
+public:
+    using SessionList = std::list<LogicalSessionId>;
+    using SessionMap =
+        stdx::unordered_map<LogicalSessionId, LogicalSessionRecord, LogicalSessionId::Hash>;
+
+    MockSessionsCollectionImpl();
+
+    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)>;
+
+    // Set custom hooks to override default behavior
+    void setFetchHook(FetchHook hook);
+    void setInsertHook(InsertHook hook);
+    void setRefreshHook(RefreshHook hook);
+    void setRemoveHook(RemoveHook hook);
+
+    // Reset all hooks to their defaults
+    void clearHooks();
+
+    // Forwarding methods from the MockSessionsCollection
+    StatusWith<LogicalSessionRecord> fetchRecord(LogicalSessionId lsid);
+    Status insertRecord(LogicalSessionRecord record);
+    SessionList refreshSessions(SessionList sessions);
+    void removeRecords(SessionList sessions);
+
+    // Test-side methods that operate on the _sessions map
+    void add(LogicalSessionRecord record);
+    void remove(LogicalSessionId lsid);
+    bool has(LogicalSessionId lsid);
+    void clearSessions();
+    const SessionMap& sessions() const;
+
+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);
+
+    stdx::mutex _mutex;
+    SessionMap _sessions;
+
+    FetchHook _fetch;
+    InsertHook _insert;
+    RefreshHook _refresh;
+    RemoveHook _remove;
+};
+
+/**
+ * To allow us to move this into the session cache while maintaining a hold
+ * on it from the test side, the MockSessionCollection will have an internal pointer
+ * to an impl that we maintain access to.
+ */
+class MockSessionsCollection : public SessionsCollection {
+public:
+    using SessionList = std::list<LogicalSessionId>;
+
+    explicit MockSessionsCollection(std::shared_ptr<MockSessionsCollectionImpl> impl)
+        : _impl(std::move(impl)) {}
+
+    StatusWith<LogicalSessionRecord> fetchRecord(LogicalSessionId lsid) override {
+        return _impl->fetchRecord(std::move(lsid));
+    }
+
+    Status insertRecord(LogicalSessionRecord record) override {
+        return _impl->insertRecord(std::move(record));
+    }
+
+    SessionList refreshSessions(SessionList sessions) override {
+        return _impl->refreshSessions(std::move(sessions));
+    }
+
+    void removeRecords(SessionList sessions) override {
+        return _impl->removeRecords(std::move(sessions));
+    }
+
+private:
+    std::shared_ptr<MockSessionsCollectionImpl> _impl;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/executor/async_timer_mock.cpp b/src/mongo/executor/async_timer_mock.cpp
index 27b09fbb620..39c9323612d 100644
--- a/src/mongo/executor/async_timer_mock.cpp
+++ b/src/mongo/executor/async_timer_mock.cpp
@@ -99,6 +99,11 @@ void AsyncTimerMockImpl::expireAfter(Milliseconds expiration) {
     }
 }
 
+int AsyncTimerMockImpl::jobs() {
+    stdx::lock_guard<stdx::mutex> lk(_mutex);
+    return _handlers.size();
+}
+
 void AsyncTimerMockImpl::_callAllHandlers(std::error_code ec) {
     std::vector<AsyncTimerInterface::Handler> tmp;
     {
@@ -153,5 +158,16 @@ Date_t AsyncTimerFactoryMock::now() {
     return Date_t::fromDurationSinceEpoch(_curTime);
 }
 
+int AsyncTimerFactoryMock::jobs() {
+    int jobs = 1;
+
+    stdx::lock_guard<stdx::recursive_mutex> lk(_timersMutex);
+    for (auto elem = _timers.begin(); elem != _timers.end(); elem++) {
+        jobs += (*elem)->jobs();
+    }
+
+    return jobs;
+}
+
 }  // namespace executor
 }  // namespace mongo
diff --git a/src/mongo/executor/async_timer_mock.h b/src/mongo/executor/async_timer_mock.h
index c0595b95364..ad77d323f5e 100644
--- a/src/mongo/executor/async_timer_mock.h
+++ b/src/mongo/executor/async_timer_mock.h
@@ -75,6 +75,11 @@ public:
      */
     void expireAfter(Milliseconds expiration);
 
+    /**
+     * Returns the number of handlers on this timer.
+     */
+    int jobs();
+
 private:
     void _callAllHandlers(std::error_code ec);
 
@@ -137,6 +142,11 @@ public:
      */
     Date_t now() override;
 
+    /**
+     * Returns the number of pending jobs across all timers.
+     */
+    int jobs();
+
 private:
     stdx::recursive_mutex _timersMutex;
     stdx::unordered_set<std::shared_ptr<AsyncTimerMockImpl>> _timers;