SERVER-39764 fix bug where $in is planned from cache incorrectly

4 files changed, 123 insertions, 4 deletions

diff --git a/src/mongo/db/query/plan_cache.cpp b/src/mongo/db/query/plan_cache.cpp
index 6422cb48174..4defd1cdd20 100644
--- a/src/mongo/db/query/plan_cache.cpp
+++ b/src/mongo/db/query/plan_cache.cpp
@@ -47,6 +47,7 @@
 #include "mongo/db/query/canonical_query_encoder.h"
 #include "mongo/db/query/collation/collator_interface.h"
 #include "mongo/db/query/plan_ranker.h"
+#include "mongo/db/query/planner_ixselect.h"
 #include "mongo/db/query/query_knobs_gen.h"
 #include "mongo/db/query/query_solution.h"
 #include "mongo/util/assert_util.h"
@@ -85,6 +86,12 @@ void encodeIndexability(const MatchExpression* tree,
 
             *keyBuilder << kEncodeDiscriminatorsEnd;
         }
+    } else if (tree->matchType() == MatchExpression::MatchType::NOT) {
+        // If the node is not compatible with any type of index, add a single '0' discriminator
+        // here. Otherwise add a '1'.
+        *keyBuilder << kEncodeDiscriminatorsBegin;
+        *keyBuilder << QueryPlannerIXSelect::logicalNodeMayBeSupportedByAnIndex(tree);
+        *keyBuilder << kEncodeDiscriminatorsEnd;
     }
 
     for (size_t i = 0; i < tree->numChildren(); ++i) {
diff --git a/src/mongo/db/query/plan_cache_test.cpp b/src/mongo/db/query/plan_cache_test.cpp
index b8903b06895..903da5cdda2 100644
--- a/src/mongo/db/query/plan_cache_test.cpp
+++ b/src/mongo/db/query/plan_cache_test.cpp
@@ -2028,4 +2028,104 @@ TEST(PlanCacheTest, StableKeyDoesNotChangeAcrossIndexCreation) {
     ASSERT_EQ(postIndexKey.getUnstablePart(), "<1>");
 }
 
+TEST(PlanCacheTest, ComputeKeyNotEqualsArray) {
+    PlanCache planCache;
+    unique_ptr<CanonicalQuery> cqNeArray(canonicalize("{a: {$ne: [1]}}"));
+    unique_ptr<CanonicalQuery> cqNeScalar(canonicalize("{a: {$ne: 123}}"));
+
+    const PlanCacheKey noIndexNeArrayKey = planCache.computeKey(*cqNeArray);
+    const PlanCacheKey noIndexNeScalarKey = planCache.computeKey(*cqNeScalar);
+    ASSERT_EQ(noIndexNeArrayKey.getUnstablePart(), "<0>");
+    ASSERT_EQ(noIndexNeScalarKey.getUnstablePart(), "<1>");
+    ASSERT_EQ(noIndexNeScalarKey.getStableKey(), noIndexNeArrayKey.getStableKey());
+
+    const auto keyPattern = BSON("a" << 1);
+    // Create a normal btree index. It will have a discriminator.
+    planCache.notifyOfIndexUpdates(
+        {CoreIndexInfo(keyPattern,
+                       IndexNames::nameToType(IndexNames::findPluginName(keyPattern)),
+                       false,                          // sparse
+                       IndexEntry::Identifier{""})});  // name
+
+    const PlanCacheKey withIndexNeArrayKey = planCache.computeKey(*cqNeArray);
+    const PlanCacheKey withIndexNeScalarKey = planCache.computeKey(*cqNeScalar);
+
+    ASSERT_NE(noIndexNeArrayKey, withIndexNeArrayKey);
+    ASSERT_EQ(noIndexNeArrayKey.getStableKey(), withIndexNeArrayKey.getStableKey());
+
+    ASSERT_EQ(noIndexNeScalarKey.getStableKey(), withIndexNeScalarKey.getStableKey());
+    // There will be one discriminator for the $not and another for the leaf node ({$eq: 123}).
+    ASSERT_EQ(withIndexNeScalarKey.getUnstablePart(), "<1><1>");
+    // There will be one discriminator for the $not and another for the leaf node ({$eq: [1]}).
+    // Since the index can support equality to an array, the second discriminator will have a value
+    // of '1'.
+    ASSERT_EQ(withIndexNeArrayKey.getUnstablePart(), "<0><1>");
+}
+
+TEST(PlanCacheTest, ComputeKeyNinArray) {
+    PlanCache planCache;
+    unique_ptr<CanonicalQuery> cqNinArray(canonicalize("{a: {$nin: [123, [1]]}}"));
+    unique_ptr<CanonicalQuery> cqNinScalar(canonicalize("{a: {$nin: [123, 456]}}"));
+
+    const PlanCacheKey noIndexNinArrayKey = planCache.computeKey(*cqNinArray);
+    const PlanCacheKey noIndexNinScalarKey = planCache.computeKey(*cqNinScalar);
+    ASSERT_EQ(noIndexNinArrayKey.getUnstablePart(), "<0>");
+    ASSERT_EQ(noIndexNinScalarKey.getUnstablePart(), "<1>");
+    ASSERT_EQ(noIndexNinScalarKey.getStableKey(), noIndexNinArrayKey.getStableKey());
+
+    const auto keyPattern = BSON("a" << 1);
+    // Create a normal btree index. It will have a discriminator.
+    planCache.notifyOfIndexUpdates(
+        {CoreIndexInfo(keyPattern,
+                       IndexNames::nameToType(IndexNames::findPluginName(keyPattern)),
+                       false,                          // sparse
+                       IndexEntry::Identifier{""})});  // name
+
+    const PlanCacheKey withIndexNinArrayKey = planCache.computeKey(*cqNinArray);
+    const PlanCacheKey withIndexNinScalarKey = planCache.computeKey(*cqNinScalar);
+
+    // The unstable part of the key for $nin: [<array>] should have changed. The stable part,
+    // however, should not.
+    ASSERT_EQ(noIndexNinArrayKey.getStableKey(), withIndexNinArrayKey.getStableKey());
+    ASSERT_NE(noIndexNinArrayKey.getUnstablePart(), withIndexNinArrayKey.getUnstablePart());
+
+    ASSERT_EQ(noIndexNinScalarKey.getStableKey(), withIndexNinScalarKey.getStableKey());
+    ASSERT_EQ(withIndexNinArrayKey.getUnstablePart(), "<0><1>");
+    ASSERT_EQ(withIndexNinScalarKey.getUnstablePart(), "<1><1>");
+}
+
+// Test for a bug which would be easy to introduce. If we only inserted discriminators for some
+// nodes, we would have a problem. For example if our "stable" key was:
+// (or[nt[eqa],nt[eqa]])
+// And there was just one discriminator:
+// <0>
+
+// Whether the discriminator referred to the first not-eq node or the second would be
+// ambiguous. This would make it possible for two queries with different shapes (and different
+// plans) to get the same plan cache key. We test that this does not happen for a simple example.
+TEST(PlanCacheTest, PlanCacheKeyCollision) {
+    PlanCache planCache;
+    unique_ptr<CanonicalQuery> cqNeA(canonicalize("{$or: [{a: {$ne: 5}}, {a: {$ne: [12]}}]}"));
+    unique_ptr<CanonicalQuery> cqNeB(canonicalize("{$or: [{a: {$ne: [12]}}, {a: {$ne: 5}}]}"));
+
+    const PlanCacheKey keyA = planCache.computeKey(*cqNeA);
+    const PlanCacheKey keyB = planCache.computeKey(*cqNeB);
+    ASSERT_EQ(keyA.getStableKey(), keyB.getStableKey());
+    ASSERT_NE(keyA.getUnstablePart(), keyB.getUnstablePart());
+
+    const auto keyPattern = BSON("a" << 1);
+    // Create a normal btree index. It will have a discriminator.
+    planCache.notifyOfIndexUpdates(
+        {CoreIndexInfo(keyPattern,
+                       IndexNames::nameToType(IndexNames::findPluginName(keyPattern)),
+                       false,                          // sparse
+                       IndexEntry::Identifier{""})});  // name
+
+    const PlanCacheKey keyAWithIndex = planCache.computeKey(*cqNeA);
+    const PlanCacheKey keyBWithIndex = planCache.computeKey(*cqNeB);
+
+    ASSERT_EQ(keyAWithIndex.getStableKey(), keyBWithIndex.getStableKey());
+    ASSERT_NE(keyAWithIndex.getUnstablePart(), keyBWithIndex.getUnstablePart());
+}
+
 }  // namespace
