SERVER-65345 Check if CanonicalQuery has pushed down stages

3 files changed, 206 insertions, 27 deletions

diff --git a/jstests/noPassthrough/plan_cache_replan_group_lookup.js b/jstests/noPassthrough/plan_cache_replan_group_lookup.js
index 403d004e0f8..4eb954eb93c 100644
--- a/jstests/noPassthrough/plan_cache_replan_group_lookup.js
+++ b/jstests/noPassthrough/plan_cache_replan_group_lookup.js
@@ -307,25 +307,6 @@ assertCacheUsage(false /*multiPlanning*/,
 // replanning the cached query.
 verifyCorrectLookupAlgorithmUsed("HashJoin", avoidReplanLookupPipeline, {allowDiskUse: true});
 
-// TODO(SERVER-65345): When the SBE plan cache is enabled, we will encode the 'allowDiskUse'
-//  option when constructing the plan cache key. As such, we will not be able to reuse the cache
-//  entry generated above to execute a HashJoin.
-if (checkSBEEnabled(db, ["featureFlagSbePlanCache"])) {
-    runLookupQuery({allowDiskUse: true});
-    assertCacheUsage(true /*multiPlanning*/,
-                     false /*activeCacheEntry*/,
-                     {b: 1} /*cachedIndex*/,
-                     avoidReplanLookupPipeline,
-                     {allowDiskUse: true});
-
-    runLookupQuery({allowDiskUse: true});
-    assertCacheUsage(true /*multiPlanning*/,
-                     true /*activeCacheEntry*/,
-                     {b: 1} /*cachedIndex*/,
-                     avoidReplanLookupPipeline,
-                     {allowDiskUse: true});
-}
-
 runLookupQuery({allowDiskUse: true});
 assertCacheUsage(false /*multiPlanning*/,
                  true /*activeCacheEntry*/,
@@ -388,15 +369,12 @@ let explain = coll.explain().aggregate(avoidReplanLookupPipeline);
 const eqLookupNodes = getAggPlanStages(explain, "EQ_LOOKUP");
 assert.eq(eqLookupNodes.length, 0, "expected no EQ_LOOKUP nodes; got " + tojson(explain));
 
-// TODO(SERVER-61507): When the SBE plan cache is enabled, we will end up creating a separate
-// plan cache entry for the non-pushed down $lookup plan. As such, we assert that we have two
-// cache entries for the same query hash.
 if (checkSBEEnabled(db, ["featureFlagSbePlanCache"])) {
     runLookupQuery();
     const profileObj = getLatestProfilerEntry(db, {op: "command", ns: coll.getFullName()});
     const matchingCacheEntries =
         coll.getPlanCache().list([{$match: {queryHash: profileObj.queryHash}}]);
-    assert.eq(2, matchingCacheEntries.length);
+    assert.eq(1, matchingCacheEntries.length);
 } else {
     // When the SBE plan cache is disabled, we will be able to reuse the same cache entry.
     runLookupQuery();
@@ -454,15 +432,12 @@ explain = coll.explain().aggregate(avoidReplanLookupPipeline);
 groupNodes = getAggPlanStages(explain, "GROUP");
 assert.eq(groupNodes.length, 0);
 
-// TODO(SERVER-61507): When the SBE plan cache is enabled, we will end up creating a separate
-// plan cache entry for the non-pushed down $lookup plan. As such, we assert that we have two
-// cache entries for the same query hash.
 if (checkSBEEnabled(db, ["featureFlagSbePlanCache"])) {
     runGroupQuery();
     const profileObj = getLatestProfilerEntry(db, {op: "command", ns: coll.getFullName()});
     const matchingCacheEntries =
         coll.getPlanCache().list([{$match: {queryHash: profileObj.queryHash}}]);
-    assert.eq(2, matchingCacheEntries.length);
+    assert.eq(1, matchingCacheEntries.length);
 } else {
     // When the SBE plan cache is disabled, we will be able to reuse the same cache entry.
     runGroupQuery();
diff --git a/jstests/noPassthrough/sbe_pipeline_plan_cache_key_reporting.js b/jstests/noPassthrough/sbe_pipeline_plan_cache_key_reporting.js
new file mode 100644
index 00000000000..5c3a0aa5401
--- /dev/null
+++ b/jstests/noPassthrough/sbe_pipeline_plan_cache_key_reporting.js
@@ -0,0 +1,144 @@
+/**
+ * Confirms that 'planCacheKey' and 'queryHash' are correctly reported when the query has $lookup
+ * and $query stages with enabled and disabled SBE Plan Cache.
+ */
+
+(function() {
+"use strict";
+
+load("jstests/libs/sbe_util.js");  // For checkSBEEnabled.
+
+const databaseName = "pipeline_plan_cache_key_reporting";
+
+function isSBEEnabled() {
+    const conn = MongoRunner.runMongod({});
+    try {
+        const db = conn.getDB(databaseName);
+        return checkSBEEnabled(db);
+    } finally {
+        MongoRunner.stopMongod(conn);
+    }
+}
+
+if (!isSBEEnabled()) {
+    jsTest.log("Skipping test because SBE is not enabled.");
+    return;
+}
+
+/**
+ * Driver function that creates mongod instances with specified parameters and run the given test
+ * cases.
+ * @param {*} params to be passed to mongod in format like { setParameter:
+ *     "featureFlagSbePlanCache=true"}
+ * @param {*} testCases a list of test cases where each test case is an object with 'setup(db)' and
+ *     'run(db, assertMessage)' functions.
+ * @returns results from 'testCase.run(db, assertMessage)'
+ */
+function runTests(params, testCases) {
+    let results = [];
+    const conn = MongoRunner.runMongod(params);
+    const db = conn.getDB(databaseName);
+
+    const assertMessage = `${tojson(params)}`;
+    try {
+        for (let testCase of testCases) {
+            testCase.setup(db);
+            results.push(testCase.run(db, assertMessage));
+        }
+    } finally {
+        MongoRunner.stopMongod(conn);
+    }
+    return results;
+}
+
+/**
+ * This function validates given explain and return and object with extracted and validated
+ * PlanCacheKey and QueryHash.
+ * @returns {planCacheKey, queryHash, explain}
+ */
+function processAndValidateExplain(explain, assertMessage) {
+    assert.neq(explain, null);
+    assert.eq(explain.explainVersion,
+              "2",
+              `[${assertMessage}] invalid explain version ${tojson(explain)}`);
+
+    const planCacheKey = explain.queryPlanner.planCacheKey;
+    validateKey(planCacheKey, `[${assertMessage}] Invalid planCacheKey: ${tojson(explain)}`);
+
+    const queryHash = explain.queryPlanner.queryHash;
+    validateKey(queryHash, `[${assertMessage}] Invalid queryHash: ${tojson(explain)}`);
+
+    return {planCacheKey, queryHash, explain};
+}
+
+/**
+ * Validates given 'key' (PlanCacheKey or QueryHash).
+ */
+function validateKey(key, assertMessage) {
+    assert.eq(typeof key, "string", assertMessage);
+    assert.gt(key.length, 0, assertMessage);
+}
+
+// 1. Create test cases for $lookup and $group stages.
+const lookupTestCase = {
+    setup: db => {
+        db.coll.drop();
+        assert.commandWorked(db.coll.createIndexes([{a: 1}, {a: 1, b: 1}]));
+
+        db.lookupColl.drop();
+        assert.commandWorked(db.lookupColl.createIndex({b: 1}));
+    },
+
+    run: (db, assertMessage) => {
+        const pipeline = [
+            {$lookup: {from: db.lookupColl.getName(), localField: "a", foreignField: "b", as: "w"}}
+        ];
+        const explain = db.coll.explain().aggregate(pipeline);
+        return processAndValidateExplain(explain, assertMessage);
+    },
+};
+
+const groupTestCase = {
+    setup: db => {
+        db.coll.drop();
+        assert.commandWorked(db.coll.insertOne({a: 1}));
+    },
+
+    run: (db, assertMessage) => {
+        const pipeline = [{
+            $group: {
+                _id: "$b",
+            }
+        }];
+        const explain = db.coll.explain().aggregate(pipeline);
+        return processAndValidateExplain(explain, assertMessage);
+    },
+};
+
+const testCases = [lookupTestCase, groupTestCase];
+
+// 2. Run the test cases with SBE Plan Cache Enabled.
+const sbeParams = {
+    setParameter: "featureFlagSbePlanCache=true"
+};
+const sbeKeys = runTests(sbeParams, testCases);
+assert.eq(testCases.length, sbeKeys.length);
+
+// 3. Run the test cases with SBE Plan Cache disabled.
+const classicParams = {
+    setParameter: "featureFlagSbePlanCache=false"
+};
+const classicKeys = runTests(classicParams, testCases);
+assert.eq(testCases.length, classicKeys.length);
+
+// 4. Validate that PlanCacheKeys and QueryHash are equal. They should be different once
+// SERVER-61507 is completed.
+for (let i = 0; i < sbeKeys.length; ++i) {
+    const sbe = sbeKeys[i];
+    const classic = classicKeys[i];
+
+    const message = `sbe=${tojson(sbe.explain)}, classic=${tojson(classic.explain)}`;
+    assert.eq(sbe.planCacheKey, classic.planCacheKey, message);
+    assert.eq(sbe.queryHash, classic.queryHash, message);
+}
+})();
diff --git a/jstests/noPassthrough/sbe_plan_cache_key_reporting.js b/jstests/noPassthrough/sbe_plan_cache_key_reporting.js
index 6dd57d377a5..41e242e872b 100644
--- a/jstests/noPassthrough/sbe_plan_cache_key_reporting.js
+++ b/jstests/noPassthrough/sbe_plan_cache_key_reporting.js
@@ -165,5 +165,65 @@ function assertQueryHashAndPlanCacheKey(sbe, classic) {
     assertQueryHashAndPlanCacheKey(sbe.attr, classic.attr);
 })();
 
+// Validate that a query with pushed down $lookup stage uses classic plan cache key encoding.
+(function validateLookupQueryHashMap() {
+    const lookupColl = db.lookupColl;
+    lookupColl.drop();
+    assert.commandWorked(lookupColl.createIndex({b: 1}));
+    const [sbe, classic] =
+        runTestAgainstSbeAndClassicEngines(
+            function(engine) {
+                const pipeline = [
+                    {
+                        $lookup:
+                        {
+                            from: lookupColl.getName(),
+                            localField: "a",
+                            foreignField: "b",
+                            as: "whatever"
+                        }
+                    }
+                ];
+                return coll.explain().aggregate(pipeline);
+            });
+
+    assert.neq(sbe, null);
+    assert.neq(classic, null);
+    assert.eq(sbe.explainVersion, "2", sbe);
+    assert.eq(classic.explainVersion, "1", classic);
+
+    // The query hashes and the plan cache keys ('the keys') are different now because
+    // 'internalQueryForceClassicEngine' flag is encoded into query shape, once this flag is removed
+    // from the query shape encoding the keys will be the same until SERVER-61507 is completed, then
+    // the keys will be different forever.
+    assertQueryHashAndPlanCacheKey(sbe.queryPlanner, classic.stages[0]["$cursor"].queryPlanner);
+})();
+
+// Validate that a query with pushed down $group stage uses classic plan cache key encoding.
+(function validateGroupQueryHashMap() {
+    const groupColl = db.groupColl;
+    groupColl.drop();
+    assert.commandWorked(groupColl.insertOne({b: 1}));
+    const [sbe, classic] = runTestAgainstSbeAndClassicEngines(function(engine) {
+        const pipeline = [{
+            $group: {
+                _id: "$b",
+            }
+        }];
+        return groupColl.explain().aggregate(pipeline);
+    });
+
+    assert.neq(sbe, null);
+    assert.neq(classic, null);
+    assert.eq(sbe.explainVersion, "2", sbe);
+    assert.eq(classic.explainVersion, "1", classic);
+
+    // The query hashes and the plan cache keys ('the keys') are different now because
+    // 'internalQueryForceClassicEngine' flag is encoded into query shape, once this flag is removed
+    // from the query shape encoding the keys will be the same until SERVER-61507 is completed, then
+    // the keys will be different forever.
+    assertQueryHashAndPlanCacheKey(sbe.queryPlanner, classic.stages[0]["$cursor"].queryPlanner);
+})();
+
 MongoRunner.stopMongod(conn);
 }());