SERVER-54205 Proactively load external keys into the keys cache

14 files changed, 171 insertions, 35 deletions

diff --git a/src/mongo/db/keys_collection_cache.cpp b/src/mongo/db/keys_collection_cache.cpp
index 46226d1c1b6..407e3168f37 100644
--- a/src/mongo/db/keys_collection_cache.cpp
+++ b/src/mongo/db/keys_collection_cache.cpp
@@ -32,7 +32,6 @@
 #include "mongo/db/keys_collection_cache.h"
 
 #include "mongo/db/keys_collection_client.h"
-#include "mongo/db/keys_collection_document_gen.h"
 #include "mongo/db/repl/replication_coordinator.h"
 #include "mongo/util/str.h"
 
@@ -204,4 +203,9 @@ void KeysCollectionCache::resetCache() {
     }
 }
 
+void KeysCollectionCache::cacheExternalKey(ExternalKeysCollectionDocument key) {
+    stdx::lock_guard<Latch> lk(_cacheMutex);
+    _externalKeysCache.emplace(key.getKeyId(), std::move(key));
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/keys_collection_cache.h b/src/mongo/db/keys_collection_cache.h
index 736ad135bc2..31103e670ee 100644
--- a/src/mongo/db/keys_collection_cache.h
+++ b/src/mongo/db/keys_collection_cache.h
@@ -83,6 +83,11 @@ public:
      */
     void resetCache();
 
+    /**
+     * Loads the given external key into the cache.
+     */
+    void cacheExternalKey(ExternalKeysCollectionDocument key);
+
 private:
     /**
      * Checks if there are new internal key documents (see definition below) with expiresAt greater
diff --git a/src/mongo/db/keys_collection_cache_test.cpp b/src/mongo/db/keys_collection_cache_test.cpp
index 82ed5670f45..c5818ba2e22 100644
--- a/src/mongo/db/keys_collection_cache_test.cpp
+++ b/src/mongo/db/keys_collection_cache_test.cpp
@@ -47,8 +47,8 @@ namespace {
 
 class CacheTest : public ConfigServerTestFixture {
 protected:
-    const UUID migrationId1 = UUID::gen();
-    const UUID migrationId2 = UUID::gen();
+    const UUID kMigrationId1 = UUID::gen();
+    const UUID kMigrationId2 = UUID::gen();
 
     void setUp() override {
         ConfigServerTestFixture::setUp();
@@ -180,13 +180,13 @@ TEST_F(CacheTest, GetKeyShouldReturnCorrectKeysAfterRefreshDirectClient) {
 
     // Use external keys with the same keyId and expiresAt as the internal key to test that the
     // cache correctly tackles key collisions.
-    ExternalKeysCollectionDocument origKey1(OID::gen(), 1, migrationId1, externalKeysTTLExpiresAt);
+    ExternalKeysCollectionDocument origKey1(OID::gen(), 1, kMigrationId1, externalKeysTTLExpiresAt);
     origKey1.setKeysCollectionDocumentBase(
         {"test", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0))});
     insertDocument(
         operationContext(), NamespaceString::kExternalKeysCollectionNamespace, origKey1.toBSON());
 
-    ExternalKeysCollectionDocument origKey2(OID::gen(), 1, migrationId2, externalKeysTTLExpiresAt);
+    ExternalKeysCollectionDocument origKey2(OID::gen(), 1, kMigrationId2, externalKeysTTLExpiresAt);
     origKey2.setKeysCollectionDocumentBase(
         {"test", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(205, 0))});
     insertDocument(
@@ -394,7 +394,7 @@ TEST_F(CacheTest, RefreshShouldNotGetExternalKeysForOtherPurpose) {
     const auto externalKeysTTLExpiresAt =
         ServiceContext().getFastClockSource()->now() + Seconds(30);
 
-    ExternalKeysCollectionDocument origKey1(OID::gen(), 1, migrationId1, externalKeysTTLExpiresAt);
+    ExternalKeysCollectionDocument origKey1(OID::gen(), 1, kMigrationId1, externalKeysTTLExpiresAt);
     origKey1.setKeysCollectionDocumentBase(
         {"dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0))});
     insertDocument(
@@ -408,7 +408,7 @@ TEST_F(CacheTest, RefreshShouldNotGetExternalKeysForOtherPurpose) {
         ASSERT_EQ(ErrorCodes::KeyNotFound, swKey.getStatus());
     }
 
-    ExternalKeysCollectionDocument origKey2(OID::gen(), 2, migrationId1, externalKeysTTLExpiresAt);
+    ExternalKeysCollectionDocument origKey2(OID::gen(), 2, kMigrationId1, externalKeysTTLExpiresAt);
     origKey2.setKeysCollectionDocumentBase(
         {"test", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(110, 0))});
     insertDocument(
@@ -493,5 +493,31 @@ TEST_F(CacheTest, RefreshCanIncrementallyGetNewKeys) {
     }
 }
 
+TEST_F(CacheTest, CacheExternalKeyBasic) {
+    const auto externalKeysTTLExpiresAt =
+        ServiceContext().getFastClockSource()->now() + Seconds(30);
+    KeysCollectionCache cache("test", catalogClient());
+
+    auto swExternalKeys = cache.getExternalKeysById(5, LogicalTime(Timestamp(10, 1)));
+    ASSERT_EQ(ErrorCodes::KeyNotFound, swExternalKeys.getStatus());
+
+    ExternalKeysCollectionDocument externalKey(
+        OID::gen(), 5, kMigrationId1, externalKeysTTLExpiresAt);
+    externalKey.setKeysCollectionDocumentBase(
+        {"test", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(100, 0))});
+
+    cache.cacheExternalKey(externalKey);
+
+    swExternalKeys = cache.getExternalKeysById(5, LogicalTime(Timestamp(10, 1)));
+    ASSERT_OK(swExternalKeys.getStatus());
+    ASSERT_EQ(1, swExternalKeys.getValue().size());
+
+    auto cachedKey = swExternalKeys.getValue().front();
+    ASSERT_EQ(externalKey.getId(), cachedKey.getId());
+    ASSERT_EQ(externalKey.getPurpose(), cachedKey.getPurpose());
+    ASSERT_EQ(externalKey.getExpiresAt().asTimestamp(), cachedKey.getExpiresAt().asTimestamp());
+    ASSERT_EQ(externalKey.getTTLExpiresAt(), cachedKey.getTTLExpiresAt());
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/keys_collection_manager.cpp b/src/mongo/db/keys_collection_manager.cpp
index e03984f4808..e3554953e95 100644
--- a/src/mongo/db/keys_collection_manager.cpp
+++ b/src/mongo/db/keys_collection_manager.cpp
@@ -201,6 +201,15 @@ void KeysCollectionManager::clearCache() {
     _keysCache.resetCache();
 }
 
+void KeysCollectionManager::cacheExternalKey(ExternalKeysCollectionDocument key) {
+    // If the refresher has been shut down, we don't cache external keys because refresh is relied
+    // on to clear expired keys. This is OK because the refresher is only shut down in cases where
+    // keys aren't needed, like on an arbiter.
+    if (!_refresher.isInShutdown()) {
+        _keysCache.cacheExternalKey(std::move(key));
+    }
+}
+
 void KeysCollectionManager::PeriodicRunner::refreshNow(OperationContext* opCtx) {
     auto refreshRequest = [this]() {
         stdx::lock_guard<Latch> lk(_mutex);
@@ -365,4 +374,9 @@ bool KeysCollectionManager::PeriodicRunner::hasSeenKeys() const noexcept {
     return _hasSeenKeys.load();
 }
 
+bool KeysCollectionManager::PeriodicRunner::isInShutdown() const {
+    stdx::lock_guard<Latch> lock(_mutex);
+    return _inShutdown;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/keys_collection_manager.h b/src/mongo/db/keys_collection_manager.h
index 73caf91fff9..50bfe04abfe 100644
--- a/src/mongo/db/keys_collection_manager.h
+++ b/src/mongo/db/keys_collection_manager.h
@@ -124,6 +124,11 @@ public:
      */
     void clearCache();
 
+    /**
+     * Loads the given external key into the keys collection cache.
+     */
+    void cacheExternalKey(ExternalKeysCollectionDocument key);
+
 private:
     /**
      * This is responsible for periodically performing refresh in the background.
@@ -174,6 +179,11 @@ private:
          */
         bool hasSeenKeys() const noexcept;
 
+        /**
+         * Returns if the periodic runner has entered shutdown.
+         */
+        bool isInShutdown() const;
+
     private:
         void _doPeriodicRefresh(ServiceContext* service,
                                 std::string threadName,
@@ -182,7 +192,7 @@ private:
         AtomicWord<bool> _hasSeenKeys{false};
 
         // protects all the member variables below.
-        Mutex _mutex = MONGO_MAKE_LATCH("PeriodicRunner::_mutex");
+        mutable Mutex _mutex = MONGO_MAKE_LATCH("PeriodicRunner::_mutex");
         std::shared_ptr<Notification<void>> _refreshRequest;
         stdx::condition_variable _refreshNeededCV;
 
diff --git a/src/mongo/db/keys_collection_manager_sharding_test.cpp b/src/mongo/db/keys_collection_manager_sharding_test.cpp
index a51ca4f7b38..71b91bb3c1f 100644
--- a/src/mongo/db/keys_collection_manager_sharding_test.cpp
+++ b/src/mongo/db/keys_collection_manager_sharding_test.cpp
@@ -50,6 +50,9 @@ public:
     }
 
 protected:
+    const UUID kMigrationId1 = UUID::gen();
+    const UUID kMigrationId2 = UUID::gen();
+
     void setUp() override {
         ConfigServerTestFixture::setUp();
 
@@ -368,6 +371,76 @@ TEST_F(KeysManagerShardedTest, HasSeenKeysIsFalseUntilKeysAreFound) {
     ASSERT_EQ(true, keyManager()->hasSeenKeys());
 }
 
+TEST_F(KeysManagerShardedTest, CacheExternalKeyBasic) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    auto externalKeysTTLExpiresAt = getServiceContext()->getFastClockSource()->now() + Seconds(30);
+    ExternalKeysCollectionDocument externalKey1(
+        OID::gen(), 1, kMigrationId1, externalKeysTTLExpiresAt);
+    externalKey1.setKeysCollectionDocumentBase(
+        {"dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(100, 0))});
+
+    keyManager()->cacheExternalKey(externalKey1);
+
+    {
+        auto keyStatus =
+            keyManager()->getKeysForValidation(operationContext(), 1, LogicalTime(Timestamp(1, 0)));
+        ASSERT_OK(keyStatus.getStatus());
+        ASSERT_EQ(1, keyStatus.getValue().size());
+
+        auto key = keyStatus.getValue().front();
+        ASSERT_EQ(1, key.getKeyId());
+
+        ASSERT_EQ(externalKey1.getKeyId(), key.getKeyId());
+        ASSERT_EQ(externalKey1.getPurpose(), key.getPurpose());
+        ASSERT_EQ(externalKey1.getExpiresAt().asTimestamp(), key.getExpiresAt().asTimestamp());
+    }
+}
+
+TEST_F(KeysManagerShardedTest, WillNotCacheExternalKeyWhenMonitoringIsStopped) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    // Insert an internal key so the key manager won't attempt to refresh after the refresher is
+    // stopped.
+    KeysCollectionDocument internalKey(1);
+    internalKey.setKeysCollectionDocumentBase(
+        {"dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(100, 0))});
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString::kKeysCollectionNamespace, internalKey.toBSON()));
+
+    auto externalKeysTTLExpiresAt = getServiceContext()->getFastClockSource()->now() + Seconds(30);
+    ExternalKeysCollectionDocument externalKey1(
+        OID::gen(), 1, kMigrationId1, externalKeysTTLExpiresAt);
+    externalKey1.setKeysCollectionDocumentBase(
+        {"dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(100, 0))});
+
+    keyManager()->cacheExternalKey(externalKey1);
+
+    {
+        auto keyStatus =
+            keyManager()->getKeysForValidation(operationContext(), 1, LogicalTime(Timestamp(1, 0)));
+        ASSERT_OK(keyStatus.getStatus());
+        ASSERT_EQ(2, keyStatus.getValue().size());
+    }
+
+    keyManager()->stopMonitoring();
+
+    ExternalKeysCollectionDocument externalKey2(
+        OID::gen(), 1, kMigrationId2, externalKeysTTLExpiresAt);
+    externalKey2.setKeysCollectionDocumentBase(
+        {"dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(100, 0))});
+
+    keyManager()->cacheExternalKey(externalKey2);
+
+    // There should still be only the first external key in the cache.
+    {
+        auto keyStatus =
+            keyManager()->getKeysForValidation(operationContext(), 1, LogicalTime(Timestamp(1, 0)));
+        ASSERT_OK(keyStatus.getStatus());
+        ASSERT_EQ(2, keyStatus.getValue().size());
+    }
+}
+
 LogicalTime addSeconds(const LogicalTime& logicalTime, const Seconds& seconds) {
     auto asTimestamp = logicalTime.asTimestamp();
     return LogicalTime(Timestamp(asTimestamp.getSecs() + seconds.count(), asTimestamp.getInc()));
diff --git a/src/mongo/db/logical_time_validator.cpp b/src/mongo/db/logical_time_validator.cpp
index b848b49cc5e..852585aff18 100644
--- a/src/mongo/db/logical_time_validator.cpp
+++ b/src/mongo/db/logical_time_validator.cpp
@@ -220,9 +220,9 @@ bool LogicalTimeValidator::shouldGossipLogicalTime() {
     return _getKeyManagerCopy()->hasSeenKeys();
 }
 
-void LogicalTimeValidator::refreshKeyManagerCache(OperationContext* opCtx) {
+void LogicalTimeValidator::cacheExternalKey(ExternalKeysCollectionDocument key) {
     invariant(_keyManager);
-    _keyManager->refreshNow(opCtx);
+    _keyManager->cacheExternalKey(std::move(key));
 }
 
 void LogicalTimeValidator::resetKeyManagerCache() {
diff --git a/src/mongo/db/logical_time_validator.h b/src/mongo/db/logical_time_validator.h
index 825b7dc70df..404d02310a1 100644
--- a/src/mongo/db/logical_time_validator.h
+++ b/src/mongo/db/logical_time_validator.h
@@ -31,6 +31,7 @@
 
 #include <memory>
 
+#include "mongo/db/keys_collection_document_gen.h"
 #include "mongo/db/signed_logical_time.h"
 #include "mongo/db/time_proof_service.h"
 #include "mongo/platform/mutex.h"
@@ -109,9 +110,9 @@ public:
     void stopKeyManager();
 
     /**
-     * Forces the key manager cache to refresh.
+     * Load the given external key into the key manager's keys cache.
      */
-    void refreshKeyManagerCache(OperationContext* opCtx);
+    void cacheExternalKey(ExternalKeysCollectionDocument key);
 
     /**
      * Reset the key manager cache of keys.
diff --git a/src/mongo/db/op_observer_impl.cpp b/src/mongo/db/op_observer_impl.cpp
index 3b31c7f1b17..ebacd067c04 100644
--- a/src/mongo/db/op_observer_impl.cpp
+++ b/src/mongo/db/op_observer_impl.cpp
@@ -45,6 +45,7 @@
 #include "mongo/db/concurrency/d_concurrency.h"
 #include "mongo/db/concurrency/write_conflict_exception.h"
 #include "mongo/db/index/index_descriptor.h"
+#include "mongo/db/keys_collection_document_gen.h"
 #include "mongo/db/logical_time_validator.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/op_observer_util.h"
@@ -513,6 +514,19 @@ void OpObserverImpl::onInserts(OperationContext* opCtx,
             ReadWriteConcernDefaults::get(opCtx).observeDirectWriteToConfigSettings(
                 opCtx, it->doc["_id"], it->doc);
         }
+    } else if (nss == NamespaceString::kExternalKeysCollectionNamespace) {
+        for (auto it = first; it != last; it++) {
+            auto externalKey = ExternalKeysCollectionDocument::parse(
+                IDLParserErrorContext("externalKey"), it->doc);
+            opCtx->recoveryUnit()->onCommit(
+                [this, opCtx, externalKey = std::move(externalKey)](
+                    boost::optional<Timestamp> unusedCommitTime) mutable {
+                    auto validator = LogicalTimeValidator::get(opCtx);
+                    if (validator) {
+                        validator->cacheExternalKey(externalKey);
+                    }
+                });
+        }
     }
 }
 
diff --git a/src/mongo/db/repl/tenant_migration_donor_service.cpp b/src/mongo/db/repl/tenant_migration_donor_service.cpp
index 8f1e7166ac4..94d6cf0a802 100644
--- a/src/mongo/db/repl/tenant_migration_donor_service.cpp
+++ b/src/mongo/db/repl/tenant_migration_donor_service.cpp
@@ -400,13 +400,12 @@ TenantMigrationDonorService::Instance::_fetchAndStoreRecipientClusterTimeKeyDocs
 
             return keyDocs;
         })
-        .then(
-            [this, self = shared_from_this(), executor, serviceToken, instanceToken](auto keyDocs) {
-                checkIfReceivedDonorAbortMigration(serviceToken, instanceToken);
+        .then([this, self = shared_from_this(), executor, serviceToken, instanceToken](
+                  auto keyDocs) {
+            checkIfReceivedDonorAbortMigration(serviceToken, instanceToken);
 
-                tenant_migration_util::storeExternalClusterTimeKeyDocsAndRefreshCache(
-                    executor, std::move(keyDocs));
-            });
+            tenant_migration_util::storeExternalClusterTimeKeyDocs(executor, std::move(keyDocs));
+        });
 }
 
 ExecutorFuture<repl::OpTime> TenantMigrationDonorService::Instance::_insertStateDoc(
diff --git a/src/mongo/db/repl/tenant_migration_recipient_service.cpp b/src/mongo/db/repl/tenant_migration_recipient_service.cpp
index 1768100e54a..e251adc76c8 100644
--- a/src/mongo/db/repl/tenant_migration_recipient_service.cpp
+++ b/src/mongo/db/repl/tenant_migration_recipient_service.cpp
@@ -1360,8 +1360,7 @@ void TenantMigrationRecipientService::Instance::_fetchAndStoreDonorClusterTimeKe
             _serviceContext, _migrationUuid, doc));
     }
 
-    tenant_migration_util::storeExternalClusterTimeKeyDocsAndRefreshCache(_scopedExecutor,
-                                                                          std::move(keyDocs));
+    tenant_migration_util::storeExternalClusterTimeKeyDocs(_scopedExecutor, std::move(keyDocs));
 }
 
 void TenantMigrationRecipientService::Instance::_compareRecipientAndDonorFCV() const {
diff --git a/src/mongo/db/repl/tenant_migration_util.cpp b/src/mongo/db/repl/tenant_migration_util.cpp
index a66f88b318f..247da7fdc89 100644
--- a/src/mongo/db/repl/tenant_migration_util.cpp
+++ b/src/mongo/db/repl/tenant_migration_util.cpp
@@ -58,9 +58,8 @@ ExternalKeysCollectionDocument makeExternalClusterTimeKeyDoc(ServiceContext* ser
     return externalKeyDoc;
 }
 
-void storeExternalClusterTimeKeyDocsAndRefreshCache(
-    std::shared_ptr<executor::ScopedTaskExecutor> executor,
-    std::vector<ExternalKeysCollectionDocument> keyDocs) {
+void storeExternalClusterTimeKeyDocs(std::shared_ptr<executor::ScopedTaskExecutor> executor,
+                                     std::vector<ExternalKeysCollectionDocument> keyDocs) {
     auto opCtxHolder = cc().makeOperationContext();
     auto opCtx = opCtxHolder.get();
     auto nss = NamespaceString::kExternalKeysCollectionNamespace;
@@ -82,14 +81,6 @@ void storeExternalClusterTimeKeyDocsAndRefreshCache(
                             /*fromMigrate=*/false);
         });
     }
-
-    auto validator = LogicalTimeValidator::get(opCtx);
-    if (validator) {
-        // Refresh the keys cache to avoid validation errors for external cluster times with
-        // a keyId that matches the keyId of an internal key since the LogicalTimeValidator
-        // only refreshes the cache when it cannot find a matching internal key.
-        validator->refreshKeyManagerCache(opCtx);
-    }
 }
 
 void createRetryableWritesView(OperationContext* opCtx, Database* db) {
diff --git a/src/mongo/db/repl/tenant_migration_util.h b/src/mongo/db/repl/tenant_migration_util.h
index a6a750884de..fe467ec5e0c 100644
--- a/src/mongo/db/repl/tenant_migration_util.h
+++ b/src/mongo/db/repl/tenant_migration_util.h
@@ -140,11 +140,9 @@ ExternalKeysCollectionDocument makeExternalClusterTimeKeyDoc(ServiceContext* ser
  * For each given ExternalKeysCollectionDocument, inserts it if there is not an existing document in
  * admin.system.external_validation_keys for it with the same keyId and replicaSetName. Otherwise,
  * updates the ttlExpiresAt of the existing document if it is less than the new ttlExpiresAt.
- * Refreshes the logical validator's cache before returning.
  */
-void storeExternalClusterTimeKeyDocsAndRefreshCache(
-    std::shared_ptr<executor::ScopedTaskExecutor> executor,
-    std::vector<ExternalKeysCollectionDocument> keyDocs);
+void storeExternalClusterTimeKeyDocs(std::shared_ptr<executor::ScopedTaskExecutor> executor,
+                                     std::vector<ExternalKeysCollectionDocument> keyDocs);
 
 /**
  * Creates a view on the oplog that allows a tenant migration recipient to fetch retryable writes
diff --git a/src/mongo/shell/replsettest.js b/src/mongo/shell/replsettest.js
index 1988f74a2ee..b8784dbd262 100644
--- a/src/mongo/shell/replsettest.js
+++ b/src/mongo/shell/replsettest.js
@@ -204,6 +204,8 @@ var ReplSetTest = function(opts) {
         }
     }
 
+    this.asCluster = asCluster;
+
     /**
      * Returns 'true' if the "conn" has been configured to run without journaling enabled.
      */