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

(cherry picked from commit 3f7016e2c2e746ff4a009750044a9bd14cf6fe6f)

5 files changed, 36 insertions, 1 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index 4d66446caf4..40d48f9170b 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -171,6 +171,9 @@ all:
     test_file: jstests/sharding/query/pipeline_length_limit.js
   - ticket: SERVER-57321
     test_file: jstests/core/mod_special_values.js
+  - ticket: SERVER-59613
+    test_file: jstests/aggregation/range.js
+
 
 suites:
 
diff --git a/jstests/aggregation/range.js b/jstests/aggregation/range.js
index 8041d0ab53c..d8d644a59dc 100644
--- a/jstests/aggregation/range.js
+++ b/jstests/aggregation/range.js
@@ -296,4 +296,11 @@ const decimalRangeResult =
         .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 e820280dc77..01a03a3d94f 100644
--- a/jstests/noPassthrough/query_knobs_validation.js
+++ b/jstests/noPassthrough/query_knobs_validation.js
@@ -38,6 +38,7 @@ const expectedParamDefaults = {
     internalPipelineLengthLimit: 2147483647,  // INT_MAX
     internalQueryMaxJsEmitBytes: 100 * 1024 * 1024,
     internalQueryMaxPushBytes: 100 * 1024 * 1024,
+    internalQueryMaxRangeBytes: 100 * 1024 * 1024,
     internalQueryMaxAddToSetBytes: 100 * 1024 * 1024,
     // Should be half the value of 'internalQueryExecYieldIterations' parameter.
     internalInsertMaxBatchSize: 500,
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index cae4fd3c854..09c49685f88 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -3796,10 +3796,23 @@ 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.idl b/src/mongo/db/query/query_knobs.idl
index 7eaba46f916..218f56af824 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -368,6 +368,17 @@ server_parameters:
     validator:
         gt: 0
 
+  internalQueryMaxRangeBytes:
+    description: "Limits the vector of values pushed into a single array while generating $range
+        result."
+    set_at: [ startup, runtime ]
+    cpp_varname: "internalQueryMaxRangeBytes"
+    cpp_vartype: AtomicWord<int>
+    default:
+      expr: 100 * 1024 * 1024
+    validator:
+        gt: 0
+
   internalQueryMaxAddToSetBytes:
     description: "Limits the vector of values pushed into a single array while grouping with the
         $addToSet accumulator."