SERVER-98650 Correctly handle StorageUnavailable exceptions during query restore (#30814)

jsteemann · MongoDB Bot · commit 62423b29f074 · 2025-04-09T20:53:41.000Z
GitOrigin-RevId: dace223539da0ae82a57a790fec9040851033687
diff --git a/jstests/noPassthrough/storage_unavailable_exception_during_query_restore.js b/jstests/noPassthrough/storage_unavailable_exception_during_query_restore.js
@@ -0,0 +1,79 @@
+/**
+ * Tests server behavior when a StorageUnavailableException is thrown when a PlanExecutor is being
+ * restored.  Specifically, tests the case where the exception is thrown when the getMore command
+ * is restoring the cursor in order to use it, as well as a case where the BatchedDeleteStage
+ * throws when restoring.
+ */
+
+const conn = MongoRunner.runMongod();
+assert.neq(null, conn, "mongod was unable to start up");
+
+const db = conn.getDB("test");
+
+function runTest() {
+    load("jstests/libs/fail_point_util.js");
+
+    let collName = jsTestName();
+    db[collName].drop();
+
+    for (let x = 0; x < 5; ++x) {
+        assert.commandWorked(db[collName].insert({_id: x, a: 1}));
+    }
+
+    assert.commandWorked(db[collName].createIndex({a: 1}));
+
+    //
+    // Test find command.
+    //
+    let res = db.runCommand({find: collName, filter: {a: 1}, batchSize: 1});
+    assert.eq(1, res.cursor.firstBatch.length, tojson(res));
+
+    // Configure the failpoint to trip once, when the getMore command restores the cursor.
+    let fp1 = configureFailPoint(db, "throwDuringIndexScanRestore", {} /* data */, {times: 1});
+
+    let getMoreRes =
+        assert.commandWorked(db.runCommand({getMore: res.cursor.id, collection: collName}));
+    assert.eq(4, getMoreRes.cursor.nextBatch.length, tojson(getMoreRes));
+
+    //
+    // Test aggregate command.
+    //
+    res =
+        db.runCommand({aggregate: collName, pipeline: [{$match: {a: 1}}], cursor: {batchSize: 1}});
+    assert.eq(1, res.cursor.firstBatch.length, tojson(res));
+
+    // Configure the failpoint to trip once, when the getMore command restores the cursor.
+    let fp2 = configureFailPoint(db, "throwDuringIndexScanRestore", {} /* data */, {times: 1});
+    getMoreRes =
+        assert.commandWorked(db.runCommand({getMore: res.cursor.id, collection: collName}));
+    assert.eq(4, getMoreRes.cursor.nextBatch.length, tojson(getMoreRes));
+
+    //
+    // Test batched delete command.
+    //
+
+    // Configure the fail point to trip once.
+    let fp3 = configureFailPoint(
+        db, "batchedDeleteStageThrowTemporarilyUnavailableException", {} /* data */, {times: 1});
+    assert.commandWorked(db[collName].remove({}));
+}
+
+// Fetch current value of 'internalQueryFrameworkControl' parameter.
+const previousQueryFrameworkControlValue =
+    assert.commandWorked(db.adminCommand({getParameter: 1, internalQueryFrameworkControl: 1}))
+        .internalQueryFrameworkControl;
+
+// Force query engine to classic for this test, as SBE does not yield for queries test here.
+assert.commandWorked(
+    db.adminCommand({setParameter: 1, internalQueryFrameworkControl: "forceClassicEngine"}));
+
+try {
+    // Run actual test.
+    runTest();
+} finally {
+    // Restore previous configuration of 'internalQueryFrameworkControl' parameter.
+    assert.commandWorked(db.adminCommand(
+        {setParameter: 1, internalQueryFrameworkControl: previousQueryFrameworkControlValue}));
+}
+
+MongoRunner.stopMongod(conn);
diff --git a/src/mongo/db/exec/batched_delete_stage.cpp b/src/mongo/db/exec/batched_delete_stage.cpp
@@ -55,6 +55,7 @@ namespace mongo {
 
 MONGO_FAIL_POINT_DEFINE(throwWriteConflictExceptionInBatchedDeleteStage);
 MONGO_FAIL_POINT_DEFINE(batchedDeleteStageSleepAfterNDocuments);
+MONGO_FAIL_POINT_DEFINE(batchedDeleteStageThrowTemporarilyUnavailableException);
 
 namespace {
 
@@ -228,15 +229,22 @@ 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;
 }
 
@@ -490,6 +498,13 @@ 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,7 +194,8 @@ 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);
+        return yieldPolicy->yieldOrInterrupt(
+            expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield);
     }
 
     return Status::OK();
