SERVER-29672 Make Config/ShardCatalogCacheLoader decorations on ServiceContext

12 files changed, 102 insertions, 65 deletions

diff --git a/src/mongo/s/catalog_cache.cpp b/src/mongo/s/catalog_cache.cpp
index a65ff175b3d..1489e0215d3 100644
--- a/src/mongo/s/catalog_cache.cpp
+++ b/src/mongo/s/catalog_cache.cpp
@@ -111,35 +111,32 @@ std::shared_ptr<ChunkManager> refreshCollectionRoutingInfo(
 
 }  // namespace
 
-CatalogCache::CatalogCache() : _cacheLoader(stdx::make_unique<ConfigServerCatalogCacheLoader>()) {}
-
-CatalogCache::CatalogCache(std::unique_ptr<CatalogCacheLoader> cacheLoader)
-    : _cacheLoader(std::move(cacheLoader)) {}
+CatalogCache::CatalogCache(CatalogCacheLoader& cacheLoader) : _cacheLoader(cacheLoader) {}
 
 CatalogCache::~CatalogCache() = default;
 
 void CatalogCache::initializeReplicaSetRole(bool isPrimary) {
-    _cacheLoader->initializeReplicaSetRole(isPrimary);
+    _cacheLoader.initializeReplicaSetRole(isPrimary);
 }
 
 void CatalogCache::onStepDown() {
-    _cacheLoader->onStepDown();
+    _cacheLoader.onStepDown();
 }
 
 void CatalogCache::onStepUp() {
-    _cacheLoader->onStepUp();
+    _cacheLoader.onStepUp();
 }
 
 void CatalogCache::notifyOfCollectionVersionUpdate(OperationContext* opCtx,
                                                    const NamespaceString& nss,
                                                    const ChunkVersion& version) {
-    _cacheLoader->notifyOfCollectionVersionUpdate(opCtx, nss, version);
+    _cacheLoader.notifyOfCollectionVersionUpdate(opCtx, nss, version);
 }
 
 Status CatalogCache::waitForCollectionVersion(OperationContext* opCtx,
                                               const NamespaceString& nss,
                                               const ChunkVersion& version) {
-    return _cacheLoader->waitForCollectionVersion(opCtx, nss, version);
+    return _cacheLoader.waitForCollectionVersion(opCtx, nss, version);
 }
 
 StatusWith<CachedDatabaseInfo> CatalogCache::getDatabase(OperationContext* opCtx,
@@ -419,7 +416,7 @@ void CatalogCache::_scheduleCollectionRefresh_inlock(
           << startingCollectionVersion;
 
     try {
-        _cacheLoader->getChunksSince(nss, startingCollectionVersion, refreshCallback);
+        _cacheLoader.getChunksSince(nss, startingCollectionVersion, refreshCallback);
     } catch (const DBException& ex) {
         const auto status = ex.toStatus();
 
diff --git a/src/mongo/s/catalog_cache.h b/src/mongo/s/catalog_cache.h
index 9b261de3a43..8b81c982645 100644
--- a/src/mongo/s/catalog_cache.h
+++ b/src/mongo/s/catalog_cache.h
@@ -34,7 +34,6 @@
 #include "mongo/s/chunk_manager.h"
 #include "mongo/s/chunk_version.h"
 #include "mongo/s/client/shard.h"
-#include "mongo/s/config_server_catalog_cache_loader.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/stdx/mutex.h"
 #include "mongo/util/concurrency/notification.h"
@@ -57,13 +56,7 @@ class CatalogCache {
     MONGO_DISALLOW_COPYING(CatalogCache);
 
 public:
-    /**
-     * Defaults to instantiating a ConfigServerCatalogCacheLoader.
-     */
-    CatalogCache();
-
-    CatalogCache(std::unique_ptr<CatalogCacheLoader> cacheLoader);
-
+    CatalogCache(CatalogCacheLoader& cacheLoader);
     ~CatalogCache();
 
     /**
@@ -222,7 +215,7 @@ private:
                                            int refreshAttempt);
 
     // Interface from which chunks will be retrieved
-    const std::unique_ptr<CatalogCacheLoader> _cacheLoader;
+    CatalogCacheLoader& _cacheLoader;
 
     // Mutex to serialize access to the structures below
     stdx::mutex _mutex;
diff --git a/src/mongo/s/catalog_cache_loader.cpp b/src/mongo/s/catalog_cache_loader.cpp
index 0f535b158ee..ba81c0fe771 100644
--- a/src/mongo/s/catalog_cache_loader.cpp
+++ b/src/mongo/s/catalog_cache_loader.cpp
@@ -31,6 +31,12 @@
 #include "mongo/s/catalog_cache_loader.h"
 
 namespace mongo {
+namespace {
+
+const auto catalogCacheLoaderDecoration =
+    ServiceContext::declareDecoration<std::unique_ptr<CatalogCacheLoader>>();
+
+}  // namespace
 
 CatalogCacheLoader::CollectionAndChangedChunks::CollectionAndChangedChunks() = default;
 
@@ -46,4 +52,30 @@ CatalogCacheLoader::CollectionAndChangedChunks::CollectionAndChangedChunks(
       shardKeyIsUnique(collShardKeyIsUnique),
       changedChunks(chunks) {}
 
+void CatalogCacheLoader::set(ServiceContext* serviceContext,
+                             std::unique_ptr<CatalogCacheLoader> loader) {
+    auto& catalogCacheLoader = catalogCacheLoaderDecoration(serviceContext);
+    invariant(!catalogCacheLoader);
+
+    catalogCacheLoader = std::move(loader);
+}
+
+void CatalogCacheLoader::clearForTests(ServiceContext* serviceContext) {
+    auto& catalogCacheLoader = catalogCacheLoaderDecoration(serviceContext);
+    invariant(catalogCacheLoader);
+
+    catalogCacheLoader.reset();
+}
+
+CatalogCacheLoader& CatalogCacheLoader::get(ServiceContext* serviceContext) {
+    auto& catalogCacheLoader = catalogCacheLoaderDecoration(serviceContext);
+    invariant(catalogCacheLoader);
+
+    return *catalogCacheLoader;
+}
+
+CatalogCacheLoader& CatalogCacheLoader::get(OperationContext* opCtx) {
+    return get(opCtx->getServiceContext());
+}
+
 }  // namespace mongo
diff --git a/src/mongo/s/catalog_cache_loader.h b/src/mongo/s/catalog_cache_loader.h
index 3257cb8fbc5..8f2b2d3067d 100644
--- a/src/mongo/s/catalog_cache_loader.h
+++ b/src/mongo/s/catalog_cache_loader.h
@@ -36,6 +36,7 @@
 #include "mongo/s/catalog/type_chunk.h"
 #include "mongo/s/catalog/type_collection.h"
 #include "mongo/s/chunk_version.h"
+#include "mongo/stdx/memory.h"
 #include "mongo/util/concurrency/notification.h"
 
 namespace mongo {
@@ -52,6 +53,16 @@ public:
     virtual ~CatalogCacheLoader() = default;
 
     /**
+     * Stores a loader on the specified service context. May only be called once for the lifetime of
+     * the service context.
+     */
+    static void set(ServiceContext* serviceContext, std::unique_ptr<CatalogCacheLoader> loader);
+
+    static CatalogCacheLoader& get(ServiceContext* serviceContext);
+    static CatalogCacheLoader& get(OperationContext* opCtx);
+
+
+    /**
      * Used as a return value for getChunksSince.
      */
     struct CollectionAndChangedChunks {
@@ -120,6 +131,12 @@ public:
         stdx::function<void(OperationContext*, StatusWith<CollectionAndChangedChunks>)>
             callbackFn) = 0;
 
+    /**
+     * Only used for unit-tests, clears a previously-created catalog cache loader from the specified
+     * service context, so that 'create' can be called again.
+     */
+    static void clearForTests(ServiceContext* serviceContext);
+
 protected:
     CatalogCacheLoader() = default;
 };
diff --git a/src/mongo/s/config_server_test_fixture.cpp b/src/mongo/s/config_server_test_fixture.cpp
index 90b815c9eaf..22bc7ac4ccb 100644
--- a/src/mongo/s/config_server_test_fixture.cpp
+++ b/src/mongo/s/config_server_test_fixture.cpp
@@ -123,6 +123,9 @@ void ConfigServerTestFixture::setUp() {
     _addShardNetworkTestEnv =
         stdx::make_unique<NetworkTestEnv>(_executorForAddShard, _mockNetworkForAddShard);
 
+    CatalogCacheLoader::set(getServiceContext(),
+                            stdx::make_unique<ConfigServerCatalogCacheLoader>());
+
     uassertStatusOK(initializeGlobalShardingStateForMongodForTest(ConnectionString::forLocal()));
 }
 
@@ -133,6 +136,8 @@ void ConfigServerTestFixture::tearDown() {
 
     ShardingCatalogManager::clearForTests(getServiceContext());
 
+    CatalogCacheLoader::clearForTests(getServiceContext());
+
     ShardingMongodTestFixture::tearDown();
 }
 
@@ -157,14 +162,8 @@ std::unique_ptr<ShardingCatalogClient> ConfigServerTestFixture::makeShardingCata
     return stdx::make_unique<ShardingCatalogClientImpl>(std::move(distLockManager));
 }
 
-std::unique_ptr<CatalogCacheLoader> ConfigServerTestFixture::makeCatalogCacheLoader() {
-    return stdx::make_unique<ConfigServerCatalogCacheLoader>();
-}
-
-std::unique_ptr<CatalogCache> ConfigServerTestFixture::makeCatalogCache(
-    std::unique_ptr<CatalogCacheLoader> catalogCacheLoader) {
-    invariant(catalogCacheLoader);
-    return stdx::make_unique<CatalogCache>(std::move(catalogCacheLoader));
+std::unique_ptr<CatalogCache> ConfigServerTestFixture::makeCatalogCache() {
+    return stdx::make_unique<CatalogCache>(CatalogCacheLoader::get(getServiceContext()));
 }
 
 std::unique_ptr<BalancerConfiguration> ConfigServerTestFixture::makeBalancerConfiguration() {
diff --git a/src/mongo/s/config_server_test_fixture.h b/src/mongo/s/config_server_test_fixture.h
index 651887617ce..37ac1f5c570 100644
--- a/src/mongo/s/config_server_test_fixture.h
+++ b/src/mongo/s/config_server_test_fixture.h
@@ -146,10 +146,7 @@ protected:
     std::unique_ptr<ShardingCatalogClient> makeShardingCatalogClient(
         std::unique_ptr<DistLockManager> distLockManager) override;
 
-    std::unique_ptr<CatalogCacheLoader> makeCatalogCacheLoader() override;
-
-    std::unique_ptr<CatalogCache> makeCatalogCache(
-        std::unique_ptr<CatalogCacheLoader> catalogCacheLoader) override;
+    std::unique_ptr<CatalogCache> makeCatalogCache() override;
 
     std::unique_ptr<ClusterCursorManager> makeClusterCursorManager() override;
 
diff --git a/src/mongo/s/server.cpp b/src/mongo/s/server.cpp
index 5e9235d8485..f388d1bf493 100644
--- a/src/mongo/s/server.cpp
+++ b/src/mongo/s/server.cpp
@@ -75,6 +75,7 @@
 #include "mongo/s/client/shard_registry.h"
 #include "mongo/s/client/shard_remote.h"
 #include "mongo/s/client/sharding_connection_hook.h"
+#include "mongo/s/config_server_catalog_cache_loader.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/is_mongos.h"
 #include "mongo/s/mongos_options.h"
@@ -158,8 +159,9 @@ static void cleanupTask() {
             }
         }
 
-        // Validator shutdown must be called after setKillAllOperations is called. Otherwise, this
-        // can deadlock.
+        // Perform all shutdown operations after setKillAllOperations is called in order to ensure
+        // that any pending threads are about to terminate
+
         if (auto validator = LogicalTimeValidator::get(serviceContext)) {
             validator->shutDown();
         }
@@ -167,9 +169,11 @@ static void cleanupTask() {
         if (auto cursorManager = Grid::get(opCtx)->getCursorManager()) {
             cursorManager->shutdown();
         }
+
         if (auto pool = Grid::get(opCtx)->getExecutorPool()) {
             pool->shutdownAndJoin();
         }
+
         if (auto catalog = Grid::get(opCtx)->catalogClient()) {
             catalog->shutDown(opCtx);
         }
@@ -209,12 +213,15 @@ static Status initializeSharding(OperationContext* opCtx) {
     auto shardFactory =
         stdx::make_unique<ShardFactory>(std::move(buildersMap), std::move(targeterFactory));
 
+    CatalogCacheLoader::set(opCtx->getServiceContext(),
+                            stdx::make_unique<ConfigServerCatalogCacheLoader>());
+
     Status status = initializeGlobalShardingState(
         opCtx,
         mongosGlobalParams.configdbs,
         generateDistLockProcessId(opCtx),
         std::move(shardFactory),
-        stdx::make_unique<CatalogCache>(),
+        stdx::make_unique<CatalogCache>(CatalogCacheLoader::get(opCtx)),
         [opCtx]() {
             auto hookList = stdx::make_unique<rpc::EgressMetadataHookList>();
             hookList->addHook(
@@ -453,6 +460,7 @@ static ExitCode initService() {
 #endif
 
 namespace {
+
 std::unique_ptr<AuthzManagerExternalState> createAuthzManagerExternalStateMongos() {
     return stdx::make_unique<AuthzManagerExternalStateMongos>();
 }
diff --git a/src/mongo/s/shard_server_test_fixture.cpp b/src/mongo/s/shard_server_test_fixture.cpp
index ec77c5a85cc..de22375398d 100644
--- a/src/mongo/s/shard_server_test_fixture.cpp
+++ b/src/mongo/s/shard_server_test_fixture.cpp
@@ -38,6 +38,7 @@
 #include "mongo/s/catalog/dist_lock_manager_mock.h"
 #include "mongo/s/catalog/sharding_catalog_client_impl.h"
 #include "mongo/s/catalog_cache.h"
+#include "mongo/s/config_server_catalog_cache_loader.h"
 #include "mongo/stdx/memory.h"
 
 namespace mongo {
@@ -81,6 +82,11 @@ void ShardServerTestFixture::setUp() {
 
     // Initialize sharding components as a shard server.
     serverGlobalParams.clusterRole = ClusterRole::ShardServer;
+
+    CatalogCacheLoader::set(getServiceContext(),
+                            stdx::make_unique<ShardServerCatalogCacheLoader>(
+                                stdx::make_unique<ConfigServerCatalogCacheLoader>()));
+
     uassertStatusOK(
         initializeGlobalShardingStateForMongodForTest(ConnectionString(kConfigHostAndPort)));
 
@@ -89,6 +95,11 @@ void ShardServerTestFixture::setUp() {
     configTargeterMock()->setFindHostReturnValue(kConfigHostAndPort);
 }
 
+void ShardServerTestFixture::tearDown() {
+    CatalogCacheLoader::clearForTests(getServiceContext());
+
+    ShardingMongodTestFixture::tearDown();
+}
 
 std::unique_ptr<DistLockCatalog> ShardServerTestFixture::makeDistLockCatalog() {
     return stdx::make_unique<DistLockCatalogMock>();
@@ -106,15 +117,8 @@ std::unique_ptr<ShardingCatalogClient> ShardServerTestFixture::makeShardingCatal
     return stdx::make_unique<ShardingCatalogClientImpl>(std::move(distLockManager));
 }
 
-std::unique_ptr<CatalogCacheLoader> ShardServerTestFixture::makeCatalogCacheLoader() {
-    return stdx::make_unique<ShardServerCatalogCacheLoader>(
-        stdx::make_unique<ConfigServerCatalogCacheLoader>());
-}
-
-std::unique_ptr<CatalogCache> ShardServerTestFixture::makeCatalogCache(
-    std::unique_ptr<CatalogCacheLoader> catalogCacheLoader) {
-    invariant(catalogCacheLoader);
-    return stdx::make_unique<CatalogCache>(std::move(catalogCacheLoader));
+std::unique_ptr<CatalogCache> ShardServerTestFixture::makeCatalogCache() {
+    return stdx::make_unique<CatalogCache>(CatalogCacheLoader::get(getServiceContext()));
 }
 
 }  // namespace mongo
diff --git a/src/mongo/s/shard_server_test_fixture.h b/src/mongo/s/shard_server_test_fixture.h
index 8d34b199168..b6f59137b97 100644
--- a/src/mongo/s/shard_server_test_fixture.h
+++ b/src/mongo/s/shard_server_test_fixture.h
@@ -67,6 +67,8 @@ protected:
      */
     void setUp() override;
 
+    void tearDown() override;
+
     /**
      * Creates a DistLockCatalogMock.
      */
@@ -84,13 +86,7 @@ protected:
     std::unique_ptr<ShardingCatalogClient> makeShardingCatalogClient(
         std::unique_ptr<DistLockManager> distLockManager) override;
 
-    /**
-     * Creates a ShardServerCatalogCacheLoader.
-     */
-    std::unique_ptr<CatalogCacheLoader> makeCatalogCacheLoader();
-
-    std::unique_ptr<CatalogCache> makeCatalogCache(
-        std::unique_ptr<CatalogCacheLoader> catalogCacheLoader);
+    std::unique_ptr<CatalogCache> makeCatalogCache() override;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/s/sharding_mongod_test_fixture.cpp b/src/mongo/s/sharding_mongod_test_fixture.cpp
index 32b50a8519d..3a9fa023b2a 100644
--- a/src/mongo/s/sharding_mongod_test_fixture.cpp
+++ b/src/mongo/s/sharding_mongod_test_fixture.cpp
@@ -236,12 +236,7 @@ std::unique_ptr<ShardingCatalogClient> ShardingMongodTestFixture::makeShardingCa
     return nullptr;
 }
 
-std::unique_ptr<CatalogCacheLoader> ShardingMongodTestFixture::makeCatalogCacheLoader() {
-    return nullptr;
-}
-
-std::unique_ptr<CatalogCache> ShardingMongodTestFixture::makeCatalogCache(
-    std::unique_ptr<CatalogCacheLoader> catalogCacheLoader) {
+std::unique_ptr<CatalogCache> ShardingMongodTestFixture::makeCatalogCache() {
     return nullptr;
 }
 
@@ -274,7 +269,7 @@ Status ShardingMongodTestFixture::initializeGlobalShardingStateForMongodForTest(
 
     auto const grid = Grid::get(operationContext());
     grid->init(makeShardingCatalogClient(std::move(distLockManagerPtr)),
-               makeCatalogCache(makeCatalogCacheLoader()),
+               makeCatalogCache(),
                makeShardRegistry(configConnStr),
                makeClusterCursorManager(),
                makeBalancerConfiguration(),
diff --git a/src/mongo/s/sharding_mongod_test_fixture.h b/src/mongo/s/sharding_mongod_test_fixture.h
index 41af161da99..6ac9a636326 100644
--- a/src/mongo/s/sharding_mongod_test_fixture.h
+++ b/src/mongo/s/sharding_mongod_test_fixture.h
@@ -221,13 +221,7 @@ protected:
     /**
      * Base class returns nullptr.
      */
-    virtual std::unique_ptr<CatalogCacheLoader> makeCatalogCacheLoader();
-
-    /**
-     * Base class returns nullptr.
-     */
-    virtual std::unique_ptr<CatalogCache> makeCatalogCache(
-        std::unique_ptr<CatalogCacheLoader> catalogCacheLoader);
+    virtual std::unique_ptr<CatalogCache> makeCatalogCache();
 
     /**
      * Base class returns nullptr.
diff --git a/src/mongo/s/sharding_test_fixture.cpp b/src/mongo/s/sharding_test_fixture.cpp
index 4db274e0863..aad4203cde3 100644
--- a/src/mongo/s/sharding_test_fixture.cpp
+++ b/src/mongo/s/sharding_test_fixture.cpp
@@ -61,6 +61,7 @@
 #include "mongo/s/client/shard_factory.h"
 #include "mongo/s/client/shard_registry.h"
 #include "mongo/s/client/shard_remote.h"
+#include "mongo/s/config_server_catalog_cache_loader.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/query/cluster_cursor_manager.h"
 #include "mongo/s/set_shard_version_request.h"
@@ -185,11 +186,13 @@ void ShardingTestFixture::setUp() {
     auto shardRegistry(stdx::make_unique<ShardRegistry>(std::move(shardFactory), configCS));
     executorPool->startup();
 
+    CatalogCacheLoader::set(serviceContext(), stdx::make_unique<ConfigServerCatalogCacheLoader>());
+
     // For now initialize the global grid object. All sharding objects will be accessible from there
     // until we get rid of it.
     Grid::get(operationContext())
         ->init(std::move(catalogClient),
-               stdx::make_unique<CatalogCache>(),
+               stdx::make_unique<CatalogCache>(CatalogCacheLoader::get(serviceContext())),
                std::move(shardRegistry),
                stdx::make_unique<ClusterCursorManager>(serviceContext()->getPreciseClockSource()),
                stdx::make_unique<BalancerConfiguration>(),
@@ -198,6 +201,8 @@ void ShardingTestFixture::setUp() {
 }
 
 void ShardingTestFixture::tearDown() {
+    CatalogCacheLoader::clearForTests(serviceContext());
+
     Grid::get(operationContext())->getExecutorPool()->shutdownAndJoin();
     Grid::get(operationContext())->catalogClient()->shutDown(_opCtx.get());
     Grid::get(operationContext())->clearForUnitTests();