SERVER-63724 ShardingDDLCoordinators should forward the WriteBlockBypass flag when running commands on participant shards

6 files changed, 40 insertions, 10 deletions

diff --git a/src/mongo/db/s/forwardable_operation_metadata.cpp b/src/mongo/db/s/forwardable_operation_metadata.cpp
index b49c73f9fc9..d3ca368cfb8 100644
--- a/src/mongo/db/s/forwardable_operation_metadata.cpp
+++ b/src/mongo/db/s/forwardable_operation_metadata.cpp
@@ -32,6 +32,7 @@
 #include "mongo/db/s/forwardable_operation_metadata.h"
 
 #include "mongo/db/auth/authorization_session.h"
+#include "mongo/db/write_block_bypass.h"
 #include "mongo/rpc/metadata/impersonated_user_metadata.h"
 
 namespace mongo {
@@ -48,6 +49,8 @@ ForwardableOperationMetadata::ForwardableOperationMetadata(OperationContext* opC
     if (const auto authMetadata = rpc::getImpersonatedUserMetadata(opCtx)) {
         setImpersonatedUserMetadata({{authMetadata->getUsers(), authMetadata->getRoles()}});
     }
+
+    setMayBypassWriteBlocking(WriteBlockBypass::get(opCtx).isWriteBlockBypassEnabled());
 }
 
 void ForwardableOperationMetadata::setOn(OperationContext* opCtx) const {
@@ -64,6 +67,8 @@ void ForwardableOperationMetadata::setOn(OperationContext* opCtx) const {
                                                                        authMetadata.getRoles());
         }
     }
+
+    WriteBlockBypass::get(opCtx).set(getMayBypassWriteBlocking());
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/s/forwardable_operation_metadata.idl b/src/mongo/db/s/forwardable_operation_metadata.idl
index 0d70bebfed6..68c8c74042c 100644
--- a/src/mongo/db/s/forwardable_operation_metadata.idl
+++ b/src/mongo/db/s/forwardable_operation_metadata.idl
@@ -59,3 +59,6 @@ structs:
                 type: AuthenticationMetadata
                 description: "The impersonated users associated with the operation."
                 optional: true
+            mayBypassWriteBlocking:
+                type: optionalBool
+                description: "Whether this operation is allowed to bypass user write blocking."
diff --git a/src/mongo/db/s/sharding_ddl_coordinator_service.cpp b/src/mongo/db/s/sharding_ddl_coordinator_service.cpp
index e684b6f4355..3d889877423 100644
--- a/src/mongo/db/s/sharding_ddl_coordinator_service.cpp
+++ b/src/mongo/db/s/sharding_ddl_coordinator_service.cpp
@@ -169,15 +169,16 @@ void ShardingDDLCoordinatorService::waitForCoordinatorsOfGivenTypeToComplete(
     });
 }
 
