SERVER-37132 Negation of $in with regex can incorrectly plan from the cache, leading to missing query results

(cherry picked from commit e786e3a313b75a1fe8aa233ed09da2d2efbaf613)

3 files changed, 96 insertions, 0 deletions

diff --git a/jstests/noPassthroughWithMongod/plan_cache_not_in_regex.js b/jstests/noPassthroughWithMongod/plan_cache_not_in_regex.js
new file mode 100644
index 00000000000..1343c40e8ec
--- /dev/null
+++ b/jstests/noPassthroughWithMongod/plan_cache_not_in_regex.js
@@ -0,0 +1,55 @@
+/**
+ * Tests that a $not-$in-$regex query, which cannot be supported by an index, cannot incorrectly
+ * hijack the cached plan for an earlier $not-$in query.
+ */
+(function() {
+    "use strict";
+
+    load('jstests/libs/analyze_plan.js');  // For isCollScan.
+
+    const coll = db.plan_cache_not_in_regex;
+    coll.drop();
+
+    // Helper function which obtains the cached plan, if any, for a given query shape.
+    function getPlanForCacheEntry(query, proj, sort) {
+        const cacheDetails =
+            coll.getPlanCache().getPlansByQuery({query: query, sort: sort, projection: proj});
+        const plans = cacheDetails.filter(plan => plan.feedback.nfeedback >= 0);
+        assert.eq(plans.length, 1, `Expected one cached plan, found: ${tojson(plans)}`);
+        return plans.shift();
+    }
+
+    // Insert a document containing a field 'a', and create two indexes that can support queries on
+    // this field. This is to ensure that the plan we choose will be cached, since if only a single
+    // index is available, the solution will not be cached.
+    assert.writeOK(coll.insert({a: "foo"}));
+    assert.commandWorked(coll.createIndex({a: 1}));
+    assert.commandWorked(coll.createIndex({a: 1, b: 1}));
+
+    // Repeat the test for query, query with projection, and query with projection and sort.
+    for (let [proj, sort] of[[{}, {}], [{_id: 0, a: 1}, {}], [{_id: 0, a: 1}, {a: 1}]]) {
+        // Perform a plain $not-$in query on 'a' and confirm that the plan is cached.
+        const queryShape = {a: {$not: {$in: [32, 33]}}};
+        assert.eq(1, coll.find(queryShape, proj).sort(sort).itcount());
+        let cacheEntry = getPlanForCacheEntry(queryShape, proj, sort);
+        assert(cacheEntry);
+
+        // If the cached plan is inactive, perform the same query to activate it.
+        if (cacheEntry.isActive === false) {
+            assert.eq(1, coll.find(queryShape, proj).sort(sort).itcount());
+            cacheEntry = getPlanForCacheEntry(queryShape, proj, sort);
+            assert(cacheEntry);
+            assert(cacheEntry.isActive);
+        }
+
+        // Now perform a $not-$in-$regex query, confirm that it obtains the correct results, and
+        // that it used a COLLSCAN rather than planning from the cache.
+        const explainOutput = assert.commandWorked(
+            coll.find({a: {$not: {$in: [34, /bar/]}}}).explain("executionStats"));
+        assert(isCollscan(explainOutput.queryPlanner.winningPlan));
+        assert.eq(1, explainOutput.executionStats.nReturned);
+
+        // Flush the plan cache before the next iteration.
+        coll.getPlanCache().clear();
+    }
+})();
\ No newline at end of file
diff --git a/src/mongo/db/query/plan_cache.cpp b/src/mongo/db/query/plan_cache.cpp
index d34df64ab88..62a3dcc33fc 100644
--- a/src/mongo/db/query/plan_cache.cpp
+++ b/src/mongo/db/query/plan_cache.cpp
@@ -505,6 +505,13 @@ void PlanCache::encodeKeyForMatch(const MatchExpression* tree, StringBuilder* ke
         encodeGeoNearMatchExpression(static_cast<const GeoNearMatchExpression*>(tree), keyBuilder);
     }
 
