SERVER-63641 Use priority queue to sort plans during multiplanning

14 files changed, 395 insertions, 331 deletions

diff --git a/src/mongo/db/exec/plan_cache_util.h b/src/mongo/db/exec/plan_cache_util.h
index 05b8e3d6a36..64d71879368 100644
--- a/src/mongo/db/exec/plan_cache_util.h
+++ b/src/mongo/db/exec/plan_cache_util.h
@@ -131,10 +131,11 @@ void updatePlanCache(
             auto&& [winnerExplainer, runnerUpExplainer] = [&]() {
                 if constexpr (std::is_same_v<PlanStageType, std::unique_ptr<sbe::PlanStage>>) {
                     return std::make_pair(
-                        plan_explainer_factory::make(
-                            winningPlan.root.get(), &winningPlan.data, winningPlan.solution.get()),
+                        plan_explainer_factory::make(winningPlan.root.get(),
+                                                     &winningPlan.data.stageData,
+                                                     winningPlan.solution.get()),
                         plan_explainer_factory::make(candidates[runnerUpIdx].root.get(),
-                                                     &candidates[runnerUpIdx].data,
+                                                     &candidates[runnerUpIdx].data.stageData,
                                                      candidates[runnerUpIdx].solution.get()));
                 } else {
                     static_assert(std::is_same_v<PlanStageType, PlanStage*>);
@@ -157,8 +158,9 @@ void updatePlanCache(
             canCache = false;
             auto winnerExplainer = [&]() {
                 if constexpr (std::is_same_v<PlanStageType, std::unique_ptr<sbe::PlanStage>>) {
-                    return plan_explainer_factory::make(
-                        winningPlan.root.get(), &winningPlan.data, winningPlan.solution.get());
+                    return plan_explainer_factory::make(winningPlan.root.get(),
+                                                        &winningPlan.data.stageData,
+                                                        winningPlan.solution.get());
                 } else {
                     static_assert(std::is_same_v<PlanStageType, PlanStage*>);
                     return plan_explainer_factory::make(winningPlan.root);
@@ -205,9 +207,9 @@ void updatePlanCache(
                         winningPlan.clonedPlan);
 
                 // Clone the winning SBE plan and its auxiliary data.
-                auto cachedPlan =
-                    std::make_unique<sbe::CachedSbePlan>(std::move(winningPlan.clonedPlan->first),
-                                                         std::move(winningPlan.clonedPlan->second));
+                auto cachedPlan = std::make_unique<sbe::CachedSbePlan>(
+                    std::move(winningPlan.clonedPlan->first),
+                    std::move(winningPlan.clonedPlan->second.stageData));
                 cachedPlan->indexFilterApplied = winningPlan.solution->indexFilterApplied;
 
                 auto buildDebugInfoFn =
diff --git a/src/mongo/db/exec/trial_run_tracker.h b/src/mongo/db/exec/trial_run_tracker.h
index ee07373a3c7..c65a722878d 100644
--- a/src/mongo/db/exec/trial_run_tracker.h
+++ b/src/mongo/db/exec/trial_run_tracker.h
@@ -122,8 +122,14 @@ public:
         return _metrics[metric];
     }
 
+    template <TrialRunMetric metric>
+    void updateMaxMetric(size_t newMaxMetric) {
+        static_assert(metric >= 0 && metric < sizeof(_metrics) / sizeof(size_t));
+        _maxMetrics[metric] = newMaxMetric;
+    }
+
 private:
-    const size_t _maxMetrics[TrialRunMetric::kLastElem];
+    size_t _maxMetrics[TrialRunMetric::kLastElem];
     size_t _metrics[TrialRunMetric::kLastElem]{0};
     bool _done{false};
     std::function<bool(TrialRunMetric)> _onMetricReached{};
diff --git a/src/mongo/db/query/plan_executor_factory.cpp b/src/mongo/db/query/plan_executor_factory.cpp
index 5b85179dea2..6d1ec882672 100644
--- a/src/mongo/db/query/plan_executor_factory.cpp
+++ b/src/mongo/db/query/plan_executor_factory.cpp
@@ -142,22 +142,23 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
                 "slots"_attr = data.debugString(),
                 "stages"_attr = sbe::DebugPrinter{}.print(*rootStage));
 
-    return {{new PlanExecutorSBE(opCtx,
-                                 std::move(cq),
-                                 std::move(optimizerData),
-                                 {makeVector<sbe::plan_ranker::CandidatePlan>(
-                                      sbe::plan_ranker::CandidatePlan{std::move(solution),
-                                                                      std::move(rootStage),
-                                                                      std::move(data),
-                                                                      false,
-                                                                      Status::OK(),
-                                                                      planIsFromCache}),
-                                  0},
-                                 plannerOptions & QueryPlannerParams::RETURN_OWNED_DATA,
-                                 std::move(nss),
-                                 false,
-                                 std::move(yieldPolicy),
-                                 generatedByBonsai),
+    return {{new PlanExecutorSBE(
+                 opCtx,
+                 std::move(cq),
+                 std::move(optimizerData),
+                 {makeVector<sbe::plan_ranker::CandidatePlan>(sbe::plan_ranker::CandidatePlan{
+                      std::move(solution),
+                      std::move(rootStage),
+                      sbe::plan_ranker::CandidatePlanData{std::move(data)},
+                      false /*exitedEarly*/,
+                      Status::OK(),
+                      planIsFromCache}),
+                  0},
+                 plannerOptions & QueryPlannerParams::RETURN_OWNED_DATA,
+                 std::move(nss),
+                 false /*isOpen*/,
+                 std::move(yieldPolicy),
+                 generatedByBonsai),
              PlanExecutor::Deleter{opCtx}}};
 }
 
@@ -172,7 +173,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     LOGV2_DEBUG(4822861,
                 5,
                 "SBE plan",
-                "slots"_attr = candidates.winner().data.debugString(),
+                "slots"_attr = candidates.winner().data.stageData.debugString(),
                 "stages"_attr = sbe::DebugPrinter{}.print(*candidates.winner().root));
 
     return {{new PlanExecutorSBE(opCtx,
diff --git a/src/mongo/db/query/plan_executor_sbe.cpp b/src/mongo/db/query/plan_executor_sbe.cpp
index f47f2a99d71..5f9c347f463 100644
--- a/src/mongo/db/query/plan_executor_sbe.cpp
+++ b/src/mongo/db/query/plan_executor_sbe.cpp
@@ -62,7 +62,7 @@ PlanExecutorSBE::PlanExecutorSBE(OperationContext* opCtx,
       _nss(std::move(nss)),
       _mustReturnOwnedBson(returnOwnedBson),
       _root{std::move(candidates.winner().root)},
-      _rootData{std::move(candidates.winner().data)},
+      _rootData{std::move(candidates.winner().data.stageData)},
       _solution{std::move(candidates.winner().solution)},
       _stash{std::move(candidates.winner().results)},
       _cq{std::move(cq)},
diff --git a/src/mongo/db/query/plan_explainer_sbe.cpp b/src/mongo/db/query/plan_explainer_sbe.cpp
index d1fbe796d2b..b4b922553d2 100644
--- a/src/mongo/db/query/plan_explainer_sbe.cpp
+++ b/src/mongo/db/query/plan_explainer_sbe.cpp
@@ -471,7 +471,8 @@ std::vector<PlanExplainer::PlanStatsDetails> PlanExplainerSBE::getRejectedPlansS
 
         auto stats = candidate.root->getStats(true /* includeDebugInfo  */);
         invariant(stats);
-        auto execPlanDebugInfo = buildExecPlanDebugInfo(candidate.root.get(), &candidate.data);
+        auto execPlanDebugInfo =
+            buildExecPlanDebugInfo(candidate.root.get(), &candidate.data.stageData);
         res.push_back(buildPlanStatsDetails(
             candidate.solution.get(), *stats, execPlanDebugInfo, boost::none, verbosity));
     }
diff --git a/src/mongo/db/query/plan_ranker_util.h b/src/mongo/db/query/plan_ranker_util.h
index 4874969db0d..76ac35fa708 100644
--- a/src/mongo/db/query/plan_ranker_util.h
+++ b/src/mongo/db/query/plan_ranker_util.h
@@ -84,8 +84,9 @@ StatusWith<std::unique_ptr<PlanRankingDecision>> pickBestPlan(
                                                     candidates[i].solution->_enumeratorExplainInfo);
             } else {
                 static_assert(std::is_same_v<PlanStageStatsType, mongo::sbe::PlanStageStats>);
-                return plan_explainer_factory::make(
-                    candidates[i].root.get(), &candidates[i].data, candidates[i].solution.get());
+                return plan_explainer_factory::make(candidates[i].root.get(),
+                                                    &candidates[i].data.stageData,
+                                                    candidates[i].solution.get());
             }
         }();
 
@@ -148,7 +149,7 @@ StatusWith<std::unique_ptr<PlanRankingDecision>> pickBestPlan(
         // Get the winning candidate's index to get the correct winning plan.
         size_t winnerIdx = scoresAndCandidateIndices[0].second;
         auto explainer = plan_explainer_factory::make(candidates[winnerIdx].root.get(),
-                                                      &candidates[winnerIdx].data,
+                                                      &candidates[winnerIdx].data.stageData,
                                                       candidates[winnerIdx].solution.get());
         auto&& [stats, _] =
             explainer->getWinningPlanStats(ExplainOptions::Verbosity::kQueryPlanner);
diff --git a/src/mongo/db/query/sbe_cached_solution_planner.cpp b/src/mongo/db/query/sbe_cached_solution_planner.cpp
index 5fbe8be2ec3..9f551b809e9 100644
--- a/src/mongo/db/query/sbe_cached_solution_planner.cpp
+++ b/src/mongo/db/query/sbe_cached_solution_planner.cpp
@@ -75,22 +75,15 @@ CandidatePlans CachedSolutionPlanner::plan(
     // If the '_decisionReads' is not present then we do not run a trial period, keeping the current
     // plan.
     if (!_decisionReads) {
-        const auto status = prepareExecutionPlan(
-            roots[0].first.get(), &roots[0].second, true /* preparingFromCache */);
-        uassertStatusOK(status);
-        bool exitedEarly;
-
-        // Discarding SlotAccessor pointers as they will be reacquired later.
-        std::tie(std::ignore, std::ignore, exitedEarly) = status.getValue();
-        tassert(
-            6693502, "TrialRunTracker is not attached therefore can not exit early", !exitedEarly);
-        return {makeVector(plan_ranker::CandidatePlan{std::move(solutions[0]),
-                                                      std::move(roots[0].first),
-                                                      std::move(roots[0].second),
-                                                      false /* exitedEarly*/,
-                                                      Status::OK(),
-                                                      true,
-                                                      /*isFromPlanCache */}),
+        prepareExecutionPlan(roots[0].first.get(), &roots[0].second, true /* preparingFromCache */);
+        roots[0].first->open(false /* reOpen */);
+        return {makeVector(plan_ranker::CandidatePlan{
+                    std::move(solutions[0]),
+                    std::move(roots[0].first),
+                    plan_ranker::CandidatePlanData{std::move(roots[0].second)},
+                    false /* exitedEarly*/,
+                    Status::OK(),
+                    true /*isFromPlanCache */}),
                 0};
     }
 
@@ -107,15 +100,15 @@ CandidatePlans CachedSolutionPlanner::plan(
 
     auto explainer = plan_explainer_factory::make(
         candidate.root.get(),
-        &candidate.data,
+        &candidate.data.stageData,
         candidate.solution.get(),
         {},    /* optimizedData */
         {},    /* rejectedCandidates */
         false, /* isMultiPlan */
         true,  /* isFromPlanCache */
-        candidate.data.debugInfo
-            ? std::make_unique<plan_cache_debug_info::DebugInfoSBE>(*candidate.data.debugInfo)
-            : nullptr);
+        candidate.data.stageData.debugInfo ? std::make_unique<plan_cache_debug_info::DebugInfoSBE>(
+                                                 *candidate.data.stageData.debugInfo)
+                                           : nullptr);
 
     if (!candidate.status.isOK()) {
         // On failure, fall back to replanning the whole query. We neither evict the existing cache
@@ -165,12 +158,11 @@ plan_ranker::CandidatePlan CachedSolutionPlanner::collectExecutionStatsForCached
 
     plan_ranker::CandidatePlan candidate{std::move(solution),
                                          std::move(root),
-                                         std::move(data),
+                                         plan_ranker::CandidatePlanData{std::move(data)},
                                          false /* exitedEarly*/,
                                          Status::OK(),
                                          true,
                                          /*is Cached plan*/};
-
     ON_BLOCK_EXIT([rootPtr = candidate.root.get()] { rootPtr->detachFromTrialRunTracker(); });
 
     // Callback for the tracker when it exceeds any of the tracked metrics. If the tracker exceeds
@@ -192,10 +184,10 @@ plan_ranker::CandidatePlan CachedSolutionPlanner::collectExecutionStatsForCached
                 MONGO_UNREACHABLE;
         }
     };
-    auto tracker = std::make_unique<TrialRunTracker>(
+    candidate.data.tracker = std::make_unique<TrialRunTracker>(
         std::move(onMetricReached), maxNumResults, maxTrialPeriodNumReads);
-    candidate.root->attachToTrialRunTracker(tracker.get());
-    executeCandidateTrial(&candidate, maxNumResults, /*isCachedPlanTrial*/ true);
+    candidate.root->attachToTrialRunTracker(candidate.data.tracker.get());
+    executeCachedCandidateTrial(&candidate, maxNumResults);
 
     return candidate;
 }
@@ -235,11 +227,8 @@ CandidatePlans CachedSolutionPlanner::replan(bool shouldCache, std::string reaso
 
         // Only one possible plan. Build the stages from the solution.
         auto [root, data] = buildExecutableTree(*solutions[0]);
-        auto status = prepareExecutionPlan(root.get(), &data);
-        uassertStatusOK(status);
-        auto [result, recordId, exitedEarly] = status.getValue();
-        tassert(
-            5323800, "cached planner unexpectedly exited early during prepare phase", !exitedEarly);
+        prepareExecutionPlan(root.get(), &data, false /*preparingFromCache*/);
+        root->open(false /* reOpen */);
 
         auto explainer = plan_explainer_factory::make(root.get(), &data, solutions[0].get());
         LOGV2_DEBUG(2058101,
@@ -267,7 +256,7 @@ CandidatePlans CachedSolutionPlanner::replan(bool shouldCache, std::string reaso
     MultiPlanner multiPlanner{_opCtx, _collections, _cq, plannerParams, cachingMode, _yieldPolicy};
     auto&& [candidates, winnerIdx] = multiPlanner.plan(std::move(solutions), std::move(roots));
     auto explainer = plan_explainer_factory::make(candidates[winnerIdx].root.get(),
-                                                  &candidates[winnerIdx].data,
+                                                  &candidates[winnerIdx].data.stageData,
                                                   candidates[winnerIdx].solution.get());
     LOGV2_DEBUG(2058201,
                 1,
diff --git a/src/mongo/db/query/sbe_multi_planner.cpp b/src/mongo/db/query/sbe_multi_planner.cpp
index 4fb0601563e..714fd4ae0d4 100644
--- a/src/mongo/db/query/sbe_multi_planner.cpp
+++ b/src/mongo/db/query/sbe_multi_planner.cpp
@@ -31,6 +31,7 @@
 
 #include "mongo/db/query/sbe_multi_planner.h"
 
+#include "mongo/db/exec/histogram_server_status_metric.h"
 #include "mongo/db/exec/sbe/expressions/expression.h"
 #include "mongo/db/exec/sbe/values/bson.h"
 #include "mongo/db/query/collection_query_info.h"
@@ -42,8 +43,45 @@
 
 #define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery
 
-
 namespace mongo::sbe {
+namespace {
+/**
+ * An element in this histogram is the number of plans in the candidate set of an invocation (of the
+ * SBE multiplanner).
+ */
+HistogramServerStatusMetric sbeNumPlansHistogram("query.multiPlanner.histograms.sbeNumPlans",
+                                                 HistogramServerStatusMetric::pow(5, 2, 2));
+
+/**
+ * Aggregation of the total number of invocations (of the SBE multiplanner).
+ */
+CounterMetric sbeCount("query.multiPlanner.sbeCount");
+
+/**
+ * Aggregation of the total number of microseconds spent (in SBE multiplanner).
+ */
+CounterMetric sbeMicrosTotal("query.multiPlanner.sbeMicros");
+
+/**
+ * Aggregation of the total number of reads done (in SBE multiplanner).
+ */
+CounterMetric sbeNumReadsTotal("query.multiPlanner.sbeNumReads");
+
+/**
+ * An element in this histogram is the number of microseconds spent in an invocation (of the SBE
+ * multiplanner).
+ */
+HistogramServerStatusMetric sbeMicrosHistogram("query.multiPlanner.histograms.sbeMicros",
+                                               HistogramServerStatusMetric::pow(11, 1024, 4));
+
+/**
+ * An element in this histogram is the number of reads performance during an invocation (of the SBE
+ * multiplanner).
+ */
+HistogramServerStatusMetric sbeNumReadsHistogram("query.multiPlanner.histograms.sbeNumReads",
+                                                 HistogramServerStatusMetric::pow(9, 128, 2));
+}  // namespace
+
 CandidatePlans MultiPlanner::plan(
     std::vector<std::unique_ptr<QuerySolution>> solutions,
     std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) {
@@ -58,6 +96,141 @@ CandidatePlans MultiPlanner::plan(
     return finalizeExecutionPlans(std::move(decision), std::move(candidates));
 }
 
+bool MultiPlanner::CandidateCmp::operator()(const plan_ranker::CandidatePlan* lhs,
+                                            const plan_ranker::CandidatePlan* rhs) const {
+    size_t lhsReads = lhs->data.tracker->getMetric<TrialRunTracker::TrialRunMetric::kNumReads>();
+    size_t rhsReads = rhs->data.tracker->getMetric<TrialRunTracker::TrialRunMetric::kNumReads>();
+    auto lhsProductivity = plan_ranker::calculateProductivity(lhs->results.size(), lhsReads);
+    auto rhsProductivity = plan_ranker::calculateProductivity(rhs->results.size(), rhsReads);
+    return lhsProductivity < rhsProductivity;
+}
+
+MultiPlanner::PlanQ MultiPlanner::preparePlans(
+    const std::vector<size_t>& planIndexes,
+    const size_t trackerResultsBudget,
+    std::vector<std::unique_ptr<QuerySolution>>& solutions,
+    std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>>& roots) {
+    PlanQ planq;
+    for (auto planIndex : planIndexes) {
+        auto&& [root, stageData] = roots[planIndex];
+        // Make a copy of the original plan. This pristine copy will be inserted into the plan
+        // cache if this candidate becomes the winner.
+        auto origPlan = std::make_pair<std::unique_ptr<PlanStage>, plan_ranker::CandidatePlanData>(
+            root->clone(), plan_ranker::CandidatePlanData{stageData});
+
+        // Attach a unique TrialRunTracker to the plan, which is configured to use at most
+        // '_maxNumReads' reads.
+        auto tracker = std::make_unique<TrialRunTracker>(trackerResultsBudget, _maxNumReads);
+        root->attachToTrialRunTracker(tracker.get());
+
+        plan_ranker::CandidatePlanData data = {std::move(stageData), std::move(tracker)};
+        _candidates.push_back({std::move(solutions[planIndex]),
+                               std::move(root),
+                               std::move(data),
+                               false /* exitedEarly */,
+                               Status::OK()});
+        auto* candidatePtr = &_candidates.back();
+        // Store the original plan in the CandidatePlan.
+        candidatePtr->clonedPlan.emplace(std::move(origPlan));
+        prepareCandidate(candidatePtr, false /*preparingFromCache*/);
+        if (fetchOneDocument(candidatePtr)) {
+            planq.push(candidatePtr);
+        }
+    }
+    return planq;
+}
+
+void MultiPlanner::trialPlans(PlanQ planq) {
+    while (!planq.empty()) {
+        plan_ranker::CandidatePlan* bestCandidate = planq.top();
+        planq.pop();
+        bestCandidate->data.tracker->updateMaxMetric<TrialRunTracker::TrialRunMetric::kNumReads>(
+            _maxNumReads);
+        if (fetchOneDocument(bestCandidate)) {
+            planq.push(bestCandidate);
+        }
+    }
+}
+
+bool MultiPlanner::fetchOneDocument(plan_ranker::CandidatePlan* candidate) {
+    if (!fetchNextDocument(candidate, _maxNumResults)) {
+        candidate->root->detachFromTrialRunTracker();
+        if (candidate->status.isOK()) {
+            _maxNumReads = std::min(
+                _maxNumReads,
+                candidate->data.tracker->getMetric<TrialRunTracker::TrialRunMetric::kNumReads>());
+        }
+        return false;
+    }
+    return true;
+}
+
+std::vector<plan_ranker::CandidatePlan> MultiPlanner::collectExecutionStats(
+    std::vector<std::unique_ptr<QuerySolution>> solutions,
+    std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots,
+    size_t maxTrialPeriodNumReads) {
+    invariant(solutions.size() == roots.size());
+
+    _maxNumResults = trial_period::getTrialPeriodNumToReturn(_cq);
+    _maxNumReads = maxTrialPeriodNumReads;
+
+    auto tickSource = _opCtx->getServiceContext()->getTickSource();
+    auto startTicks = tickSource->getTicks();
+    sbeNumPlansHistogram.increment(solutions.size());
+    sbeCount.increment();
+
+    // Determine which plans are blocking and which are non blocking. The non blocking plans will
+    // be run first in order to provide an upper bound on the number of reads allowed for the
+    // blocking plans.
+    std::vector<size_t> nonBlockingPlanIndexes;
+    std::vector<size_t> blockingPlanIndexes;
+    for (size_t index = 0; index < solutions.size(); ++index) {
+        if (solutions[index]->hasBlockingStage) {
+            blockingPlanIndexes.push_back(index);
+        } else {
+            nonBlockingPlanIndexes.push_back(index);
+        }
+    }
+
+    // If all the plans are blocking, then the trial period risks going on for too long. Because the
+    // plans are blocking, they may not provide '_maxNumResults' within the allotted budget of
+    // reads. We could end up in a situation where each plan's trial period runs for a long time,
+    // substantially slowing down the multi-planning process. For this reason, when all the plans
+    // are blocking, we pass '_maxNumResults' to the trial run tracker. This causes the sort stage
+    // to exit early as soon as it sees '_maxNumResults' _input_ values, which keeps the trial
+    // period shorter.
+    //
+    // On the other hand, if we have a mix of blocking and non-blocking plans, we don't want the
+    // sort stage to exit early based on the number of input rows it observes. This could cause the
+    // trial period for the blocking plans to run for a much shorter timeframe than the non-blocking
+    // plans. This leads to an apples-to-oranges comparison between the blocking and non-blocking
+    // plans which could artificially favor the blocking plans.
+    const size_t trackerResultsBudget = nonBlockingPlanIndexes.empty() ? _maxNumResults : 0;
+
+    // Reserve space for the candidates to avoid reallocations and have stable pointers to vector's
+    // elements.
+    _candidates.reserve(solutions.size());
+    // Run the non-blocking plans first.
+    trialPlans(preparePlans(nonBlockingPlanIndexes, trackerResultsBudget, solutions, roots));
+    // Run the blocking plans.
+    trialPlans(preparePlans(blockingPlanIndexes, trackerResultsBudget, solutions, roots));
+
+    size_t totalNumReads = 0;
+    for (const auto& candidate : _candidates) {
+        totalNumReads +=
+            candidate.data.tracker->getMetric<TrialRunTracker::TrialRunMetric::kNumReads>();
+    }
+    sbeNumReadsHistogram.increment(totalNumReads);
+    sbeNumReadsTotal.increment(totalNumReads);
+
+    auto durationMicros = durationCount<Microseconds>(
+        tickSource->ticksTo<Microseconds>(tickSource->getTicks() - startTicks));
+    sbeMicrosHistogram.increment(durationMicros);
+    sbeMicrosTotal.increment(durationMicros);
+
+    return std::move(_candidates);
+}
+
 CandidatePlans MultiPlanner::finalizeExecutionPlans(
     std::unique_ptr<mongo::plan_ranker::PlanRankingDecision> decision,
     std::vector<plan_ranker::CandidatePlan> candidates) const {
@@ -90,8 +263,8 @@ CandidatePlans MultiPlanner::finalizeExecutionPlans(
     LOGV2_DEBUG(
         4822875, 5, "Winning solution", "bestSolution"_attr = redact(winner.solution->toString()));
 
-    auto explainer =
-        plan_explainer_factory::make(winner.root.get(), &winner.data, winner.solution.get());
+    auto explainer = plan_explainer_factory::make(
+        winner.root.get(), &winner.data.stageData, winner.solution.get());
     LOGV2_DEBUG(4822876, 2, "Winning plan", "planSummary"_attr = explainer->getPlanSummary());
 
     // Close all candidate plans but the winner.
@@ -114,16 +287,17 @@ CandidatePlans MultiPlanner::finalizeExecutionPlans(
                 winner.clonedPlan);
         // Clone a new copy of the original plan to use for execution so that the 'clonedPlan' in
         // 'winner' can be inserted into the plan cache while in a clean state.
-        winner.data = stage_builder::PlanStageData(winner.clonedPlan->second);
+        winner.data.stageData = stage_builder::PlanStageData(winner.clonedPlan->second.stageData);
         // When we clone the tree below, the new tree's stats will be zeroed out. If this is an
         // explain operation, save the stats from the old tree before we discard it.
         if (_cq.getExplain()) {
-            winner.data.savedStatsOnEarlyExit = winner.root->getStats(true /* includeDebugInfo  */);
+            winner.data.stageData.savedStatsOnEarlyExit =
+                winner.root->getStats(true /* includeDebugInfo  */);
         }
         winner.root = winner.clonedPlan->first->clone();
 
         stage_builder::prepareSlotBasedExecutableTree(
-            _opCtx, winner.root.get(), &winner.data, _cq, _collections, _yieldPolicy);
+            _opCtx, winner.root.get(), &winner.data.stageData, _cq, _collections, _yieldPolicy);
         // Clear the results queue.
         winner.results = {};
         winner.root->open(false);
@@ -142,16 +316,18 @@ CandidatePlans MultiPlanner::finalizeExecutionPlans(
             _opCtx, _collections, _cq, *solution, _yieldPolicy);
         // The winner might have been replanned. So, pass through the replanning reason to the new
         // plan.
-        data.replanReason = std::move(winner.data.replanReason);
+        data.replanReason = std::move(winner.data.stageData.replanReason);
 
         // We need to clone the plan here for the plan cache to use. The clone will be stored in the
         // cache prior to preparation, whereas the original copy of the tree will be prepared and
         // used to execute this query.
-        auto clonedPlan = std::make_pair(rootStage->clone(), stage_builder::PlanStageData(data));
+        auto clonedPlan = std::make_pair(rootStage->clone(), plan_ranker::CandidatePlanData{data});
         stage_builder::prepareSlotBasedExecutableTree(
             _opCtx, rootStage.get(), &data, _cq, _collections, _yieldPolicy);
-        candidates[winnerIdx] = sbe::plan_ranker::CandidatePlan{
-            std::move(solution), std::move(rootStage), std::move(data)};
+        candidates[winnerIdx] =
+            sbe::plan_ranker::CandidatePlan{std::move(solution),
+                                            std::move(rootStage),
+                                            plan_ranker::CandidatePlanData{std::move(data)}};
         candidates[winnerIdx].clonedPlan.emplace(std::move(clonedPlan));
         candidates[winnerIdx].root->open(false);
 
diff --git a/src/mongo/db/query/sbe_multi_planner.h b/src/mongo/db/query/sbe_multi_planner.h
index 976c54581f3..f4e7ee0c85d 100644
--- a/src/mongo/db/query/sbe_multi_planner.h
+++ b/src/mongo/db/query/sbe_multi_planner.h
@@ -56,6 +56,61 @@ public:
         final;
 
 private:
+    struct CandidateCmp {
+        bool operator()(const plan_ranker::CandidatePlan* lhs,
+                        const plan_ranker::CandidatePlan* rhs) const;
+    };
+
+    using PlanQ = std::priority_queue<plan_ranker::CandidatePlan*,
+                                      std::vector<plan_ranker::CandidatePlan*>,
+                                      CandidateCmp>;
+
+    /**
+     * Moves candidates into the internal candidates vector, sets up a PlanQ entry for each,
+     * prepares each candidate's execution tree and tries to fetch the first document to initialize
+     * plan productivity stats.
+     */
+    PlanQ preparePlans(
+        const std::vector<size_t>& planIndexes,
+        size_t trackerResultsBudget,
+        std::vector<std::unique_ptr<QuerySolution>>& solutions,
+        std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>>& roots);
+
+    /**
+     * Runs the plans in the queue by fetching one document at a time. Each time the most
+     * productive plan is executed.
+     */
+    void trialPlans(PlanQ planq);
+
+    /**
+     * Tries to fetch a single document from the plan. Returns true if the trial run should continue
+     * for the given candidate and false if the trial should end.
+     */
+    bool fetchOneDocument(plan_ranker::CandidatePlan* candidate);
+
+    /**
+     * Executes each plan in to collect execution stats. Stops when all the plans have either:
+     *    * Hit EOF.
+     *    * Returned a pre-defined number of results.
+     *    * Failed
+     *    * Exited early by throwing a special signaling exception.
+     *
+     * Each plan is executed at least once. After that plans are executed by fetching one document
+     * at a time. Every time the most productive plan is executed.
+     *
+     * All documents returned by each plan are enqueued into the 'CandidatePlan->results' queue.
+     *
+     * Upon completion, returns a vector of candidate plans. Execution stats can be obtained for
+     * each of the candidate plans by calling 'CandidatePlan->root->getStats()'.
+     *
+     * After the trial period ends, plans that ran out of memory or reached EOF are closed.
+     * All other plans are open, but 'exitedEarly' plans are in an invalid state. Such plans must be
+     * re-created using the cloned copy before execution of the plan.
+     */
+    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,
+        size_t maxTrialPeriodNumReads);
     /**
      * Returns the best candidate plan selected according to the plan ranking 'decision'.
      *
@@ -68,5 +123,7 @@ private:
 
     // Describes the cases in which we should write an entry for the winning plan to the plan cache.
     const PlanCachingMode _cachingMode;
+    size_t _maxNumResults;
+    size_t _maxNumReads;
 };
 }  // namespace mongo::sbe
diff --git a/src/mongo/db/query/sbe_plan_ranker.cpp b/src/mongo/db/query/sbe_plan_ranker.cpp
index ca3e4211978..046f45c0c10 100644
--- a/src/mongo/db/query/sbe_plan_ranker.cpp
+++ b/src/mongo/db/query/sbe_plan_ranker.cpp
@@ -47,24 +47,15 @@ public:
 
 protected:
     double calculateProductivity(const mongo::sbe::PlanStageStats* root) const final {
-        auto numReads{calculateNumberOfReads(root)};
-
-        // We add one to the number of advances so that plans which returned zero documents still
-        // have a productivity of non-zero. This allows us to compare productivity scores between
-        // plans with zero advances. For example, a plan which did zero advances but examined ten
-        // documents would have a score of (0 + 1)/10, while a plan which did zero advances but
-        // examined a hundred documents would have a score of (0 + 1)/100.
-        //
-        // Similarly, we add one to the number of reads in case 0 reads were performed. This could
-        // happen if a plan encounters EOF right away, for example.
-        return static_cast<double>(root->common.advances + 1) / static_cast<double>(numReads + 1);
+        return plan_ranker::calculateProductivity(root->common.advances,
+                                                  calculateNumberOfReads(root));
     }
 
     std::string getProductivityFormula(const mongo::sbe::PlanStageStats* root) const final {
         auto numReads{calculateNumberOfReads(root)};
         StringBuilder sb;
 
-        sb << "(" << (root->common.advances + 1) << " advances)/(" << numReads << " numReads)";
+        sb << "(" << (root->common.advances) << " advances + 1)/(" << numReads << " numReads + 1)";
 
         return sb.str();
     }
@@ -90,4 +81,17 @@ std::unique_ptr<mongo::plan_ranker::PlanScorer<PlanStageStats>> makePlanScorer(
     const QuerySolution* solution) {
     return std::make_unique<DefaultPlanScorer>(solution);
 }
+
+double calculateProductivity(const size_t advances, const size_t numReads) {
+    // We add one to the number of advances so that plans which returned zero documents still
+    // have a productivity of non-zero. This allows us to compare productivity scores between
+    // plans with zero advances. For example, a plan which did zero advances but examined ten
+    // documents would have a score of (0 + 1)/10, while a plan which did zero advances but
+    // examined a hundred documents would have a score of (0 + 1)/100.
+    //
+    // Similarly, we add one to the number of reads in case 0 reads were performed. This could
+    // happen if a plan encounters EOF right away, for example.
+    return static_cast<double>(advances + 1) / static_cast<double>(numReads + 1);
+}
+
 }  // namespace mongo::sbe::plan_ranker
diff --git a/src/mongo/db/query/sbe_plan_ranker.h b/src/mongo/db/query/sbe_plan_ranker.h
index 9a08a2888fa..0eaf19fb6ed 100644
--- a/src/mongo/db/query/sbe_plan_ranker.h
+++ b/src/mongo/db/query/sbe_plan_ranker.h
@@ -34,14 +34,32 @@
 
 namespace mongo::sbe::plan_ranker {
 
+/**
+ * Structure with data needed to execute the multi-planning trial period for a single SBE candidate
+ * plan.
+ */
+struct CandidatePlanData {
+    stage_builder::PlanStageData stageData;
+    std::unique_ptr<TrialRunTracker> tracker;
+    value::SlotAccessor* resultAccessor = nullptr;
+    value::SlotAccessor* recordIdAccessor = nullptr;
+    bool open = false;
+};
+
 using CandidatePlan =
     mongo::plan_ranker::BaseCandidatePlan<std::unique_ptr<mongo::sbe::PlanStage>,
                                           std::pair<BSONObj, boost::optional<RecordId>>,
-                                          stage_builder::PlanStageData>;
+                                          CandidatePlanData>;
 
 /**
  * A factory function to create a plan ranker for an SBE plan stage stats tree.
  */
 std::unique_ptr<mongo::plan_ranker::PlanScorer<PlanStageStats>> makePlanScorer(
     const QuerySolution* solution);
+
+/**
+ * Helper to calculate productivity if you already know the advances and reads.
+ */
+double calculateProductivity(size_t advances, size_t numReads);
+
 }  // namespace mongo::sbe::plan_ranker
diff --git a/src/mongo/db/query/sbe_runtime_planner.cpp b/src/mongo/db/query/sbe_runtime_planner.cpp
index 0dd7481ec64..ed4527fbaeb 100644
--- a/src/mongo/db/query/sbe_runtime_planner.cpp
+++ b/src/mongo/db/query/sbe_runtime_planner.cpp
@@ -37,67 +37,19 @@
 #include "mongo/db/query/plan_executor_sbe.h"
 
 namespace mongo::sbe {
-namespace {
-
-/**
- * Aggregation of the total number of microseconds spent (in SBE multiplanner).
- */
-CounterMetric sbeMicrosTotal("query.multiPlanner.sbeMicros");
-
-/**
- * Aggregation of the total number of reads done (in SBE multiplanner).
- */
-CounterMetric sbeNumReadsTotal("query.multiPlanner.sbeNumReads");
-
-/**
- * Aggregation of the total number of invocations (of the SBE multiplanner).
- */
-CounterMetric sbeCount("query.multiPlanner.sbeCount");
-
-/**
- * An element in this histogram is the number of microseconds spent in an invocation (of the SBE
- * multiplanner).
- */
-HistogramServerStatusMetric sbeMicrosHistogram("query.multiPlanner.histograms.sbeMicros",
-                                               HistogramServerStatusMetric::pow(11, 1024, 4));
-
-/**
- * An element in this histogram is the number of reads performance during an invocation (of the SBE
- * multiplanner).
- */
-HistogramServerStatusMetric sbeNumReadsHistogram("query.multiPlanner.histograms.sbeNumReads",
-                                                 HistogramServerStatusMetric::pow(9, 128, 2));
-
-/**
- * An element in this histogram is the number of plans in the candidate set of an invocation (of the
- * SBE multiplanner).
- */
-HistogramServerStatusMetric sbeNumPlansHistogram("query.multiPlanner.histograms.sbeNumPlans",
-                                                 HistogramServerStatusMetric::pow(5, 2, 2));
-
-/**
- * Fetches a next document form the given plan stage tree and returns 'true' if the plan stage
- * returns EOF, or throws 'TrialRunTracker::EarlyExitException' exception. Otherwise, the
- * loaded document is placed into the candidate's plan result queue.
- *
- * If the plan stage throws a 'QueryExceededMemoryLimitNoDiskUseAllowed', it will be caught and the
- * 'candidate->failed' flag will be set to 'true', and the 'numFailures' parameter incremented by 1.
- * This failure is considered recoverable, as another candidate plan may require less memory, or may
- * not contain a stage requiring spilling to disk at all.
- */
-enum class FetchDocStatus {
-    done = 0,
-    exitedEarly,
-    inProgress,
-};
-FetchDocStatus fetchNextDocument(
-    plan_ranker::CandidatePlan* candidate,
-    const std::pair<value::SlotAccessor*, value::SlotAccessor*>& slots) {
+bool BaseRuntimePlanner::fetchNextDocument(plan_ranker::CandidatePlan* candidate,
+                                           size_t maxNumResults) {
+    auto* resultSlot = candidate->data.resultAccessor;
+    auto* recordIdSlot = candidate->data.recordIdAccessor;
     try {
+        if (!candidate->data.open) {
+            candidate->root->open(false);
+            candidate->data.open = true;
+        }
+
         BSONObj obj;
         RecordId recordId;
 
-        auto [resultSlot, recordIdSlot] = slots;
         auto state = fetchNext(candidate->root.get(),
                                resultSlot,
                                recordIdSlot,
@@ -106,26 +58,28 @@ FetchDocStatus fetchNextDocument(
                                true /* must return owned BSON */);
         if (state == PlanState::IS_EOF) {
             candidate->root->close();
-            return FetchDocStatus::done;
+            return false;
         }
 
         invariant(state == PlanState::ADVANCED);
         invariant(obj.isOwned());
-        candidate->results.push_back({obj, {recordIdSlot != nullptr, recordId}});
+        candidate->results.push_back({std::move(obj), {recordIdSlot != nullptr, recordId}});
+        if (candidate->results.size() >= maxNumResults) {
+            return false;
+        }
     } catch (const ExceptionFor<ErrorCodes::QueryTrialRunCompleted>&) {
-        return FetchDocStatus::exitedEarly;
+        candidate->exitedEarly = true;
+        return false;
     } catch (const ExceptionFor<ErrorCodes::QueryExceededMemoryLimitNoDiskUseAllowed>& ex) {
         candidate->root->close();
         candidate->status = ex.toStatus();
+        return false;
     }
-    return FetchDocStatus::inProgress;
+    return true;
 }
-}  // namespace
 
-StatusWith<std::tuple<value::SlotAccessor*, value::SlotAccessor*, bool>>
-BaseRuntimePlanner::prepareExecutionPlan(PlanStage* root,
-                                         stage_builder::PlanStageData* data,
-                                         const bool preparingFromCache) const {
+std::pair<value::SlotAccessor*, value::SlotAccessor*> BaseRuntimePlanner::prepareExecutionPlan(
+    PlanStage* root, stage_builder::PlanStageData* data, const bool preparingFromCache) const {
     invariant(root);
     invariant(data);
 
@@ -144,141 +98,20 @@ BaseRuntimePlanner::prepareExecutionPlan(PlanStage* root,
         tassert(4822872, "Query does not have a recordId slot.", recordIdSlot);
     }
 
-    auto exitedEarly{false};
-    try {
-        root->open(false);
-    } catch (const ExceptionFor<ErrorCodes::QueryTrialRunCompleted>&) {
-        exitedEarly = true;
-    } catch (const ExceptionFor<ErrorCodes::QueryExceededMemoryLimitNoDiskUseAllowed>& ex) {
-        root->close();
-        return ex.toStatus();
-    }
-
-    return std::make_tuple(resultSlot, recordIdSlot, exitedEarly);
+    return std::make_pair(resultSlot, recordIdSlot);
 }
 
-void BaseRuntimePlanner::executeCandidateTrial(plan_ranker::CandidatePlan* candidate,
-                                               size_t maxNumResults,
-                                               const bool isCachedPlanTrial) {
+void BaseRuntimePlanner::prepareCandidate(plan_ranker::CandidatePlan* candidate,
+                                          bool preparingFromCache) {
     _indexExistenceChecker.check(_opCtx, _collections);
-
-    auto status = prepareExecutionPlan(candidate->root.get(), &candidate->data, isCachedPlanTrial);
-    if (!status.isOK()) {
-        candidate->status = status.getStatus();
-        return;
-    }
-
-    auto [resultAccessor, recordIdAccessor, exitedEarly] = status.getValue();
-    if (exitedEarly) {
-        candidate->exitedEarly = true;
-        return;
-    }
-
-    for (size_t i = 0; i < maxNumResults && candidate->status.isOK(); ++i) {
-        FetchDocStatus fetch =
-            fetchNextDocument(candidate, std::make_pair(resultAccessor, recordIdAccessor));
-        if (fetch == FetchDocStatus::done || fetch == FetchDocStatus::exitedEarly) {
-            candidate->exitedEarly = (fetch == FetchDocStatus::exitedEarly);
-            return;
-        }
-    }
+    std::tie(candidate->data.resultAccessor, candidate->data.recordIdAccessor) =
+        prepareExecutionPlan(candidate->root.get(), &candidate->data.stageData, preparingFromCache);
 }
 
-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,
-    size_t maxTrialPeriodNumReads) {
-    invariant(solutions.size() == roots.size());
-
-    std::vector<plan_ranker::CandidatePlan> candidates;
-    std::vector<std::pair<value::SlotAccessor*, value::SlotAccessor*>> accessors;
-
-    const auto maxNumResults{trial_period::getTrialPeriodNumToReturn(_cq)};
-
-    auto tickSource = _opCtx->getServiceContext()->getTickSource();
-    auto startTicks = tickSource->getTicks();
-    sbeNumPlansHistogram.increment(solutions.size());
-    sbeCount.increment();
-
-    // Determine which plans are blocking and which are non blocking. The non blocking plans will
-    // be run first in order to provide an upper bound on the number of reads allowed for the
-    // blocking plans.
-    std::vector<size_t> nonBlockingPlanIndexes;
-    std::vector<size_t> blockingPlanIndexes;
-    for (size_t index = 0; index < solutions.size(); ++index) {
-        if (solutions[index]->hasBlockingStage) {
-            blockingPlanIndexes.push_back(index);
-        } else {
-            nonBlockingPlanIndexes.push_back(index);
-        }
+void BaseRuntimePlanner::executeCachedCandidateTrial(plan_ranker::CandidatePlan* candidate,
+                                                     size_t maxNumResults) {
+    prepareCandidate(candidate, true /*preparingFromCache*/);
+    while (fetchNextDocument(candidate, maxNumResults)) {
     }
-
-    // If all the plans are blocking, then the trial period risks going on for too long. Because the
-    // plans are blocking, they may not provide 'maxNumResults' within the allotted budget of reads.
-    // We could end up in a situation where each plan's trial period runs for a long time,
-    // substantially slowing down the multi-planning process. For this reason, when all the plans
-    // are blocking, we pass 'maxNumResults' to the trial run tracker. This causes the sort stage to
-    // exit early as soon as it sees 'maxNumResults' _input_ values, which keeps the trial period
-    // shorter.
-    //
-    // On the other hand, if we have a mix of blocking and non-blocking plans, we don't want the
-    // sort stage to exit early based on the number of input rows it observes. This could cause the
-    // trial period for the blocking plans to run for a much shorter timeframe than the non-blocking
-    // plans. This leads to an apples-to-oranges comparison between the blocking and non-blocking
-    // plans which could artificially favor the blocking plans.
-    const size_t trackerResultsBudget = nonBlockingPlanIndexes.empty() ? maxNumResults : 0;
-
-    uint64_t totalNumReads = 0;
-
-    auto runPlans = [&](const std::vector<size_t>& planIndexes, size_t& maxNumReads) -> void {
-        for (auto planIndex : planIndexes) {
-            // Prepare the plan.
-            auto&& [root, data] = roots[planIndex];
-            // Make a copy of the original plan. This pristine copy will be inserted into the plan
-            // cache if this candidate becomes the winner.
-            auto origPlan =
-                std::make_pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>(
-                    root->clone(), stage_builder::PlanStageData(data));
-
-            // Attach a unique TrialRunTracker to the plan, which is configured to use at most
-            // 'maxNumReads' reads.
-            auto tracker = std::make_unique<TrialRunTracker>(trackerResultsBudget, maxNumReads);
-            ON_BLOCK_EXIT([rootPtr = root.get()] { rootPtr->detachFromTrialRunTracker(); });
-            root->attachToTrialRunTracker(tracker.get());
-
-            candidates.push_back({std::move(solutions[planIndex]),
-                                  std::move(root),
-                                  std::move(data),
-                                  false /* exitedEarly */,
-                                  Status::OK()});
-            auto& currentCandidate = candidates.back();
-            // Store the original plan in the CandidatePlan.
-            currentCandidate.clonedPlan.emplace(std::move(origPlan));
-            executeCandidateTrial(&currentCandidate, maxNumResults, /*isCachedPlanTrial*/ false);
-
-            auto reads = tracker->getMetric<TrialRunTracker::TrialRunMetric::kNumReads>();
-            // We intentionally increment the metrics outside of the isOk/existedEarly check.
-            totalNumReads += reads;
-
-            // Reduce the number of reads the next candidates are allocated if this candidate is
-            // more efficient than the current bound.
-            if (currentCandidate.status.isOK() && !currentCandidate.exitedEarly) {
-                maxNumReads = std::min(maxNumReads, reads);
-            }
-        }
-    };
-
-    runPlans(nonBlockingPlanIndexes, maxTrialPeriodNumReads);
-    runPlans(blockingPlanIndexes, maxTrialPeriodNumReads);
-
-    sbeNumReadsHistogram.increment(totalNumReads);
-    sbeNumReadsTotal.increment(totalNumReads);
-
-    auto durationMicros = durationCount<Microseconds>(
-        tickSource->ticksTo<Microseconds>(tickSource->getTicks() - startTicks));
-    sbeMicrosHistogram.increment(durationMicros);
-    sbeMicrosTotal.increment(durationMicros);
-
-    return candidates;
 }
 }  // namespace mongo::sbe
diff --git a/src/mongo/db/query/sbe_runtime_planner.h b/src/mongo/db/query/sbe_runtime_planner.h
index 9078f79e70d..44e781335c6 100644
--- a/src/mongo/db/query/sbe_runtime_planner.h
+++ b/src/mongo/db/query/sbe_runtime_planner.h
@@ -93,55 +93,34 @@ public:
 
 protected:
     /**
-     * Prepares the given plan stage tree for execution, attaches it to the operation context and
-     * returns two slot accessors for the result and recordId slots, and a boolean value indicating
-     * if the plan has exited early from the trial period. If the plan has failed in a recoverable
-     * fashion, it will return a non-OK status.
+     * Fetches a next document from the given plan stage tree and the loaded document is placed into
+     * the candidate's plan result queue.
+     *
+     * Returns true if a document was fetched, and false if the plan stage tree reached EOF, an
+     * exception was thrown or the plan stage tree returned maxNumResults documents.
      *
-     * The caller should pass true for 'preparingFromCache' if the SBE plan being prepared is being
-     * recovered from the SBE plan cache.
+     * If the plan stage throws a 'QueryExceededMemoryLimitNoDiskUseAllowed', it will be caught and
+     * the 'candidate->status' will be set. This failure is considered recoverable, as another
+     * candidate plan may require less memory, or may not contain a stage requiring spilling to disk
+     * at all.
      */
-    StatusWith<std::tuple<sbe::value::SlotAccessor*, sbe::value::SlotAccessor*, bool>>
-    prepareExecutionPlan(PlanStage* root,
-                         stage_builder::PlanStageData* data,
-                         bool preparingFromCache = false) const;
+    static bool fetchNextDocument(plan_ranker::CandidatePlan* candidate, size_t maxNumResults);
 
     /**
-     * Executes a candidate plan until it
-     *   - reaches EOF, or
-     *   - reaches the 'maxNumResults' limit, or
-     *   - early exits via the TrialRunTracker, or
-     *   - returns a failure Status.
-     *
-     * The execution process populates the 'results' array of the 'candidate' plan with any results
-     * from execution the plan. This function also sets the 'status' and 'exitedEarly' fields of the
-     * input 'candidate' object when applicable.
-     *
-     * If we are running the trial period for a plan recovered from the plan cache, then the caller
-     * must pass true for 'isCachedPlanTrial'.
+     * Prepares the given plan stage tree for execution, attaches it to the operation context and
+     * returns two slot accessors for the result and recordId slots. The caller should pass true
+     * for 'preparingFromCache' if the SBE plan being prepared is being recovered from the SBE plan
+     * cache.
      */
-    void executeCandidateTrial(plan_ranker::CandidatePlan* candidate,
-                               size_t maxNumResults,
-                               bool isCachedPlanTrial);
+    std::pair<sbe::value::SlotAccessor*, sbe::value::SlotAccessor*> prepareExecutionPlan(
+        PlanStage* root, stage_builder::PlanStageData* data, bool preparingFromCache) const;
 
     /**
-     * Executes each plan in a round-robin fashion to collect execution stats. Stops when:
-     *    * Any plan hits EOF.
-     *    * Or returns a pre-defined number of results.
-     *    * Or all candidate plans fail or exit early by throwing a special signaling exception.
-     *
-     * All documents returned by each plan are enqueued into the 'CandidatePlan->results' queue.
-     *
-     * Upon completion returns a vector of candidate plans. Execution stats can be obtained for each
-     * of the candidate plans by calling 'CandidatePlan->root->getStats()'.
-     *
-     * After the trial period ends, all plans remain open, but 'exitedEarly' plans are in an invalid
-     * state. Any 'exitedEarly' plans must be closed and reopened before they can be executed.
+     * Wraps prepareExecutionPlan(), checks index validity, and caches outputAccessors.
      */
-    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,
-        size_t maxTrialPeriodNumReads);
+    void prepareCandidate(plan_ranker::CandidatePlan* candidate, bool preparingFromCache);
+
+    void executeCachedCandidateTrial(plan_ranker::CandidatePlan* candidate, size_t maxNumResults);
 
     OperationContext* const _opCtx;
     const MultipleCollectionAccessor& _collections;
@@ -149,5 +128,7 @@ protected:
     const QueryPlannerParams _queryParams;
     PlanYieldPolicySBE* const _yieldPolicy;
     const AllIndicesRequiredChecker _indexExistenceChecker;
+
+    std::vector<plan_ranker::CandidatePlan> _candidates;
 };
 }  // namespace mongo::sbe
diff --git a/src/mongo/db/query/sbe_sub_planner.cpp b/src/mongo/db/query/sbe_sub_planner.cpp
index e3da52fb0d9..4082e7bec39 100644
--- a/src/mongo/db/query/sbe_sub_planner.cpp
+++ b/src/mongo/db/query/sbe_sub_planner.cpp
@@ -118,10 +118,8 @@ CandidatePlans SubPlanner::plan(
     // range deletion).
     plan_cache_util::updatePlanCache(_opCtx, _collections, _cq, *compositeSolution, *root, data);
 
-    auto status = prepareExecutionPlan(root.get(), &data);
-    uassertStatusOK(status);
-    auto [result, recordId, exitedEarly] = status.getValue();
-    tassert(5323804, "sub-planner unexpectedly exited early during prepare phase", !exitedEarly);
+    prepareExecutionPlan(root.get(), &data, false /*preparingFromCache*/);
+    root->open(false);
 
     return {makeVector(plan_ranker::CandidatePlan{
                 std::move(compositeSolution), std::move(root), std::move(data)}),
@@ -143,11 +141,8 @@ CandidatePlans SubPlanner::planWholeQuery() const {
 
         auto&& [root, data] = stage_builder::buildSlotBasedExecutableTree(
             _opCtx, _collections, _cq, *solutions[0], _yieldPolicy);
-        auto status = prepareExecutionPlan(root.get(), &data);
-        uassertStatusOK(status);
-        auto [result, recordId, exitedEarly] = status.getValue();
-        tassert(
-            5323805, "sub-planner unexpectedly exited early during prepare phase", !exitedEarly);
+        prepareExecutionPlan(root.get(), &data, false /*preparingFromCache*/);
+        root->open(false);
         return {makeVector(plan_ranker::CandidatePlan{
                     std::move(solutions[0]), std::move(root), std::move(data)}),
                 0};