SERVER-49976 Cleanup database cache on shard removal

11 files changed, 152 insertions, 40 deletions

diff --git a/src/mongo/db/auth/authorization_manager_impl.cpp b/src/mongo/db/auth/authorization_manager_impl.cpp
index c4bca48c570..96ad1aa0e9d 100644
--- a/src/mongo/db/auth/authorization_manager_impl.cpp
+++ b/src/mongo/db/auth/authorization_manager_impl.cpp
@@ -608,7 +608,7 @@ void AuthorizationManagerImpl::invalidateUsersFromDB(OperationContext* opCtx, St
     LOGV2_DEBUG(20236, 2, "Invalidating all users from database", "database"_attr = dbname);
     _updateCacheGeneration();
     _authSchemaVersionCache.invalidateAll();
-    _userCache.invalidateIfKey(
+    _userCache.invalidateKeyIf(
         [&](const UserRequest& userRequest) { return userRequest.name.getDB() == dbname; });
 }
 
diff --git a/src/mongo/db/s/SConscript b/src/mongo/db/s/SConscript
index 5ebc2fd9211..1299619916c 100644
--- a/src/mongo/db/s/SConscript
+++ b/src/mongo/db/s/SConscript
@@ -392,7 +392,6 @@ env.CppUnitTest(
         'active_migrations_registry_test.cpp',
         'active_move_primaries_registry_test.cpp',
         'active_shard_collection_registry_test.cpp',
-        'catalog_cache_loader_mock.cpp',
         'chunk_split_state_driver_test.cpp',
         'migration_chunk_cloner_source_legacy_test.cpp',
         'migration_destination_manager_test.cpp',
diff --git a/src/mongo/db/s/migration_util_test.cpp b/src/mongo/db/s/migration_util_test.cpp
index c39871ab2a1..0d06c8175d8 100644
--- a/src/mongo/db/s/migration_util_test.cpp
+++ b/src/mongo/db/s/migration_util_test.cpp
@@ -32,7 +32,6 @@
 #include "mongo/db/catalog_raii.h"
 #include "mongo/db/persistent_task_store.h"
 #include "mongo/db/repl/wait_for_majority_service.h"
-#include "mongo/db/s/catalog_cache_loader_mock.h"
 #include "mongo/db/s/collection_sharding_runtime.h"
 #include "mongo/db/s/collection_sharding_state.h"
 #include "mongo/db/s/migration_util.h"
@@ -42,6 +41,7 @@
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/s/catalog/sharding_catalog_client_mock.h"
 #include "mongo/s/catalog/type_shard.h"
+#include "mongo/s/catalog_cache_loader_mock.h"
 #include "mongo/s/database_version_helpers.h"
 #include "mongo/util/future.h"
 
diff --git a/src/mongo/db/s/shard_server_catalog_cache_loader_test.cpp b/src/mongo/db/s/shard_server_catalog_cache_loader_test.cpp
index 66590127863..8c215b8b664 100644
--- a/src/mongo/db/s/shard_server_catalog_cache_loader_test.cpp
+++ b/src/mongo/db/s/shard_server_catalog_cache_loader_test.cpp
@@ -29,11 +29,11 @@
 
 #include "mongo/platform/basic.h"
 
-#include "mongo/db/s/catalog_cache_loader_mock.h"
 #include "mongo/db/s/shard_server_catalog_cache_loader.h"
 #include "mongo/db/s/shard_server_test_fixture.h"
 #include "mongo/s/catalog/type_chunk.h"
 #include "mongo/s/catalog/type_collection.h"
+#include "mongo/s/catalog_cache_loader_mock.h"
 
 namespace mongo {
 namespace {
diff --git a/src/mongo/s/SConscript b/src/mongo/s/SConscript
index 3a364a924c2..14f85e390a1 100644
--- a/src/mongo/s/SConscript
+++ b/src/mongo/s/SConscript
@@ -196,6 +196,7 @@ env.Library(
     target='sharding_test_fixture_common',
     source=[
         'sharding_test_fixture_common.cpp',
+        'catalog_cache_loader_mock.cpp',
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/client/remote_command_targeter_mock',
@@ -565,6 +566,7 @@ env.CppUnitTest(
         'append_raw_responses_test.cpp',
         'balancer_configuration_test.cpp',
         'build_versioned_requests_for_targeted_shards_test.cpp',
+        'catalog_cache_test.cpp',
         'catalog_cache_refresh_test.cpp',
         'catalog/type_changelog_test.cpp',
         'catalog/type_chunk_test.cpp',
diff --git a/src/mongo/s/catalog_cache.cpp b/src/mongo/s/catalog_cache.cpp
index acad122d913..7e2cfb9501b 100644
--- a/src/mongo/s/catalog_cache.cpp
+++ b/src/mongo/s/catalog_cache.cpp
@@ -162,23 +162,6 @@ StatusWith<CachedDatabaseInfo> CatalogCache::getDatabase(OperationContext* opCtx
             return {ErrorCodes::NamespaceNotFound,
                     str::stream() << "database " << dbName << " not found"};
         }
-        const auto primaryShardExists =
-            Grid::get(opCtx)->shardRegistry()->getShard(opCtx, dbEntry->getPrimary()).isOK();
-
-        if (!primaryShardExists) {
-            LOGV2_FOR_CATALOG_REFRESH(
-                4947103,
-                2,
-                "Invalidating cached database entry because its primary shard hasn't been found",
-                "db"_attr = dbName);
-            _databaseCache.invalidate(dbName.toString());
-            dbEntry = _databaseCache.acquire(opCtx, dbName.toString());
-            if (!dbEntry) {
-                return {ErrorCodes::NamespaceNotFound,
-                        str::stream() << "database " << dbName << " not found"};
-            }
-        }
-
         const auto primaryShard = uassertStatusOKWithContext(
             Grid::get(opCtx)->shardRegistry()->getShard(opCtx, dbEntry->getPrimary()),
             str::stream() << "could not find the primary shard for database " << dbName);
@@ -436,24 +419,19 @@ void CatalogCache::invalidateShardForShardedCollection(const NamespaceString& ns
 }
 
 void CatalogCache::invalidateEntriesThatReferenceShard(const ShardId& shardId) {
-    stdx::lock_guard<Latch> lg(_mutex);
+    LOGV2_DEBUG(4997600,
+                1,
+                "Invalidating databases and collections referencing a specific shard",
+                "shardId"_attr = shardId);
 
-    LOGV2(22643,
-          "Starting to invalidate collections with data on shard {shardId}",
-          "Starting to invalidate collections referencing a specific shard",
-          "shardId"_attr = shardId);
+    _databaseCache.invalidateCachedValueIf(
+        [&](const DatabaseType& dbt) { return dbt.getPrimary() == shardId; });
+
+    stdx::lock_guard<Latch> lg(_mutex);
 
     // Invalidate collections which contain data on this shard.
     for (const auto& [db, collInfoMap] : _collectionsByDb) {
         for (const auto& [collNs, collRoutingInfoEntry] : collInfoMap) {
-
-            LOGV2_DEBUG(22646,
-                        3,
-                        "Checking if collection {namespace} has data on shard {shardId}",
-                        "Checking if collection has data on specific shard",
-                        "namespace"_attr = collNs,
-                        "shardId"_attr = shardId);
-
             if (!collRoutingInfoEntry->needsRefresh && collRoutingInfoEntry->routingInfo) {
                 // The set of shards on which this collection contains chunks.
                 std::set<ShardId> shardsOwningDataForCollection;
diff --git a/src/mongo/db/s/catalog_cache_loader_mock.cpp b/src/mongo/s/catalog_cache_loader_mock.cpp
index 1711358e0c8..7ed96e720a3 100644
--- a/src/mongo/db/s/catalog_cache_loader_mock.cpp
+++ b/src/mongo/s/catalog_cache_loader_mock.cpp
@@ -29,7 +29,7 @@
 
 #include "mongo/platform/basic.h"
 
-#include "mongo/db/s/catalog_cache_loader_mock.h"
+#include "mongo/s/catalog_cache_loader_mock.h"
 
 #include "mongo/db/operation_context.h"
 #include "mongo/s/catalog/type_chunk.h"
diff --git a/src/mongo/db/s/catalog_cache_loader_mock.h b/src/mongo/s/catalog_cache_loader_mock.h
index 1bb13a46285..1bb13a46285 100644
--- a/src/mongo/db/s/catalog_cache_loader_mock.h
+++ b/src/mongo/s/catalog_cache_loader_mock.h
diff --git a/src/mongo/s/catalog_cache_test.cpp b/src/mongo/s/catalog_cache_test.cpp
new file mode 100644
index 00000000000..fce177bdd4f
--- /dev/null
+++ b/src/mongo/s/catalog_cache_test.cpp
@@ -0,0 +1,133 @@
+/**
+ *    Copyright (C) 2018-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    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
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    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 Server Side 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.
+ */
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kDefault
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/s/catalog/type_database.h"
+#include "mongo/s/catalog_cache.h"
+#include "mongo/s/catalog_cache_loader_mock.h"
+#include "mongo/s/sharding_router_test_fixture.h"
+
+namespace mongo {
+namespace {
+
+class CatalogCacheTest : public ShardingTestFixture {
+protected:
+    void setUp() override {
+        ShardingTestFixture::setUp();
+
+        // Setup dummy config server
+        setRemote(kConfigHostAndPort);
+        configTargeter()->setFindHostReturnValue(kConfigHostAndPort);
+
+        // Setup catalogCache with mock loader
+        _catalogCacheLoader = std::make_unique<CatalogCacheLoaderMock>();
+        _catalogCache = std::make_unique<CatalogCache>(getServiceContext(), *_catalogCacheLoader);
+
+        // Populate the shardRegistry with the shards from kShards vector
+        std::vector<std::tuple<ShardId, HostAndPort>> shardInfos;
+        for (const auto& shardId : kShards) {
+            shardInfos.emplace_back(
+                std::make_tuple(shardId, HostAndPort(shardId.toString(), kDummyPort)));
+        }
+        addRemoteShards(shardInfos);
+    };
+
+    void loadDatabases(const std::vector<DatabaseType>& databases) {
+        for (const auto& db : databases) {
+            _catalogCacheLoader->setDatabaseRefreshReturnValue(db);
+            const auto swDatabase = _catalogCache->getDatabase(operationContext(), db.getName());
+            ASSERT_OK(swDatabase.getStatus());
+        }
+
+        // Reset the database return value to avoid false positive results
+        _catalogCacheLoader->setDatabaseRefreshReturnValue(kErrorStatus);
+    }
+
+    const NamespaceString kNss{"catalgoCacheTestDB.foo"};
+    const int kDummyPort{12345};
+    const HostAndPort kConfigHostAndPort{"DummyConfig", kDummyPort};
+    const std::vector<ShardId> kShards{{"0"}, {"1"}};
+    const Status kErrorStatus{ErrorCodes::InternalError,
+                              "Received an unexpected CatalogCacheLoader request"};
+
+    std::unique_ptr<CatalogCacheLoaderMock> _catalogCacheLoader;
+    std::unique_ptr<CatalogCache> _catalogCache;
+};
+
+TEST_F(CatalogCacheTest, GetDatabase) {
+    const auto dbName = "testDB";
+    const auto dbVersion = DatabaseVersion(UUID::gen(), 1);
+    _catalogCacheLoader->setDatabaseRefreshReturnValue(
+        DatabaseType(dbName, kShards[0], true, dbVersion));
+
+    const auto swDatabase = _catalogCache->getDatabase(operationContext(), dbName);
+
+    ASSERT_OK(swDatabase.getStatus());
+    const auto cachedDb = swDatabase.getValue();
+    ASSERT_TRUE(cachedDb.shardingEnabled());
+    ASSERT_EQ(cachedDb.primaryId(), kShards[0]);
+    ASSERT_EQ(cachedDb.databaseVersion().getUuid(), dbVersion.getUuid());
+    ASSERT_EQ(cachedDb.databaseVersion().getLastMod(), dbVersion.getLastMod());
+}
+
+TEST_F(CatalogCacheTest, GetCachedDatabase) {
+    const auto dbName = "testDB";
+    const auto dbVersion = DatabaseVersion(UUID::gen(), 1);
+    loadDatabases({DatabaseType(dbName, kShards[0], true, dbVersion)});
+
+    const auto swDatabase = _catalogCache->getDatabase(operationContext(), dbName);
+
+    ASSERT_OK(swDatabase.getStatus());
+    const auto cachedDb = swDatabase.getValue();
+    ASSERT_TRUE(cachedDb.shardingEnabled());
+    ASSERT_EQ(cachedDb.primaryId(), kShards[0]);
+    ASSERT_EQ(cachedDb.databaseVersion().getUuid(), dbVersion.getUuid());
+    ASSERT_EQ(cachedDb.databaseVersion().getLastMod(), dbVersion.getLastMod());
+}
+
+TEST_F(CatalogCacheTest, InvalidateSingleDbOnShardRemoval) {
+    const auto dbName = "testDB";
+    const auto dbVersion = DatabaseVersion(UUID::gen(), 1);
+    loadDatabases({DatabaseType(dbName, kShards[0], true, dbVersion)});
+
+    _catalogCache->invalidateEntriesThatReferenceShard(kShards[0]);
+    _catalogCacheLoader->setDatabaseRefreshReturnValue(
+        DatabaseType(dbName, kShards[1], true, dbVersion));
+    const auto swDatabase = _catalogCache->getDatabase(operationContext(), dbName);
+
+    ASSERT_OK(swDatabase.getStatus());
+    auto cachedDb = swDatabase.getValue();
+    ASSERT_EQ(cachedDb.primaryId(), kShards[1]);
+}
+
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/util/read_through_cache.h b/src/mongo/util/read_through_cache.h
index ffa4d7295e2..7d2e24b9b20 100644
--- a/src/mongo/util/read_through_cache.h
+++ b/src/mongo/util/read_through_cache.h
@@ -335,7 +335,7 @@ public:
      * Invalidates all cached entries and in progress lookups with keys that matches the preidcate.
      */
     template <typename Pred>
-    void invalidateIfKey(const Pred& predicate) {
+    void invalidateKeyIf(const Pred& predicate) {
         stdx::lock_guard lg(_mutex);
         for (auto& entry : _inProgressLookups) {
             if (predicate(entry.first))
@@ -348,14 +348,14 @@ public:
      * Invalidates all cached entries with stored values that matches the preidcate.
      */
     template <typename Pred>
-    void invalidateIfCachedValue(const Pred& predicate) {
+    void invalidateCachedValueIf(const Pred& predicate) {
         stdx::lock_guard lg(_mutex);
         _cache.invalidateIf(
             [&](const Key&, const StoredValue* value) { return predicate(value->value); });
     }
 
     void invalidateAll() {
-        invalidateIfKey([](const Key&) { return true; });
+        invalidateKeyIf([](const Key&) { return true; });
     }
 
     /**
diff --git a/src/mongo/util/read_through_cache_test.cpp b/src/mongo/util/read_through_cache_test.cpp
index 0093f6aa8dc..70d452790e7 100644
--- a/src/mongo/util/read_through_cache_test.cpp
+++ b/src/mongo/util/read_through_cache_test.cpp
@@ -183,7 +183,7 @@ TEST_F(ReadThroughCacheTest, FetchInvalidateKeyAndRefetch) {
             ASSERT(cache.acquire(_opCtx, "TestKey"));
             ASSERT_EQ(i, cache.countLookups);
 
-            cache.invalidateIfKey([](const std::string& key) { return key == "TestKey"; });
+            cache.invalidateKeyIf([](const std::string& key) { return key == "TestKey"; });
         }
     };
 
@@ -221,7 +221,7 @@ TEST_F(ReadThroughCacheTest, FetchInvalidateValueAndRefetch) {
             ASSERT(cache.acquire(_opCtx, "TestKey"));
             ASSERT_EQ(i, cache.countLookups);
 
-            cache.invalidateIfCachedValue(
+            cache.invalidateCachedValueIf(
                 [i](const CachedValue& value) { return value.counter == 100 * i; });
         }
     };