SERVER-29672 Make Config/ShardCatalogCacheLoader decorations on ServiceContext

1 files changed, 32 insertions, 0 deletions

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