SERVER-33773 Add 'waitForMovePrimaryCriticalSection' flag to OperationShardingState

3 files changed, 50 insertions, 3 deletions

diff --git a/src/mongo/db/s/database_sharding_state.cpp b/src/mongo/db/s/database_sharding_state.cpp
index f2841d5a391..8624bb0b7bc 100644
--- a/src/mongo/db/s/database_sharding_state.cpp
+++ b/src/mongo/db/s/database_sharding_state.cpp
@@ -102,9 +102,8 @@ void DatabaseShardingState::checkDbVersion(OperationContext* opCtx) const {
                                                         ? ShardingMigrationCriticalSection::kWrite
                                                         : ShardingMigrationCriticalSection::kRead);
     if (criticalSectionSignal) {
-        // TODO (SERVER-33773): Set movePrimary critical section signal on the
-        // OperationShardingState (so that the operation can wait outside the DBLock for the
-        // movePrimary critical section to end before returning to the client).
+        OperationShardingState::get(opCtx).setMovePrimaryCriticalSectionSignal(
+            criticalSectionSignal);
 
         uasserted(StaleDbRoutingVersion(dbName, *clientDbVersion, boost::none),
                   "movePrimary critical section active");
diff --git a/src/mongo/db/s/operation_sharding_state.cpp b/src/mongo/db/s/operation_sharding_state.cpp
index 4a2e46ecb76..7b54e18c408 100644
--- a/src/mongo/db/s/operation_sharding_state.cpp
+++ b/src/mongo/db/s/operation_sharding_state.cpp
@@ -42,6 +42,9 @@ const OperationContext::Decoration<OperationShardingState> shardingMetadataDecor
 // Max time to wait for the migration critical section to complete
 const Milliseconds kMaxWaitForMigrationCriticalSection = Minutes(5);
 
+// Max time to wait for the movePrimary critical section to complete
+const Milliseconds kMaxWaitForMovePrimaryCriticalSection = Minutes(5);
+
 // The name of the field in which the client attaches its database version.
 constexpr auto kDbVersionField = "databaseVersion"_sd;
 }  // namespace
@@ -161,4 +164,27 @@ void OperationShardingState::setMigrationCriticalSectionSignal(
     _migrationCriticalSectionSignal = std::move(critSecSignal);
 }
 
+bool OperationShardingState::waitForMovePrimaryCriticalSectionSignal(OperationContext* opCtx) {
+    // Must not block while holding a lock
+    invariant(!opCtx->lockState()->isLocked());
+
+    if (_movePrimaryCriticalSectionSignal) {
+        _movePrimaryCriticalSectionSignal->waitFor(
+            opCtx,
+            opCtx->hasDeadline() ? std::min(opCtx->getRemainingMaxTimeMillis(),
+                                            kMaxWaitForMovePrimaryCriticalSection)
+                                 : kMaxWaitForMovePrimaryCriticalSection);
+        _movePrimaryCriticalSectionSignal = nullptr;
+        return true;
+    }
+
+    return false;
+}
+
+void OperationShardingState::setMovePrimaryCriticalSectionSignal(
+    std::shared_ptr<Notification<void>> critSecSignal) {
+    invariant(critSecSignal);
+    _movePrimaryCriticalSectionSignal = std::move(critSecSignal);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/operation_sharding_state.h b/src/mongo/db/s/operation_sharding_state.h
index 995738f3ebb..c274c6935f3 100644
--- a/src/mongo/db/s/operation_sharding_state.h
+++ b/src/mongo/db/s/operation_sharding_state.h
@@ -135,6 +135,23 @@ public:
      */
     void setMigrationCriticalSectionSignal(std::shared_ptr<Notification<void>> critSecSignal);
 
+    /**
+     * This call is a no op if there isn't a currently active movePrimary critical section.
+     * Otherwise it will wait for the critical section to complete up to the remaining operation
+     * time.
+     *
+     * Returns true if the call actually waited because of movePrimary critical section (regardless
+     * whether it timed out or not), false if there was no active movePrimary critical section.
+     */
+    bool waitForMovePrimaryCriticalSectionSignal(OperationContext* opCtx);
+
+    /**
+     * Setting this value indicates that when the version check failed, there was an active
+     * movePrimary for the namespace and that it would be prudent to wait for the critical section
+     * to complete before retrying so the router doesn't make wasteful requests.
+     */
+    void setMovePrimaryCriticalSectionSignal(std::shared_ptr<Notification<void>> critSecSignal);
+
 private:
     // Specifies whether the request is allowed to create database/collection implicitly
     bool _allowImplicitCollectionCreation{true};
@@ -147,6 +164,11 @@ private:
     // This value will only be non-null if version check during the operation execution failed due
     // to stale version and there was a migration for that namespace, which was in critical section.
     std::shared_ptr<Notification<void>> _migrationCriticalSectionSignal;
+
+    // This value will only be non-null if version check during the operation execution failed due
+    // to stale version and there was a movePrimary for that namespace, which was in critical
+    // section.
+    std::shared_ptr<Notification<void>> _movePrimaryCriticalSectionSignal;
 };
 
 }  // namespace mongo