SERVER-56114 [SBE] Fix bounding number of reads before triggering a replan

7 files changed, 32 insertions, 15 deletions

diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
index 294a437fe22..80f0861f26c 100644
--- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp
+++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
@@ -324,14 +324,16 @@ PlanState IndexScanStage::getNext() {
             _nextRecord->keyString, *_ordering, &_valuesBuffer, &_accessors, _indexKeysToInclude);
     }
 
+    ++_specificStats.numReads;
     if (_tracker && _tracker->trackProgress<TrialRunTracker::kNumReads>(1)) {
         // If we're collecting execution stats during multi-planning and reached the end of the
-        // trial period (trackProgress() will return 'true' in this case), then we can reset the
-        // tracker. Note that a trial period is executed only once per a PlanStge tree, and once
-        // completed never run again on the same tree.
+        // trial period because we've performed enough physical reads, bail out from the trial run
+        // by raising a special exception to signal a runtime planner that this candidate plan has
+        // completed its trial run early. Note that a trial period is executed only once per a
+        // PlanStage tree, and once completed never run again on the same tree.
         _tracker = nullptr;
+        uasserted(ErrorCodes::QueryTrialRunCompleted, "Trial run early exit in ixscan");
     }
-    ++_specificStats.numReads;
     return trackPlanState(PlanState::ADVANCED);
 }
 
diff --git a/src/mongo/db/exec/sbe/stages/scan.cpp b/src/mongo/db/exec/sbe/stages/scan.cpp
index 2d7e6fabe65..bee851362c4 100644
--- a/src/mongo/db/exec/sbe/stages/scan.cpp
+++ b/src/mongo/db/exec/sbe/stages/scan.cpp
@@ -323,14 +323,16 @@ PlanState ScanStage::getNext() {
         }
     }
 
+    ++_specificStats.numReads;
     if (_tracker && _tracker->trackProgress<TrialRunTracker::kNumReads>(1)) {
         // If we're collecting execution stats during multi-planning and reached the end of the
-        // trial period (trackProgress() will return 'true' in this case), then we can reset the
-        // tracker. Note that a trial period is executed only once per a PlanStge tree, and once
-        // completed never run again on the same tree.
+        // trial period because we've performed enough physical reads, bail out from the trial run
+        // by raising a special exception to signal a runtime planner that this candidate plan has
+        // completed its trial run early. Note that a trial period is executed only once per a
+        // PlanStage tree, and once completed never run again on the same tree.
         _tracker = nullptr;
+        uasserted(ErrorCodes::QueryTrialRunCompleted, "Trial run early exit in scan");
     }
-    ++_specificStats.numReads;
     return trackPlanState(PlanState::ADVANCED);
 }
 
diff --git a/src/mongo/db/exec/sbe/stages/sort.cpp b/src/mongo/db/exec/sbe/stages/sort.cpp
index ec931cd9c67..647f0dd346a 100644
--- a/src/mongo/db/exec/sbe/stages/sort.cpp
+++ b/src/mongo/db/exec/sbe/stages/sort.cpp
@@ -194,7 +194,7 @@ void SortStage::open(bool reOpen) {
             // higher level stages.
             _tracker = nullptr;
             _children[0]->close();
-            uasserted(ErrorCodes::QueryTrialRunCompleted, "Trial run early exit");
+            uasserted(ErrorCodes::QueryTrialRunCompleted, "Trial run early exit in sort");
         }
     }
 
