SERVER-96350 $push, $addToSet, $setUnion, $concatArrays window functions respect memory limits (#35136)

leozzx · MongoDB Bot · commit dabc243744d8 · 2025-04-22T18:12:53.000Z
GitOrigin-RevId: 2aeb354c71e4d3aa76d9eb403b936736969703d3
diff --git a/jstests/noPassthrough/query/agg/agg_configurable_memory_limits.js b/jstests/noPassthrough/query/agg/agg_configurable_memory_limits.js
@@ -1,18 +1,16 @@
-// @tags: [featureFlagAccuratePercentiles]
-
-import {
-    testLargeUniformDataset,
-    testLargeUniformDataset_WithInfinities,
-    testWithMultipleGroups,
-    testWithSingleGroup
-} from "jstests/aggregation/libs/percentiles_util.js";
-
 // Tests that certain aggregation operators have configurable memory limits.
 const conn = MongoRunner.runMongod();
 assert.neq(null, conn, "mongod was unable to start up");
 const db = conn.getDB("test");
 const coll = db.agg_configurable_memory_limit;
 
+// Function to change the parameter value and return the previous value.
+function setParam(param, val) {
+    const res = db.adminCommand({setParameter: 1, [param]: val});
+    assert.commandWorked(res);
+    return res.was;
+}
+
 // The approximate size of the strings below is 22-25 bytes, so configure one memory limit such that
 // 100 of these strings will surely exceed it but 24 of them won't, and another such that 24 will
 // exceed as well.
@@ -36,10 +34,23 @@ assert.commandWorked(bulk.execute());
     assert.doesNotThrow(() => coll.aggregate([{$group: {_id: null, strings: {$push: "$y"}}}]));
 
     // Now lower the limit to test that its configuration is obeyed.
-    assert.commandWorked(
-        db.adminCommand({setParameter: 1, internalQueryMaxPushBytes: memLimitArray}));
+    const originalVal = setParam('internalQueryMaxPushBytes', memLimitArray);
     assert.throwsWithCode(() => coll.aggregate([{$group: {_id: null, strings: {$push: "$y"}}}]),
                           ErrorCodes.ExceededMemoryLimit);
+    setParam('internalQueryMaxPushBytes', originalVal);
+}());
+
+(function testInternalQueryMaxPushBytesSettingWindowFunc() {
+    let pipeline = [
+        {$setWindowFields: {sortBy: {_id: 1}, output: {v: {$push: '$y'}}}},
+    ];
+    // Test that the default 100MB memory limit isn't reached with our data.
+    assert.doesNotThrow(() => coll.aggregate(pipeline));
+
+    // Now lower the limit to test that its configuration is obeyed.
+    const originalVal = setParam('internalQueryMaxPushBytes', memLimitArray);
+    assert.throwsWithCode(() => coll.aggregate(pipeline), ErrorCodes.ExceededMemoryLimit);
+    setParam('internalQueryMaxPushBytes', originalVal);
 }());
 
 (function testInternalQueryMaxAddToSetBytesSetting() {
@@ -48,14 +59,31 @@ assert.commandWorked(bulk.execute());
 
     // Test that $addToSet needs a tighter limit than $push (because some of the strings are the
     // same).
-    assert.commandWorked(
-        db.adminCommand({setParameter: 1, internalQueryMaxAddToSetBytes: memLimitArray}));
+    const originalVal = setParam('internalQueryMaxAddToSetBytes', memLimitArray);
     assert.doesNotThrow(() => coll.aggregate([{$group: {_id: null, strings: {$addToSet: "$y"}}}]));
 
-    assert.commandWorked(
-        db.adminCommand({setParameter: 1, internalQueryMaxAddToSetBytes: memLimitSet}));
+    setParam('internalQueryMaxAddToSetBytes', memLimitSet);
     assert.throwsWithCode(() => coll.aggregate([{$group: {_id: null, strings: {$addToSet: "$y"}}}]),
                           ErrorCodes.ExceededMemoryLimit);
+    setParam('internalQueryMaxAddToSetBytes', originalVal);
+}());
+
+(function testInternalQueryMaxAddToSetBytesSettingWindowFunc() {
+    let pipeline = [
+        {$setWindowFields: {sortBy: {_id: 1}, output: {v: {$addToSet: '$y'}}}},
+    ];
+    // Test that the default 100MB memory limit isn't reached with our data.
+    assert.doesNotThrow(() => coll.aggregate(pipeline));
+
+    // Test that $addToSet needs a tighter limit than $concatArrays (because some of the strings are
+    // the same).
+    const originalVal = setParam('internalQueryMaxAddToSetBytes', memLimitArray);
+    assert.doesNotThrow(() => coll.aggregate(pipeline));
+
+    // Test that a tighter limit for $addToSet is obeyed.
+    setParam('internalQueryMaxAddToSetBytes', memLimitSet);
+    assert.throwsWithCode(() => coll.aggregate(pipeline), ErrorCodes.ExceededMemoryLimit);
+    setParam('internalQueryMaxAddToSetBytes', originalVal);
 }());
 
 (function testInternalQueryTopNAccumulatorBytesSetting() {
diff --git a/jstests/noPassthrough/query/agg/agg_configurable_memory_limits_array_accumulators.js b/jstests/noPassthrough/query/agg/agg_configurable_memory_limits_array_accumulators.js
@@ -1,5 +1,5 @@
 /**
- * Tests that $concatArrays and $setUnion accumulators respect configurable memory limits.
+ * Tests that $concatArrays and $setUnion accumulators/operators respect configurable memory limits.
  * @tags: [requires_fcv_81]
  */
 
@@ -32,17 +32,40 @@ for (let i = 0; i < nSets; i++) {
 }
 assert.commandWorked(bulk.execute());
 
+// Function to change the parameter value and return the previous value.
+function setParam(param, val) {
+    const res = db.adminCommand({setParameter: 1, [param]: val});
+    assert.commandWorked(res);
+    return res.was;
+}
+
 (function testInternalQueryMaxConcatArraysBytesSetting() {
     // Test that the default 100MB memory limit isn't reached with our data.
     assert.doesNotThrow(
         () => coll.aggregate([{$group: {_id: null, strings: {$concatArrays: "$y"}}}]));
 
     // Now lower the limit to test that its configuration is obeyed.
-    assert.commandWorked(
-        db.adminCommand({setParameter: 1, internalQueryMaxConcatArraysBytes: memLimitArray}));
+    const originalVal = setParam('internalQueryMaxConcatArraysBytes', memLimitArray);
     assert.throwsWithCode(
         () => coll.aggregate([{$group: {_id: null, strings: {$concatArrays: "$y"}}}]),
         ErrorCodes.ExceededMemoryLimit);
+    setParam('internalQueryMaxConcatArraysBytes', originalVal);
+}());
+
+(function testInternalQueryMaxConcatArraysBytesSettingWindowFunc() {
+    let pipeline = [
+        {
+            $setWindowFields:
+                {partitionBy: null, sortBy: {_id: 1}, output: {allStr: {$concatArrays: '$y'}}}
+        },
+    ];
+    // Test that the default 100MB memory limit isn't reached with our data.
+    assert.doesNotThrow(() => coll.aggregate(pipeline));
+
+    // Now lower the limit to test that its configuration is obeyed.
+    const originalVal = setParam('internalQueryMaxConcatArraysBytes', memLimitArray);
+    assert.throwsWithCode(() => coll.aggregate(pipeline), ErrorCodes.ExceededMemoryLimit);
+    setParam('internalQueryMaxConcatArraysBytes', originalVal);
 }());
 
 (function testInternalQueryMaxSetUnionBytesSetting() {
@@ -51,15 +74,32 @@ assert.commandWorked(bulk.execute());
 
     // Test that $setUnion needs a tighter limit than $concatArrays (because some of the strings are
     // the same).
-    assert.commandWorked(
-        db.adminCommand({setParameter: 1, internalQueryMaxSetUnionBytes: memLimitArray}));
+    const originalVal = setParam('internalQueryMaxSetUnionBytes', memLimitArray);
     assert.doesNotThrow(() => coll.aggregate([{$group: {_id: null, strings: {$setUnion: "$y"}}}]));
 
     // Test that a tighter limit for $setUnion is obeyed.
-    assert.commandWorked(
-        db.adminCommand({setParameter: 1, internalQueryMaxSetUnionBytes: memLimitSet}));
+    setParam('internalQueryMaxSetUnionBytes', memLimitSet);
     assert.throwsWithCode(() => coll.aggregate([{$group: {_id: null, strings: {$setUnion: "$y"}}}]),
                           ErrorCodes.ExceededMemoryLimit);
+    setParam('internalQueryMaxSetUnionBytes', originalVal);
+}());
+
+(function testInternalQueryMaxSetUnionBytesSettingWindowFunc() {
+    let pipeline = [
+        {$setWindowFields: {sortBy: {_id: 1}, output: {v: {$setUnion: '$y'}}}},
+    ];
+    // Test that the default 100MB memory limit isn't reached with our data.
+    assert.doesNotThrow(() => coll.aggregate(pipeline));
+
+    // Test that $setUnion needs a tighter limit than $concatArrays (because some of the strings are
+    // the same).
+    const originalVal = setParam('internalQueryMaxSetUnionBytes', memLimitArray);
+    assert.doesNotThrow(() => coll.aggregate(pipeline));
+
+    // Test that a tighter limit for $setUnion is obeyed.
+    setParam('internalQueryMaxSetUnionBytes', memLimitSet);
+    assert.throwsWithCode(() => coll.aggregate(pipeline), ErrorCodes.ExceededMemoryLimit);
+    setParam('internalQueryMaxSetUnionBytes', originalVal);
 }());
 
 MongoRunner.stopMongod(conn);
diff --git a/src/mongo/db/pipeline/window_function/window_function_add_to_set.h b/src/mongo/db/pipeline/window_function/window_function_add_to_set.h
@@ -53,14 +53,19 @@ class WindowFunctionAddToSet final : public WindowFunctionState {
     }
 
     explicit WindowFunctionAddToSet(ExpressionContext* const expCtx)
-        : WindowFunctionState(expCtx),
+        : WindowFunctionState(expCtx, internalQueryMaxAddToSetBytes.load()),
           _values(MemoryTokenValueComparator(&_expCtx->getValueComparator())) {
         _memUsageTracker.set(sizeof(*this));
     }
 
     void add(Value value) override {
         _values.emplace(SimpleMemoryUsageToken{value.getApproximateSize(), &_memUsageTracker},
                         std::move(value));
+        uassert(ErrorCodes::ExceededMemoryLimit,
+                str::stream() << "$addToSet used too much memory and cannot spill to disk. Used: "
+                              << _memUsageTracker.currentMemoryBytes() << " bytes. Memory limit: "
+                              << _memUsageTracker.maxAllowedMemoryUsageBytes() << " bytes",
+                _memUsageTracker.withinMemoryLimit());
     }
 
     /**
diff --git a/src/mongo/db/pipeline/window_function/window_function_concat_arrays.h b/src/mongo/db/pipeline/window_function/window_function_concat_arrays.h
@@ -44,7 +44,7 @@ class WindowFunctionConcatArrays final : public WindowFunctionState {
     }
 
     explicit WindowFunctionConcatArrays(ExpressionContext* const expCtx)
-        : WindowFunctionState(expCtx) {
+        : WindowFunctionState(expCtx, internalQueryMaxConcatArraysBytes.load()) {
         _memUsageTracker.set(sizeof(*this));
     }
 
@@ -62,6 +62,12 @@ class WindowFunctionConcatArrays final : public WindowFunctionState {
         _count += value.getArrayLength();
         _values.emplace_back(SimpleMemoryUsageToken{value.getApproximateSize(), &_memUsageTracker},
                              std::move(value));
+        uassert(ErrorCodes::ExceededMemoryLimit,
+                str::stream() << "$concatArrays used too much memory and spilling to disk will not "
+                                 "reduce memory usage. Used: "
+                              << _memUsageTracker.currentMemoryBytes() << " bytes. Memory limit: "
+                              << _memUsageTracker.maxAllowedMemoryUsageBytes() << " bytes",
+                _memUsageTracker.withinMemoryLimit());
     }
 
     void remove(Value value) override {
diff --git a/src/mongo/db/pipeline/window_function/window_function_push.h b/src/mongo/db/pipeline/window_function/window_function_push.h
@@ -55,7 +55,8 @@ class WindowFunctionPush final : public WindowFunctionState {
         return std::make_unique<WindowFunctionPush>(expCtx);
     }
 
-    explicit WindowFunctionPush(ExpressionContext* const expCtx) : WindowFunctionState(expCtx) {
+    explicit WindowFunctionPush(ExpressionContext* const expCtx)
+        : WindowFunctionState(expCtx, internalQueryMaxPushBytes.load()) {
         _memUsageTracker.set(sizeof(*this));
     }
 
@@ -65,6 +66,11 @@ class WindowFunctionPush final : public WindowFunctionState {
         }
         _values.emplace_back(SimpleMemoryUsageToken{value.getApproximateSize(), &_memUsageTracker},
                              std::move(value));
+        uassert(ErrorCodes::ExceededMemoryLimit,
+                str::stream() << "$push used too much memory and cannot spill to disk. Used: "
+                              << _memUsageTracker.currentMemoryBytes() << "bytes. Memory limit: "
+                              << _memUsageTracker.maxAllowedMemoryUsageBytes() << " bytes",
+                _memUsageTracker.withinMemoryLimit());
     }
 
     /**
diff --git a/src/mongo/db/pipeline/window_function/window_function_set_union.h b/src/mongo/db/pipeline/window_function/window_function_set_union.h
@@ -43,7 +43,7 @@ class WindowFunctionSetUnion final : public WindowFunctionState {
     }
 
     explicit WindowFunctionSetUnion(ExpressionContext* const expCtx)
-        : WindowFunctionState(expCtx),
+        : WindowFunctionState(expCtx, internalQueryMaxSetUnionBytes.load()),
           _values(MemoryTokenValueComparator(&_expCtx->getValueComparator())) {
         _memUsageTracker.set(sizeof(*this));
     }
@@ -62,6 +62,13 @@ class WindowFunctionSetUnion final : public WindowFunctionState {
         for (const auto& val : value.getArray()) {
             _values.emplace(SimpleMemoryUsageToken{val.getApproximateSize(), &_memUsageTracker},
                             val);
+            uassert(ErrorCodes::ExceededMemoryLimit,
+                    str::stream() << "$setUnion used too much memory and spilling to disk will not "
+                                     "reduce memory usage. Used: "
+                                  << _memUsageTracker.currentMemoryBytes()
+                                  << "bytes. Memory limit: "
+                                  << _memUsageTracker.maxAllowedMemoryUsageBytes() << " bytes",
+                    _memUsageTracker.withinMemoryLimit());
         }
     }
 

Original file line number	Diff line number	Diff line change
`@@ -53,14 +53,19 @@ class WindowFunctionAddToSet final : public WindowFunctionState {`
`53`	`53`	`}`
`54`	`54`
`55`	`55`	`explicit WindowFunctionAddToSet(ExpressionContext* const expCtx)`
`56`		`- : WindowFunctionState(expCtx),`
	`56`	`+ : WindowFunctionState(expCtx, internalQueryMaxAddToSetBytes.load()),`
`57`	`57`	`_values(MemoryTokenValueComparator(&_expCtx->getValueComparator())) {`
`58`	`58`	`_memUsageTracker.set(sizeof(*this));`
`59`	`59`	`}`
`60`	`60`
`61`	`61`	`void add(Value value) override {`
`62`	`62`	`_values.emplace(SimpleMemoryUsageToken{value.getApproximateSize(), &_memUsageTracker},`
`63`	`63`	`std::move(value));`
	`64`	`+ uassert(ErrorCodes::ExceededMemoryLimit,`
	`65`	`+ str::stream() << "$addToSet used too much memory and cannot spill to disk. Used: "`
	`66`	`+ << _memUsageTracker.currentMemoryBytes() << " bytes. Memory limit: "`
	`67`	`+ << _memUsageTracker.maxAllowedMemoryUsageBytes() << " bytes",`
	`68`	`+ _memUsageTracker.withinMemoryLimit());`
`64`	`69`	`}`
`65`	`70`
`66`	`71`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ class WindowFunctionConcatArrays final : public WindowFunctionState {`
`44`	`44`	`}`
`45`	`45`
`46`	`46`	`explicit WindowFunctionConcatArrays(ExpressionContext* const expCtx)`
`47`		`- : WindowFunctionState(expCtx) {`
	`47`	`+ : WindowFunctionState(expCtx, internalQueryMaxConcatArraysBytes.load()) {`
`48`	`48`	`_memUsageTracker.set(sizeof(*this));`
`49`	`49`	`}`
`50`	`50`
`@@ -62,6 +62,12 @@ class WindowFunctionConcatArrays final : public WindowFunctionState {`
`62`	`62`	`_count += value.getArrayLength();`
`63`	`63`	`_values.emplace_back(SimpleMemoryUsageToken{value.getApproximateSize(), &_memUsageTracker},`
`64`	`64`	`std::move(value));`
	`65`	`+ uassert(ErrorCodes::ExceededMemoryLimit,`
	`66`	`+ str::stream() << "$concatArrays used too much memory and spilling to disk will not "`
	`67`	`+ "reduce memory usage. Used: "`
	`68`	`+ << _memUsageTracker.currentMemoryBytes() << " bytes. Memory limit: "`
	`69`	`+ << _memUsageTracker.maxAllowedMemoryUsageBytes() << " bytes",`
	`70`	`+ _memUsageTracker.withinMemoryLimit());`
`65`	`71`	`}`
`66`	`72`
`67`	`73`	`void remove(Value value) override {`
Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,8 @@ class WindowFunctionPush final : public WindowFunctionState {`
`55`	`55`	`return std::make_unique<WindowFunctionPush>(expCtx);`
`56`	`56`	`}`
`57`	`57`
`58`		`- explicit WindowFunctionPush(ExpressionContext* const expCtx) : WindowFunctionState(expCtx) {`
	`58`	`+ explicit WindowFunctionPush(ExpressionContext* const expCtx)`
	`59`	`+ : WindowFunctionState(expCtx, internalQueryMaxPushBytes.load()) {`
`59`	`60`	`_memUsageTracker.set(sizeof(*this));`
`60`	`61`	`}`
`61`	`62`
`@@ -65,6 +66,11 @@ class WindowFunctionPush final : public WindowFunctionState {`
`65`	`66`	`}`
`66`	`67`	`_values.emplace_back(SimpleMemoryUsageToken{value.getApproximateSize(), &_memUsageTracker},`
`67`	`68`	`std::move(value));`
	`69`	`+ uassert(ErrorCodes::ExceededMemoryLimit,`
	`70`	`+ str::stream() << "$push used too much memory and cannot spill to disk. Used: "`
	`71`	`+ << _memUsageTracker.currentMemoryBytes() << "bytes. Memory limit: "`
	`72`	`+ << _memUsageTracker.maxAllowedMemoryUsageBytes() << " bytes",`
	`73`	`+ _memUsageTracker.withinMemoryLimit());`
`68`	`74`	`}`
`69`	`75`
`70`	`76`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ class WindowFunctionSetUnion final : public WindowFunctionState {`
`43`	`43`	`}`
`44`	`44`
`45`	`45`	`explicit WindowFunctionSetUnion(ExpressionContext* const expCtx)`
`46`		`- : WindowFunctionState(expCtx),`
	`46`	`+ : WindowFunctionState(expCtx, internalQueryMaxSetUnionBytes.load()),`
`47`	`47`	`_values(MemoryTokenValueComparator(&_expCtx->getValueComparator())) {`
`48`	`48`	`_memUsageTracker.set(sizeof(*this));`
`49`	`49`	`}`
`@@ -62,6 +62,13 @@ class WindowFunctionSetUnion final : public WindowFunctionState {`
`62`	`62`	`for (const auto& val : value.getArray()) {`
`63`	`63`	`_values.emplace(SimpleMemoryUsageToken{val.getApproximateSize(), &_memUsageTracker},`
`64`	`64`	`val);`
	`65`	`+ uassert(ErrorCodes::ExceededMemoryLimit,`
	`66`	`+ str::stream() << "$setUnion used too much memory and spilling to disk will not "`
	`67`	`+ "reduce memory usage. Used: "`
	`68`	`+ << _memUsageTracker.currentMemoryBytes()`
	`69`	`+ << "bytes. Memory limit: "`
	`70`	`+ << _memUsageTracker.maxAllowedMemoryUsageBytes() << " bytes",`
	`71`	`+ _memUsageTracker.withinMemoryLimit());`
`65`	`72`	`}`
`66`	`73`	`}`
`67`	`74`