SERVER-64870 Make sure $lookup query is cached when SBE plan cache is enabled

4 files changed, 81 insertions, 12 deletions

diff --git a/jstests/noPassthrough/lookup_pushdown_cache.js b/jstests/noPassthrough/lookup_pushdown_cache.js
new file mode 100644
index 00000000000..24fd84959f8
--- /dev/null
+++ b/jstests/noPassthrough/lookup_pushdown_cache.js
@@ -0,0 +1,66 @@
+/**
+ * Tests basic functionality of integrating plan cache with lowered $lookup. Currently only the
+ * stages below group/lookup get cached in the classic cache.
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/profiler.js");  // For 'getLatestProfilerEntry'.
+
+const conn = MongoRunner.runMongod({setParameter: {featureFlagSBELookupPushdown: true}});
+assert.neq(null, conn, "mongod was unable to start up");
+const name = "lookup_pushdown";
+const foreignCollName = "foreign_lookup_pushdown";
+let db = conn.getDB(name);
+let coll = db[name];
+let foreignColl = db[foreignCollName];
+
+function verifyPlanCache({query, isActive, planCacheKey}) {
+    const cacheEntries = coll.aggregate([{$planCacheStats: {}}]).toArray();
+    assert.eq(cacheEntries.length, 1);
+    const cacheEntry = cacheEntries[0];
+    // TODO(SERVER-61507): Convert the assertion to SBE cache once lowered $lookup integrates
+    // with SBE plan cache.
+    assert.eq(cacheEntry.version, 1);
+    assert.docEq(cacheEntry.createdFromQuery.query, query);
+    assert.eq(cacheEntry.isActive, isActive);
+    if (planCacheKey) {
+        assert.eq(cacheEntry.planCacheKey, planCacheKey);
+    }
+    return cacheEntry;
+}
+
+// Create two indices to make sure the query gets multi-planned, so that the query subtree will
+// be saved in the classic cache.
+assert.commandWorked(coll.createIndexes([{a: 1, b: 1}, {a: 1, c: 1}]));
+assert.commandWorked(coll.insert([{a: 1}, {a: 2}]));
+assert.commandWorked(foreignColl.insert([{c: 1}, {c: 2}]));
+const query = {
+    a: {$gt: 1}
+};
+const pipeline = [
+    {$match: query},
+    {$lookup: {from: foreignCollName, localField: "a", foreignField: "c", as: "c_out"}}
+];
+
+// First run should create an inactive cache entry.
+assert.eq(1, coll.aggregate(pipeline).itcount());
+const cacheEntry = verifyPlanCache({query, isActive: false});
+const planCacheKey = cacheEntry.planCacheKey;
+
+// Second run should mark the cache entry active.
+assert.eq(1, coll.aggregate(pipeline).itcount());
+verifyPlanCache({query, planCacheKey, isActive: true});
+
+// Third run should use the active cached entry.
+assert.commandWorked(db.setProfilingLevel(2));
+assert.eq(1, coll.aggregate(pipeline).itcount());
+const profileEntry = getLatestProfilerEntry(db, {});
+assert.eq(planCacheKey, profileEntry.planCacheKey);
+
+// Explain output should show the same plan cache key.
+const explain = coll.explain().aggregate(pipeline);
+assert.eq(planCacheKey, explain.queryPlanner.planCacheKey);
+
+MongoRunner.stopMongod(conn);
+}());
diff --git a/src/mongo/db/exec/plan_cache_util.h b/src/mongo/db/exec/plan_cache_util.h
index 4c5249cb1e6..5540c94c161 100644
--- a/src/mongo/db/exec/plan_cache_util.h
+++ b/src/mongo/db/exec/plan_cache_util.h
@@ -107,13 +107,6 @@ void updatePlanCache(
     invariant(winnerIdx >= 0 && winnerIdx < candidates.size());
     auto& winningPlan = candidates[winnerIdx];
 
-    // TODO SERVER-61507: Integration between lowering parts of aggregation pipeline into the find
-    // subsystem and the new SBE cache isn't implemented yet.
-    if (!query.pipeline().empty() &&
-        feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV()) {
-        return;
-    }
-
     // Even if the query is of a cacheable shape, the caller might have indicated that we shouldn't
     // write to the plan cache.
     //
@@ -200,7 +193,10 @@ void updatePlanCache(
 
         if (winningPlan.solution->cacheData != nullptr) {
             if constexpr (std::is_same_v<PlanStageType, std::unique_ptr<sbe::PlanStage>>) {
-                if (feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV()) {
+                // TODO SERVER-61507: Integration between lowering parts of aggregation pipeline
+                // into the find subsystem and the new SBE cache isn't implemented yet.
+                if (feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV() &&
+                    query.pipeline().empty()) {
                     tassert(6142201,
                             "The winning CandidatePlan should contain the original plan",
                             winningPlan.clonedPlan);
@@ -222,8 +218,9 @@ void updatePlanCache(
                         &callbacks,
                         boost::none /* worksGrowthCoefficient */));
                 } else {
-                    // Fall back to use the classic plan cache. Remove this branch after
-                    // "gFeatureFlagSbePlanCache" is removed.
+                    // TODO(SERVER-61507, SERVER-64882): Fall back to use the classic plan cache.
+                    // Remove this branch after "gFeatureFlagSbePlanCache" is removed and lowering
+                    // parts of pipeline is integrated with SBE cache.
                     cacheClassicPlan();
                 }
             } else {
diff --git a/src/mongo/db/query/explain.cpp b/src/mongo/db/query/explain.cpp
index 82f1c6e50eb..f887c93beb8 100644
--- a/src/mongo/db/query/explain.cpp
+++ b/src/mongo/db/query/explain.cpp
@@ -96,7 +96,10 @@ void generatePlannerInfo(PlanExecutor* exec,
             QuerySettingsDecoration::get(collection->getSharedDecorations());
         if (exec->getCanonicalQuery()->isSbeCompatible() &&
             feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV() &&
-            !exec->getCanonicalQuery()->getForceClassicEngine()) {
+            !exec->getCanonicalQuery()->getForceClassicEngine() &&
+            // TODO(SERVER-61507): Remove pipeline check once lowered pipelines are integrated with
+            // SBE plan cache.
+            exec->getCanonicalQuery()->pipeline().empty()) {
             const auto planCacheKeyInfo = plan_cache_key_factory::make<sbe::PlanCacheKey>(
                 *exec->getCanonicalQuery(), collection);
             planCacheKeyHash = planCacheKeyInfo.planCacheKeyHash();
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index f53e9f1debb..8cefa6efded 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -1083,7 +1083,10 @@ protected:
     std::unique_ptr<SlotBasedPrepareExecutionResult> buildCachedPlan(
         const sbe::PlanCacheKey& planCacheKey) final {
         if (shouldCacheQuery(*_cq)) {
-            if (!feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV()) {
+            // TODO SERVER-61507: remove _cq->pipeline().empty() check when $group pushdown is
+            // integrated with SBE plan cache.
+            if (!feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV() ||
+                !_cq->pipeline().empty()) {
                 // If the feature flag is off, we first try to build an "id hack" plan because the
                 // id hack plans are not cached in the classic cache. We then fall back to use the
                 // classic plan cache.