SERVER-44869 Add query knobs for $push and $addToSet memory limits

(cherry picked from commit f7d0d2f3d6fe9e4c472bd1d8893c8b6dc96881b6)
(cherry picked from commit f69a356120ae8edd6fec5ab6efaebccc1427dd28)

7 files changed, 92 insertions, 13 deletions

diff --git a/jstests/noPassthrough/agg_configurable_memory_limits.js b/jstests/noPassthrough/agg_configurable_memory_limits.js
new file mode 100644
index 00000000000..6ffd61d6990
--- /dev/null
+++ b/jstests/noPassthrough/agg_configurable_memory_limits.js
@@ -0,0 +1,49 @@
+// Tests that certain aggregation operators have configurable memory limits.
+(function() {
+    "use strict";
+
+    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;
+
+    const bulk = coll.initializeUnorderedBulkOp();
+    for (let i = 0; i < 100; i++) {
+        bulk.insert(
+            {_id: i, x: i, y: ["string 1", "string 2", "string 3", "string 4", "string " + i]});
+    }
+    assert.commandWorked(bulk.execute());
+
+    // Test that pushing a bunch of strings to an array does not exceed the default 100MB memory
+    // limit.
+    assert.doesNotThrow(
+        () => coll.aggregate([{$unwind: "$y"}, {$group: {_id: null, strings: {$push: "$y"}}}]));
+
+    // Now lower the limit to test that it's configuration is obeyed.
+    assert.commandWorked(db.adminCommand({setParameter: 1, internalQueryMaxPushBytes: 100}));
+    let error = assert.throws(
+        () => coll.aggregate([{$unwind: "$y"}, {$group: {_id: null, strings: {$push: "$y"}}}]));
+    assert.eq(error.code, ErrorCodes.ExceededMemoryLimit);
+
+    // Test that using $addToSet behaves similarly.
+    assert.doesNotThrow(
+        () => coll.aggregate([{$unwind: "$y"}, {$group: {_id: null, strings: {$addToSet: "$y"}}}]));
+
+    assert.commandWorked(db.adminCommand({setParameter: 1, internalQueryMaxAddToSetBytes: 100}));
+    error = assert.throws(
+        () => coll.aggregate([{$unwind: "$y"}, {$group: {_id: null, strings: {$addToSet: "$y"}}}]));
+    assert.eq(error.code, ErrorCodes.ExceededMemoryLimit);
+
+    // Test that you cannot set the parameter to a negative value or 0.
+    assert.commandFailedWithCode(db.adminCommand({setParameter: 1, internalQueryMaxPushBytes: -10}),
+                                 ErrorCodes.BadValue);
+    assert.commandFailedWithCode(db.adminCommand({setParameter: 1, internalQueryMaxPushBytes: 0}),
+                                 ErrorCodes.BadValue);
+    assert.commandFailedWithCode(
+        db.adminCommand({setParameter: 1, internalQueryMaxAddToSetBytes: -10}),
+        ErrorCodes.BadValue);
+    assert.commandFailedWithCode(
+        db.adminCommand({setParameter: 1, internalQueryMaxAddToSetBytes: 0}), ErrorCodes.BadValue);
+
+    MongoRunner.stopMongod(conn);
+}());
diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index a8945ea2f14..fb911760111 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -161,6 +161,7 @@ env.Library(
         ],
     LIBDEPS=[
         'document_value',
+        '$BUILD_DIR/mongo/db/query/query_knobs',
         '$BUILD_DIR/mongo/util/summation',
         'expression',
         'field_path',
diff --git a/src/mongo/db/pipeline/accumulator.h b/src/mongo/db/pipeline/accumulator.h
index 4d10fc051b5..8f89a13b616 100644
--- a/src/mongo/db/pipeline/accumulator.h
+++ b/src/mongo/db/pipeline/accumulator.h
@@ -105,9 +105,12 @@ private:
 
 class AccumulatorAddToSet final : public Accumulator {
 public:
-    static constexpr int kDefaultMaxMemoryUsageBytes = 100 * 1024 * 1024;
-    explicit AccumulatorAddToSet(const boost::intrusive_ptr<ExpressionContext>& expCtx,
-                                 int maxMemoryUsageBytes = kDefaultMaxMemoryUsageBytes);
+    /**
+     * Creates a new $addToSet accumulator. If no memory limit is given, defaults to the value of
+     * the server parameter 'internalQueryMaxAddToSetBytes'.
+     */
+    AccumulatorAddToSet(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                        boost::optional<int> maxMemoryUsageBytes = boost::none);
 
     void processInternal(const Value& input, bool merging) final;
     Value getValue(bool toBeMerged) final;
@@ -239,9 +242,12 @@ public:
 
 class AccumulatorPush final : public Accumulator {
 public:
-    static constexpr int kDefaultMaxMemoryUsageBytes = 100 * 1024 * 1024;
-    explicit AccumulatorPush(const boost::intrusive_ptr<ExpressionContext>& expCtx,
-                             int maxMemoryUsageBytes = kDefaultMaxMemoryUsageBytes);
+    /**
+     * Creates a new $push accumulator. If no memory limit is given, defaults to the value of the
+     * server parameter 'internalQueryMaxPushBytes'.
+     */
+    AccumulatorPush(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                    boost::optional<int> maxMemoryUsageBytes = boost::none);
 
     void processInternal(const Value& input, bool merging) final;
     Value getValue(bool toBeMerged) final;
diff --git a/src/mongo/db/pipeline/accumulator_add_to_set.cpp b/src/mongo/db/pipeline/accumulator_add_to_set.cpp
index fb591247866..a10cd2b45f8 100644
--- a/src/mongo/db/pipeline/accumulator_add_to_set.cpp
+++ b/src/mongo/db/pipeline/accumulator_add_to_set.cpp
@@ -35,6 +35,7 @@
 #include "mongo/db/pipeline/accumulation_statement.h"
 #include "mongo/db/pipeline/expression_context.h"
 #include "mongo/db/pipeline/value.h"
+#include "mongo/db/query/query_knobs.h"
 
 namespace mongo {
 
@@ -81,10 +82,10 @@ Value AccumulatorAddToSet::getValue(bool toBeMerged) {
 }
 
 AccumulatorAddToSet::AccumulatorAddToSet(const boost::intrusive_ptr<ExpressionContext>& expCtx,
-                                         int maxMemoryUsageBytes)
+                                         boost::optional<int> maxMemoryUsageBytes)
     : Accumulator(expCtx),
       _set(expCtx->getValueComparator().makeUnorderedValueSet()),
-      _maxMemUsageBytes(maxMemoryUsageBytes) {
+      _maxMemUsageBytes(maxMemoryUsageBytes.value_or(internalQueryMaxAddToSetBytes.load())) {
     _memUsageBytes = sizeof(*this);
 }
 
@@ -95,7 +96,7 @@ void AccumulatorAddToSet::reset() {
 
 intrusive_ptr<Accumulator> AccumulatorAddToSet::create(
     const boost::intrusive_ptr<ExpressionContext>& expCtx) {
-    return new AccumulatorAddToSet(expCtx);
+    return new AccumulatorAddToSet(expCtx, boost::none);
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/accumulator_push.cpp b/src/mongo/db/pipeline/accumulator_push.cpp
index 7e42d29546d..d45b37138d5 100644
--- a/src/mongo/db/pipeline/accumulator_push.cpp
+++ b/src/mongo/db/pipeline/accumulator_push.cpp
@@ -35,6 +35,7 @@
 #include "mongo/db/pipeline/accumulation_statement.h"
 #include "mongo/db/pipeline/expression_context.h"
 #include "mongo/db/pipeline/value.h"
+#include "mongo/db/query/query_knobs.h"
 
 namespace mongo {
 
@@ -84,8 +85,9 @@ Value AccumulatorPush::getValue(bool toBeMerged) {
 }
 
 AccumulatorPush::AccumulatorPush(const boost::intrusive_ptr<ExpressionContext>& expCtx,
-                                 int maxMemoryUsageBytes)
-    : Accumulator(expCtx), _maxMemUsageBytes(maxMemoryUsageBytes) {
+                                 boost::optional<int> maxMemoryUsageBytes)
+    : Accumulator(expCtx),
+      _maxMemUsageBytes(maxMemoryUsageBytes.value_or(internalQueryMaxPushBytes.load())) {
     _memUsageBytes = sizeof(*this);
 }
 
@@ -96,6 +98,6 @@ void AccumulatorPush::reset() {
 
 intrusive_ptr<Accumulator> AccumulatorPush::create(
     const boost::intrusive_ptr<ExpressionContext>& expCtx) {
-    return new AccumulatorPush(expCtx);
-}
+    return new AccumulatorPush(expCtx, boost::none);
 }
+}  // namespace mongo
diff --git a/src/mongo/db/query/query_knobs.cpp b/src/mongo/db/query/query_knobs.cpp
index 741b5293865..bdde620c87f 100644
--- a/src/mongo/db/query/query_knobs.cpp
+++ b/src/mongo/db/query/query_knobs.cpp
@@ -97,4 +97,20 @@ MONGO_EXPORT_SERVER_PARAMETER(internalQueryPlannerGenerateCoveredWholeIndexScans
 MONGO_EXPORT_SERVER_PARAMETER(internalQueryIgnoreUnknownJSONSchemaKeywords, bool, false);
 
 MONGO_EXPORT_SERVER_PARAMETER(internalQueryProhibitBlockingMergeOnMongoS, bool, false);
+
+MONGO_EXPORT_SERVER_PARAMETER(internalQueryMaxPushBytes, int, 100 * 1024 * 1024)
+    ->withValidator([](const int& newVal) {
+        if (newVal <= 0) {
+            return Status(ErrorCodes::BadValue, "internalQueryMaxPushBytes must be positive");
+        }
+        return Status::OK();
+    });
+
+MONGO_EXPORT_SERVER_PARAMETER(internalQueryMaxAddToSetBytes, int, 100 * 1024 * 1024)
+    ->withValidator([](const int& newVal) {
+        if (newVal <= 0) {
+            return Status(ErrorCodes::BadValue, "internalQueryMaxAddToSetBytes must be positive");
+        }
+        return Status::OK();
+    });
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_knobs.h b/src/mongo/db/query/query_knobs.h
index f2fe5046296..035915267e8 100644
--- a/src/mongo/db/query/query_knobs.h
+++ b/src/mongo/db/query/query_knobs.h
@@ -127,4 +127,8 @@ extern AtomicInt32 internalDocumentSourceCursorBatchSizeBytes;
 extern AtomicInt32 internalDocumentSourceLookupCacheSizeBytes;
 
 extern AtomicBool internalQueryProhibitBlockingMergeOnMongoS;
+
+extern AtomicInt32 internalQueryMaxPushBytes;
+
+extern AtomicInt32 internalQueryMaxAddToSetBytes;
 }  // namespace mongo