SERVER-122449 Disable $bit* expressions in SBE and add knob to limit their memory use (#50442)

GitOrigin-RevId: b2ff61d350fe6abd50f9c51e3fea40bc5d6579f0

jstests/core/query/sbe/sbe_plan_cache_autoparameterize_collscan.js

@@ -230,52 +230,6 @@ runTest(
     true,
 );
 
-// Test basic auto-parameterization of $bitsAllClear.
-runTest(
-    {query: {a: {$bitsAllClear: [0, 3]}}, projection: {_id: 1}},
-    [{_id: 1}, {_id: 3}, {_id: 4}, {_id: 5}, {_id: 16}],
-    {query: {a: {$bitsAllClear: [0, 2, 65]}}, projection: {_id: 1}},
-    [{_id: 1}, {_id: 6}, {_id: 16}],
-    true,
-);
-
-// Test basic auto-parameterization of $bitsAllSet.
-runTest(
-    {query: {a: {$bitsAllSet: [0, 2]}}, projection: {_id: 1}},
-    [{_id: 5}, {_id: 6}],
-    {query: {a: {$bitsAllSet: [0, 1]}}, projection: {_id: 1}},
-    [{_id: 2}, {_id: 5}, {_id: 6}],
-    true,
-);
-
-// Test basic auto-parameterization of $bitsAnyClear.
-runTest(
-    {query: {a: {$bitsAnyClear: 1}}, projection: {_id: 1}},
-    [{_id: 1}, {_id: 3}, {_id: 4}, {_id: 5}, {_id: 6}, {_id: 16}],
-    {query: {a: {$bitsAnyClear: 3}}, projection: {_id: 1}},
-    [{_id: 0}, {_id: 1}, {_id: 3}, {_id: 4}, {_id: 5}, {_id: 6}, {_id: 16}],
-    true,
-);
-
-// Test basic auto-parameterization of $bitsAnySet.
-runTest(
-    {query: {a: {$bitsAnySet: 1}}, projection: {_id: 1}},
-    [{_id: 0}, {_id: 2}, {_id: 5}, {_id: 6}],
-    {query: {a: {$bitsAnySet: 3}}, projection: {_id: 1}},
-    [{_id: 0}, {_id: 1}, {_id: 2}, {_id: 5}, {_id: 6}],
-    true,
-);
-
-// Auto-parameterization of bit-test operators should work even if looking past 64 bits is required
-// in order to match against binary data.
-runTest(
-    {query: {a: {$bitsAllSet: [0, 94]}}, projection: {_id: 1}},
-    [],
-    {query: {a: {$bitsAllSet: [88, 89, 90, 91, 92, 93]}}, projection: {_id: 1}},
-    [{_id: 16}],
-    true,
-);
-
 // Test auto-parameterization of $elemMatch object.
 runTest(
     {query: {a: {$elemMatch: {b: {$gt: 3, $lt: 5}}}}, projection: {_id: 1}},

jstests/noPassthrough/query/bit_test_disabled_in_sbe.js

@@ -0,0 +1,29 @@
+/**
+ * Tests that $bit* match expressions ($bitsAllSet, $bitsAllClear, $bitsAnySet, $bitsAnyClear) are
+ * executed using the SBE engine when in trySbeEngine or featureFlagSbeFull mode, and use the
+ * classic engine otherwise (forceClassicEngine or trySbeRestricted).
+ */
+import {getEngine} from "jstests/libs/query/analyze_plan.js";
+import {checkSbeFullyEnabled} from "jstests/libs/query/sbe_util.js";
+
+const conn = MongoRunner.runMongod();
+assert.neq(conn, null, "mongod failed to start up");
+const db = conn.getDB(jsTestName());
+const coll = db.bit_test_not_sbe;
+coll.drop();
+
+assert.commandWorked(coll.insert({x: 7}));
+
+const expectSbe = checkSbeFullyEnabled(db);
+
+for (const op of ["$bitsAllSet", "$bitsAllClear", "$bitsAnySet", "$bitsAnyClear"]) {
+    const explain = coll.find({x: {[op]: [0, 1, 2]}}).explain();
+    const expectedEngine = expectSbe ? "sbe" : "classic";
+    assert.eq(
+        getEngine(explain),
+        expectedEngine,
+        `expected ${op} to use ${expectedEngine} engine but got: ${tojson(explain)}`,
+    );
+}
+
+MongoRunner.stopMongod(conn);

jstests/noPassthrough/query/bit_test_max_positions.js

@@ -0,0 +1,51 @@
+/**
+ * Tests that internalQueryMaxBitTestIntermediatePositions limits the number of set bits accepted
+ * by a $bitTest expression ($bitsAllSet, $bitsAllClear, $bitsAnySet, $bitsAnyClear) when the
+ * bitmask is supplied as BinData.
+ */
+const kLimitExceededCode = 12244901;
+
+const conn = MongoRunner.runMongod();
+assert.neq(conn, null, "mongod failed to start up");
+const db = conn.getDB(jsTestName());
+const coll = db.bit_test_max_positions;
+coll.drop();
+
+assert.commandWorked(coll.insert({x: 7}));
+
+// Set a limit of 64 set bits (the minimum allowed value).
+assert.commandWorked(db.adminCommand({setParameter: 1, internalQueryMaxBitTestIntermediatePositions: 64}));
+
+// 8 bytes of 0xFF = exactly 64 set bits. Should succeed for all operators.
+const sixtyFourBits = BinData(0, "//////////8=");
+for (const op of ["$bitsAllSet", "$bitsAllClear", "$bitsAnySet", "$bitsAnyClear"]) {
+    assert.doesNotThrow(
+        () => coll.find({x: {[op]: sixtyFourBits}}).itcount(),
+        [],
+        `expected ${op} with 64-bit mask to succeed`,
+    );
+}
+
+// 8 bytes of 0xFF + 0x01 = 65 set bits, exceeds the limit of 64. Should fail.
+const sixtyFiveBits = BinData(0, "//////////8B");
+for (const op of ["$bitsAllSet", "$bitsAllClear", "$bitsAnySet", "$bitsAnyClear"]) {
+    assert.commandFailedWithCode(
+        db.runCommand({find: coll.getName(), filter: {x: {[op]: sixtyFiveBits}}}),
+        kLimitExceededCode,
+        `expected ${op} with 65-bit mask to fail`,
+    );
+}
+
+// Values below 64 are invalid since a 64-bit integer can have up to 64 set bits.
+assert.commandFailedWithCode(
+    db.adminCommand({setParameter: 1, internalQueryMaxBitTestIntermediatePositions: 63}),
+    ErrorCodes.BadValue,
+    "expected setting internalQueryMaxBitTestIntermediatePositions below 64 to fail",
+);
+assert.commandFailedWithCode(
+    db.adminCommand({setParameter: 1, internalQueryMaxBitTestIntermediatePositions: 0}),
+    ErrorCodes.BadValue,
+    "expected setting internalQueryMaxBitTestIntermediatePositions to 0 to fail",
+);
+
+MongoRunner.stopMongod(conn);

src/mongo/db/matcher/expression_leaf.cpp

@@ -46,6 +46,7 @@
 #include "mongo/db/query/query_execution_knobs_gen.h"
 #include "mongo/db/query/query_integration_knobs_gen.h"
 #include "mongo/db/query/query_optimization_knobs_gen.h"
+#include "mongo/logv2/log.h"
 #include "mongo/util/errno_util.h"
 #include "mongo/util/pcre.h"
 #include "mongo/util/pcre_util.h"
@@ -55,6 +56,8 @@
 #include <cmath>
 #include <memory>
 
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery
+
 namespace mongo {
 
 template <typename T>
@@ -475,6 +478,7 @@ BitTestMatchExpression::BitTestMatchExpression(MatchType type,
                                                uint32_t bitMaskLen,
                                                clonable_ptr<ErrorAnnotation> annotation)
     : LeafMatchExpression(type, path, std::move(annotation)) {
+    const auto maxPositions = internalQueryMaxBitTestIntermediatePositions.load();
     for (uint32_t byte = 0; byte < bitMaskLen; byte++) {
         char byteAt = bitMaskBinary[byte];
         if (!byteAt) {
@@ -493,7 +497,18 @@ BitTestMatchExpression::BitTestMatchExpression(MatchType type,
 
         for (int bit = 0; bit < 8; bit++) {
             if (byteAt & (1 << bit)) {
+                uassert(12244901,
+                        str::stream() << "BinData bitmask for " << name()
+                                      << " has too many set bits; maximum is " << maxPositions
+                                      << " (controlled by "
+                                         "internalQueryMaxBitTestIntermediatePositions)",
+                        _bitPositions.size() < static_cast<size_t>(maxPositions));
                 _bitPositions.push_back(8 * byte + bit);
+                if (_bitPositions.size() == 100001) {
+                    LOGV2(12244900,
+                          "Creating large bitPosition vector",
+                          "size"_attr = _bitPositions.size());
+                }
             }
         }
     }

src/mongo/db/query/compiler/parsers/matcher/expression_parser.cpp

@@ -867,6 +867,7 @@ StatusWithMatchExpression parseBitTest(boost::optional<StringData> name,
                           << name << " takes an Array, a number, or a BinData but received: " << e);
     }
 
+    expCtx->setSbeCompatibility(SbeCompatibility::requiresTrySbe);
     return {std::move(bitTestMatchExpression)};
 }
 

src/mongo/db/query/query_execution_knobs.idl

@@ -962,6 +962,20 @@ server_parameters:
             gte: 0
         redact: false
 
+    internalQueryMaxBitTestIntermediatePositions:
+        description: >-
+            Maximum number of bit positions that a $bitTest expression (bitsAllSet, bitsAllClear,
+            bitsAnySet, bitsAnyClear) may accumulate when constructed from a BinData mask. This
+            limits peak memory usage when parsing queries with very large binary bitmasks.
+        set_at: [startup, runtime]
+        cpp_varname: "internalQueryMaxBitTestIntermediatePositions"
+        cpp_vartype: AtomicWord<int>
+        default:
+            expr: std::numeric_limits<int>::max()
+        validator:
+            gte: 64
+        redact: false
+
     internalQueryEnableWriteConflictBackoffWithoutTicket:
         description: >-
             Enables behavior where the query engine will release resources before sleeping for write