SERVER-61760 Stop migrations in sharded collections before executing a collMod command

4 files changed, 77 insertions, 28 deletions

diff --git a/src/mongo/db/s/collmod_coordinator.cpp b/src/mongo/db/s/collmod_coordinator.cpp
index 9c13e2756a5..1b379b0d4cf 100644
--- a/src/mongo/db/s/collmod_coordinator.cpp
+++ b/src/mongo/db/s/collmod_coordinator.cpp
@@ -167,10 +167,15 @@ ExecutorFuture<void> CollModCoordinator::_runImpl(
                                 << "Cannot update granularity of a sharded time-series collection.",
                             !hasTimeSeriesGranularityUpdate(_doc.getCollModRequest()));
                     }
+                    _doc.setCollUUID(
+                        sharding_ddl_util::getCollectionUUID(opCtx, nss(), true /* allowViews */));
 
-                    if (_recoveredFromDisk) {
+                    sharding_ddl_util::stopMigrations(opCtx, nss(), _doc.getCollUUID());
+
+                    if (!_firstExecution) {
                         _performNoopRetryableWriteOnParticipants(opCtx, **executor);
                     }
+
                     _doc = _updateSession(opCtx, _doc);
                     const OperationSessionInfo osi = getCurrentSession(_doc);
 
@@ -201,6 +206,7 @@ ExecutorFuture<void> CollModCoordinator::_runImpl(
                         CommandHelpers::appendSimpleCommandStatus(builder, ok, errmsg);
                     }
                     _result = builder.obj();
+                    sharding_ddl_util::resumeMigrations(opCtx, nss(), _doc.getCollUUID());
                 } else {
                     CollMod cmd(nss());
                     cmd.setCollModRequest(_doc.getCollModRequest());
@@ -228,6 +234,15 @@ ExecutorFuture<void> CollModCoordinator::_runImpl(
                             "Error running collMod",
                             "namespace"_attr = nss(),
                             "error"_attr = redact(status));
+                // If we have the collection UUID set, this error happened in a sharded collection,
+                // we should restore the migrations.
+                if (_doc.getCollUUID()) {
+                    auto opCtxHolder = cc().makeOperationContext();
+                    auto* opCtx = opCtxHolder.get();
+                    getForwardableOpMetadata().setOn(opCtx);
+
+                    sharding_ddl_util::resumeMigrations(opCtx, nss(), _doc.getCollUUID());
+                }
             }
             return status;
         });
diff --git a/src/mongo/db/s/collmod_coordinator_document.idl b/src/mongo/db/s/collmod_coordinator_document.idl
index 8ff37dc6308..9c0aaacbf98 100644
--- a/src/mongo/db/s/collmod_coordinator_document.idl
+++ b/src/mongo/db/s/collmod_coordinator_document.idl
@@ -61,3 +61,7 @@ structs:
             collModRequest:
                 type: CollModRequest
                 description: "Initial collMod request."
+            collUUID:
+                type: uuid
+                description: "Collection uuid."
+                optional: true
diff --git a/src/mongo/db/s/sharding_ddl_util.cpp b/src/mongo/db/s/sharding_ddl_util.cpp
index 2d4dc68a02b..4d6eb5e8016 100644
--- a/src/mongo/db/s/sharding_ddl_util.cpp
+++ b/src/mongo/db/s/sharding_ddl_util.cpp
@@ -146,6 +146,35 @@ write_ops::UpdateCommandRequest buildNoopWriteRequestCommand() {
     return updateOp;
 }
 
