From a5d69928f116c7cfe1f9441f0bc0c13a82ae2442 Mon Sep 17 00:00:00 2001
From: Gregory Noma <gregory.noma@gmail.com>
Date: Fri, 11 Nov 2022 19:33:14 +0000
Subject: SERVER-71191 Unlock and relock RSTL during index build setup

(cherry picked from commit f1203e555321e326b0621479bcc6b607c5854391)
---
 .../replication_state_transition_lock_guard.cpp    | 10 ++++++
 .../replication_state_transition_lock_guard.h      |  3 +-
 src/mongo/db/index_builds_coordinator.cpp          | 37 ++++++++++++++++++++--
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/mongo/db/concurrency/replication_state_transition_lock_guard.cpp b/src/mongo/db/concurrency/replication_state_transition_lock_guard.cpp
index 649dbce681d..57187a87224 100644
--- a/src/mongo/db/concurrency/replication_state_transition_lock_guard.cpp
+++ b/src/mongo/db/concurrency/replication_state_transition_lock_guard.cpp
@@ -50,6 +50,12 @@ ReplicationStateTransitionLockGuard::ReplicationStateTransitionLockGuard(Operati
     _enqueueLock();
 }
 
+ReplicationStateTransitionLockGuard::ReplicationStateTransitionLockGuard(
+    ReplicationStateTransitionLockGuard&& other)
+    : _opCtx(other._opCtx), _mode(other._mode), _result(other._result) {
+    other._result = LockResult::LOCK_INVALID;
+}
+
 ReplicationStateTransitionLockGuard::~ReplicationStateTransitionLockGuard() {
     _unlock();
 }
@@ -81,6 +87,10 @@ void ReplicationStateTransitionLockGuard::_enqueueLock() {
 }
 
 void ReplicationStateTransitionLockGuard::_unlock() {
+    if (_result == LockResult::LOCK_INVALID) {
+        return;
+    }
+
     // If ReplicationStateTransitionLockGuard is called in a WriteUnitOfWork, we won't accept
     // any exceptions to be thrown between _enqueueLock and waitForLockUntil because that would
     // delay cleaning up any failed RSTL lock attempt state from lock manager.
diff --git a/src/mongo/db/concurrency/replication_state_transition_lock_guard.h b/src/mongo/db/concurrency/replication_state_transition_lock_guard.h
index 64107a74251..cf94dc77c28 100644
--- a/src/mongo/db/concurrency/replication_state_transition_lock_guard.h
+++ b/src/mongo/db/concurrency/replication_state_transition_lock_guard.h
@@ -63,8 +63,7 @@ public:
      */
     ReplicationStateTransitionLockGuard(OperationContext* opCtx, LockMode mode, EnqueueOnly);
 
-    ReplicationStateTransitionLockGuard(ReplicationStateTransitionLockGuard&&) = delete;
-    ReplicationStateTransitionLockGuard& operator=(ReplicationStateTransitionLockGuard&&) = delete;
+    ReplicationStateTransitionLockGuard(ReplicationStateTransitionLockGuard&&);
 
     ~ReplicationStateTransitionLockGuard();
 
diff --git a/src/mongo/db/index_builds_coordinator.cpp b/src/mongo/db/index_builds_coordinator.cpp
index 8fdc338f836..dd6252eaa28 100644
--- a/src/mongo/db/index_builds_coordinator.cpp
+++ b/src/mongo/db/index_builds_coordinator.cpp
@@ -1876,10 +1876,41 @@ IndexBuildsCoordinator::PostSetupAction IndexBuildsCoordinator::_setUpIndexBuild
     std::shared_ptr<ReplIndexBuildState> replState,
     Timestamp startTimestamp,
     const IndexBuildOptions& indexBuildOptions) {
-    const NamespaceStringOrUUID nssOrUuid{replState->dbName, replState->collectionUUID};
+    auto [dbLock, collLock, rstl] = [&] {
+        while (true) {
+            Lock::DBLock dbLock{opCtx, replState->dbName, MODE_IX};
 
-    AutoGetCollection coll(opCtx, nssOrUuid, MODE_X);
-    CollectionWriter collection(coll);
+            // Unlock the RSTL to avoid deadlocks with prepared transactions and replication state
+            // transitions. See SERVER-71191.
+            unlockRSTL(opCtx);
+
+            Lock::CollectionLock collLock{
+                opCtx, {replState->dbName, replState->collectionUUID}, MODE_X};
+            repl::ReplicationStateTransitionLockGuard rstl{
+                opCtx, MODE_IX, repl::ReplicationStateTransitionLockGuard::EnqueueOnly{}};
+
+            try {
+                // Since this thread is not killable by state transitions, this deadline is
+                // effectively the longest period of time we can block a state transition. State
+                // transitions are infrequent, but need to happen quickly. It should be okay to set
+                // this to a low value because the RSTL is rarely contended and, if this does time
+                // out, we will retry and reacquire the RSTL again without a deadline.
+                rstl.waitForLockUntil(Date_t::now() + Milliseconds{10});
+            } catch (const ExceptionFor<ErrorCodes::LockTimeout>&) {
+                // We weren't able to re-acquire the RSTL within the timeout, which means there is
+                // an active state transition. Release our locks and try again from the beginning.
+                LOGV2(7119100,
+                      "Unable to acquire RSTL for index build setup within deadline, releasing "
+                      "locks and trying again",
+                      "buildUUID"_attr = replState->buildUUID);
+                continue;
+            }
+
+            return std::make_tuple(std::move(dbLock), std::move(collLock), std::move(rstl));
+        }
+    }();
+
+    CollectionWriter collection(opCtx, replState->collectionUUID);
     CollectionShardingState::get(opCtx, collection->ns())->checkShardVersionOrThrow(opCtx);
 
     auto replCoord = repl::ReplicationCoordinator::get(opCtx);
-- 
cgit v1.2.1