SERVER-53406 Store external keys once per migration

10 files changed, 94 insertions, 74 deletions

diff --git a/jstests/replsets/libs/tenant_migration_test.js b/jstests/replsets/libs/tenant_migration_test.js
index 70d6c24c571..d49fa3bc342 100644
--- a/jstests/replsets/libs/tenant_migration_test.js
+++ b/jstests/replsets/libs/tenant_migration_test.js
@@ -272,9 +272,6 @@ function TenantMigrationTest({
             this.checkTenantDBHashes(tenantId);
         }
 
-        this.assertNoDuplicatedExternalKeyDocs(this.getDonorRst());
-        this.assertNoDuplicatedExternalKeyDocs(this.getRecipientRst());
-
         return stateRes;
     };
 
@@ -578,22 +575,6 @@ function TenantMigrationTest({
     };
 
     /**
-     * Asserts that there are no admin.system.external_validation_keys docs with the same keyId
-     * and replicaSetName.
-     */
-    this.assertNoDuplicatedExternalKeyDocs = function(rst) {
-        grantFindExternalClusterTimeKeysPrivilegeIfNeeded(rst);
-
-        const docs = rst.getPrimary().getDB("admin").system.external_validation_keys.find();
-        let aggRes = rst.getPrimary().getDB("admin").system.external_validation_keys.aggregate([
-            {$group: {_id: {keyId: "$keyId", replicaSetName: "$replicaSetName"}, count: {$sum: 1}}}
-        ]);
-        aggRes.forEach(doc => {
-            assert.eq(1, doc.count, "all external key docs: " + tojson(docs.toArray()));
-        });
-    };
-
-    /**
      * Shuts down the donor and recipient sets, only if they were not passed in as parameters.
      * If they were passed in, the test that initialized them should be responsible for shutting
      * them down.
diff --git a/jstests/replsets/tenant_migration_cluster_time_keys_cloning.js b/jstests/replsets/tenant_migration_cluster_time_keys_cloning.js
index 4642934e4b7..c911af0417a 100644
--- a/jstests/replsets/tenant_migration_cluster_time_keys_cloning.js
+++ b/jstests/replsets/tenant_migration_cluster_time_keys_cloning.js
@@ -20,7 +20,7 @@ const kExternalKeysNs = "admin.system.external_validation_keys";
  * Asserts that the donor and recipient have copied each other's cluster time keys into
  * admin.system.external_validation_keys.
  */
-function assertCopiedExternalKeys(tenantMigrationTest) {
+function assertCopiedExternalKeys(tenantMigrationTest, migrationId) {
     const donorPrimary = tenantMigrationTest.getDonorPrimary();
     const recipientPrimary = tenantMigrationTest.getRecipientPrimary();
 
@@ -29,7 +29,7 @@ function assertCopiedExternalKeys(tenantMigrationTest) {
             keyId: internalKeyDoc._id,
             key: internalKeyDoc.key,
             expiresAt: internalKeyDoc.expiresAt,
-            replicaSetName: tenantMigrationTest.getRecipientRst().name
+            migrationId,
         }));
     });
 
@@ -38,12 +38,23 @@ function assertCopiedExternalKeys(tenantMigrationTest) {
             keyId: internalKeyDoc._id,
             key: internalKeyDoc.key,
             expiresAt: internalKeyDoc.expiresAt,
-            replicaSetName: tenantMigrationTest.getDonorRst().name
+            migrationId,
         }));
     });
 }
 
-const kTenantId = "testTenantId";
+function runMigrationAndAssertExternalKeysCopied(tenantMigrationTest, tenantId) {
+    const migrationId = UUID();
+    const migrationOpts = {
+        migrationIdString: extractUUIDFromObject(migrationId),
+        tenantId: tenantId,
+    };
+    assert.commandWorked(tenantMigrationTest.runMigration(migrationOpts));
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
+}
+
+const kTenantId1 = "testTenantId1";
+const kTenantId2 = "testTenantId2";
 const migrationX509Options = TenantMigrationUtil.makeX509OptionsForTest();
 
 (() => {
@@ -59,10 +70,14 @@ const migrationX509Options = TenantMigrationUtil.makeX509OptionsForTest();
     const migrationId = UUID();
     const migrationOpts = {
         migrationIdString: extractUUIDFromObject(migrationId),
-        tenantId: kTenantId,
+        tenantId: kTenantId1,
     };
     assert.commandWorked(tenantMigrationTest.runMigration(migrationOpts));
-    assertCopiedExternalKeys(tenantMigrationTest);
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
+
+    // After another migration, the first's keys should still exist.
+    runMigrationAndAssertExternalKeysCopied(tenantMigrationTest, kTenantId2);
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
 
     tenantMigrationTest.stop();
 })();