-void ShardingDDLCoordinatorService::waitForOngoingCoordinatorsToFinish(OperationContext* opCtx) {
+void ShardingDDLCoordinatorService::waitForOngoingCoordinatorsToFinish(
+    OperationContext* opCtx, std::function<bool(const ShardingDDLCoordinator&)> pred) {
     std::vector<SharedSemiFuture<void>> futuresToWait;
 
     const auto instances = getAllInstances(opCtx);
     for (const auto& instance : instances) {
         auto typedInstance = checked_pointer_cast<ShardingDDLCoordinator>(instance);
-        // TODO: SERVER-63724 Wait only for coordinators that don't have the user-write-blocking
-        // bypass enabled.
-        futuresToWait.emplace_back(typedInstance->getCompletionFuture());
+        if (pred(*typedInstance)) {
+            futuresToWait.emplace_back(typedInstance->getCompletionFuture());
+        }
     }
 
     for (auto&& future : futuresToWait) {
diff --git a/src/mongo/db/s/sharding_ddl_coordinator_service.h b/src/mongo/db/s/sharding_ddl_coordinator_service.h
index a2c4aae17cc..d0fb1d023f9 100644
--- a/src/mongo/db/s/sharding_ddl_coordinator_service.h
+++ b/src/mongo/db/s/sharding_ddl_coordinator_service.h
@@ -35,6 +35,8 @@
 
 namespace mongo {
 
+class ShardingDDLCoordinator;
+
 class ShardingDDLCoordinatorService final : public repl::PrimaryOnlyService {
 public:
     static constexpr StringData kServiceName = "ShardingDDLCoordinator"_sd;
@@ -76,10 +78,13 @@ public:
                                                   DDLCoordinatorTypeEnum type) const;
 
     /**
-     * Waits for all currently running coordinators to finish. While waiting here, new coordinators
-     * may start, but they will not be waited for.
+     * Waits for all currently running coordinators matching the predicate 'pred' to finish. While
+     * waiting here, new coordinators may start, but they will not be waited for.
      */
-    void waitForOngoingCoordinatorsToFinish(OperationContext* opCtx);
+    void waitForOngoingCoordinatorsToFinish(
+        OperationContext* opCtx,
+        std::function<bool(const ShardingDDLCoordinator&)> pred = {
+            [](const ShardingDDLCoordinator&) { return true; }});
 
 private:
     ExecutorFuture<void> _rebuildService(std::shared_ptr<executor::ScopedTaskExecutor> executor,
diff --git a/src/mongo/db/s/sharding_ddl_util.cpp b/src/mongo/db/s/sharding_ddl_util.cpp
index c66375b6c06..8a0c739d051 100644
--- a/src/mongo/db/s/sharding_ddl_util.cpp
+++ b/src/mongo/db/s/sharding_ddl_util.cpp
@@ -42,7 +42,9 @@
 #include "mongo/db/s/shard_filtering_metadata_refresh.h"
 #include "mongo/db/s/sharding_logging.h"
 #include "mongo/db/s/sharding_util.h"
+#include "mongo/db/server_feature_flags_gen.h"
 #include "mongo/db/vector_clock.h"
+#include "mongo/db/write_block_bypass.h"
 #include "mongo/logv2/log.h"
 #include "mongo/rpc/metadata/impersonated_user_metadata.h"
 #include "mongo/s/catalog/type_chunk.h"
@@ -209,6 +211,9 @@ std::vector<AsyncRequestsSender::Response> sendAuthenticatedCommandToShards(
     // the command
     BSONObjBuilder bob(command);
     rpc::writeAuthDataToImpersonatedUserMetadata(opCtx, &bob);
+    if (gFeatureFlagUserWriteBlocking.isEnabled(serverGlobalParams.featureCompatibility)) {
+        WriteBlockBypass::get(opCtx).writeAsMetadata(&bob);
+    }
     auto authenticatedCommand = bob.obj();
     return sharding_util::sendCommandToShards(
         opCtx, dbName, authenticatedCommand, shardIds, executor);
diff --git a/src/mongo/db/s/shardsvr_set_user_write_block_mode_command.cpp b/src/mongo/db/s/shardsvr_set_user_write_block_mode_command.cpp
index 71053ef77a6..e5b7515398f 100644
--- a/src/mongo/db/s/shardsvr_set_user_write_block_mode_command.cpp
+++ b/src/mongo/db/s/shardsvr_set_user_write_block_mode_command.cpp
@@ -33,6 +33,7 @@
 
 #include "mongo/db/auth/authorization_session.h"
 #include "mongo/db/commands.h"
+#include "mongo/db/s/sharding_ddl_coordinator.h"
 #include "mongo/db/s/sharding_ddl_coordinator_service.h"
 #include "mongo/db/s/user_writes_recoverable_critical_section_service.h"
 #include "mongo/logv2/log.h"
@@ -76,9 +77,19 @@ public:
                         // Wait for ongoing ShardingDDLCoordinators to finish. This ensures that all
                         // coordinators that started before enabling blocking have finish, and that
                         // any new coordinator that is started after this point will see the
-                        // blocking is enabled.
-                        ShardingDDLCoordinatorService::getService(opCtx)
-                            ->waitForOngoingCoordinatorsToFinish(opCtx);
+                        // blocking is enabled. Wait only for coordinators that don't have
+                        // the user-write-blocking bypass enabled -- the ones allowed to bypass user
+                        // write blocking don't care about the write blocking state.
+                        {
+                            const auto mayNotBypassUserWriteBlockPred =
+                                [](const ShardingDDLCoordinator& coordinatorInstance) -> bool {
+                                return !coordinatorInstance.getForwardableOpMetadata()
+                                            .getMayBypassWriteBlocking();
+                            };
+                            ShardingDDLCoordinatorService::getService(opCtx)
+                                ->waitForOngoingCoordinatorsToFinish(
+                                    opCtx, mayNotBypassUserWriteBlockPred);
+                        }
                         break;
                     case ShardsvrSetUserWriteBlockModePhaseEnum::kComplete:
                         UserWritesRecoverableCriticalSectionService::get(opCtx)