SERVER-78009 Commit/Abort Resharding is Retryable on Shutdownr7.0.3-rc0

(cherry picked from commit 7caff9c00c6bc56b55107b1ef1b10a14c2f50516)

6 files changed, 57 insertions, 19 deletions

diff --git a/src/mongo/db/repl/replica_set_aware_service.h b/src/mongo/db/repl/replica_set_aware_service.h
index 5192bb79482..91c5b1aad4c 100644
--- a/src/mongo/db/repl/replica_set_aware_service.h
+++ b/src/mongo/db/repl/replica_set_aware_service.h
@@ -142,6 +142,8 @@ public:
 
     /**
      * Called as part of ReplicationCoordinator shutdown.
+     * Note that it is possible that we are still a writable primary after onShutdown() has been
+     * called (see SERVER-81115).
      */
     virtual void onShutdown() = 0;
 
diff --git a/src/mongo/db/s/config/configsvr_abort_reshard_collection_command.cpp b/src/mongo/db/s/config/configsvr_abort_reshard_collection_command.cpp
index 9388b0e735a..83225f43a00 100644
--- a/src/mongo/db/s/config/configsvr_abort_reshard_collection_command.cpp
+++ b/src/mongo/db/s/config/configsvr_abort_reshard_collection_command.cpp
@@ -78,10 +78,10 @@ void assertExistsReshardingDocument(OperationContext* opCtx, UUID reshardingUUID
 }
 
 auto assertGetReshardingMachine(OperationContext* opCtx, UUID reshardingUUID) {
-    auto machine = resharding::tryGetReshardingStateMachine<ReshardingCoordinatorService,
-                                                            ReshardingCoordinator,
-                                                            ReshardingCoordinatorDocument>(
-        opCtx, reshardingUUID);
+    auto machine = resharding::tryGetReshardingStateMachineAndThrowIfShuttingDown<
+        ReshardingCoordinatorService,
+        ReshardingCoordinator,
+        ReshardingCoordinatorDocument>(opCtx, reshardingUUID);
 
     uassert(ErrorCodes::NoSuchReshardCollection,
             "Could not find in-progress resharding operation to abort",
diff --git a/src/mongo/db/s/config/configsvr_commit_reshard_collection_command.cpp b/src/mongo/db/s/config/configsvr_commit_reshard_collection_command.cpp
index 0e7fe758552..4c67b8430be 100644
--- a/src/mongo/db/s/config/configsvr_commit_reshard_collection_command.cpp
+++ b/src/mongo/db/s/config/configsvr_commit_reshard_collection_command.cpp
@@ -87,10 +87,10 @@ public:
 
             UUID reshardingUUID = retrieveReshardingUUID(opCtx, ns());
 
-            auto machine = resharding::tryGetReshardingStateMachine<ReshardingCoordinatorService,
-                                                                    ReshardingCoordinator,
-                                                                    ReshardingCoordinatorDocument>(
-                opCtx, reshardingUUID);
+            auto machine = resharding::tryGetReshardingStateMachineAndThrowIfShuttingDown<
+                ReshardingCoordinatorService,
+                ReshardingCoordinator,
+                ReshardingCoordinatorDocument>(opCtx, reshardingUUID);
 
             uassert(ErrorCodes::NoSuchReshardCollection,
                     "Could not find in-progress resharding operation to commit",
diff --git a/src/mongo/db/s/resharding/resharding_donor_recipient_common.h b/src/mongo/db/s/resharding/resharding_donor_recipient_common.h
index 834982a75c4..a624ccbc18d 100644
--- a/src/mongo/db/s/resharding/resharding_donor_recipient_common.h
+++ b/src/mongo/db/s/resharding/resharding_donor_recipient_common.h
@@ -38,15 +38,50 @@ namespace resharding {
 using ReshardingFields = TypeCollectionReshardingFields;
 
 /**
- * Looks up the StateMachine by the 'reshardingUUID'. If it does not exist, returns boost::none.
+ * Looks up the StateMachine by the 'reshardingUUID'. Returns boost::none in the following cases:
+ * 1. The state machine does not exist.
+ * 2. In certain cases when the node is shutting down.
+ * Additionally returns a bool indicating if the node is stepping or shutting down to disambiguate
+ * the two.
  */
 template <class Service, class StateMachine, class ReshardingDocument>
-boost::optional<std::shared_ptr<StateMachine>> tryGetReshardingStateMachine(
-    OperationContext* opCtx, const UUID& reshardingUUID) {
+std::pair<boost::optional<std::shared_ptr<StateMachine>>, bool>
+tryGetReshardingStateMachineAndShutdownState(OperationContext* opCtx, const UUID& reshardingUUID) {
     auto instanceId = BSON(ReshardingDocument::kReshardingUUIDFieldName << reshardingUUID);
     auto registry = repl::PrimaryOnlyServiceRegistry::get(opCtx->getServiceContext());
     auto service = registry->lookupServiceByName(Service::kServiceName);
-    auto [instance, _] = StateMachine::lookup(opCtx, service, instanceId);
+    return StateMachine::lookup(opCtx, service, instanceId);
+}
+
+/**
+ * Same as tryGetReshardingStateMachineAndShutdownState, except does not return the shutdown state.
+ */
+template <class Service, class StateMachine, class ReshardingDocument>
+boost::optional<std::shared_ptr<StateMachine>> tryGetReshardingStateMachine(
+    OperationContext* opCtx, const UUID& reshardingUUID) {
+    auto [instance, _] =
+        tryGetReshardingStateMachineAndShutdownState<Service, StateMachine, ReshardingDocument>(
+            opCtx, reshardingUUID);
+    return instance;
+}
+
+/**
+ * Same as tryGetReshardingStateMachine, except throws if we were stepping or shutting down when we
+ * tried to access the PrimaryOnlyService. Use this function in situations where you need to
+ * guarantee that a return of boost::none means that there is no state document on disk for the
+ * associated state machine.
+ */
+template <class Service, class StateMachine, class ReshardingDocument>
+boost::optional<std::shared_ptr<StateMachine>> tryGetReshardingStateMachineAndThrowIfShuttingDown(
+    OperationContext* opCtx, const UUID& reshardingUUID) {
+    auto [instance, steppingOrShuttingDown] =
+        tryGetReshardingStateMachineAndShutdownState<Service, StateMachine, ReshardingDocument>(
+            opCtx, reshardingUUID);
+
+    uassert(ErrorCodes::InterruptedDueToReplStateChange,
+            "Unable to get resharding state machine, if it exists, because the node is "
+            "stepping or shutting down.",
+            !steppingOrShuttingDown);
 
     return instance;
 }
diff --git a/src/mongo/db/s/shardsvr_abort_reshard_collection_command.cpp b/src/mongo/db/s/shardsvr_abort_reshard_collection_command.cpp
index d4c3195f3da..72b174dfcaa 100644
--- a/src/mongo/db/s/shardsvr_abort_reshard_collection_command.cpp
+++ b/src/mongo/db/s/shardsvr_abort_reshard_collection_command.cpp
@@ -68,7 +68,7 @@ public:
 
             std::vector<SharedSemiFuture<void>> futuresToWait;
 
-            if (auto machine = resharding::tryGetReshardingStateMachine<
+            if (auto machine = resharding::tryGetReshardingStateMachineAndThrowIfShuttingDown<
                     ReshardingRecipientService,
                     ReshardingRecipientService::RecipientStateMachine,
                     ReshardingRecipientDocument>(opCtx, uuid())) {
@@ -80,7 +80,7 @@ public:
                 (*machine)->abort(isUserCanceled());
             }
 
-            if (auto machine = resharding::tryGetReshardingStateMachine<
+            if (auto machine = resharding::tryGetReshardingStateMachineAndThrowIfShuttingDown<
                     ReshardingDonorService,
                     ReshardingDonorService::DonorStateMachine,
                     ReshardingDonorDocument>(opCtx, uuid())) {
diff --git a/src/mongo/db/s/shardsvr_commit_reshard_collection_command.cpp b/src/mongo/db/s/shardsvr_commit_reshard_collection_command.cpp
index 9f719315717..e18f7488080 100644
--- a/src/mongo/db/s/shardsvr_commit_reshard_collection_command.cpp
+++ b/src/mongo/db/s/shardsvr_commit_reshard_collection_command.cpp
@@ -69,12 +69,13 @@ public:
             std::vector<SharedSemiFuture<void>> futuresToWait;
 
             {
-                auto recipientMachine = resharding::tryGetReshardingStateMachine<
-                    ReshardingRecipientService,
-                    ReshardingRecipientService::RecipientStateMachine,
-                    ReshardingRecipientDocument>(opCtx, uuid());
+                auto recipientMachine =
+                    resharding::tryGetReshardingStateMachineAndThrowIfShuttingDown<
+                        ReshardingRecipientService,
+                        ReshardingRecipientService::RecipientStateMachine,
+                        ReshardingRecipientDocument>(opCtx, uuid());
 
-                auto donorMachine = resharding::tryGetReshardingStateMachine<
+                auto donorMachine = resharding::tryGetReshardingStateMachineAndThrowIfShuttingDown<
                     ReshardingDonorService,
                     ReshardingDonorService::DonorStateMachine,
                     ReshardingDonorDocument>(opCtx, uuid());