SERVER-61835 Fix how SBE plan cache deals with ShardFilterer

6 files changed, 82 insertions, 23 deletions

diff --git a/src/mongo/db/exec/plan_cache_util.cpp b/src/mongo/db/exec/plan_cache_util.cpp
index 0a4aab74886..7a74e4683bf 100644
--- a/src/mongo/db/exec/plan_cache_util.cpp
+++ b/src/mongo/db/exec/plan_cache_util.cpp
@@ -160,5 +160,21 @@ plan_cache_debug_info::DebugInfoSBE buildDebugInfo(const QuerySolution* solution
 
     return debugInfo;
 }
+
+void resetRuntimeEnvironmentBeforeCaching(stage_builder::PlanStageData* data) {
+    tassert(6183501, "PlanStageData should not be null", data);
+
+    // Manually reset "shardFilterer" to "Nothing" because we should not store
+    // "shardFilterer" which holds a "ScopedCollectionFilter" preventing data that
+    // may have been migrated from being deleted.
+    if (auto shardFiltererSlot = data->env->getSlotIfExists("shardFilterer"_sd)) {
+        data->env->resetSlot(*shardFiltererSlot, sbe::value::TypeTags::Nothing, 0, true);
+    }
+
+    // Reset all the parameterized slots.
+    for (auto [paramId, slotId] : data->inputParamToSlotMap) {
+        data->env->resetSlot(slotId, sbe::value::TypeTags::Nothing, 0, true);
+    }
+}
 }  // namespace plan_cache_util
 }  // namespace mongo
diff --git a/src/mongo/db/exec/plan_cache_util.h b/src/mongo/db/exec/plan_cache_util.h
index 9b8851fac22..06c1a4ddbf0 100644
--- a/src/mongo/db/exec/plan_cache_util.h
+++ b/src/mongo/db/exec/plan_cache_util.h
@@ -38,6 +38,7 @@
 #include "mongo/db/query/query_solution.h"
 #include "mongo/db/query/sbe_plan_cache.h"
 #include "mongo/db/query/sbe_plan_ranker.h"
