SERVER-66389 Fix $where parameter bind-in optimization

(cherry picked from commit 48acc21bc952810c7028f79773905f9cdcce44af)

10 files changed, 102 insertions, 29 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index bddd06f3357..a0aa3bcafd1 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -216,6 +216,8 @@ last-continuous:
       test_file: jstests/core/wildcard_index_cached_plans.js
     - ticket: SERVER-65821
       test_file: jstests/replsets/dbhash_lock_acquisition.js
+    - ticket: SERVER-66389
+      test_file: jstests/core/where_multiple_plans.js
 
   # Tests that should only be excluded from particular suites should be listed under that suite.
   suites:
@@ -567,6 +569,8 @@ last-lts:
       test_file: jstests/core/wildcard_index_cached_plans.js
     - ticket: SERVER-65821
       test_file: jstests/replsets/dbhash_lock_acquisition.js
+    - ticket: SERVER-66389
+      test_file: jstests/core/where_multiple_plans.js
 
   # Tests that should only be excluded from particular suites should be listed under that suite.
   suites:
diff --git a/jstests/core/where_multiple_plans.js b/jstests/core/where_multiple_plans.js
new file mode 100644
index 00000000000..86f2e050747
--- /dev/null
+++ b/jstests/core/where_multiple_plans.js
@@ -0,0 +1,33 @@
+/**
+ * Tests that a $where predicate works as expected when there are multiple candidate plans.
+ *
+ * @tags: [
+ *   requires_fcv_60,
+ *   requires_scripting,
+ * ]
+ */
+(function() {
+"use strict";
+
+const coll = db.where_multiple_plans;
+coll.drop();
+
+assert.commandWorked(coll.createIndex({a: 1}));
+assert.commandWorked(coll.createIndex({b: 1}));
+
+assert.commandWorked(coll.insert({_id: 0, a: 1, b: 3, c: 4}));
+assert.commandWorked(coll.insert({_id: 1, a: 2, b: 1, c: 4}));
+assert.commandWorked(coll.insert({_id: 2, a: 2, b: 3, c: 1}));
+assert.commandWorked(coll.insert({_id: 3, a: 2, b: 3, c: 4}));
+assert.commandWorked(coll.insert({_id: 4, a: 2, b: 3, c: 4}));
+
+assert.eq(2,
+          coll.find({
+                  a: {$eq: 2},
+                  b: {$eq: 3},
+                  $where: function() {
+                      return this.c == 4;
+                  }
+              })
+              .itcount());
+}());
diff --git a/src/mongo/db/query/bind_input_params.cpp b/src/mongo/db/query/bind_input_params.cpp
index c6276a5d60b..89a1b469375 100644
--- a/src/mongo/db/query/bind_input_params.cpp
+++ b/src/mongo/db/query/bind_input_params.cpp
@@ -42,8 +42,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);
     }
 
@@ -171,18 +174,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);
+        }
     }
 
     /**
@@ -299,6 +308,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 {
@@ -322,8 +335,10 @@ private:
 
 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 82092c1ab6f..c86af36c872 100644
--- a/src/mongo/db/query/bind_input_params.h
+++ b/src/mongo/db/query/bind_input_params.h
@@ -40,8 +40,12 @@ 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);
 }  // namespace mongo::input_params
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index 07b2e440e64..8b801dcb33b 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -1350,7 +1350,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 8aa5c2105a0..c594c70f22e 100644
--- a/src/mongo/db/query/sbe_cached_solution_planner.cpp
+++ b/src/mongo/db/query/sbe_cached_solution_planner.cpp
@@ -161,7 +161,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 74e09a76ecd..b3f35938fdf 100644
--- a/src/mongo/db/query/sbe_runtime_planner.cpp
+++ b/src/mongo/db/query/sbe_runtime_planner.cpp
@@ -146,12 +146,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) {
@@ -179,10 +180,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;
@@ -274,7 +276,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 ebfc50387eb..9ace5b5e150 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -530,7 +530,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);
@@ -591,7 +592,7 @@ void prepareSlotBasedExecutableTree(OperationContext* opCtx,
 
     // 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) {
         bindIndexBoundsParams(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;