diff --git a/src/mongo/db/query/sbe_cached_solution_planner.cpp b/src/mongo/db/query/sbe_cached_solution_planner.cpp
index ea3fe9f562d..01bc55c9653 100644
--- a/src/mongo/db/query/sbe_cached_solution_planner.cpp
+++ b/src/mongo/db/query/sbe_cached_solution_planner.cpp
@@ -48,7 +48,12 @@ CandidatePlans CachedSolutionPlanner::plan(
     invariant(solutions.size() == roots.size());
 
     auto candidate = [&]() {
-        auto candidates = collectExecutionStats(std::move(solutions), std::move(roots));
+        // In cached solution planning we collect execution stats with an upper bound on reads
+        // allowed per trial run computed based on previous decision reads.
+        auto candidates = collectExecutionStats(
+            std::move(solutions),
+            std::move(roots),
+            static_cast<size_t>(internalQueryCacheEvictionRatio * _decisionReads));
         invariant(candidates.size() == 1);
         return std::move(candidates[0]);
     }();
diff --git a/src/mongo/db/query/sbe_multi_planner.cpp b/src/mongo/db/query/sbe_multi_planner.cpp
index 42857f3d39d..28b78a6c0d2 100644
--- a/src/mongo/db/query/sbe_multi_planner.cpp
+++ b/src/mongo/db/query/sbe_multi_planner.cpp
@@ -46,7 +46,10 @@ namespace mongo::sbe {
 CandidatePlans MultiPlanner::plan(
     std::vector<std::unique_ptr<QuerySolution>> solutions,
     std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) {
-    auto candidates = collectExecutionStats(std::move(solutions), std::move(roots));
+    auto candidates =
+        collectExecutionStats(std::move(solutions),
+                              std::move(roots),
+                              trial_period::getTrialPeriodMaxWorks(_opCtx, _collection));
     auto decision = uassertStatusOK(mongo::plan_ranker::pickBestPlan<PlanStageStats>(candidates));
     return finalizeExecutionPlans(std::move(decision), std::move(candidates));
 }
diff --git a/src/mongo/db/query/sbe_runtime_planner.cpp b/src/mongo/db/query/sbe_runtime_planner.cpp
index 0acf0a10134..0dc280b0efc 100644
--- a/src/mongo/db/query/sbe_runtime_planner.cpp
+++ b/src/mongo/db/query/sbe_runtime_planner.cpp
@@ -115,7 +115,8 @@ BaseRuntimePlanner::prepareExecutionPlan(PlanStage* root,
 
 std::vector<plan_ranker::CandidatePlan> BaseRuntimePlanner::collectExecutionStats(
     std::vector<std::unique_ptr<QuerySolution>> solutions,
-    std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) {
+    std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots,
+    size_t maxTrialPeriodNumReads) {
     invariant(solutions.size() == roots.size());
 
     std::vector<plan_ranker::CandidatePlan> candidates;
@@ -131,12 +132,12 @@ std::vector<plan_ranker::CandidatePlan> BaseRuntimePlanner::collectExecutionStat
     });
 
     const auto maxNumResults{trial_period::getTrialPeriodNumToReturn(_cq)};
-    const auto maxNumReads{trial_period::getTrialPeriodMaxWorks(_opCtx, _collection)};
+
     for (size_t ix = 0; ix < roots.size(); ++ix) {
         auto&& [root, data] = roots[ix];
 
         // Attach a unique TrialRunTracker to each SBE plan.
-        auto tracker = std::make_unique<TrialRunTracker>(maxNumResults, maxNumReads);
+        auto tracker = std::make_unique<TrialRunTracker>(maxNumResults, maxTrialPeriodNumReads);
         root->attachToTrialRunTracker(tracker.get());
         trialRunTrackers.emplace_back(root.get(), std::move(tracker));
 
diff --git a/src/mongo/db/query/sbe_runtime_planner.h b/src/mongo/db/query/sbe_runtime_planner.h
index c622a863756..90ad01acc68 100644
--- a/src/mongo/db/query/sbe_runtime_planner.h
+++ b/src/mongo/db/query/sbe_runtime_planner.h
@@ -104,10 +104,14 @@ protected:
      * of the candidate plans by calling 'CandidatePlan->root->getStats()'.
      *
      * After the trial period ends, all plans remain open.
+     *
+     * The number of reads allowed for a trial execution period is bounded by
+     * 'maxTrialPeriodNumReads'.
      */
     std::vector<plan_ranker::CandidatePlan> collectExecutionStats(
         std::vector<std::unique_ptr<QuerySolution>> solutions,
-        std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots);
+        std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots,
+        size_t maxTrialPeriodNumReads);
 
     OperationContext* const _opCtx;
     const CollectionPtr& _collection;