@@ -86,11 +101,15 @@ const migrationX509Options = TenantMigrationUtil.makeX509OptionsForTest();
     const migrationId = UUID();
     const migrationOpts = {
         migrationIdString: extractUUIDFromObject(migrationId),
-        tenantId: kTenantId,
+        tenantId: kTenantId1,
         readPreference: {mode: "secondary"}
     };
     assert.commandWorked(tenantMigrationTest.runMigration(migrationOpts));
-    assertCopiedExternalKeys(tenantMigrationTest);
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
+
+    // After another migration, the first's keys should still exist.
+    runMigrationAndAssertExternalKeysCopied(tenantMigrationTest, kTenantId2);
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
 
     recipientRst.stopSet();
     tenantMigrationTest.stop();
@@ -119,7 +138,7 @@ const migrationX509Options = TenantMigrationUtil.makeX509OptionsForTest();
     const migrationId = UUID();
     const migrationOpts = {
         migrationIdString: extractUUIDFromObject(migrationId),
-        tenantId: kTenantId,
+        tenantId: kTenantId1,
     };
     assert.commandWorked(tenantMigrationTest.startMigration(migrationOpts));
     fp.wait();
@@ -132,7 +151,11 @@ const migrationX509Options = TenantMigrationUtil.makeX509OptionsForTest();
     assert.commandWorked(tenantMigrationTest.waitForMigrationToComplete(
         migrationOpts, true /* retryOnRetryableErrors */));
 
-    assertCopiedExternalKeys(tenantMigrationTest);
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
+
+    // After another migration, the first's keys should still exist.
+    runMigrationAndAssertExternalKeysCopied(tenantMigrationTest, kTenantId2);
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
 
     donorRst.stopSet();
     tenantMigrationTest.stop();
@@ -162,7 +185,7 @@ const migrationX509Options = TenantMigrationUtil.makeX509OptionsForTest();
     const migrationId = UUID();
     const migrationOpts = {
         migrationIdString: extractUUIDFromObject(migrationId),
-        tenantId: kTenantId,
+        tenantId: kTenantId1,
     };
     assert.commandWorked(tenantMigrationTest.startMigration(migrationOpts));
     fp.wait();
@@ -175,7 +198,11 @@ const migrationX509Options = TenantMigrationUtil.makeX509OptionsForTest();
     assert.commandWorked(tenantMigrationTest.waitForMigrationToComplete(
         migrationOpts, true /* retryOnRetryableErrors */));
 
-    assertCopiedExternalKeys(tenantMigrationTest);
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
+
+    // After another migration, the first's keys should still exist.
+    runMigrationAndAssertExternalKeysCopied(tenantMigrationTest, kTenantId2);
+    assertCopiedExternalKeys(tenantMigrationTest, migrationId);
 
     recipientRst.stopSet();
     tenantMigrationTest.stop();
diff --git a/src/mongo/db/keys_collection_cache.cpp b/src/mongo/db/keys_collection_cache.cpp
index 5e3533ea5f9..46226d1c1b6 100644
--- a/src/mongo/db/keys_collection_cache.cpp
+++ b/src/mongo/db/keys_collection_cache.cpp
@@ -129,7 +129,7 @@ Status KeysCollectionCache::_refreshExternalKeys(OperationContext* opCtx) {
     }
 
     for (auto&& key : newKeys) {
-        _externalKeysCache[key.getKeyId()].emplace(key.getReplicaSetName(), std::move(key));
+        _externalKeysCache.emplace(key.getKeyId(), std::move(key));
     }
 
     return Status::OK();
@@ -157,23 +157,27 @@ StatusWith<std::vector<ExternalKeysCollectionDocument>> KeysCollectionCache::get
     stdx::lock_guard<Latch> lk(_cacheMutex);
     std::vector<ExternalKeysCollectionDocument> keys;
 
