1 files changed, 31 insertions, 46 deletions

diff --git a/src/mongo/db/query/plan_yield_policy.cpp b/src/mongo/db/query/plan_yield_policy.cpp
index 79816505026..58064f76d6e 100644
--- a/src/mongo/db/query/plan_yield_policy.cpp
+++ b/src/mongo/db/query/plan_yield_policy.cpp
@@ -32,7 +32,6 @@
 #include "mongo/db/query/plan_yield_policy.h"
 
 #include "mongo/db/catalog/collection.h"
-#include "mongo/db/concurrency/exception_util.h"
 #include "mongo/db/concurrency/write_conflict_exception.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/util/scopeguard.h"
@@ -40,43 +39,17 @@
 
 namespace mongo {
 
-PlanYieldPolicy::PlanYieldPolicy(OperationContext* opCtx,
-                                 YieldPolicy policy,
+PlanYieldPolicy::PlanYieldPolicy(YieldPolicy policy,
                                  ClockSource* cs,
                                  int yieldIterations,
                                  Milliseconds yieldPeriod,
                                  const Yieldable* yieldable,
                                  std::unique_ptr<const YieldPolicyCallbacks> callbacks)
-    : _policy(getPolicyOverrideForOperation(opCtx, policy)),
+    : _policy(policy),
       _yieldable(yieldable),
       _callbacks(std::move(callbacks)),
       _elapsedTracker(cs, yieldIterations, yieldPeriod) {}
 
-PlanYieldPolicy::YieldPolicy PlanYieldPolicy::getPolicyOverrideForOperation(
-    OperationContext* opCtx, PlanYieldPolicy::YieldPolicy desired) {
-    // We may have a null opCtx in testing.
-    if (MONGO_unlikely(!opCtx)) {
-        return desired;
-    }
-    // Multi-document transactions cannot yield locks or snapshots. We convert to a non-yielding
-    // interruptible plan.
-    if (opCtx->inMultiDocumentTransaction() &&
-        (desired == YieldPolicy::YIELD_AUTO || desired == YieldPolicy::YIELD_MANUAL ||
-         desired == YieldPolicy::WRITE_CONFLICT_RETRY_ONLY)) {
-        return YieldPolicy::INTERRUPT_ONLY;
-    }
-
-    // If the state of our locks held is not yieldable at all, we will assume this is an internal
-    // operation that should not be interrupted or yielded.
-    // TODO: SERVER-76238 Evaluate if we can make everything INTERRUPT_ONLY instead.
-    if (!opCtx->lockState()->canSaveLockState() &&
-        (desired == YieldPolicy::YIELD_AUTO || desired == YieldPolicy::YIELD_MANUAL)) {
-        return YieldPolicy::NO_YIELD;
-    }
-
-    return desired;
-}
-
 bool PlanYieldPolicy::shouldYieldOrInterrupt(OperationContext* opCtx) {
     if (_policy == YieldPolicy::INTERRUPT_ONLY) {
         return _elapsedTracker.intervalHasElapsed();
@@ -157,7 +130,7 @@ Status PlanYieldPolicy::yieldOrInterrupt(OperationContext* opCtx,
             if (_callbacks) {
                 _callbacks->handledWriteConflict(opCtx);
             }
-            logWriteConflictAndBackoff(attempt, "query yield", ""_sd);
+            WriteConflictException::logAndBackoff(attempt, "query yield", ""_sd);
             // Retry the yielding process.
         } catch (...) {
             // Errors other than write conflicts don't get retried, and should instead result in
@@ -174,34 +147,46 @@ void PlanYieldPolicy::performYield(OperationContext* opCtx,
                                    std::function<void()> whileYieldingFn) {
     // Things have to happen here in a specific order:
     //   * Release 'yieldable'.
-    //   * Abandon the current storage engine snapshot.
+    //   * Release lock mgr locks.
     //   * Check for interrupt if the yield policy requires.
-    //   * Release lock manager locks.
-    //   * Reacquire lock manager locks.
+    //   * Abondon the query's current storage engine snapshot.
+    //   * Reacquire lock mgr locks.
     //   * Restore 'yieldable'.
-    invariant(_policy == YieldPolicy::YIELD_AUTO || _policy == YieldPolicy::YIELD_MANUAL);
+    Locker* locker = opCtx->lockState();
+
+    if (locker->isGlobalLockedRecursively()) {
+        // No purpose in yielding if the locks are recursively held and cannot be released.
+        return;
+    }
 
-    // If we are here, the caller has guaranteed locks are not recursively held. This is a top level
-    // operation and we can safely clear the 'yieldable' state before unlocking and then
-    // re-establish it after re-locking.
+    // Since the locks are not recursively held, this is a top level operation and we can safely
+    // clear the 'yieldable' state before unlocking and then re-establish it after re-locking.
     if (yieldable) {
         yieldable->yield();
     }
 
-    // Release any storage engine resources. This requires holding a global lock to correctly
-    // synchronize with states such as shutdown and rollback.
-    opCtx->recoveryUnit()->abandonSnapshot();
+    Locker::LockSnapshot snapshot;
+    auto unlocked = locker->saveLockStateAndUnlock(&snapshot);
 
-    // Check for interrupt before releasing locks. This avoids the complexities of having to
-    // re-acquire locks to clean up when we are interrupted. This is the main interrupt check during
-    // query execution. Yield points and interrupt points are one and the same.
+    // After all steps to relinquish locks and save the execution plan have been taken, check
+    // for interrupt. This is the main interrupt check during query execution. Yield points and
+    // interrupt points are one and the same.
     if (getPolicy() == PlanYieldPolicy::YieldPolicy::YIELD_AUTO) {
         opCtx->checkForInterrupt();  // throws
     }
 
-    Locker* locker = opCtx->lockState();
-    Locker::LockSnapshot snapshot;
-    locker->saveLockStateAndUnlock(&snapshot);
+    if (!unlocked) {
+        // Nothing was unlocked. Recursively held locks are not the only reason locks cannot be
+        // released. Restore the 'yieldable' state before returning.
+        if (yieldable) {
+            yieldable->restore();
+        }
+        return;
+    }
+
+    // Top-level locks are freed, release any potential low-level (storage engine-specific
+    // locks). If we are yielding, we are at a safe place to do so.
+    opCtx->recoveryUnit()->abandonSnapshot();
 
     if (_callbacks) {
         _callbacks->duringYield(opCtx);