diff --git a/src/mongo/db/query/planner_ixselect.cpp b/src/mongo/db/query/planner_ixselect.cpp
index e0e8188d4ed..7383659e7f6 100644
--- a/src/mongo/db/query/planner_ixselect.cpp
+++ b/src/mongo/db/query/planner_ixselect.cpp
@@ -55,7 +55,9 @@ namespace {
 
 namespace wcp = ::mongo::wildcard_planning;
 
-bool isNotEqualsArrayOrNotInArray(const MatchExpression* me) {
+// Can't index negations of {$eq: <Array>} or {$in: [<Array>, ...]}. Note that we could
+// use the index in principle, though we would need to generate special bounds.
+bool isEqualsArrayOrNotInArray(const MatchExpression* me) {
     const auto type = me->matchType();
     if (type == MatchExpression::EQ) {
         return static_cast<const ComparisonMatchExpression*>(me)->getData().type() ==
@@ -433,9 +435,7 @@ bool QueryPlannerIXSelect::_compatible(const BSONElement& keyPatternElt,
                 return false;
             }
 
-            // Can't index negations of {$eq: <Array>} or {$in: [<Array>, ...]}. Note that we could
-            // use the index in principle, though we would need to generate special bounds.
-            if (isNotEqualsArrayOrNotInArray(child)) {
+            if (isEqualsArrayOrNotInArray(child)) {
                 return false;
             }
 
@@ -610,6 +610,11 @@ bool QueryPlannerIXSelect::nodeIsSupportedBySparseIndex(const MatchExpression* q
     return true;
 }
 
+bool QueryPlannerIXSelect::logicalNodeMayBeSupportedByAnIndex(const MatchExpression* queryExpr) {
+    return !(queryExpr->matchType() == MatchExpression::NOT &&
+             isEqualsArrayOrNotInArray(queryExpr->getChild(0)));
+}
+
 bool QueryPlannerIXSelect::nodeIsSupportedByWildcardIndex(const MatchExpression* queryExpr) {
     // Wildcard indexes only store index keys for "leaf" nodes in an object. That is, they do not
     // store keys for nested objects, meaning that any kind of comparison to an object or array
diff --git a/src/mongo/db/query/planner_ixselect.h b/src/mongo/db/query/planner_ixselect.h
index 93f4c4b9cd6..b4269f9b64d 100644
--- a/src/mongo/db/query/planner_ixselect.h
+++ b/src/mongo/db/query/planner_ixselect.h
@@ -146,6 +146,13 @@ public:
      */
     static bool nodeIsSupportedBySparseIndex(const MatchExpression* queryExpr, bool isInElemMatch);
 
+    /**
+     * Some types of matches are not supported by any type of index. If this function returns
+     * false, then 'queryExpr' is definitely not supported for any type of index. If the function
+     * returns true then 'queryExpr' may (or may not) be supported by some index.
+     */
+    static bool logicalNodeMayBeSupportedByAnIndex(const MatchExpression* queryExpr);
+
 private:
     /**
      * Used to keep track of if any $elemMatch predicates were encountered when walking a