SERVER-59698 Make subplanning work with the SBE plan cache

5 files changed, 164 insertions, 110 deletions

diff --git a/jstests/core/sbe/plan_cache_sbe_with_or_queries.js b/jstests/core/sbe/plan_cache_sbe_with_or_queries.js
new file mode 100644
index 00000000000..7bf52c462b4
--- /dev/null
+++ b/jstests/core/sbe/plan_cache_sbe_with_or_queries.js
@@ -0,0 +1,57 @@
+// Tests that rooted $OR queries are added to the SBE plan cache.
+// @tags: [
+//   # This test attempts to perform queries and introspect the server's plan cache entries. The
+//   # former operation may be routed to a secondary in the replica set, whereas the latter must be
+//   # routed to the primary.
+//   assumes_read_preference_unchanged,
+//   assumes_read_concern_unchanged,
+//   assumes_unsharded_collection,
+//   does_not_support_stepdowns,
+// ]
+
+(function() {
+"use strict";
+
+load("jstests/libs/analyze_plan.js");
+load("jstests/libs/sbe_util.js");  // For checkSBEEnabled.
+
+const isSBEEnabled = checkSBEEnabled(db, ["featureFlagSbePlanCache"]);
+if (!isSBEEnabled) {
+    jsTest.log("Skip running the test because featureFlagSbePlanCache is not enabled");
+    return;
+}
+
+function getPlanCacheEntries(query, collection, db) {
+    const planCacheKey = getPlanCacheKeyFromShape({query, collection, db});
+    return coll.aggregate([{$planCacheStats: {}}, {$match: {planCacheKey}}]).toArray();
+}
+
+function assertRootedOrPlan(query, collection, db) {
+    const explain = collection.find(query).explain();
+    const winningPlan = getWinningPlan(explain.queryPlanner);
+    assert(planHasStage(db, winningPlan, "OR"), explain);
+}
+
+const coll = db.sbe_subplan;
+coll.drop();
+
+assert.commandWorked(coll.createIndexes([{a: 1, b: -1}, {b: 1}]));
+assert.commandWorked(coll.insertMany([{a: 1, b: 1}, {a: 1, b: 2}, {a: 2, b: 2}, {a: 2, b: 3}]));
+
+const query = {
+    $or: [{a: 1}, {b: 2}]
+};
+assert.eq(3, coll.find(query).itcount());
+
+assertRootedOrPlan(query, coll, db);
+
+const planCacheEntries = getPlanCacheEntries(query, coll, db);
+// A subplan query adds an entry to the plan cache.
+assert.eq(1, planCacheEntries.length, planCacheEntries);
+
+// And this entry must be pinned and active.
+assert.eq(true, planCacheEntries[0].isPinned, planCacheEntries);
+assert.eq(true, planCacheEntries[0].isActive, planCacheEntries);
+// Works is always 0 for pinned plan cache entries.
+assert.eq(0, planCacheEntries[0].works, planCacheEntries);
+}());
diff --git a/src/mongo/db/exec/subplan.cpp b/src/mongo/db/exec/subplan.cpp
index 40d6a4d6252..0cf6e4e0eb4 100644
--- a/src/mongo/db/exec/subplan.cpp
+++ b/src/mongo/db/exec/subplan.cpp
@@ -171,19 +171,26 @@ Status SubplanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {
     // Dismiss the requirement that no indices can be dropped when this method returns.
     ON_BLOCK_EXIT([this] { releaseAllIndicesRequirement(); });
 
-    std::function<PlanCacheKey(const CanonicalQuery& cq, const CollectionPtr& coll)>
-        createPlanCacheKey = [](const CanonicalQuery& cq, const CollectionPtr& coll) {
-            return plan_cache_key_factory::make<PlanCacheKey>(cq, coll);
-        };
+    std::function<std::unique_ptr<SolutionCacheData>(const CanonicalQuery& cq,
+                                                     const CollectionPtr& coll)>
+        getSolutionCachedData =
+            [](const CanonicalQuery& cq,
+               const CollectionPtr& coll) -> std::unique_ptr<SolutionCacheData> {
+        auto planCache = CollectionQueryInfo::get(coll).getPlanCache();
+        tassert(5969800, "Classic Plan Cache not found", planCache);
+        if (shouldCacheQuery(cq)) {
+            auto planCacheKey = plan_cache_key_factory::make<PlanCacheKey>(cq, coll);
+            if (auto cachedSol = planCache->getCacheEntryIfActive(planCacheKey)) {
+                return std::move(cachedSol->cachedPlan);
+            }
+        }
+
+        return nullptr;
+    };
 
     // Plan each branch of the $or.
-    auto subplanningStatus =
-        QueryPlanner::planSubqueries(expCtx()->opCtx,
-                                     CollectionQueryInfo::get(collection()).getPlanCache(),
-                                     createPlanCacheKey,
-                                     collection(),
-                                     *_query,
-                                     _plannerParams);
+    auto subplanningStatus = QueryPlanner::planSubqueries(
+        expCtx()->opCtx, getSolutionCachedData, collection(), *_query, _plannerParams);
     if (!subplanningStatus.isOK()) {
         return choosePlanWholeQuery(yieldPolicy);
     }
@@ -192,7 +199,7 @@ Status SubplanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {
     auto subplanningResult = std::move(subplanningStatus.getValue());
     _branchPlannedFromCache.clear();
     for (auto&& branch : subplanningResult.branches) {
-        _branchPlannedFromCache.push_back(branch->cachedSolution != nullptr);
+        _branchPlannedFromCache.push_back(branch->cachedData != nullptr);
     }
 
     // Use the multi plan stage to select a winning plan for each branch, and then construct
diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp
index 848e8ba6aef..63cf94f6376 100644
--- a/src/mongo/db/query/query_planner.cpp
+++ b/src/mongo/db/query/query_planner.cpp
@@ -1309,13 +1309,10 @@ StatusWith<std::unique_ptr<QuerySolution>> QueryPlanner::choosePlanForSubqueries
         auto orChild = planningResult.orExpression->getChild(i);
         auto branchResult = planningResult.branches[i].get();
 
-        if (branchResult->cachedSolution.get()) {
+        if (branchResult->cachedData.get()) {
             // We can get the index tags we need out of the cache.
-            Status tagStatus =
-                tagOrChildAccordingToCache(cacheData.get(),
-                                           branchResult->cachedSolution->cachedPlan.get(),
-                                           orChild,
-                                           planningResult.indexMap);
+            Status tagStatus = tagOrChildAccordingToCache(
+                cacheData.get(), branchResult->cachedData.get(), orChild, planningResult.indexMap);
             if (!tagStatus.isOK()) {
                 return tagStatus;
             }
@@ -1399,4 +1396,76 @@ StatusWith<std::unique_ptr<QuerySolution>> QueryPlanner::choosePlanForSubqueries
 
     return std::move(compositeSolution);
 }
+
+StatusWith<QueryPlanner::SubqueriesPlanningResult> QueryPlanner::planSubqueries(
+    OperationContext* opCtx,
+    std::function<std::unique_ptr<SolutionCacheData>(
+        const CanonicalQuery& cq, const CollectionPtr& coll)> getSolutionCachedData,
+    const CollectionPtr& collection,
+    const CanonicalQuery& query,
+    const QueryPlannerParams& params) {
+    invariant(query.root()->matchType() == MatchExpression::OR);
+    invariant(query.root()->numChildren(), "Cannot plan subqueries for an $or with no children");
+
+    SubqueriesPlanningResult planningResult{query.root()->shallowClone()};
+    for (size_t i = 0; i < params.indices.size(); ++i) {
+        const IndexEntry& ie = params.indices[i];
+        const auto insertionRes = planningResult.indexMap.insert(std::make_pair(ie.identifier, i));
+        // Be sure the key was not already in the map.
+        invariant(insertionRes.second);
+        log_detail::logSubplannerIndexEntry(ie, i);
+    }
+
+    for (size_t i = 0; i < planningResult.orExpression->numChildren(); ++i) {
+        // We need a place to shove the results from planning this branch.
+        planningResult.branches.push_back(
+            std::make_unique<SubqueriesPlanningResult::BranchPlanningResult>());
+        auto branchResult = planningResult.branches.back().get();
+        auto orChild = planningResult.orExpression->getChild(i);
+
+        // Turn the i-th child into its own query.
+        auto statusWithCQ = CanonicalQuery::canonicalize(opCtx, query, orChild);
+        if (!statusWithCQ.isOK()) {
+            str::stream ss;
+            ss << "Can't canonicalize subchild " << orChild->debugString() << " "
+               << statusWithCQ.getStatus().reason();
+            return Status(ErrorCodes::BadValue, ss);
+        }
+
+        branchResult->canonicalQuery = std::move(statusWithCQ.getValue());
+
+        // Plan the i-th child. We might be able to find a plan for the i-th child in the plan
+        // cache. If there's no cached plan, then we generate and rank plans using the MPS.
+
+        // Populate branchResult->cachedData if an active cachedData entry exists.
+        if (getSolutionCachedData) {
+            branchResult->cachedData =
+                getSolutionCachedData(*branchResult->canonicalQuery, collection);
+        }
+
+        if (branchResult->cachedData) {
+            log_detail::logCachedPlanFound(planningResult.orExpression->numChildren(), i);
+        } else {
+            // No CachedSolution found. We'll have to plan from scratch.
+            log_detail::logCachedPlanNotFound(planningResult.orExpression->numChildren(), i);
+
+            // We don't set NO_TABLE_SCAN because peeking at the cache data will keep us from
+            // considering any plan that's a collscan.
+            invariant(branchResult->solutions.empty());
+            auto statusWithMultiPlanSolns =
+                QueryPlanner::plan(*branchResult->canonicalQuery, params);
+            if (!statusWithMultiPlanSolns.isOK()) {
+                str::stream ss;
+                ss << "Can't plan for subchild " << branchResult->canonicalQuery->toString() << " "
+                   << statusWithMultiPlanSolns.getStatus().reason();
+                return Status(ErrorCodes::BadValue, ss);
+            }
+            branchResult->solutions = std::move(statusWithMultiPlanSolns.getValue());
+
+            log_detail::logNumberOfSolutions(branchResult->solutions.size());
+        }
+    }
+
+    return std::move(planningResult);
+}
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_planner.h b/src/mongo/db/query/query_planner.h
index 170083d523b..142a6f6d7c7 100644
--- a/src/mongo/db/query/query_planner.h
+++ b/src/mongo/db/query/query_planner.h
@@ -73,7 +73,7 @@ public:
             // a set of alternate plans for the branch. The index tags from the cache data
             // can be applied directly to the parent $or MatchExpression when generating the
             // composite solution.
-            std::unique_ptr<CachedSolution> cachedSolution;
+            std::unique_ptr<SolutionCacheData> cachedData;
 
             // Query solutions resulting from planning the $or branch.
             std::vector<std::unique_ptr<QuerySolution>> solutions;
@@ -122,12 +122,10 @@ public:
      * The 'createPlanCacheKey' callback is used to create a plan cache key of the specified
      * 'KeyType' for each of the branches to look up the plan in the 'planCache'.
      */
-    template <typename KeyType, typename... Args>
     static StatusWith<SubqueriesPlanningResult> planSubqueries(
         OperationContext* opCtx,
-        const PlanCacheBase<KeyType, Args...>* planCache,
-        std::function<KeyType(const CanonicalQuery& cq, const CollectionPtr& coll)>
-            createPlanCacheKey,
+        std::function<std::unique_ptr<SolutionCacheData>(
+            const CanonicalQuery& cq, const CollectionPtr& coll)> getSolutionCachedData,
         const CollectionPtr& collection,
         const CanonicalQuery& query,
         const QueryPlannerParams& params);
@@ -179,79 +177,4 @@ public:
             CanonicalQuery* cq, std::vector<std::unique_ptr<QuerySolution>>)> multiplanCallback);
 };
 
