SERVER-66389 Fix $where parameter bind-in optimization

8 files changed, 65 insertions, 29 deletions

diff --git a/src/mongo/db/query/bind_input_params.cpp b/src/mongo/db/query/bind_input_params.cpp
index 7d489d275ff..d96aa4294bf 100644
--- a/src/mongo/db/query/bind_input_params.cpp
+++ b/src/mongo/db/query/bind_input_params.cpp
@@ -44,8 +44,11 @@ class MatchExpressionParameterBindingVisitor final : public MatchExpressionConst
 public:
     MatchExpressionParameterBindingVisitor(
         const stage_builder::InputParamToSlotMap& inputParamToSlotMap,
-        sbe::RuntimeEnvironment* runtimeEnvironment)
-        : _inputParamToSlotMap(inputParamToSlotMap), _runtimeEnvironment(runtimeEnvironment) {
+        sbe::RuntimeEnvironment* runtimeEnvironment,
+        bool bindingCachedPlan)
+        : _inputParamToSlotMap(inputParamToSlotMap),
+          _runtimeEnvironment(runtimeEnvironment),
+          _bindingCachedPlan(bindingCachedPlan) {
         invariant(_runtimeEnvironment);
     }
 
@@ -173,18 +176,24 @@ public:
             return;
         }
 
-        // Generally speaking, this visitor is non-destructive and does not mutate the
-        // MatchExpression tree. However, in order to apply an optimization to avoid making a copy
-        // of the 'JsFunction' object stored within 'WhereMatchExpression', we can transfer its
-        // ownership from the match expression node into the SBE runtime environment. Hence, we need
-        // to drop the const qualifier. This should be a safe operation, given that the match
-        // expression tree is allocated on the heap, and this visitor has exclusive access to this
-        // tree (after we have bound in all input parameters, it's no longer used).
-        bindParam(*inputParam,
-                  true /*owned*/,
-                  sbe::value::TypeTags::jsFunction,
-                  sbe::value::bitcastFrom<JsFunction*>(
-                      const_cast<WhereMatchExpression*>(expr)->extractPredicate().release()));
+        if (_bindingCachedPlan) {
+            // Generally speaking, this visitor is non-destructive and does not mutate the
+            // MatchExpression tree. However, in order to apply an optimization to avoid making a
+            // copy of the 'JsFunction' object stored within 'WhereMatchExpression', we can transfer
+            // its ownership from the match expression node into the SBE runtime environment. Hence,
+            // we need to drop the const qualifier. This is a safe operation only when the plan is
+            // being recovered from the SBE plan cache -- in this case, the visitor has exclusive
+            // access to this match expression tree. Furthermore, after all input parameters are
+            // bound the match expression tree is no longer used.
+            bindParam(*inputParam,
+                      true /*owned*/,
+                      sbe::value::TypeTags::jsFunction,
+                      sbe::value::bitcastFrom<JsFunction*>(
+                          const_cast<WhereMatchExpression*>(expr)->extractPredicate().release()));
+        } else {
+            auto [typeTag, value] = sbe::value::makeCopyJsFunction(expr->getPredicate());
+            bindParam(*inputParam, true /*owned*/, typeTag, value);
+        }
     }
 
     /**
@@ -298,6 +307,10 @@ private:
     const stage_builder::InputParamToSlotMap& _inputParamToSlotMap;
 
     sbe::RuntimeEnvironment* const _runtimeEnvironment;
+
+    // True if the plan for which we are binding parameter values is being recovered from the SBE
+    // plan cache.
+    const bool _bindingCachedPlan;
 };
 
 class MatchExpressionParameterBindingWalker {
@@ -426,8 +439,10 @@ void bindGenericPlanSlots(const stage_builder::IndexBoundsEvaluationInfo& indexB
 
 void bind(const CanonicalQuery& canonicalQuery,
           const stage_builder::InputParamToSlotMap& inputParamToSlotMap,
-          sbe::RuntimeEnvironment* runtimeEnvironment) {
-    MatchExpressionParameterBindingVisitor visitor{inputParamToSlotMap, runtimeEnvironment};
+          sbe::RuntimeEnvironment* runtimeEnvironment,
+          const bool bindingCachedPlan) {
+    MatchExpressionParameterBindingVisitor visitor{
+        inputParamToSlotMap, runtimeEnvironment, bindingCachedPlan};
     MatchExpressionParameterBindingWalker walker{&visitor};
     tree_walker::walk<true, MatchExpression>(canonicalQuery.root(), &walker);
 }
diff --git a/src/mongo/db/query/bind_input_params.h b/src/mongo/db/query/bind_input_params.h
index 910ce991878..9384bd3b5f1 100644
--- a/src/mongo/db/query/bind_input_params.h
+++ b/src/mongo/db/query/bind_input_params.h
@@ -40,10 +40,14 @@ namespace mongo::input_params {
  * parameter id, along with a constant associated with that parameter id. For each such match
  * expression node, looks up the corresponding slot in the 'RuntimeEnvironment' using the
  * 'InputParamToSlotMap' and sets the value of that slot.
+ *
+ * The caller should pass true for 'bindingCachedPlan' if we are binding-in new parameter values for
+ * a plan that was recovered from the SBE plan cache.
  */
 void bind(const CanonicalQuery&,
           const stage_builder::InputParamToSlotMap&,
-          sbe::RuntimeEnvironment*);
+          sbe::RuntimeEnvironment*,
+          bool bindingCachedPlan);
 
 /**
  * Binds index bounds evaluated from IETs to index bounds slots for the given query.
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index c7f66647d9a..9eba8febf1f 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -1362,7 +1362,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getSlotBasedExe
 
     // Prepare the SBE tree for execution.
     stage_builder::prepareSlotBasedExecutableTree(
-        opCtx, root.get(), &data, *cq, collections, yieldPolicy.get());
+        opCtx, root.get(), &data, *cq, collections, yieldPolicy.get(), true);
 
     return plan_executor_factory::make(opCtx,
                                        std::move(cq),
diff --git a/src/mongo/db/query/sbe_cached_solution_planner.cpp b/src/mongo/db/query/sbe_cached_solution_planner.cpp
index 6c243294011..5f1b8f008d6 100644
--- a/src/mongo/db/query/sbe_cached_solution_planner.cpp
+++ b/src/mongo/db/query/sbe_cached_solution_planner.cpp
@@ -163,7 +163,7 @@ plan_ranker::CandidatePlan CachedSolutionPlanner::collectExecutionStatsForCached
     auto tracker = std::make_unique<TrialRunTracker>(
         std::move(onMetricReached), maxNumResults, maxTrialPeriodNumReads);
     candidate.root->attachToTrialRunTracker(tracker.get());
-    executeCandidateTrial(&candidate, maxNumResults);
+    executeCandidateTrial(&candidate, maxNumResults, /*isCachedPlanTrial*/ true);
 
     return candidate;
 }