+    // We must encode the presence of any regular expressions within MATCH_IN to ensure correctness.
+    if (MatchExpression::MATCH_IN == tree->matchType()) {
+        if (!static_cast<const InMatchExpression*>(tree)->getRegexes().empty()) {
+            encodeUserString("_re"_sd, keyBuilder);
+        }
+    }
+
     // Encode indexability.
     const IndexToDiscriminatorMap& discriminators =
         _indexabilityState.getDiscriminators(tree->path());
diff --git a/src/mongo/db/query/plan_cache_test.cpp b/src/mongo/db/query/plan_cache_test.cpp
index bac9704e0cf..5be65b1b27d 100644
--- a/src/mongo/db/query/plan_cache_test.cpp
+++ b/src/mongo/db/query/plan_cache_test.cpp
@@ -1392,6 +1392,40 @@ TEST(PlanCacheTest, ComputeKeyGeoNear) {
         "gnanrsp");
 }
 
+TEST(PlanCacheTest, ComputeKeyMatchInDependsOnPresenceOfRegex) {
+    // Test that an $in containing a single regex is unwrapped to $regex.
+    testComputeKey("{a: {$in: [/foo/]}}", "{}", "{}", "rea");
+    testComputeKey("{a: {$in: [/foo/i]}}", "{}", "{}", "rea");
+
+    // Test that an $in with no regexes does not include any regex information.
+    testComputeKey("{a: {$in: [1, 'foo']}}", "{}", "{}", "ina");
+
+    // Test that an $in with a regex encodes the presence of the regex.
+    testComputeKey("{a: {$in: [1, /foo/]}}", "{}", "{}", "ina_re");
+
+    // Test that an $in with a regex and flags encodes the presence of the regex.
+    testComputeKey("{a: {$in: [1, /foo/is]}}", "{}", "{}", "ina_re");
+    testComputeKey("{a: {$in: [1, /foo/i]}}", "{}", "{}", "ina_re");
+
+    // Test that an $in with multiple regexes encodes their presence only once.
+    testComputeKey("{a: {$in: [1, /foo/i, /bar/m, /baz/s]}}", "{}", "{}", "ina_re");
+    testComputeKey(
+        "{a: {$in: [1, /foo/i, /bar/m, /baz/s, /qux/i, /quux/s]}}", "{}", "{}", "ina_re");
+    testComputeKey(
+        "{a: {$in: [1, /foo/ism, /bar/msi, /baz/im, /qux/si, /quux/im]}}", "{}", "{}", "ina_re");
+    testComputeKey(
+        "{a: {$in: [1, /foo/msi, /bar/ism, /baz/is, /qux/mi, /quux/im]}}", "{}", "{}", "ina_re");
+
+    // Test that $not-$in-$regex similarly records the presence of any regexes.
+    testComputeKey("{a: {$not: {$in: [1, 'foo']}}}", "{}", "{}", "nt[ina]");
+    testComputeKey("{a: {$not: {$in: [1, /foo/]}}}", "{}", "{}", "nt[ina_re]");
+    testComputeKey("{a: {$not: {$in: [1, /foo/i, /bar/i, /baz/msi]}}}", "{}", "{}", "nt[ina_re]");
+
+    // Test that a $not-$in containing a single regex is unwrapped to $not-$regex.
+    testComputeKey("{a: {$not: {$in: [/foo/]}}}", "{}", "{}", "nt[rea]");
+    testComputeKey("{a: {$not: {$in: [/foo/i]}}}", "{}", "{}", "nt[rea]");
+}
+
 // When a sparse index is present, computeKey() should generate different keys depending on
 // whether or not the predicates in the given query can use the index.
 TEST(PlanCacheTest, ComputeKeySparseIndex) {