diff --git a/src/mongo/db/exec/index_scan.cpp b/src/mongo/db/exec/index_scan.cpp
@@ -60,6 +60,8 @@ int sgn(int i) {
 
 namespace mongo {
 
+MONGO_FAIL_POINT_DEFINE(throwDuringIndexScanRestore);
+
 // static
 const char* IndexScan::kStageType = "IXSCAN";
 
@@ -284,8 +286,14 @@ void IndexScan::doSaveStateRequiresIndex() {
 }
 
 void IndexScan::doRestoreStateRequiresIndex() {
-    if (_indexCursor)
+    if (_indexCursor) {
+        if (MONGO_unlikely(throwDuringIndexScanRestore.shouldFail())) {
+            throwTemporarilyUnavailableException(str::stream()
+                                                 << "Hit failpoint '"
+                                                 << throwDuringIndexScanRestore.getName() << "'.");
+        }
         _indexCursor->restore();
+    }
 }
 
 void IndexScan::doDetachFromOperationContext() {
diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp
@@ -189,7 +189,8 @@ 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));
+        uassertStatusOK(yieldPolicy->yieldOrInterrupt(
+            expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield));
     }
 }
 
diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
@@ -521,7 +521,8 @@ class CanInterrupt {
                 opCtx->checkForInterrupt();
             }
         } else if (_yieldPolicy->shouldYieldOrInterrupt(opCtx)) {
-            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(opCtx));
+            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(
+                opCtx, nullptr, RestoreContext::RestoreType::kYield));
         }
     }
 
diff --git a/src/mongo/db/exec/trial_stage.cpp b/src/mongo/db/exec/trial_stage.cpp
@@ -83,7 +83,8 @@ Status TrialStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {
                 throwWriteConflictException(
                     "Write conflict during TrialStage plan selection and yielding is disabled.");
             }
-            auto yieldStatus = yieldPolicy->yieldOrInterrupt(expCtx()->opCtx);
+            auto yieldStatus = yieldPolicy->yieldOrInterrupt(
+                expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield);
             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,7 +49,9 @@ class MockYieldPolicy : public PlanYieldPolicy {
         MONGO_UNREACHABLE;
     }
 
-    void restoreState(OperationContext* opCtx, const Yieldable* yieldable) override final {
+    void restoreState(OperationContext* opCtx,
+                      const Yieldable* yieldable,
+                      RestoreContext::RestoreType restoreType) final {
         MONGO_UNREACHABLE;
     }
 };
@@ -67,7 +69,9 @@ class AlwaysTimeOutYieldPolicy final : public MockYieldPolicy {
         return true;
     }
 
-    Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {
+    Status yieldOrInterrupt(OperationContext*,
+                            std::function<void()> whileYieldingFn,
+                            RestoreContext::RestoreType restoreType) override {
         return {ErrorCodes::ExceededTimeLimit, "Using AlwaysTimeOutYieldPolicy"};
     }
 };
@@ -85,7 +89,9 @@ class AlwaysPlanKilledYieldPolicy final : public MockYieldPolicy {
         return true;
     }
 
-    Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {
+    Status yieldOrInterrupt(OperationContext*,
+                            std::function<void()> whileYieldingFn,
+                            RestoreContext::RestoreType restoreType) override {
         return {ErrorCodes::QueryPlanKilled, "Using AlwaysPlanKilledYieldPolicy"};
     }
 };
@@ -103,7 +109,9 @@ class NoopYieldPolicy final : public MockYieldPolicy {
         return false;
     }
 
-    Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {
+    Status yieldOrInterrupt(OperationContext*,
+                            std::function<void()> whileYieldingFn,
+                            RestoreContext::RestoreType restoreType) override {
         MONGO_UNREACHABLE;
     }
 };
diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp
@@ -265,7 +265,8 @@ void PlanExecutorImpl::restoreState(const RestoreContext& context) {
             throw;
 
         // Handles retries by calling restoreStateWithoutRetrying() in a loop.
-        uassertStatusOK(_yieldPolicy->yieldOrInterrupt(getOpCtx()));
+        uassertStatusOK(_yieldPolicy->yieldOrInterrupt(
+            getOpCtx(), nullptr /* whileYieldingFn */, context.type()));
     }
 }
 
@@ -383,7 +384,8 @@ PlanExecutor::ExecState PlanExecutorImpl::_getNextImpl(Snapshotted<Document>* ob
         };
 
         if (_yieldPolicy->shouldYieldOrInterrupt(_opCtx)) {
-            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx, whileYieldingFn));
+            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(
+                _opCtx, whileYieldingFn, RestoreContext::RestoreType::kYield));
         }
 
         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,16 +114,20 @@ 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();
