SERVER-59613 $range expression should error if it exceeds memory limit

(cherry picked from commit ca2053d172076c106c994564b92109d1a973ed81)

6 files changed, 39 insertions, 1 deletions

diff --git a/jstests/aggregation/range.js b/jstests/aggregation/range.js
index 7ac445c16df..ea7ffdfa0b6 100644
--- a/jstests/aggregation/range.js
+++ b/jstests/aggregation/range.js
@@ -311,4 +311,12 @@
             .toArray();
 
     assert(arrayEq(decimalRangeExpectedResult, decimalRangeResult));
+
+    assert(coll.drop());
+    assert.commandWorked(coll.insertOne({_id: 1}));
+    assertErrorCode(
+        coll, [{$project: {result: {$range: [0, 1073741924]}}}], ErrorCodes.ExceededMemoryLimit);
+    assert(
+        arrayEq([{_id: 1, result: []}],
+                coll.aggregate([{$project: {result: {$range: [0, 1073741924, -1]}}}]).toArray()));
 }());
diff --git a/jstests/noPassthrough/query_knobs_validation.js b/jstests/noPassthrough/query_knobs_validation.js
index abe7869d666..edc5ce81d74 100644
--- a/jstests/noPassthrough/query_knobs_validation.js
+++ b/jstests/noPassthrough/query_knobs_validation.js
@@ -34,6 +34,7 @@
         internalLookupStageIntermediateDocumentMaxSizeBytes: 100 * 1024 * 1024,
         internalQueryMaxPushBytes: 100 * 1024 * 1024,
         internalQueryMaxAddToSetBytes: 100 * 1024 * 1024,
+        internalQueryMaxRangeBytes: 100 * 1024 * 1024,
         // Should be half the value of 'internalQueryExecYieldIterations' parameter.
         internalInsertMaxBatchSize: 64,
         internalQueryPlannerGenerateCoveredWholeIndexScans: false,
diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index 0e5744cce40..79179655cdd 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -142,6 +142,9 @@ env.Library(
         'dependencies',
         'document_value',
         'expression_context',
+    ],
+    LIBDEPS_PRIVATE=[
+        '$BUILD_DIR/mongo/db/query/query_knobs',
     ]
 )
 
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 8daf8c24f2a..0aabae477c7 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -44,6 +44,7 @@
 #include "mongo/db/pipeline/expression_context.h"
 #include "mongo/db/pipeline/value.h"
 #include "mongo/db/query/datetime/date_time_support.h"
+#include "mongo/db/query/query_knobs.h"
 #include "mongo/platform/bits.h"
 #include "mongo/platform/decimal128.h"
 #include "mongo/util/mongoutils/str.h"
@@ -3640,10 +3641,25 @@ Value ExpressionRange::evaluate(const Document& root, Variables* variables) cons
         uassert(34449, "$range requires a non-zero step value", step != 0);
     }
 
+    // Calculate how much memory is needed to generate the array and avoid going over the memLimit.
+    auto steps = (end - current) / step;
+    // If steps not positive then no amount of steps can get you from start to end. For example
+    // with start=5, end=7, step=-1 steps would be negative and in this case we would return an
+    // empty array.
+    auto length = steps >= 0 ? 1 + steps : 0;
+    int64_t memNeeded = sizeof(std::vector<Value>) + length * startVal.getApproximateSize();
+    auto memLimit = internalQueryMaxRangeBytes.load();
+    uassert(ErrorCodes::ExceededMemoryLimit,
+            str::stream() << "$range would use too much memory (" << memNeeded << " bytes) "
+                          << "and cannot spill to disk. Memory limit: "
+                          << memLimit
+                          << " bytes",
+            memNeeded < memLimit);
+
     std::vector<Value> output;
 
     while ((step > 0 ? current < end : current > end)) {
-        output.push_back(Value(static_cast<int>(current)));
+        output.emplace_back(static_cast<int>(current));
         current += step;
     }
 
diff --git a/src/mongo/db/query/query_knobs.cpp b/src/mongo/db/query/query_knobs.cpp
index 2cd2c61f022..207d1f89114 100644
--- a/src/mongo/db/query/query_knobs.cpp
+++ b/src/mongo/db/query/query_knobs.cpp
@@ -137,6 +137,14 @@ MONGO_EXPORT_SERVER_PARAMETER(internalQueryMaxAddToSetBytes, int, 100 * 1024 * 1
         return Status::OK();
     });
 
+MONGO_EXPORT_SERVER_PARAMETER(internalQueryMaxRangeBytes, int, 100 * 1024 * 1024)
+    ->withValidator([](const int& newVal) {
+        if (newVal <= 0) {
+            return Status(ErrorCodes::BadValue, "internalQueryMaxRangeBytes must be positive");
+        }
+        return Status::OK();
+    });
+
 MONGO_EXPORT_SERVER_PARAMETER(internalQueryExplainSizeThresholdBytes, int, 10 * 1024 * 1024)
     ->withValidator([](const int& newVal) {
         if (newVal <= 0) {
diff --git a/src/mongo/db/query/query_knobs.h b/src/mongo/db/query/query_knobs.h
index 6165be149f8..8cc699012d8 100644
--- a/src/mongo/db/query/query_knobs.h
+++ b/src/mongo/db/query/query_knobs.h
@@ -157,6 +157,8 @@ extern AtomicInt32 internalQueryMaxPushBytes;
 
 extern AtomicInt32 internalQueryMaxAddToSetBytes;
 
+extern AtomicInt32 internalQueryMaxRangeBytes;
+
 // The number of bytes after which explain should start truncating portions of its output.
 extern AtomicInt32 internalQueryExplainSizeThresholdBytes;
 }  // namespace mongo