-    auto keysIter = _externalKeysCache.find(keyId);
-
-    if (keysIter == _externalKeysCache.end()) {
+    if (_externalKeysCache.empty()) {
         return {ErrorCodes::KeyNotFound,
                 str::stream() << "Cache Reader No external keys found for " << _purpose
                               << " with id: " << keyId};
     }
 
-    invariant(!keysIter->second.empty());
-
-    for (auto keyIter = keysIter->second.begin(); keyIter != keysIter->second.end(); keyIter++) {
+    auto keysRange = _externalKeysCache.equal_range(keyId);
+    for (auto keyIter = keysRange.first; keyIter != keysRange.second; keyIter++) {
         auto key = keyIter->second;
         if (key.getExpiresAt() > forThisTime) {
             keys.push_back(key);
         }
     }
 
+    if (keys.empty()) {
+        return {ErrorCodes::KeyNotFound,
+                str::stream() << "Cache Reader No external keys found for " << _purpose
+                              << " that is valid for time: " << forThisTime.toString()
+                              << " with id: " << keyId};
+    }
+
     return std::move(keys);
 }
 
diff --git a/src/mongo/db/keys_collection_cache.h b/src/mongo/db/keys_collection_cache.h
index a37b4e45c65..736ad135bc2 100644
--- a/src/mongo/db/keys_collection_cache.h
+++ b/src/mongo/db/keys_collection_cache.h
@@ -109,8 +109,7 @@ private:
     // Stores keys for validating cluster times created by other clusters. These key documents
     // cannot be stored in a regular map like _internalKeysCache since expiresAt and keyId are not
     // necessarily unique across clusters so there is chance of collision.
-    stdx::unordered_map<long long, StringMap<ExternalKeysCollectionDocument>>
-        _externalKeysCache;  // keyId -> (replicaSetName -> ExternalKeysDocument)
+    std::multimap<long long, ExternalKeysCollectionDocument> _externalKeysCache;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/keys_collection_cache_test.cpp b/src/mongo/db/keys_collection_cache_test.cpp
index bb42c3f46c3..82ed5670f45 100644
--- a/src/mongo/db/keys_collection_cache_test.cpp
+++ b/src/mongo/db/keys_collection_cache_test.cpp
@@ -47,6 +47,9 @@ namespace {
 
 class CacheTest : public ConfigServerTestFixture {
 protected:
+    const UUID migrationId1 = UUID::gen();
+    const UUID migrationId2 = UUID::gen();
+
     void setUp() override {
         ConfigServerTestFixture::setUp();
 
@@ -177,17 +180,15 @@ 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, "replicaSetName1", externalKeysTTLExpiresAt);
+    ExternalKeysCollectionDocument origKey1(OID::gen(), 1, migrationId1, externalKeysTTLExpiresAt);
     origKey1.setKeysCollectionDocumentBase(
         {"test", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0))});
     insertDocument(
         operationContext(), NamespaceString::kExternalKeysCollectionNamespace, origKey1.toBSON());
 
-    ExternalKeysCollectionDocument origKey2(
-        OID::gen(), 1, "replicaSetName2", externalKeysTTLExpiresAt);
+    ExternalKeysCollectionDocument origKey2(OID::gen(), 1, migrationId2, externalKeysTTLExpiresAt);
     origKey2.setKeysCollectionDocumentBase(
-        {"test", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0))});
+        {"test", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(205, 0))});
     insertDocument(
         operationContext(), NamespaceString::kExternalKeysCollectionNamespace, origKey2.toBSON());
 
@@ -248,6 +249,23 @@ TEST_F(CacheTest, GetKeyShouldReturnCorrectKeysAfterRefreshDirectClient) {
             ASSERT_EQ(expectedKey.getTTLExpiresAt(), key.getTTLExpiresAt());
         }
     }