-        }
-    });
+    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);
     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,7 +99,8 @@ void PlanYieldPolicy::resetTimer() {
 }
 
 Status PlanYieldPolicy::yieldOrInterrupt(OperationContext* opCtx,
-                                         std::function<void()> whileYieldingFn) {
+                                         std::function<void()> whileYieldingFn,
+                                         RestoreContext::RestoreType restoreType) {
     invariant(opCtx);
 
     if (_policy == YieldPolicy::INTERRUPT_ONLY) {
@@ -156,7 +157,8 @@ Status PlanYieldPolicy::yieldOrInterrupt(OperationContext* opCtx,
                 performYield(opCtx, *yieldable, whileYieldingFn);
             }
 
-            restoreState(opCtx, yieldable);
+            // This copies 'yieldable' back to '_yieldable' where needed.
+            restoreState(opCtx, yieldable, restoreType);
             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,6 +32,9 @@
 #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"
@@ -224,7 +227,8 @@ class PlanYieldPolicy {
      * been relinquished.
      */
     virtual Status yieldOrInterrupt(OperationContext* opCtx,
-                                    std::function<void()> whileYieldingFn = nullptr);
+                                    std::function<void()> whileYieldingFn,
+                                    RestoreContext::RestoreType restoreType);
 
     /**
      * All calls to shouldYieldOrInterrupt() will return true until the next call to
@@ -293,7 +297,9 @@ class PlanYieldPolicy {
      * specific query execution engines.
      */
     virtual void saveState(OperationContext* opCtx) = 0;
-    virtual void restoreState(OperationContext* opCtx, const Yieldable* yieldable) = 0;
+    virtual void restoreState(OperationContext* opCtx,
+                              const Yieldable* yieldable,
+                              RestoreContext::RestoreType restoreType) = 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,9 +52,13 @@ void PlanYieldPolicyImpl::saveState(OperationContext* opCtx) {
     _planYielding->saveState();
 }
 
-void PlanYieldPolicyImpl::restoreState(OperationContext* opCtx, const Yieldable* yieldable) {
-    _planYielding->restoreStateWithoutRetrying({RestoreContext::RestoreType::kYield, nullptr},
-                                               yieldable);
+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);
 }
 
 }  // 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,7 +45,9 @@ class PlanYieldPolicyImpl final : public PlanYieldPolicy {
 private:
     void saveState(OperationContext* opCtx) override;
 
-    void restoreState(OperationContext* opCtx, const Yieldable* yieldable) override;
+    void restoreState(OperationContext* opCtx,
+                      const Yieldable* yieldable,
+                      RestoreContext::RestoreType restoreType) 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,7 +38,9 @@ void PlanYieldPolicySBE::saveState(OperationContext* opCtx) {
     }
 }
 
