SERVER-48555 Allow yielding during SBE runtime planning

6 files changed, 40 insertions, 7 deletions

diff --git a/src/mongo/db/query/plan_executor_sbe.cpp b/src/mongo/db/query/plan_executor_sbe.cpp
index c02828cc39c..d4237cb65fc 100644
--- a/src/mongo/db/query/plan_executor_sbe.cpp
+++ b/src/mongo/db/query/plan_executor_sbe.cpp
@@ -91,7 +91,11 @@ PlanExecutorSBE::PlanExecutorSBE(
 
     // Callers are allowed to disable yielding for this plan by passing a null yield policy.
     if (_yieldPolicy) {
-        _yieldPolicy->setRootStage(_root.get());
+        // Clear any formerly registered plans and register '_root' to yield. This is needed because
+        // multiple candidate plans may have been registered during runtime planning, before the
+        // PlanExecutor was created. All but one candidate plan ('_root') have since been discarded.
+        _yieldPolicy->clearRegisteredPlans();
+        _yieldPolicy->registerPlan(_root.get());
     }
 
     // We may still need to initialize _nss from either collection or _cq.
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();
     }
diff --git a/src/mongo/db/query/plan_yield_policy_sbe.h b/src/mongo/db/query/plan_yield_policy_sbe.h
index 1b9041bcc0a..0c1fcb161f7 100644
--- a/src/mongo/db/query/plan_yield_policy_sbe.h
+++ b/src/mongo/db/query/plan_yield_policy_sbe.h
@@ -46,14 +46,26 @@ public:
                 policy != YieldPolicy::WRITE_CONFLICT_RETRY_ONLY);
     }
 
-    void setRootStage(sbe::PlanStage* rootStage) {
-        _rootStage = rootStage;
+    /**
+     * Registers the tree rooted at 'plan' to yield, in addition to all other plans that have been
+     * previously registered with this yield policy.
+     */
+    void registerPlan(sbe::PlanStage* plan) {
+        _yieldingPlans.push_back(plan);
+    }
+
+    /**
+     * Clears the list of plans currently registered to yield.
+     */
+    void clearRegisteredPlans() {
+        _yieldingPlans.clear();
     }
 
 private:
     Status yield(OperationContext* opCtx, std::function<void()> whileYieldingFn = nullptr) override;
 
-    sbe::PlanStage* _rootStage = nullptr;
+    // The list of plans registered to yield when the configured policy triggers a yield.
+    std::vector<sbe::PlanStage*> _yieldingPlans;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/sbe_cached_solution_planner.cpp b/src/mongo/db/query/sbe_cached_solution_planner.cpp
index abc3027ce35..17980a75de2 100644
--- a/src/mongo/db/query/sbe_cached_solution_planner.cpp
+++ b/src/mongo/db/query/sbe_cached_solution_planner.cpp
@@ -101,6 +101,10 @@ plan_ranker::CandidatePlan CachedSolutionPlanner::finalizeExecutionPlan(
 }
 
 plan_ranker::CandidatePlan CachedSolutionPlanner::replan(bool shouldCache) const {
+    // The plan drawn from the cache is being discarded, and should no longer be registered with the
+    // yield policy.
+    _yieldPolicy->clearRegisteredPlans();
+
     if (shouldCache) {
         // Deactivate the current cache entry.
         auto cache = CollectionQueryInfo::get(_collection).getPlanCache();
diff --git a/src/mongo/db/query/sbe_sub_planner.cpp b/src/mongo/db/query/sbe_sub_planner.cpp
index 4ecdc93344a..05dadabbf5c 100644
--- a/src/mongo/db/query/sbe_sub_planner.cpp
+++ b/src/mongo/db/query/sbe_sub_planner.cpp
@@ -53,6 +53,11 @@ plan_ranker::CandidatePlan SubPlanner::plan(
     auto multiplanCallback = [&](CanonicalQuery* cq,
                                  std::vector<std::unique_ptr<QuerySolution>> solutions)
         -> StatusWith<std::unique_ptr<QuerySolution>> {
+        // Before planning a new branch of the $or, we need to clear any plans previously registered
+        // to yield. Each branch of the $or gets multi-planned independently, so we don't want to
+        // try to yield plans leftover from a different branch.
+        _yieldPolicy->clearRegisteredPlans();
+
         std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots;
         for (auto&& solution : solutions) {
             roots.push_back(stage_builder::buildSlotBasedExecutableTree(
diff --git a/src/mongo/db/query/stage_builder_util.cpp b/src/mongo/db/query/stage_builder_util.cpp
index 5f7c6c462e3..733675aa0b7 100644
--- a/src/mongo/db/query/stage_builder_util.cpp
+++ b/src/mongo/db/query/stage_builder_util.cpp
@@ -73,6 +73,10 @@ buildSlotBasedExecutableTree(OperationContext* opCtx,
         opCtx, collection, cq, solution, sbeYieldPolicy, needsTrialRunProgressTracker);
     auto root = builder->build(solution.root.get());
     auto data = builder->getPlanStageData();
+
+    // Register this plan to yield according to the configured policy.
+    sbeYieldPolicy->registerPlan(root.get());
+
     return {std::move(root), std::move(data)};
 }
 }  // namespace mongo::stage_builder