+
+    swExternalKeys = cache.getExternalKeysById(1, LogicalTime(Timestamp(150, 0)));
+    ASSERT_OK(swExternalKeys.getStatus());
+
+    {
+        auto keys = swExternalKeys.getValue();
+        ASSERT_EQ(1, keys.size());
+        auto key = keys.front();
+
+        ASSERT_EQ(origKey2.getId(), key.getId());
+        ASSERT_EQ(origKey2.getPurpose(), key.getPurpose());
+        ASSERT_EQ(origKey2.getExpiresAt().asTimestamp(), key.getExpiresAt().asTimestamp());
+        ASSERT_EQ(origKey2.getTTLExpiresAt(), key.getTTLExpiresAt());
+    }
+
+    swExternalKeys = cache.getExternalKeysById(1, LogicalTime(Timestamp(300, 0)));
+    ASSERT_EQ(ErrorCodes::KeyNotFound, swExternalKeys.getStatus());
 }
 
 TEST_F(CacheTest, GetInternalKeyShouldReturnErrorIfNoKeyIsValidForGivenTime) {
@@ -376,8 +394,7 @@ TEST_F(CacheTest, RefreshShouldNotGetExternalKeysForOtherPurpose) {
     const auto externalKeysTTLExpiresAt =
         ServiceContext().getFastClockSource()->now() + Seconds(30);
 
-    ExternalKeysCollectionDocument origKey1(
-        OID::gen(), 1, "replicaSetName1", externalKeysTTLExpiresAt);
+    ExternalKeysCollectionDocument origKey1(OID::gen(), 1, migrationId1, externalKeysTTLExpiresAt);
     origKey1.setKeysCollectionDocumentBase(
         {"dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0))});
     insertDocument(
@@ -391,8 +408,7 @@ TEST_F(CacheTest, RefreshShouldNotGetExternalKeysForOtherPurpose) {
         ASSERT_EQ(ErrorCodes::KeyNotFound, swKey.getStatus());
     }
 
-    ExternalKeysCollectionDocument origKey2(
-        OID::gen(), 2, "replicaSetName1", externalKeysTTLExpiresAt);
+    ExternalKeysCollectionDocument origKey2(OID::gen(), 2, migrationId1, externalKeysTTLExpiresAt);
     origKey2.setKeysCollectionDocumentBase(
         {"test", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(110, 0))});
     insertDocument(
diff --git a/src/mongo/db/keys_collection_document.idl b/src/mongo/db/keys_collection_document.idl
index df546b7eff4..e2914519c83 100644
--- a/src/mongo/db/keys_collection_document.idl
+++ b/src/mongo/db/keys_collection_document.idl
@@ -81,7 +81,6 @@ structs:
         chained_structs:
             keysCollectionDocumentBase: keysCollectionDocumentBase
         fields:
-            # TODO (SERVER-53406): Add ttlExpiresAt field.
             _id:
                 type: objectid
                 description: "Unique identifier for this key document."
@@ -92,9 +91,9 @@ structs:
                     NumberLong representation of the cluster time at which the key was created.
                     Corresponds to the _id of the admin.system.keys document for this key.
                 cpp_name: keyId
-            replicaSetName:
-                type: string
-                description: "The name of the replica set that created this key."
+            migrationId:
+                type: uuid
+                description: "The id of the tenant migration that inserted this key."
             ttlExpiresAt:
                 type: date
                 description: >-
diff --git a/src/mongo/db/repl/tenant_migration_donor_service.cpp b/src/mongo/db/repl/tenant_migration_donor_service.cpp
index af57ada365c..618b859f348 100644
--- a/src/mongo/db/repl/tenant_migration_donor_service.cpp
+++ b/src/mongo/db/repl/tenant_migration_donor_service.cpp
@@ -351,7 +351,7 @@ TenantMigrationDonorService::Instance::_fetchAndStoreRecipientClusterTimeKeyDocs
                 const auto& data = dataStatus.getValue();
                 for (const BSONObj& doc : data.documents) {
                     keyDocs.push_back(tenant_migration_util::makeExternalClusterTimeKeyDoc(
-                        _serviceContext, _recipientUri.getSetName(), doc.getOwned()));
+                        _serviceContext, _stateDoc.getId(), doc.getOwned()));
                 }
                 fetchStatus = Status::OK();
 
@@ -401,7 +401,7 @@ TenantMigrationDonorService::Instance::_fetchAndStoreRecipientClusterTimeKeyDocs
                 checkIfReceivedDonorAbortMigration(serviceToken, instanceToken);
 
                 tenant_migration_util::storeExternalClusterTimeKeyDocsAndRefreshCache(
-                    executor, std::move(keyDocs), instanceToken);
+                    executor, std::move(keyDocs));
             });
 }
 
diff --git a/src/mongo/db/repl/tenant_migration_recipient_service.cpp b/src/mongo/db/repl/tenant_migration_recipient_service.cpp
index 8d569280538..19473158f53 100644
--- a/src/mongo/db/repl/tenant_migration_recipient_service.cpp
+++ b/src/mongo/db/repl/tenant_migration_recipient_service.cpp
@@ -1258,11 +1258,11 @@ void TenantMigrationRecipientService::Instance::_fetchAndStoreDonorClusterTimeKe
     while (cursor->more()) {
         const auto doc = cursor->nextSafe().getOwned();
         keyDocs.push_back(tenant_migration_util::makeExternalClusterTimeKeyDoc(
-            _serviceContext, _donorUri.getSetName(), doc));
+            _serviceContext, _migrationUuid, doc));
     }
 
