SERVER-57204 Extend the _configsvrSetAllowMigrations command to take the collection uuid to make it idempotent

3 files changed, 32 insertions, 11 deletions

diff --git a/src/mongo/db/s/config/configsvr_set_allow_migrations_command.cpp b/src/mongo/db/s/config/configsvr_set_allow_migrations_command.cpp
index 8ab105f72e6..e467753f864 100644
--- a/src/mongo/db/s/config/configsvr_set_allow_migrations_command.cpp
+++ b/src/mongo/db/s/config/configsvr_set_allow_migrations_command.cpp
@@ -64,9 +64,10 @@ public:
                 repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern);
 
             const auto allowMigrations = request().getAllowMigrations();
+            const auto& collectionUUID = request().getCollectionUUID();
 
             ShardingCatalogManager::get(opCtx)->setAllowMigrationsAndBumpOneChunk(
-                opCtx, nss, allowMigrations);
+                opCtx, nss, collectionUUID, allowMigrations);
         }
 
     private:
diff --git a/src/mongo/db/s/config/sharding_catalog_manager.h b/src/mongo/db/s/config/sharding_catalog_manager.h
index 4e7f4f1ef35..47e01831332 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager.h
+++ b/src/mongo/db/s/config/sharding_catalog_manager.h
@@ -312,6 +312,7 @@ public:
      */
     void setAllowMigrationsAndBumpOneChunk(OperationContext* opCtx,
                                            const NamespaceString& nss,
+                                           const boost::optional<UUID>& collectionUUID,
                                            bool allowMigrations);
 
     //
diff --git a/src/mongo/db/s/config/sharding_catalog_manager_chunk_operations.cpp b/src/mongo/db/s/config/sharding_catalog_manager_chunk_operations.cpp
index 7dc2a338a33..97aaf4cb686 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager_chunk_operations.cpp
+++ b/src/mongo/db/s/config/sharding_catalog_manager_chunk_operations.cpp
@@ -1519,9 +1519,11 @@ void ShardingCatalogManager::splitOrMarkJumbo(OperationContext* opCtx,
     }
 }
 
-void ShardingCatalogManager::setAllowMigrationsAndBumpOneChunk(OperationContext* opCtx,
-                                                               const NamespaceString& nss,
-                                                               bool allowMigrations) {
+void ShardingCatalogManager::setAllowMigrationsAndBumpOneChunk(
+    OperationContext* opCtx,
+    const NamespaceString& nss,
+    const boost::optional<UUID>& collectionUUID,
+    bool allowMigrations) {
     std::set<ShardId> shardsIds;
     {
         // Take _kChunkOpLock in exclusive mode to prevent concurrent chunk splits, merges, and
@@ -1531,6 +1533,13 @@ void ShardingCatalogManager::setAllowMigrationsAndBumpOneChunk(OperationContext*
         const auto cm = uassertStatusOK(
             Grid::get(opCtx)->catalogCache()->getShardedCollectionRoutingInfoWithRefresh(opCtx,
                                                                                          nss));
+
+        uassert(ErrorCodes::InvalidUUID,
+                str::stream() << "Collection uuid " << collectionUUID
+                              << " in the request does not match the current uuid " << cm.getUUID()
+                              << " for ns " << nss,
+                !collectionUUID || collectionUUID == cm.getUUID());
+
         cm.getAllShardIds(&shardsIds);
         withTransaction(
             opCtx, CollectionType::ConfigNS, [&](OperationContext* opCtx, TxnNumber txnNumber) {
@@ -1541,16 +1550,26 @@ void ShardingCatalogManager::setAllowMigrationsAndBumpOneChunk(OperationContext*
                     ? BSON("$unset" << BSON(CollectionType::kAllowMigrationsFieldName << ""))
                     : BSON("$set" << BSON(CollectionType::kAllowMigrationsFieldName << false));
 
-                writeToConfigDocumentInTxn(
+                BSONObj query = BSON(CollectionType::kNssFieldName << nss.ns());
+                if (collectionUUID) {
+                    query =
+                        query.addFields(BSON(CollectionType::kUuidFieldName << *collectionUUID));
+                }
+
+                const auto res = writeToConfigDocumentInTxn(
                     opCtx,
                     CollectionType::ConfigNS,
-                    BatchedCommandRequest::buildUpdateOp(
-                        CollectionType::ConfigNS,
-                        BSON(CollectionType::kNssFieldName << nss.ns()) /* query */,
-                        update /* update */,
-                        false /* upsert */,
-                        false /* multi */),
+                    BatchedCommandRequest::buildUpdateOp(CollectionType::ConfigNS,
+                                                         query,
+                                                         update /* update */,
+                                                         false /* upsert */,
+                                                         false /* multi */),
                     txnNumber);
+                const auto numDocsModified = UpdateOp::parseResponse(res).getN();
+                uassert(5720400,
+                        str::stream() << "Expected to match one doc for query " << query
+                                      << " but matched " << numDocsModified,
+                        numDocsModified == 1);
 
                 // Bump the chunk version for one single chunk
                 invariant(!shardsIds.empty());