Revert "SERVER-98650 Correctly handle StorageUnavailable exceptions during query restore (#32463)

jsteemann · MongoDB Bot · commit 6951d91380da · 2025-04-09T20:54:46.000Z
GitOrigin-RevId: f099987179b0e9919aa2fcba25afe48f35e53ae9
diff --git a/jstests/noPassthrough/storage_unavailable_exception_during_query_restore.js b/jstests/noPassthrough/storage_unavailable_exception_during_query_restore.js
diff --git a/src/mongo/db/exec/batched_delete_stage.cpp b/src/mongo/db/exec/batched_delete_stage.cpp
@@ -55,7 +55,6 @@ namespace mongo {
 
 MONGO_FAIL_POINT_DEFINE(throwWriteConflictExceptionInBatchedDeleteStage);
 MONGO_FAIL_POINT_DEFINE(batchedDeleteStageSleepAfterNDocuments);
-MONGO_FAIL_POINT_DEFINE(batchedDeleteStageThrowTemporarilyUnavailableException);
 
 namespace {
 
@@ -229,22 +228,15 @@ PlanStage::StageState BatchedDeleteStage::doWork(WorkingSetID* out) {
         _commitStagedDeletes = _passStagingComplete || !_stagedDeletesBuffer.empty();
     }
 
-    // We need to check if 'planStateStage' is 'NEED_YIELD' earlier than we check 'isEOF()'.
-    // The reason is that 'isEOF()' returns true if we have completed staging and have an empty
-    // buffer of staged deletes. However, even if 'isEOF()' returns true, the 'planStateStage' can
-    // be 'NEED_YIELD' here if inside 'restoreState()' in '_deleteBatch()' a
-    // 'StorageEngineException' was thrown and caught. In this case, the 'planStateStage' is
-    // 'NEED_YIELD' and 'isEOF()' already returns true.
-    if (planStageState == PlanStage::NEED_YIELD) {
-        *out = idToReturn;
-        return PlanStage::NEED_YIELD;
-    }
-
     if (isEOF()) {
         invariant(planStageState != PlanStage::NEED_YIELD);
         return PlanStage::IS_EOF;
     }
 
+    if (planStageState == PlanStage::NEED_YIELD) {
+        *out = idToReturn;
+    }
+
     return planStageState;
 }
 
@@ -498,13 +490,6 @@ PlanStage::StageState BatchedDeleteStage::_tryRestoreState(WorkingSetID* out) {
         "BatchedDeleteStage::_tryRestoreState",
         collection()->ns().ns(),
         [&] {
-            if (MONGO_unlikely(
-                    batchedDeleteStageThrowTemporarilyUnavailableException.shouldFail())) {
-                throwTemporarilyUnavailableException(
-                    str::stream()
-                    << "Hit failpoint '"
-                    << batchedDeleteStageThrowTemporarilyUnavailableException.getName() << "'.");
-            }
             child()->restoreState(&collection());
             return PlanStage::NEED_TIME;
         },
diff --git a/src/mongo/db/exec/cached_plan.cpp b/src/mongo/db/exec/cached_plan.cpp
@@ -194,8 +194,7 @@ Status CachedPlanStage::tryYield(PlanYieldPolicy* yieldPolicy) {
     // In all cases, the actual yielding happens here.
     if (yieldPolicy->shouldYieldOrInterrupt(expCtx()->opCtx)) {
         // Here's where we yield.
-        return yieldPolicy->yieldOrInterrupt(
-            expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield);
+        return yieldPolicy->yieldOrInterrupt(expCtx()->opCtx);
     }
 
     return Status::OK();
diff --git a/src/mongo/db/exec/index_scan.cpp b/src/mongo/db/exec/index_scan.cpp
@@ -60,8 +60,6 @@ int sgn(int i) {
 
 namespace mongo {
 
-MONGO_FAIL_POINT_DEFINE(throwDuringIndexScanRestore);
-
 // static
 const char* IndexScan::kStageType = "IXSCAN";
 
@@ -286,14 +284,8 @@ void IndexScan::doSaveStateRequiresIndex() {
 }
 
 void IndexScan::doRestoreStateRequiresIndex() {
-    if (_indexCursor) {
-        if (MONGO_unlikely(throwDuringIndexScanRestore.shouldFail())) {
-            throwTemporarilyUnavailableException(str::stream()
-                                                 << "Hit failpoint '"
-                                                 << throwDuringIndexScanRestore.getName() << "'.");
-        }
+    if (_indexCursor)
         _indexCursor->restore();
-    }
 }
 
 void IndexScan::doDetachFromOperationContext() {
diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp
@@ -189,8 +189,7 @@ void MultiPlanStage::tryYield(PlanYieldPolicy* yieldPolicy) {
     //   3) we need to yield and retry due to a WriteConflictException.
     // In all cases, the actual yielding happens here.
     if (yieldPolicy->shouldYieldOrInterrupt(expCtx()->opCtx)) {
-        uassertStatusOK(yieldPolicy->yieldOrInterrupt(
-            expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield));
+        uassertStatusOK(yieldPolicy->yieldOrInterrupt(expCtx()->opCtx));
     }
 }
 
diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
@@ -521,8 +521,7 @@ class CanInterrupt {
                 opCtx->checkForInterrupt();
             }
         } else if (_yieldPolicy->shouldYieldOrInterrupt(opCtx)) {
-            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(
-                opCtx, nullptr, RestoreContext::RestoreType::kYield));
+            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(opCtx));
         }
     }
 
diff --git a/src/mongo/db/exec/trial_stage.cpp b/src/mongo/db/exec/trial_stage.cpp
@@ -83,8 +83,7 @@ Status TrialStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {
                 throwWriteConflictException(
                     "Write conflict during TrialStage plan selection and yielding is disabled.");
             }
-            auto yieldStatus = yieldPolicy->yieldOrInterrupt(
-                expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield);
+            auto yieldStatus = yieldPolicy->yieldOrInterrupt(expCtx()->opCtx);
             if (!yieldStatus.isOK()) {
                 return yieldStatus;
             }
diff --git a/src/mongo/db/query/mock_yield_policies.h b/src/mongo/db/query/mock_yield_policies.h
@@ -49,9 +49,7 @@ class MockYieldPolicy : public PlanYieldPolicy {
         MONGO_UNREACHABLE;
     }
 
-    void restoreState(OperationContext* opCtx,
-                      const Yieldable* yieldable,
-                      RestoreContext::RestoreType restoreType) final {
+    void restoreState(OperationContext* opCtx, const Yieldable* yieldable) override final {
         MONGO_UNREACHABLE;
     }
 };
@@ -69,9 +67,7 @@ class AlwaysTimeOutYieldPolicy final : public MockYieldPolicy {
         return true;
     }
 
-    Status yieldOrInterrupt(OperationContext*,
-                            std::function<void()> whileYieldingFn,
-                            RestoreContext::RestoreType restoreType) override {
+    Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {
         return {ErrorCodes::ExceededTimeLimit, "Using AlwaysTimeOutYieldPolicy"};
     }
 };
@@ -89,9 +85,7 @@ class AlwaysPlanKilledYieldPolicy final : public MockYieldPolicy {
         return true;
     }
 
-    Status yieldOrInterrupt(OperationContext*,
-                            std::function<void()> whileYieldingFn,
-                            RestoreContext::RestoreType restoreType) override {
+    Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {
         return {ErrorCodes::QueryPlanKilled, "Using AlwaysPlanKilledYieldPolicy"};
     }
 };
@@ -109,9 +103,7 @@ class NoopYieldPolicy final : public MockYieldPolicy {
         return false;
     }
 
-    Status yieldOrInterrupt(OperationContext*,
-                            std::function<void()> whileYieldingFn,
-                            RestoreContext::RestoreType restoreType) override {
+    Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {
         MONGO_UNREACHABLE;
     }
 };
diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp
@@ -265,8 +265,7 @@ void PlanExecutorImpl::restoreState(const RestoreContext& context) {
             throw;
 
         // Handles retries by calling restoreStateWithoutRetrying() in a loop.
-        uassertStatusOK(_yieldPolicy->yieldOrInterrupt(
-            getOpCtx(), nullptr /* whileYieldingFn */, context.type()));
+        uassertStatusOK(_yieldPolicy->yieldOrInterrupt(getOpCtx()));
     }
 }
 
@@ -384,8 +383,7 @@ PlanExecutor::ExecState PlanExecutorImpl::_getNextImpl(Snapshotted<Document>* ob
         };
 
         if (_yieldPolicy->shouldYieldOrInterrupt(_opCtx)) {
-            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(
-                _opCtx, whileYieldingFn, RestoreContext::RestoreType::kYield));
+            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx, whileYieldingFn));
         }
 
         WorkingSetID id = WorkingSet::INVALID_ID;