-template <typename KeyType, typename... Args>
-StatusWith<QueryPlanner::SubqueriesPlanningResult> QueryPlanner::planSubqueries(
-    OperationContext* opCtx,
-    const PlanCacheBase<KeyType, Args...>* planCache,
-    std::function<KeyType(const CanonicalQuery& cq, const CollectionPtr& coll)> createPlanCacheKey,
-    const CollectionPtr& collection,
-    const CanonicalQuery& query,
-    const QueryPlannerParams& params) {
-    invariant(query.root()->matchType() == MatchExpression::OR);
-    invariant(query.root()->numChildren(), "Cannot plan subqueries for an $or with no children");
-
-    SubqueriesPlanningResult planningResult{query.root()->shallowClone()};
-    for (size_t i = 0; i < params.indices.size(); ++i) {
-        const IndexEntry& ie = params.indices[i];
-        const auto insertionRes = planningResult.indexMap.insert(std::make_pair(ie.identifier, i));
-        // Be sure the key was not already in the map.
-        invariant(insertionRes.second);
-        log_detail::logSubplannerIndexEntry(ie, i);
-    }
-
-    for (size_t i = 0; i < planningResult.orExpression->numChildren(); ++i) {
-        // We need a place to shove the results from planning this branch.
-        planningResult.branches.push_back(
-            std::make_unique<SubqueriesPlanningResult::BranchPlanningResult>());
-        auto branchResult = planningResult.branches.back().get();
-        auto orChild = planningResult.orExpression->getChild(i);
-
-        // Turn the i-th child into its own query.
-        auto statusWithCQ = CanonicalQuery::canonicalize(opCtx, query, orChild);
-        if (!statusWithCQ.isOK()) {
-            str::stream ss;
-            ss << "Can't canonicalize subchild " << orChild->debugString() << " "
-               << statusWithCQ.getStatus().reason();
-            return Status(ErrorCodes::BadValue, ss);
-        }
-
-        branchResult->canonicalQuery = std::move(statusWithCQ.getValue());
-
-        // Plan the i-th child. We might be able to find a plan for the i-th child in the plan
-        // cache. If there's no cached plan, then we generate and rank plans using the MPS.
-
-        // Populate branchResult->cachedSolution if an active cachedSolution entry exists.
-        if (planCache && shouldCacheQuery(*branchResult->canonicalQuery)) {
-            if (auto cachedSol = planCache->getCacheEntryIfActive(
-                    createPlanCacheKey(*branchResult->canonicalQuery, collection))) {
-                // We have a CachedSolution. Store it for later.
-                log_detail::logCachedPlanFound(planningResult.orExpression->numChildren(), i);
-
-                branchResult->cachedSolution = std::move(cachedSol);
-            }
-        }
-
-        if (!branchResult->cachedSolution) {
-            // No CachedSolution found. We'll have to plan from scratch.
-            log_detail::logCachedPlanNotFound(planningResult.orExpression->numChildren(), i);
-
-            // We don't set NO_TABLE_SCAN because peeking at the cache data will keep us from
-            // considering any plan that's a collscan.
-            invariant(branchResult->solutions.empty());
-            auto statusWithMultiPlanSolns =
-                QueryPlanner::plan(*branchResult->canonicalQuery, params);
-            if (!statusWithMultiPlanSolns.isOK()) {
-                str::stream ss;
-                ss << "Can't plan for subchild " << branchResult->canonicalQuery->toString() << " "
-                   << statusWithMultiPlanSolns.getStatus().reason();
-                return Status(ErrorCodes::BadValue, ss);
-            }
-            branchResult->solutions = std::move(statusWithMultiPlanSolns.getValue());
-
-            log_detail::logNumberOfSolutions(branchResult->solutions.size());
-        }
-    }
-
-    return std::move(planningResult);
-}
 }  // namespace mongo