diff --git a/src/mongo/db/query/sbe_runtime_planner.cpp b/src/mongo/db/query/sbe_runtime_planner.cpp
index a0c2309e49b..5edf89539bf 100644
--- a/src/mongo/db/query/sbe_runtime_planner.cpp
+++ b/src/mongo/db/query/sbe_runtime_planner.cpp
@@ -131,12 +131,13 @@ FetchDocStatus fetchNextDocument(
 
 StatusWith<std::tuple<value::SlotAccessor*, value::SlotAccessor*, bool>>
 BaseRuntimePlanner::prepareExecutionPlan(PlanStage* root,
-                                         stage_builder::PlanStageData* data) const {
+                                         stage_builder::PlanStageData* data,
+                                         const bool preparingFromCache) const {
     invariant(root);
     invariant(data);
 
     stage_builder::prepareSlotBasedExecutableTree(
-        _opCtx, root, data, _cq, _collections, _yieldPolicy);
+        _opCtx, root, data, _cq, _collections, _yieldPolicy, preparingFromCache);
 
     value::SlotAccessor* resultSlot{nullptr};
     if (auto slot = data->outputs.getIfExists(stage_builder::PlanStageSlots::kResult); slot) {
@@ -164,10 +165,11 @@ BaseRuntimePlanner::prepareExecutionPlan(PlanStage* root,
 }
 
 void BaseRuntimePlanner::executeCandidateTrial(plan_ranker::CandidatePlan* candidate,
-                                               size_t maxNumResults) {
+                                               size_t maxNumResults,
+                                               const bool isCachedPlanTrial) {
     _indexExistenceChecker.check();
 
-    auto status = prepareExecutionPlan(candidate->root.get(), &candidate->data);
+    auto status = prepareExecutionPlan(candidate->root.get(), &candidate->data, isCachedPlanTrial);
     if (!status.isOK()) {
         candidate->status = status.getStatus();
         return;
@@ -259,7 +261,7 @@ std::vector<plan_ranker::CandidatePlan> BaseRuntimePlanner::collectExecutionStat
             auto& currentCandidate = candidates.back();
             // Store the original plan in the CandidatePlan.
             currentCandidate.clonedPlan.emplace(std::move(origPlan));
-            executeCandidateTrial(&currentCandidate, maxNumResults);
+            executeCandidateTrial(&currentCandidate, maxNumResults, /*isCachedPlanTrial*/ false);
 
             auto reads = tracker->getMetric<TrialRunTracker::TrialRunMetric::kNumReads>();
             // We intentionally increment the metrics outside of the isOk/existedEarly check.
diff --git a/src/mongo/db/query/sbe_runtime_planner.h b/src/mongo/db/query/sbe_runtime_planner.h
index 6dc49449165..7d28e268ae3 100644
--- a/src/mongo/db/query/sbe_runtime_planner.h
+++ b/src/mongo/db/query/sbe_runtime_planner.h
@@ -97,9 +97,14 @@ protected:
      * 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.
+     *
+     * The caller should pass true for 'preparingFromCache' if the SBE plan being prepared is being
+     * recovered from the SBE plan cache.
      */
     StatusWith<std::tuple<sbe::value::SlotAccessor*, sbe::value::SlotAccessor*, bool>>
-    prepareExecutionPlan(PlanStage* root, stage_builder::PlanStageData* data) const;
+    prepareExecutionPlan(PlanStage* root,
+                         stage_builder::PlanStageData* data,
+                         bool preparingFromCache = false) const;
 
     /**
      * Executes a candidate plan until it
@@ -111,8 +116,13 @@ protected:
      * 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'.
      */
-    void executeCandidateTrial(plan_ranker::CandidatePlan* candidate, size_t maxNumResults);
+    void executeCandidateTrial(plan_ranker::CandidatePlan* candidate,
+                               size_t maxNumResults,
+                               bool isCachedPlanTrial);
 
     /**
      * Executes each plan in a round-robin fashion to collect execution stats. Stops when:
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index f1e1e3d6e74..c79beaa9dfa 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -402,7 +402,8 @@ void prepareSlotBasedExecutableTree(OperationContext* opCtx,
                                     PlanStageData* data,
                                     const CanonicalQuery& cq,
                                     const MultipleCollectionAccessor& collections,
-                                    PlanYieldPolicySBE* yieldPolicy) {
+                                    PlanYieldPolicySBE* yieldPolicy,
+                                    const bool preparingFromCache) {
     tassert(6183502, "PlanStage cannot be null", root);
     tassert(6142205, "PlanStageData cannot be null", data);
     tassert(6142206, "yieldPolicy cannot be null", yieldPolicy);
@@ -464,7 +465,7 @@ void prepareSlotBasedExecutableTree(OperationContext* opCtx,
     // TODO SERVER-66039: Add dassert on sbe::validateInputParamsBindings().
     // If the cached plan is parameterized, bind new values for the parameters into the runtime
     // environment.
-    input_params::bind(cq, data->inputParamToSlotMap, env);
+    input_params::bind(cq, data->inputParamToSlotMap, env, preparingFromCache);
 
     for (auto&& indexBoundsInfo : data->indexBoundsEvaluationInfos) {
         input_params::bindIndexBounds(cq, indexBoundsInfo, env);
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index 0f1813437ea..7abd0e2fa46 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -55,13 +55,17 @@ std::unique_ptr<sbe::RuntimeEnvironment> makeRuntimeEnvironment(
  * This function prepares the SBE tree for execution, such as attaching the OperationContext,
  * ensuring that the SBE tree is registered with the PlanYieldPolicySBE and populating the
  * "RuntimeEnvironment".
+ *
+ * The caller should pass true for 'preparingFromCache' if the SBE plan being prepared is being
+ * recovered from the SBE plan cache.
  */
 void prepareSlotBasedExecutableTree(OperationContext* opCtx,
                                     sbe::PlanStage* root,
                                     PlanStageData* data,
                                     const CanonicalQuery& cq,
                                     const MultipleCollectionAccessor& collections,
-                                    PlanYieldPolicySBE* yieldPolicy);
+                                    PlanYieldPolicySBE* yieldPolicy,
+                                    bool preparingFromCache = false);
 
 class PlanStageReqs;