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

6 files changed, 39 insertions, 1 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index 568adaaa903..36553dae3a8 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -118,6 +118,8 @@ all:
     test_file: jstests/replsets/sessions_collection_reaping.js
   - ticket: SERVER-37904
     test_file: jstests/replsets/cluster_chaining_override.js
+  - ticket: SERVER-59613
+    test_file: jstests/aggregation/range.js
 
 # Tests that should only be excluded from particular suites should be listed under that suite.
 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 60fdb57bae7..f578153460f 100644
--- a/jstests/noPassthrough/query_knobs_validation.js
+++ b/jstests/noPassthrough/query_knobs_validation.js
@@ -40,6 +40,7 @@ const expectedParamDefaults = {
     internalDocumentSourceGroupMaxMemoryBytes: 100 * 1024 * 1024,
     internalPipelineLengthLimit: 2147483647,  // INT_MAX
     internalQueryMaxPushBytes: 100 * 1024 * 1024,
+    internalQueryMaxRangeBytes: 100 * 1024 * 1024,
     internalQueryMaxAddToSetBytes: 100 * 1024 * 1024,
     // Should be half the value of 'internalQueryExecYieldIterations' parameter.
     internalInsertMaxBatchSize: 64,
diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index 73bfdbdb404..01b96b37101 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -145,6 +145,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 66b9bdae9c2..51c0275efba 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -46,6 +46,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_gen.h"
 #include "mongo/platform/bits.h"
 #include "mongo/platform/decimal128.h"
 #include "mongo/util/regex_util.h"
@@ -3639,10 +3640,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 49550b175dd..6cea0503012 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -381,6 +381,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."