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

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

6 files changed, 43 insertions, 13 deletions

diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index 70783b7b1ac..eb8a5d1d046 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -208,6 +208,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