-void PlanYieldPolicySBE::restoreState(OperationContext* opCtx, const Yieldable*) {
+void PlanYieldPolicySBE::restoreState(OperationContext* opCtx,
+                                      const Yieldable*,
+                                      RestoreContext::RestoreType restoreType) {
     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,7 +75,9 @@ class PlanYieldPolicySBE final : public PlanYieldPolicy {
 private:
     void saveState(OperationContext* opCtx) override;
 
-    void restoreState(OperationContext* opCtx, const Yieldable* yieldable) override;
+    void restoreState(OperationContext* opCtx,
+                      const Yieldable* yieldable,
+                      RestoreContext::RestoreType restoreType) override;
 
     // TODO SERVER-59620: Remove this.
     bool useExperimentalCommitTxnBehavior() const override {

Original file line number	Diff line number	Diff line change
`@@ -194,7 +194,8 @@ 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(expCtx()->opCtx);`
	`197`	`+ return yieldPolicy->yieldOrInterrupt(`
	`198`	`+ expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield);`
`198`	`199`	`}`
`199`	`200`
`200`	`201`	`return Status::OK();`
Original file line number	Diff line number	Diff line change
`@@ -189,7 +189,8 @@ 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(expCtx()->opCtx));`
	`192`	`+ uassertStatusOK(yieldPolicy->yieldOrInterrupt(`
	`193`	`+ expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield));`
`193`	`194`	`}`
`194`	`195`	`}`
`195`	`196`
Original file line number	Diff line number	Diff line change
`@@ -521,7 +521,8 @@ class CanInterrupt {`
`521`	`521`	`opCtx->checkForInterrupt();`
`522`	`522`	`}`
`523`	`523`	`} else if (_yieldPolicy->shouldYieldOrInterrupt(opCtx)) {`
`524`		`- uassertStatusOK(_yieldPolicy->yieldOrInterrupt(opCtx));`
	`524`	`+ uassertStatusOK(_yieldPolicy->yieldOrInterrupt(`
	`525`	`+ opCtx, nullptr, RestoreContext::RestoreType::kYield));`
`525`	`526`	`}`
`526`	`527`	`}`
`527`	`528`
Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,8 @@ 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(expCtx()->opCtx);`
	`86`	`+ auto yieldStatus = yieldPolicy->yieldOrInterrupt(`
	`87`	`+ expCtx()->opCtx, nullptr, RestoreContext::RestoreType::kYield);`
`87`	`88`	`if (!yieldStatus.isOK()) {`
`88`	`89`	`return yieldStatus;`
`89`	`90`	`}`
Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,9 @@ class MockYieldPolicy : public PlanYieldPolicy {`
`49`	`49`	`MONGO_UNREACHABLE;`
`50`	`50`	`}`
`51`	`51`
`52`		`- void restoreState(OperationContext* opCtx, const Yieldable* yieldable) override final {`
	`52`	`+ void restoreState(OperationContext* opCtx,`
	`53`	`+ const Yieldable* yieldable,`
	`54`	`+ RestoreContext::RestoreType restoreType) final {`
`53`	`55`	`MONGO_UNREACHABLE;`
`54`	`56`	`}`
`55`	`57`	`};`
`@@ -67,7 +69,9 @@ class AlwaysTimeOutYieldPolicy final : public MockYieldPolicy {`
`67`	`69`	`return true;`
`68`	`70`	`}`
`69`	`71`
`70`		`- Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {`
	`72`	`+ Status yieldOrInterrupt(OperationContext*,`
	`73`	`+ std::function<void()> whileYieldingFn,`
	`74`	`+ RestoreContext::RestoreType restoreType) override {`
`71`	`75`	`return {ErrorCodes::ExceededTimeLimit, "Using AlwaysTimeOutYieldPolicy"};`
`72`	`76`	`}`
`73`	`77`	`};`
`@@ -85,7 +89,9 @@ class AlwaysPlanKilledYieldPolicy final : public MockYieldPolicy {`
`85`	`89`	`return true;`
`86`	`90`	`}`
`87`	`91`
`88`		`- Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {`
	`92`	`+ Status yieldOrInterrupt(OperationContext*,`
	`93`	`+ std::function<void()> whileYieldingFn,`
	`94`	`+ RestoreContext::RestoreType restoreType) override {`
`89`	`95`	`return {ErrorCodes::QueryPlanKilled, "Using AlwaysPlanKilledYieldPolicy"};`
`90`	`96`	`}`
`91`	`97`	`};`
`@@ -103,7 +109,9 @@ class NoopYieldPolicy final : public MockYieldPolicy {`
`103`	`109`	`return false;`
`104`	`110`	`}`
`105`	`111`
`106`		`- Status yieldOrInterrupt(OperationContext*, std::function<void()> whileYieldingFn) override {`
	`112`	`+ Status yieldOrInterrupt(OperationContext*,`
	`113`	`+ std::function<void()> whileYieldingFn,`
	`114`	`+ RestoreContext::RestoreType restoreType) override {`
`107`	`115`	`MONGO_UNREACHABLE;`
`108`	`116`	`}`
`109`	`117`	`};`
Original file line number	Diff line number	Diff line change
`@@ -265,7 +265,8 @@ 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(getOpCtx()));`
	`268`	`+ uassertStatusOK(_yieldPolicy->yieldOrInterrupt(`
	`269`	`+ getOpCtx(), nullptr /* whileYieldingFn */, context.type()));`
`269`	`270`	`}`
`270`	`271`	`}`
`271`	`272`
`@@ -383,7 +384,8 @@ PlanExecutor::ExecState PlanExecutorImpl::_getNextImpl(Snapshotted<Document>* ob`
`383`	`384`	`};`
`384`	`385`
`385`	`386`	`if (_yieldPolicy->shouldYieldOrInterrupt(_opCtx)) {`
`386`		`- uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx, whileYieldingFn));`
	`387`	`+ uassertStatusOK(_yieldPolicy->yieldOrInterrupt(`
	`388`	`+ _opCtx, whileYieldingFn, RestoreContext::RestoreType::kYield));`
`387`	`389`	`}`
`388`	`390`
`389`	`391`	`WorkingSetID id = WorkingSet::INVALID_ID;`