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

(cherry picked from commit 348a869208b3be1ca9e04f0f998c7ff7e8586873)

6 files changed, 60 insertions, 22 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index 42c927a7497..efad6127b50 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -96,6 +96,8 @@ last-continuous:
       test_file: jstests/sharding/change_stream_show_migration_events.js
     - ticket: SERVER-58636
       test_file: jstests/replsets/initial_sync_replicate_drop_mid_secondary_batch.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:
@@ -352,6 +354,8 @@ last-lts:
       test_file: jstests/sharding/change_stream_show_migration_events.js
     - ticket: SERVER-58636
       test_file: jstests/replsets/initial_sync_replicate_drop_mid_secondary_batch.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 af5d3dc57f3..0859979f3b1 100644
--- a/jstests/aggregation/range.js
+++ b/jstests/aggregation/range.js
@@ -297,4 +297,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 f7c9f52183d..a647bc24df1 100644
--- a/jstests/noPassthrough/query_knobs_validation.js
+++ b/jstests/noPassthrough/query_knobs_validation.js
@@ -40,6 +40,7 @@ const expectedParamDefaults = {
     internalDocumentSourceSetWindowFieldsMaxMemoryBytes: 100 * 1024 * 1024,
     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/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index 5b46789312d..54e0af4a73c 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -48,6 +48,7 @@
 #include "mongo/db/index/btree_key_generator.h"
 #include "mongo/db/query/collation/collation_index_key.h"
 #include "mongo/db/query/datetime/date_time_support.h"
+#include "mongo/db/query/query_knobs_gen.h"
 #include "mongo/db/storage/key_string.h"
 #include "mongo/logv2/log.h"
 #include "mongo/util/fail_point.h"
@@ -1396,22 +1397,23 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinNewArrayFromRan
         return {false, value::TypeTags::Nothing, 0};
     }
 
-    auto isPositiveStep = stepVal > 0;
-
-    if (isPositiveStep) {
-        if (startVal < endVal) {
-            arr->reserve(1 + (endVal - startVal) / stepVal);
-            for (auto i = startVal; i < endVal; i += stepVal) {
-                arr->push_back(value::TypeTags::NumberInt32, value::bitcastTo<int32_t>(i));
-            }
-        }
-    } else {
-        if (startVal > endVal) {
-            arr->reserve(1 + (startVal - endVal) / (-stepVal));
-            for (auto i = startVal; i > endVal; i += stepVal) {
-                arr->push_back(value::TypeTags::NumberInt32, value::bitcastTo<int32_t>(i));
-            }
-        }
+    // Calculate how much memory is needed to generate the array and avoid going over the memLimit.
+    auto steps = (endVal - startVal) / stepVal;
+    // 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(value::Array) + length * value::getApproximateSize(startTag, start);
+    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);
+
+    arr->reserve(length);
+    for (auto i = startVal; stepVal > 0 ? i < endVal : i > endVal; i += stepVal) {
+        arr->push_back(value::TypeTags::NumberInt32, value::bitcastTo<int32_t>(i));
     }
 
     guard.reset();
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 8d8c1d27d98..f35cd8c81b3 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -4145,10 +4145,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 23caa71f114..3de5c9d773e 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -142,7 +142,7 @@ server_parameters:
 
   planCacheSize:
     description: "The maximum amount of memory that the system will allocate for the plan cache.
-      It takes value value in one of the two formats: 
+      It takes value value in one of the two formats:
       1. <number>% indicates a percentage of the physical memory available to the process. E.g.: 15%.
       2. <number>(MB|GB), indicates the amount of memory in MB or GB. E.g.: 1.5GB, 100MB.
       The defualt value is 5% which means 5% of the physical memory available to the process."
@@ -151,7 +151,7 @@ server_parameters:
     cpp_vartype: synchronized_value<std::string>
     default: "5%"
     on_update: plan_cache_util::onPlanCacheSizeUpdate
-    
+
   #
   # Parsing
   #
@@ -405,6 +405,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."
@@ -546,9 +557,9 @@ server_parameters:
         expr: 100 * 1024 * 1024
     validator:
         gt: 0
-        
+
   enableSearchMeta:
-    description: "Exists for backwards compatibility in startup parameters, 
+    description: "Exists for backwards compatibility in startup parameters,
       enabling this was required on 4.4 to access SEARCH_META variables. Does not do anything."
     set_at: [ startup, runtime ]
     cpp_varname: "enableSearchMeta"
@@ -570,7 +581,7 @@ server_parameters:
     default:  500000
     validator:
         gt: 0
- 
+
   internalDocumentSourceDensifyMaxMemoryBytes:
     description: "Limits the number of bytes densification can use to store partition information."
     set_at: [ startup, runtime ]