SERVER-48555 Allow yielding during SBE runtime planning

1 files changed, 7 insertions, 3 deletions

diff --git a/src/mongo/db/query/plan_yield_policy_sbe.cpp b/src/mongo/db/query/plan_yield_policy_sbe.cpp
index c7a76154f3c..113389c6a7b 100644
--- a/src/mongo/db/query/plan_yield_policy_sbe.cpp
+++ b/src/mongo/db/query/plan_yield_policy_sbe.cpp
@@ -37,13 +37,15 @@ namespace mongo {
 // so that SBE does not depend on CurOp. But we should still expose that statistic and keep it fresh
 // as the operation executes.
 Status PlanYieldPolicySBE::yield(OperationContext* opCtx, std::function<void()> whileYieldingFn) {
-    if (!_rootStage) {
+    if (!_yieldingPlans.empty()) {
         // This yield policy isn't bound to an execution tree yet.
         return Status::OK();
     }
 
     try {
-        _rootStage->saveState();
+        for (auto&& root : _yieldingPlans) {
+            root->saveState();
+        }
 
         opCtx->recoveryUnit()->abandonSnapshot();
 
@@ -51,7 +53,9 @@ Status PlanYieldPolicySBE::yield(OperationContext* opCtx, std::function<void()>
             whileYieldingFn();
         }
 
-        _rootStage->restoreState();
+        for (auto&& root : _yieldingPlans) {
+            root->restoreState();
+        }
     } catch (...) {
         return exceptionToStatus();
     }