diff --git a/src/mongo/db/query/plan_insert_listener.cpp b/src/mongo/db/query/plan_insert_listener.cpp
@@ -114,20 +114,16 @@ void waitForInserts(OperationContext* opCtx,
     ON_BLOCK_EXIT([curOp] { curOp->resumeTimer(); });
 
     notifier->prepareForWait(opCtx);
-    auto yieldResult = yieldPolicy->yieldOrInterrupt(
-        opCtx,
-        [opCtx, &notifier] {
-            const auto deadline = awaitDataState(opCtx).waitForInsertsDeadline;
-            notifier->waitUntil(opCtx, deadline);
-            if (MONGO_unlikely(planExecutorHangWhileYieldedInWaitForInserts.shouldFail())) {
-                LOGV2(4452903,
-                      "PlanExecutor - planExecutorHangWhileYieldedInWaitForInserts fail point "
-                      "enabled. "
-                      "Blocking until fail point is disabled");
-                planExecutorHangWhileYieldedInWaitForInserts.pauseWhileSet();
-            }
-        },
-        RestoreContext::RestoreType::kYield);
+    auto yieldResult = yieldPolicy->yieldOrInterrupt(opCtx, [opCtx, &notifier] {
+        const auto deadline = awaitDataState(opCtx).waitForInsertsDeadline;
+        notifier->waitUntil(opCtx, deadline);
+        if (MONGO_unlikely(planExecutorHangWhileYieldedInWaitForInserts.shouldFail())) {
+            LOGV2(4452903,
+                  "PlanExecutor - planExecutorHangWhileYieldedInWaitForInserts fail point enabled. "
+                  "Blocking until fail point is disabled");
+            planExecutorHangWhileYieldedInWaitForInserts.pauseWhileSet();
+        }
+    });
     notifier->doneWaiting(opCtx);
 
     uassertStatusOK(yieldResult);
diff --git a/src/mongo/db/query/plan_yield_policy.cpp b/src/mongo/db/query/plan_yield_policy.cpp
@@ -99,8 +99,7 @@ void PlanYieldPolicy::resetTimer() {
 }
 
 Status PlanYieldPolicy::yieldOrInterrupt(OperationContext* opCtx,
-                                         std::function<void()> whileYieldingFn,
-                                         RestoreContext::RestoreType restoreType) {
+                                         std::function<void()> whileYieldingFn) {
     invariant(opCtx);
 
     if (_policy == YieldPolicy::INTERRUPT_ONLY) {
@@ -157,8 +156,7 @@ Status PlanYieldPolicy::yieldOrInterrupt(OperationContext* opCtx,
                 performYield(opCtx, *yieldable, whileYieldingFn);
             }
 
-            // This copies 'yieldable' back to '_yieldable' where needed.
-            restoreState(opCtx, yieldable, restoreType);
+            restoreState(opCtx, yieldable);
             return Status::OK();
         } catch (const StorageUnavailableException&) {
             if (_callbacks) {
diff --git a/src/mongo/db/query/plan_yield_policy.h b/src/mongo/db/query/plan_yield_policy.h
@@ -32,9 +32,6 @@
 #include <functional>
 
 #include "mongo/db/operation_context.h"
-#include "mongo/db/query/restore_context.h"
-#include "mongo/util/assert_util.h"
-#include "mongo/util/clock_source.h"
 #include "mongo/util/duration.h"
 #include "mongo/util/elapsed_tracker.h"
 #include "mongo/util/uuid.h"
@@ -227,8 +224,7 @@ class PlanYieldPolicy {
      * been relinquished.
      */
     virtual Status yieldOrInterrupt(OperationContext* opCtx,
-                                    std::function<void()> whileYieldingFn,
-                                    RestoreContext::RestoreType restoreType);
+                                    std::function<void()> whileYieldingFn = nullptr);
 
     /**
      * All calls to shouldYieldOrInterrupt() will return true until the next call to
@@ -297,9 +293,7 @@ class PlanYieldPolicy {
      * specific query execution engines.
      */
     virtual void saveState(OperationContext* opCtx) = 0;
-    virtual void restoreState(OperationContext* opCtx,
-                              const Yieldable* yieldable,
-                              RestoreContext::RestoreType restoreType) = 0;
+    virtual void restoreState(OperationContext* opCtx, const Yieldable* yieldable) = 0;
 
     /**
      * TODO SERVER-59620: Remove this.
diff --git a/src/mongo/db/query/plan_yield_policy_impl.cpp b/src/mongo/db/query/plan_yield_policy_impl.cpp
@@ -52,13 +52,9 @@ void PlanYieldPolicyImpl::saveState(OperationContext* opCtx) {
     _planYielding->saveState();
 }
 
-void PlanYieldPolicyImpl::restoreState(OperationContext* opCtx,
-                                       const Yieldable* yieldable,
-                                       RestoreContext::RestoreType restoreType) {
-    // We expect 'yieldable' to either be null or a CollectionPtr. TODO (SERVER-98914), consider
-    // refactoring this code to avoid the need for a runtime cast here.
-    auto collectionPtr = checked_cast<const CollectionPtr*>(yieldable);
-    _planYielding->restoreStateWithoutRetrying({restoreType, collectionPtr}, yieldable);
+void PlanYieldPolicyImpl::restoreState(OperationContext* opCtx, const Yieldable* yieldable) {
+    _planYielding->restoreStateWithoutRetrying({RestoreContext::RestoreType::kYield, nullptr},
+                                               yieldable);
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/plan_yield_policy_impl.h b/src/mongo/db/query/plan_yield_policy_impl.h
@@ -45,9 +45,7 @@ class PlanYieldPolicyImpl final : public PlanYieldPolicy {
 private:
     void saveState(OperationContext* opCtx) override;
 
-    void restoreState(OperationContext* opCtx,
-                      const Yieldable* yieldable,
-                      RestoreContext::RestoreType restoreType) override;
+    void restoreState(OperationContext* opCtx, const Yieldable* yieldable) override;
 
     // The plan executor which this yield policy is responsible for yielding. Must not outlive the
     // plan executor.
diff --git a/src/mongo/db/query/plan_yield_policy_sbe.cpp b/src/mongo/db/query/plan_yield_policy_sbe.cpp
@@ -38,9 +38,7 @@ void PlanYieldPolicySBE::saveState(OperationContext* opCtx) {
     }
 }
 
-void PlanYieldPolicySBE::restoreState(OperationContext* opCtx,
-                                      const Yieldable*,
-                                      RestoreContext::RestoreType restoreType) {
+void PlanYieldPolicySBE::restoreState(OperationContext* opCtx, const Yieldable*) {
     for (auto&& root : _yieldingPlans) {
         root->restoreState(!_useExperimentalCommitTxnBehavior /* relinquish cursor */);
     }
diff --git a/src/mongo/db/query/plan_yield_policy_sbe.h b/src/mongo/db/query/plan_yield_policy_sbe.h
@@ -75,9 +75,7 @@ class PlanYieldPolicySBE final : public PlanYieldPolicy {
 private:
     void saveState(OperationContext* opCtx) override;
 
-    void restoreState(OperationContext* opCtx,
-                      const Yieldable* yieldable,
-                      RestoreContext::RestoreType restoreType) override;
+    void restoreState(OperationContext* opCtx, const Yieldable* yieldable) override;
 
     // TODO SERVER-59620: Remove this.
     bool useExperimentalCommitTxnBehavior() const override {

Original file line number	Diff line number	Diff line change
`@@ -194,8 +194,7 @@ Status CachedPlanStage::tryYield(PlanYieldPolicy* yieldPolicy) {`
`194`	`194`	`// In all cases, the actual yielding happens here.`
`195`	`195`	`if (yieldPolicy->shouldYieldOrInterrupt(expCtx()->opCtx)) {`
`196`	`196`	`// Here's where we yield.`
`197`		`- return yieldPolicy->yieldOrInterrupt(`
`198`		`- expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield);`
	`197`	`+ return yieldPolicy->yieldOrInterrupt(expCtx()->opCtx);`
`199`	`198`	`}`
`200`	`199`
`201`	`200`	`return Status::OK();`
Original file line number	Diff line number	Diff line change
`@@ -189,8 +189,7 @@ void MultiPlanStage::tryYield(PlanYieldPolicy* yieldPolicy) {`
`189`	`189`	`// 3) we need to yield and retry due to a WriteConflictException.`
`190`	`190`	`// In all cases, the actual yielding happens here.`
`191`	`191`	`if (yieldPolicy->shouldYieldOrInterrupt(expCtx()->opCtx)) {`
`192`		`- uassertStatusOK(yieldPolicy->yieldOrInterrupt(`
`193`		`- expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield));`
	`192`	`+ uassertStatusOK(yieldPolicy->yieldOrInterrupt(expCtx()->opCtx));`
`194`	`193`	`}`
`195`	`194`	`}`
`196`	`195`
Original file line number	Diff line number	Diff line change
`@@ -521,8 +521,7 @@ class CanInterrupt {`
`521`	`521`	`opCtx->checkForInterrupt();`
`522`	`522`	`}`
`523`	`523`	`} else if (_yieldPolicy->shouldYieldOrInterrupt(opCtx)) {`
`524`		`- uassertStatusOK(_yieldPolicy->yieldOrInterrupt(`
`525`		`- opCtx, nullptr, RestoreContext::RestoreType::kYield));`
	`524`	`+ uassertStatusOK(_yieldPolicy->yieldOrInterrupt(opCtx));`
`526`	`525`	`}`
`527`	`526`	`}`
`528`	`527`
Original file line number	Diff line number	Diff line change
`@@ -83,8 +83,7 @@ Status TrialStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {`
`83`	`83`	`throwWriteConflictException(`
`84`	`84`	`"Write conflict during TrialStage plan selection and yielding is disabled.");`
`85`	`85`	`}`
`86`		`- auto yieldStatus = yieldPolicy->yieldOrInterrupt(`
`87`		`- expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield);`
	`86`	`+ auto yieldStatus = yieldPolicy->yieldOrInterrupt(expCtx()->opCtx);`
`88`	`87`	`if (!yieldStatus.isOK()) {`
`89`	`88`	`return yieldStatus;`
`90`	`89`	`}`
Original file line number	Diff line number	Diff line change
`@@ -49,9 +49,7 @@ class MockYieldPolicy : public PlanYieldPolicy {`
`49`	`49`	`MONGO_UNREACHABLE;`
`50`	`50`	`}`
`51`	`51`
`52`		`- void restoreState(OperationContext* opCtx,`
`53`		`- const Yieldable* yieldable,`
`54`		`- RestoreContext::RestoreType restoreType) final {`
	`52`	`+ void restoreState(OperationContext* opCtx, const Yieldable* yieldable) override final {`
`55`	`53`	`MONGO_UNREACHABLE;`
`56`	`54`	`}`
`57`	`55`	`};`
`@@ -69,9 +67,7 @@ class AlwaysTimeOutYieldPolicy final : public MockYieldPolicy {`
`69`	`67`	`return true;`
`70`	`68`	`}`
`71`	`69`
`72`		`- Status yieldOrInterrupt(OperationContext*,`
`73`		`- std::function<void()> whileYieldingFn,`
`74`		`- RestoreContext::RestoreType restoreType) override {`
	`70`	`+ Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {`
`75`	`71`	`return {ErrorCodes::ExceededTimeLimit, "Using AlwaysTimeOutYieldPolicy"};`
`76`	`72`	`}`
`77`	`73`	`};`
`@@ -89,9 +85,7 @@ class AlwaysPlanKilledYieldPolicy final : public MockYieldPolicy {`
`89`	`85`	`return true;`
`90`	`86`	`}`
`91`	`87`
`92`		`- Status yieldOrInterrupt(OperationContext*,`
`93`		`- std::function<void()> whileYieldingFn,`
`94`		`- RestoreContext::RestoreType restoreType) override {`
	`88`	`+ Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {`
`95`	`89`	`return {ErrorCodes::QueryPlanKilled, "Using AlwaysPlanKilledYieldPolicy"};`
`96`	`90`	`}`
`97`	`91`	`};`
`@@ -109,9 +103,7 @@ class NoopYieldPolicy final : public MockYieldPolicy {`
`109`	`103`	`return false;`
`110`	`104`	`}`
`111`	`105`
`112`		`- Status yieldOrInterrupt(OperationContext*,`
`113`		`- std::function<void()> whileYieldingFn,`
`114`		`- RestoreContext::RestoreType restoreType) override {`
	`106`	`+ Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {`
`115`	`107`	`MONGO_UNREACHABLE;`
`116`	`108`	`}`
`117`	`109`	`};`
Original file line number	Diff line number	Diff line change
`@@ -265,8 +265,7 @@ void PlanExecutorImpl::restoreState(const RestoreContext& context) {`
`265`	`265`	`throw;`
`266`	`266`
`267`	`267`	`// Handles retries by calling restoreStateWithoutRetrying() in a loop.`
`268`		`- uassertStatusOK(_yieldPolicy->yieldOrInterrupt(`
`269`		`- getOpCtx(), nullptr /* whileYieldingFn */, context.type()));`
	`268`	`+ uassertStatusOK(_yieldPolicy->yieldOrInterrupt(getOpCtx()));`
`270`	`269`	`}`
`271`	`270`	`}`
`272`	`271`
`@@ -384,8 +383,7 @@ PlanExecutor::ExecState PlanExecutorImpl::_getNextImpl(Snapshotted<Document>* ob`
`384`	`383`	`};`
`385`	`384`
`386`	`385`	`if (_yieldPolicy->shouldYieldOrInterrupt(_opCtx)) {`
`387`		`- uassertStatusOK(_yieldPolicy->yieldOrInterrupt(`
`388`		`- _opCtx, whileYieldingFn, RestoreContext::RestoreType::kYield));`
	`386`	`+ uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx, whileYieldingFn));`
`389`	`387`	`}`
`390`	`388`
`391`	`389`	`WorkingSetID id = WorkingSet::INVALID_ID;`