SERVER-54302 Write tenant migration test when the state doc collection is dropped

6 files changed, 51 insertions, 9 deletions

diff --git a/src/mongo/db/repl/tenant_migration_access_blocker_registry.cpp b/src/mongo/db/repl/tenant_migration_access_blocker_registry.cpp
index 6117435218a..fd4afafdc24 100644
--- a/src/mongo/db/repl/tenant_migration_access_blocker_registry.cpp
+++ b/src/mongo/db/repl/tenant_migration_access_blocker_registry.cpp
@@ -62,9 +62,7 @@ void TenantMigrationAccessBlockerRegistry::add(StringData tenantId,
     _tenantMigrationAccessBlockers.emplace(tenantId, mtabPair);
 }
 
-void TenantMigrationAccessBlockerRegistry::remove(StringData tenantId, MtabType type) {
-    stdx::lock_guard<Latch> lg(_mutex);
-
+void TenantMigrationAccessBlockerRegistry::_remove(WithLock, StringData tenantId, MtabType type) {
     auto it = _tenantMigrationAccessBlockers.find(tenantId);
     invariant(it != _tenantMigrationAccessBlockers.end());
     auto mtabPair = it->second;
@@ -75,6 +73,19 @@ void TenantMigrationAccessBlockerRegistry::remove(StringData tenantId, MtabType
     }
 }
 
+void TenantMigrationAccessBlockerRegistry::remove(StringData tenantId, MtabType type) {
+    stdx::lock_guard<Latch> lg(_mutex);
+    _remove(lg, tenantId, type);
+}
+
+void TenantMigrationAccessBlockerRegistry::removeAll(MtabType type) {
+    stdx::lock_guard<Latch> lg(_mutex);
+
+    for (auto& [tenantId, _] : _tenantMigrationAccessBlockers) {
+        _remove(lg, tenantId, type);
+    }
+}
+
 boost::optional<MtabPair>
 TenantMigrationAccessBlockerRegistry::getTenantMigrationAccessBlockerForDbName(StringData dbName) {
     stdx::lock_guard<Latch> lg(_mutex);
diff --git a/src/mongo/db/repl/tenant_migration_access_blocker_registry.h b/src/mongo/db/repl/tenant_migration_access_blocker_registry.h
index 51f06185868..2d6c588a883 100644
--- a/src/mongo/db/repl/tenant_migration_access_blocker_registry.h
+++ b/src/mongo/db/repl/tenant_migration_access_blocker_registry.h
@@ -95,6 +95,12 @@ public:
      * Invariants that an entry for tenantId exists, and then removes the entry for (tenantId, mtab)
      */
     void remove(StringData tenantId, TenantMigrationAccessBlocker::BlockerType type);
+    void _remove(WithLock, StringData tenantId, TenantMigrationAccessBlocker::BlockerType type);
+
+    /**
+     * Removes all mtabs of the given type.
+     */
+    void removeAll(TenantMigrationAccessBlocker::BlockerType type);
 
     /**
      * Iterates through each of the TenantMigrationAccessBlockers and
diff --git a/src/mongo/db/repl/tenant_migration_donor_op_observer.cpp b/src/mongo/db/repl/tenant_migration_donor_op_observer.cpp
index 2d8a242507f..8c8fca5a31a 100644
--- a/src/mongo/db/repl/tenant_migration_donor_op_observer.cpp
+++ b/src/mongo/db/repl/tenant_migration_donor_op_observer.cpp
@@ -269,6 +269,20 @@ void TenantMigrationDonorOpObserver::onDelete(OperationContext* opCtx,
     }
 }
 
+repl::OpTime TenantMigrationDonorOpObserver::onDropCollection(OperationContext* opCtx,
+                                                              const NamespaceString& collectionName,
+                                                              OptionalCollectionUUID uuid,
+                                                              std::uint64_t numRecords,
+                                                              const CollectionDropType dropType) {
+    if (collectionName == NamespaceString::kTenantMigrationDonorsNamespace) {
+        opCtx->recoveryUnit()->onCommit([opCtx](boost::optional<Timestamp>) {
+            TenantMigrationAccessBlockerRegistry::get(opCtx->getServiceContext())
+                .removeAll(TenantMigrationAccessBlocker::BlockerType::kDonor);
+        });
+    }
+    return {};
+}
+
 void TenantMigrationDonorOpObserver::onMajorityCommitPointUpdate(
     ServiceContext* service, const repl::OpTime& newCommitPoint) {
     TenantMigrationAccessBlockerRegistry::get(service).onMajorityCommitPointUpdate(newCommitPoint);
diff --git a/src/mongo/db/repl/tenant_migration_donor_op_observer.h b/src/mongo/db/repl/tenant_migration_donor_op_observer.h
index 3e7cfc169ea..4cdd5d30c3f 100644
--- a/src/mongo/db/repl/tenant_migration_donor_op_observer.h
+++ b/src/mongo/db/repl/tenant_migration_donor_op_observer.h
@@ -124,9 +124,7 @@ public:
                                   const NamespaceString& collectionName,
                                   OptionalCollectionUUID uuid,
                                   std::uint64_t numRecords,
-                                  CollectionDropType dropType) final {
-        return repl::OpTime();
-    }
+                                  CollectionDropType dropType) final;
 
     void onDropIndex(OperationContext* opCtx,
                      const NamespaceString& nss,
diff --git a/src/mongo/db/repl/tenant_migration_recipient_op_observer.cpp b/src/mongo/db/repl/tenant_migration_recipient_op_observer.cpp
index 0312203101d..8dfaab1cbe8 100644
--- a/src/mongo/db/repl/tenant_migration_recipient_op_observer.cpp
+++ b/src/mongo/db/repl/tenant_migration_recipient_op_observer.cpp
@@ -165,5 +165,20 @@ void TenantMigrationRecipientOpObserver::onDelete(OperationContext* opCtx,
     }
 }
 
+repl::OpTime TenantMigrationRecipientOpObserver::onDropCollection(
+    OperationContext* opCtx,
+    const NamespaceString& collectionName,
+    OptionalCollectionUUID uuid,
+    std::uint64_t numRecords,
+    const CollectionDropType dropType) {
+    if (collectionName == NamespaceString::kTenantMigrationRecipientsNamespace) {
+        opCtx->recoveryUnit()->onCommit([opCtx](boost::optional<Timestamp>) {
+            TenantMigrationAccessBlockerRegistry::get(opCtx->getServiceContext())
+                .removeAll(TenantMigrationAccessBlocker::BlockerType::kRecipient);
+        });
+    }
+    return {};
+}
+
 }  // namespace repl
 }  // namespace mongo
diff --git a/src/mongo/db/repl/tenant_migration_recipient_op_observer.h b/src/mongo/db/repl/tenant_migration_recipient_op_observer.h
index 4a6fae354ab..ccf3e6db38f 100644
--- a/src/mongo/db/repl/tenant_migration_recipient_op_observer.h
+++ b/src/mongo/db/repl/tenant_migration_recipient_op_observer.h
@@ -125,9 +125,7 @@ public:
                                   const NamespaceString& collectionName,
                                   OptionalCollectionUUID uuid,
                                   std::uint64_t numRecords,
-                                  CollectionDropType dropType) final {
-        return repl::OpTime();
-    }
+                                  CollectionDropType dropType) final;
 
     void onDropIndex(OperationContext* opCtx,
                      const NamespaceString& nss,