-    tenant_migration_util::storeExternalClusterTimeKeyDocsAndRefreshCache(
-        _scopedExecutor, std::move(keyDocs), token);
+    tenant_migration_util::storeExternalClusterTimeKeyDocsAndRefreshCache(_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 f92baa8b1e7..3c6e4a5878d 100644
--- a/src/mongo/db/repl/tenant_migration_util.cpp
+++ b/src/mongo/db/repl/tenant_migration_util.cpp
@@ -43,14 +43,14 @@ namespace mongo {
 namespace tenant_migration_util {
 
 ExternalKeysCollectionDocument makeExternalClusterTimeKeyDoc(ServiceContext* serviceContext,
-                                                             std::string rsName,
+                                                             UUID migrationId,
                                                              BSONObj keyDoc) {
     auto originalKeyDoc = KeysCollectionDocument::parse(IDLParserErrorContext("keyDoc"), keyDoc);
 
     ExternalKeysCollectionDocument externalKeyDoc(
         OID::gen(),
         originalKeyDoc.getKeyId(),
-        rsName,
+        migrationId,
         serviceContext->getFastClockSource()->now() +
             Seconds{repl::tenantMigrationExternalKeysRemovalDelaySecs.load()});
     externalKeyDoc.setKeysCollectionDocumentBase(originalKeyDoc.getKeysCollectionDocumentBase());
@@ -60,8 +60,7 @@ ExternalKeysCollectionDocument makeExternalClusterTimeKeyDoc(ServiceContext* ser
 
 void storeExternalClusterTimeKeyDocsAndRefreshCache(
     std::shared_ptr<executor::ScopedTaskExecutor> executor,
-    std::vector<ExternalKeysCollectionDocument> keyDocs,
-    const CancelationToken& token) {
+    std::vector<ExternalKeysCollectionDocument> keyDocs) {
     auto opCtxHolder = cc().makeOperationContext();
     auto opCtx = opCtxHolder.get();
     auto nss = NamespaceString::kExternalKeysCollectionNamespace;
@@ -70,15 +69,11 @@ void storeExternalClusterTimeKeyDocsAndRefreshCache(
         AutoGetCollection collection(opCtx, nss, MODE_IX);
 
         writeConflictRetry(opCtx, "CloneExternalKeyDocs", nss.ns(), [&] {
-            const auto filter = BSON(ExternalKeysCollectionDocument::kKeyIdFieldName
-                                     << keyDoc.getKeyId()
-                                     << ExternalKeysCollectionDocument::kReplicaSetNameFieldName
-                                     << keyDoc.getReplicaSetName()
-                                     << ExternalKeysCollectionDocument::kTTLExpiresAtFieldName
-                                     << BSON("$lte" << keyDoc.getTTLExpiresAt()));
-
-            // Remove _id since updating _id is not allowed.
-            const auto updateMod = keyDoc.toBSON().removeField("_id");
+            // Note that each external key's _id is generated by the migration, so this upsert can
+            // only insert.
+            const auto filter =
+                BSON(ExternalKeysCollectionDocument::kIdFieldName << keyDoc.getId());
+            const auto updateMod = keyDoc.toBSON();
 
             Helpers::upsert(opCtx,
                             nss.ns(),
diff --git a/src/mongo/db/repl/tenant_migration_util.h b/src/mongo/db/repl/tenant_migration_util.h
index e76e8464ad0..a6a750884de 100644
--- a/src/mongo/db/repl/tenant_migration_util.h
+++ b/src/mongo/db/repl/tenant_migration_util.h
@@ -133,7 +133,7 @@ inline Status validatePrivateKeyPEMPayload(const StringData& payload) {
  * document from the given the admin.system.keys document BSONObj.
  */
 ExternalKeysCollectionDocument makeExternalClusterTimeKeyDoc(ServiceContext* serviceContext,
-                                                             std::string rsName,
+                                                             UUID migrationId,
                                                              BSONObj keyDoc);
 
 /*
@@ -144,8 +144,7 @@ ExternalKeysCollectionDocument makeExternalClusterTimeKeyDoc(ServiceContext* ser
  */
 void storeExternalClusterTimeKeyDocsAndRefreshCache(
     std::shared_ptr<executor::ScopedTaskExecutor> executor,
-    std::vector<ExternalKeysCollectionDocument> keyDocs,
-    const CancelationToken& token);
+    std::vector<ExternalKeysCollectionDocument> keyDocs);
 
 /**
  * Creates a view on the oplog that allows a tenant migration recipient to fetch retryable writes