+void setAllowMigrations(OperationContext* opCtx,
+                        const NamespaceString& nss,
+                        const boost::optional<UUID>& expectedCollectionUUID,
+                        bool allowMigrations) {
+    ConfigsvrSetAllowMigrations configsvrSetAllowMigrationsCmd(nss, allowMigrations);
+    configsvrSetAllowMigrationsCmd.setCollectionUUID(expectedCollectionUUID);
+
+    const auto swSetAllowMigrationsResult =
+        Grid::get(opCtx)->shardRegistry()->getConfigShard()->runCommandWithFixedRetryAttempts(
+            opCtx,
+            ReadPreferenceSetting{ReadPreference::PrimaryOnly},
+            NamespaceString::kAdminDb.toString(),
+            CommandHelpers::appendMajorityWriteConcern(configsvrSetAllowMigrationsCmd.toBSON({})),
+            Shard::RetryPolicy::kIdempotent  // Although ConfigsvrSetAllowMigrations is not really
+                                             // idempotent (because it will cause the collection
+                                             // version to be bumped), it is safe to be retried.
+        );
+    try {
+        uassertStatusOKWithContext(
+            Shard::CommandResponse::getEffectiveStatus(std::move(swSetAllowMigrationsResult)),
+            str::stream() << "Error setting allowMigrations to " << allowMigrations
+                          << " for collection " << nss.toString());
+    } catch (const ExceptionFor<ErrorCodes::NamespaceNotSharded>&) {
+        // Collection no longer exists
+    } catch (const ExceptionFor<ErrorCodes::ConflictingOperationInProgress>&) {
+        // Collection metadata was concurrently dropped
+    }
+}
+
 }  // namespace
 
 void linearizeCSRSReads(OperationContext* opCtx) {
@@ -404,34 +433,23 @@ boost::optional<CreateCollectionResponse> checkIfCollectionAlreadySharded(
 void stopMigrations(OperationContext* opCtx,
                     const NamespaceString& nss,
                     const boost::optional<UUID>& expectedCollectionUUID) {
-    ConfigsvrSetAllowMigrations configsvrSetAllowMigrationsCmd(nss, false /* allowMigrations */);
-    configsvrSetAllowMigrationsCmd.setCollectionUUID(expectedCollectionUUID);
-
-    const auto swSetAllowMigrationsResult =
-        Grid::get(opCtx)->shardRegistry()->getConfigShard()->runCommandWithFixedRetryAttempts(
-            opCtx,
-            ReadPreferenceSetting{ReadPreference::PrimaryOnly},
-            NamespaceString::kAdminDb.toString(),
-            CommandHelpers::appendMajorityWriteConcern(configsvrSetAllowMigrationsCmd.toBSON({})),
-            Shard::RetryPolicy::kIdempotent  // Although ConfigsvrSetAllowMigrations is not really
-                                             // idempotent (because it will cause the collection
-                                             // version to be bumped), it is safe to be retried.
-        );
+    setAllowMigrations(opCtx, nss, expectedCollectionUUID, false);
+}
 
-    try {
-        uassertStatusOKWithContext(
-            Shard::CommandResponse::getEffectiveStatus(std::move(swSetAllowMigrationsResult)),
-            str::stream() << "Error setting allowMigrations to false for collection "
-                          << nss.toString());
-    } catch (const ExceptionFor<ErrorCodes::NamespaceNotSharded>&) {
-        // Collection no longer exists
-    } catch (const ExceptionFor<ErrorCodes::ConflictingOperationInProgress>&) {
-        // Collection metadata was concurrently dropped
-    }
+void resumeMigrations(OperationContext* opCtx,
+                      const NamespaceString& nss,
+                      const boost::optional<UUID>& expectedCollectionUUID) {
+    setAllowMigrations(opCtx, nss, expectedCollectionUUID, true);
 }
 
-boost::optional<UUID> getCollectionUUID(OperationContext* opCtx, const NamespaceString& nss) {
-    AutoGetCollection autoColl(opCtx, nss, MODE_IS, AutoGetCollectionViewMode::kViewsForbidden);
+boost::optional<UUID> getCollectionUUID(OperationContext* opCtx,
+                                        const NamespaceString& nss,
+                                        bool allowViews) {
+    AutoGetCollection autoColl(opCtx,
+                               nss,
+                               MODE_IS,
+                               allowViews ? AutoGetCollectionViewMode::kViewsPermitted
+                                          : AutoGetCollectionViewMode::kViewsForbidden);
     return autoColl ? boost::make_optional(autoColl->uuid()) : boost::none;
 }
 
diff --git a/src/mongo/db/s/sharding_ddl_util.h b/src/mongo/db/s/sharding_ddl_util.h
index ec198cc7ad6..1cfba12be74 100644
--- a/src/mongo/db/s/sharding_ddl_util.h
+++ b/src/mongo/db/s/sharding_ddl_util.h
@@ -138,16 +138,28 @@ boost::optional<CreateCollectionResponse> checkIfCollectionAlreadySharded(
 /**
  * Stops ongoing migrations and prevents future ones to start for the given nss.
  * If expectedCollectionUUID is set and doesn't match that of that collection, then this is a no-op.
+ * If expectedCollectionUUID is not set, no UUID check will be performed before stopping migrations.
  */
 void stopMigrations(OperationContext* opCtx,
                     const NamespaceString& nss,
                     const boost::optional<UUID>& expectedCollectionUUID);
 
+/**
+ * Resume migrations and balancing rounds for the given nss.
+ * If expectedCollectionUUID is set and doesn't match that of the collection, then this is a no-op.
+ * If expectedCollectionUUID is not set, no UUID check will be performed before resuming migrations.
+ */
+void resumeMigrations(OperationContext* opCtx,
+                      const NamespaceString& nss,
+                      const boost::optional<UUID>& expectedCollectionUUID);
+
 /*
  * Returns the UUID of the collection (if exists) using the catalog. It does not provide any locking
- *guarantees.
+ * guarantees after the call.
  **/
-boost::optional<UUID> getCollectionUUID(OperationContext* opCtx, const NamespaceString& nss);
+boost::optional<UUID> getCollectionUUID(OperationContext* opCtx,
+                                        const NamespaceString& nss,
+                                        bool allowViews = false);
 
 /*
  * Performs a noop retryable write on the given shards using the session and txNumber specified in