diff --git a/src/mongo/db/query/sbe_sub_planner.cpp b/src/mongo/db/query/sbe_sub_planner.cpp
index e6830dc5622..e5e714ad3aa 100644
--- a/src/mongo/db/query/sbe_sub_planner.cpp
+++ b/src/mongo/db/query/sbe_sub_planner.cpp
@@ -39,22 +39,13 @@
 
 namespace mongo::sbe {
 CandidatePlans SubPlanner::plan(
-    std::vector<std::unique_ptr<QuerySolution>>,
+    std::vector<std::unique_ptr<QuerySolution>> solutions,
     std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>>) {
-    std::function<mongo::PlanCacheKey(const CanonicalQuery& cq, const CollectionPtr& coll)>
-        createPlanCacheKey = [](const CanonicalQuery& cq, const CollectionPtr& coll) {
-            return plan_cache_key_factory::make<mongo::PlanCacheKey>(cq, coll);
-        };
 
     const auto& mainColl = _collections.getMainCollection();
     // Plan each branch of the $or.
-    auto subplanningStatus =
-        QueryPlanner::planSubqueries(_opCtx,
-                                     CollectionQueryInfo::get(mainColl).getPlanCache(),
-                                     createPlanCacheKey,
-                                     mainColl,
-                                     _cq,
-                                     _queryParams);
+    auto subplanningStatus = QueryPlanner::planSubqueries(
+        _opCtx, {} /* getSolutionCachedData */, mainColl, _cq, _queryParams);
     if (!subplanningStatus.isOK()) {
         return planWholeQuery();
     }
@@ -122,6 +113,13 @@ CandidatePlans SubPlanner::plan(
     uassertStatusOK(status);
     auto [result, recordId, exitedEarly] = status.getValue();
     tassert(5323804, "sub-planner unexpectedly exited early during prepare phase", !exitedEarly);
+
+    // TODO SERVER-61507: do it unconditionally when $group pushdown is integrated with the SBE plan
+    // cache.
+    if (_cq.pipeline().empty()) {
+        plan_cache_util::updatePlanCache(_opCtx, mainColl, _cq, *compositeSolution, *root, data);
+    }
+
     return {makeVector(plan_ranker::CandidatePlan{
                 std::move(compositeSolution), std::move(root), std::move(data)}),
             0};