+#include "mongo/db/query/sbe_stage_builder.h"
 
 namespace mongo {
 /**
@@ -73,19 +74,24 @@ void logNotCachingZeroResults(std::string&& query, double score, std::string win
 void logNotCachingNoData(std::string&& solution);
 }  // namespace log_detail
 
-/*
+/**
  * Builds "DebugInfo" for storing in the classic plan cache.
  */
 plan_cache_debug_info::DebugInfo buildDebugInfo(
     const CanonicalQuery& query, std::unique_ptr<const plan_ranker::PlanRankingDecision> decision);
 
-/*
+/**
  * Builds "DebugInfoSBE" for storing in the SBE plan cache. Pre-computes necessary debugging
  * information to build "PlanExplainerSBE" when recoverying the cached SBE plan from the cache.
  */
 plan_cache_debug_info::DebugInfoSBE buildDebugInfo(const QuerySolution* solution);
 
 /**
+ * Resets all the parameterized slots in the RuntimeEnvironment of 'data'.
+ */
+void resetRuntimeEnvironmentBeforeCaching(stage_builder::PlanStageData* data);
+
+/**
  * Caches the best candidate plan, chosen from the given 'candidates' based on the 'ranking'
  * decision, if the 'query' is of a type that can be cached. Otherwise, does nothing.
  *
@@ -203,6 +209,8 @@ void updatePlanCache(
                     auto cachedPlan = std::make_unique<sbe::CachedSbePlan>(
                         winningPlan.root->clone(), winningPlan.data);
 
+                    resetRuntimeEnvironmentBeforeCaching(&cachedPlan->planStageData);
+
                     PlanCacheLoggingCallbacks<sbe::PlanCacheKey,
                                               sbe::CachedSbePlan,
                                               plan_cache_debug_info::DebugInfoSBE>
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index 160ff81fd8a..1f6f08b06e5 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -89,6 +89,7 @@
 #include "mongo/db/query/sbe_multi_planner.h"
 #include "mongo/db/query/sbe_sub_planner.h"
 #include "mongo/db/query/sbe_utils.h"
+#include "mongo/db/query/shard_filterer_factory_impl.h"
 #include "mongo/db/query/stage_builder_util.h"
 #include "mongo/db/query/util/make_data_structure.h"
 #include "mongo/db/query/wildcard_multikey_paths.h"
@@ -444,6 +445,29 @@ bool shouldWaitForOplogVisibility(OperationContext* opCtx,
     return repl::ReplicationCoordinator::get(opCtx)->canAcceptWritesForDatabase(opCtx, "admin");
 }
 
+void prepareExecutionTree(OperationContext* opCtx,
+                          const CollectionPtr& collection,
+                          const CanonicalQuery& cq,
+                          stage_builder::PlanStageData* stageData) {
+    tassert(6183502, "PlanStageData should not be null", stageData);
+    // Populate/renew "shardFilterer" if there exists a "shardFilterer" slot. The slot value
+    // should have been reset to "Nothing" on caching.
+    //
+    // TODO SERVER-61422: Populate the "shardFilterer" when preparing SBE plan.
+    if (auto shardFiltererSlot = stageData->env->getSlotIfExists("shardFilterer"_sd)) {
+        auto shardFiltererFactory = std::make_unique<ShardFiltererFactoryImpl>(collection);
+        auto shardFilterer = shardFiltererFactory->makeShardFilterer(opCtx);
+        stageData->env->resetSlot(*shardFiltererSlot,
+                                  sbe::value::TypeTags::shardFilterer,
+                                  sbe::value::bitcastFrom<ShardFilterer*>(shardFilterer.release()),
+                                  true);
+    }
+
+    // If the cached plan is parameterized, bind new values for the parameters into the runtime
+    // environment.
+    input_params::bind(cq, stageData->inputParamToSlotMap, stageData->env);
+}
+
 namespace {
 /**
  * A class to hold the result of preparation of the query to be executed using classic engine. This
@@ -1074,9 +1098,7 @@ protected:
         invariant(sbeYieldPolicy);
         sbeYieldPolicy->registerPlan(root.get());
 
-        // If the cached plan is parameterized, bind new values for the parameters into the runtime
-        // environment.
-        input_params::bind(*_cq, stageData.inputParamToSlotMap, stageData.env);
+        prepareExecutionTree(_opCtx, _collections.getMainCollection(), *_cq, &stageData);
 
         auto result = makeResult();
         result->setDecisionWorks(cacheEntry->decisionWorks);
diff --git a/src/mongo/db/query/get_executor.h b/src/mongo/db/query/get_executor.h
index 4e038a646c9..09dd1722412 100644
--- a/src/mongo/db/query/get_executor.h
+++ b/src/mongo/db/query/get_executor.h
@@ -108,6 +108,15 @@ void fillOutPlannerParams(OperationContext* opCtx,
                           QueryPlannerParams* plannerParams);
 
 /**
+ * Prepare the "PlanStageData" for execution by populating necessary value to "RuntimeEnvironment"
+ * and binding parameterized values with corresponding slots.
+ */
+void prepareExecutionTree(OperationContext* opCtx,
+                          const CollectionPtr& collection,
+                          const CanonicalQuery& cq,
+                          stage_builder::PlanStageData* data);
+
+/**
  * Return whether or not any component of the path 'path' is multikey given an index key pattern
  * and multikeypaths. If no multikey metdata is available for the index, and the index is marked
  * multikey, conservatively assumes that a component of 'path' _is_ multikey. The 'isMultikey'
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index 33d066d3143..ab2618d08e1 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -2738,18 +2738,16 @@ namespace {
 /**
  * Given an SBE subtree 'childStage' which computes the shard key and puts it into the given
  * 'shardKeySlot', augments the SBE plan to actually perform shard filtering. Namely, a FilterStage
- * is added at the root of the tree whose filter expression uses 'shardFilterer' to determine
+ * is added at the root of the tree whose filter expression uses 'shardFiltererSlot' to determine
  * whether the shard key value in 'shardKeySlot' belongs to an owned range or not.
  */
 auto buildShardFilterGivenShardKeySlot(sbe::value::SlotId shardKeySlot,
                                        std::unique_ptr<sbe::PlanStage> childStage,
-                                       std::unique_ptr<ShardFilterer> shardFilterer,
+                                       sbe::value::SlotId shardFiltererSlot,
                                        PlanNodeId nodeId) {
-    auto shardFilterFn =
-        makeFunction("shardFilter",
-                     makeConstant(sbe::value::TypeTags::shardFilterer,
-                                  sbe::value::bitcastFrom<ShardFilterer*>(shardFilterer.release())),
-                     sbe::makeE<sbe::EVariable>(shardKeySlot));
+    auto shardFilterFn = makeFunction("shardFilter",
+                                      sbe::makeE<sbe::EVariable>(shardFiltererSlot),
+                                      sbe::makeE<sbe::EVariable>(shardKeySlot));
 
     return sbe::makeS<sbe::FilterStage<false>>(
         std::move(childStage), std::move(shardFilterFn), nodeId);
@@ -2758,7 +2756,7 @@ auto buildShardFilterGivenShardKeySlot(sbe::value::SlotId shardKeySlot,
 
 std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots>
 SlotBasedStageBuilder::buildShardFilterCovered(const ShardingFilterNode* filterNode,
-                                               std::unique_ptr<ShardFilterer> shardFilterer,
+                                               sbe::value::SlotId shardFiltererSlot,
                                                BSONObj shardKeyPattern,
                                                BSONObj indexKeyPattern,
                                                const QuerySolutionNode* child,
@@ -2865,7 +2863,7 @@ SlotBasedStageBuilder::buildShardFilterCovered(const ShardingFilterNode* filterN
                                                         filterNode->nodeId());
 
     auto filterStage = buildShardFilterGivenShardKeySlot(
-        shardKeySlot, std::move(mkObjStage), std::move(shardFilterer), filterNode->nodeId());
+        shardKeySlot, std::move(mkObjStage), shardFiltererSlot, filterNode->nodeId());
 
     outputs.setIndexKeySlots(!parentIndexKeyReqs ? boost::none
                                                  : boost::optional<sbe::value::SlotVector>{
@@ -2887,8 +2885,15 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     // there are orphaned documents from aborted migrations. To check if the document is owned by
     // the shard, we need to own a 'ShardFilterer', and extract the document's shard key as a
     // BSONObj.
+    // TODO SERVER-61422: Should not construct the "ShardFilterer" here.
     auto shardFilterer = _shardFiltererFactory->makeShardFilterer(_opCtx);
     auto shardKeyPattern = shardFilterer->getKeyPattern().toBSON();
+    auto shardFiltererSlot =
+        _data.env->registerSlot("shardFilterer"_sd,
+                                sbe::value::TypeTags::shardFilterer,
+                                sbe::value::bitcastFrom<ShardFilterer*>(shardFilterer.release()),
+                                true,
+                                &_slotIdGenerator);
 
     // Determine if our child is an index scan and extract it's key pattern, or empty BSONObj if our
     // child is not an IXSCAN node.
@@ -2915,8 +2920,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     auto childReqs = reqs.copy().setIf(kResult, indexKeyPattern.isEmpty());
     if (!childReqs.has(kResult)) {
         return buildShardFilterCovered(filterNode,
-                                       std::move(shardFilterer),
-                                       std::move(shardKeyPattern),
+                                       shardFiltererSlot,
+                                       shardKeyPattern,
                                        std::move(indexKeyPattern),
                                        filterNode->children[0],
                                        std::move(childReqs));
@@ -2984,11 +2989,10 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     auto finalShardKeyObjStage = makeProjectStage(
         std::move(shardKeyObjStage), root->nodeId(), finalShardKeySlot, std::move(arrayChecks));
 
-    return {buildShardFilterGivenShardKeySlot(finalShardKeySlot,
-                                              std::move(finalShardKeyObjStage),
-                                              std::move(shardFilterer),
-                                              root->nodeId()),
-            std::move(outputs)};
+    return {
+        buildShardFilterGivenShardKeySlot(
+            finalShardKeySlot, std::move(finalShardKeyObjStage), shardFiltererSlot, root->nodeId()),
+        std::move(outputs)};
 }
 
 // Returns a non-null pointer to the root of a plan tree, or a non-OK status if the PlanStage tree
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index c004f614046..8d5d14f20e8 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -420,11 +420,11 @@ private:
      * 'shardKeyPattern' are provided by 'childIxscan'. In this case, the SBE plan for the child
      * index scan node will fill out slots for the necessary components of the index key. These
      * slots can be read directly in order to determine the shard key that should be passed to the
-     * 'shardFilterer'.
+     * 'shardFiltererSlot'.
      */
     std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> buildShardFilterCovered(
         const ShardingFilterNode* filterNode,
-        std::unique_ptr<ShardFilterer> shardFilterer,
+        sbe::value::SlotId shardFiltererSlot,
         BSONObj shardKeyPattern,
         BSONObj indexKeyPattern,
         const QuerySolutionNode* child,