SERVER-58979 Add limit to number of documents generated by $densify

5 files changed, 202 insertions, 43 deletions

diff --git a/jstests/aggregation/sources/densify/generated_limit.js b/jstests/aggregation/sources/densify/generated_limit.js
new file mode 100644
index 00000000000..704de96dac2
--- /dev/null
+++ b/jstests/aggregation/sources/densify/generated_limit.js
@@ -0,0 +1,72 @@
+/**
+ * Test that densify correctly limits the number of generated documents.
+ * @tags: [
+ *   # Needed as $densify is a 51 feature.
+ *   requires_fcv_51,
+ * ]
+ */
+
+(function() {
+"use strict";
+
+load("jstests/noPassthrough/libs/server_parameter_helpers.js");  // For setParameterOnAllHosts.
+load("jstests/libs/discover_topology.js");                       // For findNonConfigNodes.
+
+const featureEnabled =
+    assert.commandWorked(db.adminCommand({getParameter: 1, featureFlagDensify: 1}))
+        .featureFlagDensify.value;
+if (!featureEnabled) {
+    jsTestLog("Skipping test because the densify feature flag is disabled");
+    return;
+}
+const paramName = "internalQueryMaxAllowedDensifyDocs";
+const origParamValue = assert.commandWorked(
+    db.adminCommand({getParameter: 1, internalQueryMaxAllowedDensifyDocs: 1}))[paramName];
+function setMaxDocs(max) {
+    setParameterOnAllHosts(DiscoverTopology.findNonConfigNodes(db.getMongo()), paramName, max);
+}
+const coll = db[jsTestName()];
+coll.drop();
+
+function runAggregate(densifyStage, failCode = null) {
+    if (failCode === null) {
+        return db.runCommand({aggregate: coll.getName(), pipeline: [densifyStage], cursor: {}});
+    } else {
+        assert.commandFailedWithCode(
+            db.runCommand({aggregate: coll.getName(), pipeline: [densifyStage], cursor: {}}),
+            failCode);
+    }
+}
+
+// Test that with explicit range and no documents we can't generate more than the limit.
+setMaxDocs(10);
+runAggregate({$densify: {field: "val", range: {step: 1, bounds: [0, 11]}}}, 5897900);
+runAggregate({$densify: {field: "val", range: {step: 4, bounds: [0, 45]}}}, 5897900);
+// Exactly ten documents should pass.
+runAggregate({$densify: {field: "val", range: {step: 1, bounds: [0, 10]}}});
+
+// Full fails if there are enough points between min and max to pass limit
+assert.commandWorked(coll.insert({val: 0, part: 1}));
+assert.commandWorked(coll.insert({val: 12, part: 1}));
+runAggregate({$densify: {field: "val", range: {step: 1, bounds: "full"}}}, 5897900);
+
+// Test that count is shared across partitions.
+setMaxDocs(20);
+assert.commandWorked(coll.insert({val: 0, part: 2}));
+assert.commandWorked(coll.insert({val: 12, part: 2}));
+runAggregate(
+    {$densify: {field: "val", partitionByFields: ["part"], range: {step: 1, bounds: "partition"}}},
+    5897900);
+
+// Test that already existing documents don't count towards the limit.
+coll.drop();
+setMaxDocs(10);
+assert.commandWorked(coll.insert({val: 0, part: 1}));
+assert.commandWorked(coll.insert({val: 12, part: 1}));
+runAggregate({$densify: {field: "val", range: {step: 1, bounds: "full"}}}, 5897900);
+assert.commandWorked(coll.insert({val: 5, part: 1}));
+runAggregate({$densify: {field: "val", range: {step: 1, bounds: "full"}}});
+
+// Reset parameter.
+setMaxDocs(origParamValue);
+})();
diff --git a/src/mongo/db/pipeline/document_source_densify.cpp b/src/mongo/db/pipeline/document_source_densify.cpp
index 34a8eafa214..38cd2080284 100644
--- a/src/mongo/db/pipeline/document_source_densify.cpp
+++ b/src/mongo/db/pipeline/document_source_densify.cpp
@@ -230,12 +230,14 @@ DocumentSourceInternalDensify::DocGenerator::DocGenerator(
     FieldPath fieldName,
     boost::optional<Document> includeFields,
     boost::optional<Document> finalDoc,
-    ValueComparator comp)
+    ValueComparator comp,
+    size_t* counter)
     : _comp(std::move(comp)),
       _range(std::move(range)),
       _path(std::move(fieldName.fullPath())),
       _finalDoc(std::move(finalDoc)),
-      _min(std::move(min)) {
+      _min(std::move(min)),
+      _counter(counter) {
 
     if (includeFields) {
         _includeFields = *includeFields;
@@ -338,6 +340,7 @@ Document DocumentSourceInternalDensify::DocGenerator::getNextDocument() {
 
     MutableDocument retDoc(_includeFields);
     retDoc.setNestedField(_path, valueToAdd);
+    ++(*_counter);
     return retDoc.freeze();
 }
 
@@ -358,16 +361,14 @@ DocumentSource::GetNextResult DocumentSourceInternalDensify::densifyAfterEOF(Num
                 stdx::holds_alternative<NumericBounds>(bounds));
         auto lowerBound = stdx::get<NumericBounds>(bounds).first;
         _current = lowerBound;
-        _docGenerator = DocGenerator(
-            lowerBound, _range, _field, boost::none, boost::none, pExpCtx->getValueComparator());
+        createDocGenerator(lowerBound, _range);
     } else if (compareValues(addValues(stdx::get<Value>(*_current), _range.getStep()) >=
                              bounds.second)) {
         _densifyState = DensifyState::kDensifyDone;
         return DocumentSource::GetNextResult::makeEOF();
     } else {
         auto lowerBound = addValues(stdx::get<Value>(*_current), _range.getStep());
-        _docGenerator = DocGenerator(
-            lowerBound, _range, _field, boost::none, boost::none, pExpCtx->getValueComparator());
+        createDocGenerator(lowerBound, _range);
     }
     _densifyState = DensifyState::kHaveGenerator;
     auto generatedDoc = _docGenerator->getNextDocument();
@@ -400,14 +401,12 @@ DocumentSource::GetNextResult DocumentSourceInternalDensify::processDocAboveMinB
         val = subtractValues(val, _range.getStep());
     }
     Value upperBound = (compareValues(val <= bounds.second)) ? val : bounds.second;
-    _docGenerator = DocGenerator(
+    createDocGenerator(
         lowerBound,
         RangeStatement(_range.getStep(), NumericBounds(lowerBound, upperBound), _range.getUnit()),
-        _field,
         _partitionExpr ? boost::make_optional<Document>(getDensifyPartition(doc).getDocument())
                        : boost::none,
-        doc,
-        pExpCtx->getValueComparator());
+        doc);
     Document nextFromGen = _docGenerator->getNextDocument();
     _current = getDensifyValue(nextFromGen);
     _densifyState = DensifyState::kHaveGenerator;
@@ -492,13 +491,12 @@ DocumentSource::GetNextResult DocumentSourceInternalDensify::finishDensifyingPar
         if (minOverride && compareValues(valToGenerate, *minOverride) < 0) {
             valToGenerate = *minOverride;
         }
-        _docGenerator = DocGenerator(
+        createDocGenerator(
             valToGenerate,
             RangeStatement(_range.getStep(), NumericBounds(valToGenerate, max), _range.getUnit()),
-            _field,
             firstPartition.getDocument(),
-            boost::none,  // final doc.
-            pExpCtx->getValueComparator());
+            boost::none  // final doc.
+        );
         // Remove this partition from the table, we're done with it.
         _partitionTable.erase(firstPartitionKeyVal);
         _densifyState = DensifyState::kHaveGenerator;
@@ -612,15 +610,13 @@ DocumentSource::GetNextResult DocumentSourceInternalDensify::handleNeedGen(Docum
         compareValues(offsetFromStep == Value(0)) ? subtractValues(max, _range.getStep()) : max;
 
     Value newCurrent = addValues(stdx::get<Value>(*_current), _range.getStep());
-    _docGenerator = DocumentSourceInternalDensify::DocGenerator(
+    createDocGenerator(
         DensifyValueType(newCurrent),
         RangeStatement(_range.getStep(), NumericBounds(newCurrent, maxAdjusted), _range.getUnit()),
-        _field,
         _partitionExpr
             ? boost::make_optional<Document>(getDensifyPartition(currentDoc).getDocument())
             : boost::none,
-        currentDoc,
-        pExpCtx->getValueComparator());
+        currentDoc);
 
     _densifyState = DensifyState::kHaveGenerator;
     auto nextDoc = _docGenerator->getNextDocument();
@@ -715,6 +711,14 @@ void DocumentSourceInternalDensify::initializePartitionState(Document initialDoc
 }
 
 DocumentSource::GetNextResult DocumentSourceInternalDensify::doGetNext() {
+    // When we return a generated document '_docsGenerated' is incremented. Check that the last
+    // document didn't put us over the limit.
+    uassert(5897900,
+            str::stream() << "Generated " << _docsGenerated
+                          << " documents in $densify, which is over the limit of " << _maxDocs
+                          << ". Increase the 'internalQueryMaxAllowedDensifyDocs' parameter to "
+                             "allow more generated documents",
+            _docsGenerated <= _maxDocs);
     switch (_densifyState) {
         case DensifyState::kUninitializedOrBelowRange: {
             // This state represents the first run of doGetNext() or that the value that we last
diff --git a/src/mongo/db/pipeline/document_source_densify.h b/src/mongo/db/pipeline/document_source_densify.h
index a00cd33f478..b48992357fd 100644
--- a/src/mongo/db/pipeline/document_source_densify.h
+++ b/src/mongo/db/pipeline/document_source_densify.h
@@ -143,8 +143,9 @@ public:
           _field(std::move(field)),
           _partitions(std::move(partitions)),
           _range(std::move(range)),
-          _partitionTable(pExpCtx->getValueComparator().makeUnorderedValueMap<Value>()){};
-
+          _partitionTable(pExpCtx->getValueComparator().makeUnorderedValueMap<Value>()) {
+        _maxDocs = internalQueryMaxAllowedDensifyDocs.load();
+    };
     class DocGenerator {
     public:
         DocGenerator(DensifyValueType current,
@@ -152,7 +153,8 @@ public:
                      FieldPath fieldName,
                      boost::optional<Document> includeFields,
                      boost::optional<Document> finalDoc,
-                     ValueComparator comp);
+                     ValueComparator comp,
+                     size_t* counter);
         Document getNextDocument();
         bool done() const;
 
@@ -179,6 +181,10 @@ public:
         };
 
         GeneratorState _state = GeneratorState::kGeneratingDocuments;
+        // Value to increment when returning a generated document. This is a pointer to the counter
+        // that keeps track of the total number of documents generated by the owning stage across
+        // all generators.
+        size_t* _counter;
     };
 
     DocumentSourceInternalDensify(
@@ -343,6 +349,26 @@ private:
             _partitionTable[getDensifyPartition(doc)] = getDensifyValue(doc);
         }
     }
+
+    /**
+     * Helpers to create doc generators. Sets _docGenerator to the created generator.
+     */
+    void createDocGenerator(DensifyValueType min,
+                            RangeStatement range,
+                            boost::optional<Document> includeFields,
+                            boost::optional<Document> finalDoc) {
+        _docGenerator = DocGenerator(min,
+                                     range,
+                                     _field,
+                                     includeFields,
+                                     finalDoc,
+                                     pExpCtx->getValueComparator(),
+                                     &_docsGenerated);
+    }
+    void createDocGenerator(DensifyValueType min, RangeStatement range) {
+        createDocGenerator(min, range, boost::none, boost::none);
+    }
+
     boost::optional<DocGenerator> _docGenerator = boost::none;
 
     /**
@@ -373,5 +399,9 @@ private:
     RangeStatement _range;
     // Store of the value we've seen for each partition.
     ValueUnorderedMap<Value> _partitionTable;
+
+    // Keep track of documents generated, error if it goes above the limit.
+    size_t _docsGenerated = 0;
+    size_t _maxDocs = 0;
 };
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_densify_test.cpp b/src/mongo/db/pipeline/document_source_densify_test.cpp
index 309421c4acf..e29a207de69 100644
--- a/src/mongo/db/pipeline/document_source_densify_test.cpp
+++ b/src/mongo/db/pipeline/document_source_densify_test.cpp
@@ -75,52 +75,60 @@ Date_t makeDate(std::string dateStr) {
 
 DEATH_TEST(DensifyGeneratorTest, ErrorsIfMinOverMax, "lower or equal to max") {
     Document doc{{"a", 1}};
+    size_t counter = 0;
     ASSERT_THROWS_CODE(
         GenClass(Value(1),
                  RangeStatement(Value(1), NumericBounds(Value(1), Value(0)), boost::none),
                  "path",
                  doc,
                  doc,
-                 ValueComparator()),
+                 ValueComparator(),
+                 &counter),
         AssertionException,
         5733303);
 }
 
 DEATH_TEST(DensifyGeneratorTest, ErrorsIfStepIsZero, "be positive") {
     Document doc{{"a", 1}};
+    size_t counter = 0;
     ASSERT_THROWS_CODE(
         GenClass(Value(0),
                  RangeStatement(Value(0), NumericBounds(Value(0), Value(1)), boost::none),
                  "path",
                  doc,
                  doc,
-                 ValueComparator()),
+                 ValueComparator(),
+                 &counter),
         AssertionException,
         5733305);
 }
 
 DEATH_TEST(DensifyGeneratorTest, ErrorsOnMixedValues, "same type") {
     Document doc{{"a", 1}};
+    size_t counter = 0;
     ASSERT_THROWS_CODE(
         GenClass(Date_t::max(),
                  RangeStatement(Value(1), NumericBounds(Value(0), Value(1)), boost::none),
                  "path",
                  doc,
                  doc,
-                 ValueComparator()),
+                 ValueComparator(),
+                 &counter),
         AssertionException,
         5733300);
 }
 
 DEATH_TEST(DensifyGeneratorTest, ErrorsIfFieldExistsInDocument, "cannot include field") {
     Document doc{{"path", 1}};
+    size_t counter = 0;
     ASSERT_THROWS_CODE(
         GenClass(Value(0),
                  RangeStatement(Value(1), NumericBounds(Value(0), Value(1)), boost::none),
                  "path",
                  doc,
                  doc,
-                 ValueComparator()),
+                 ValueComparator(),
+                 &counter),
         AssertionException,
         5733306);
 }
@@ -131,13 +139,15 @@ DEATH_TEST(DensifyGeneratorTest, ErrorsIfFieldExistsButIsArray, "cannot include
     docArray.push_back(doc);
     docArray.push_back(doc);
     Document preservedFields{{"arr", Value(docArray)}};
+    size_t counter = 0;
     ASSERT_THROWS_CODE(
         GenClass(Value(0),
                  RangeStatement(Value(1), NumericBounds(Value(0), Value(1)), boost::none),
                  "arr",
                  preservedFields,
                  doc,
-                 ValueComparator()),
+                 ValueComparator(),
+                 &counter),
         AssertionException,
         5733306);
 }
@@ -148,39 +158,45 @@ TEST(DensifyGeneratorTest, ErrorsIfFieldIsInArray) {
     docArray.push_back(doc);
     docArray.push_back(doc);
     Document preservedFields{{"arr", Value(docArray)}};
+    size_t counter = 0;
     ASSERT_THROWS_CODE(
         GenClass(Value(0),
                  RangeStatement(Value(1), NumericBounds(Value(0), Value(1)), boost::none),
                  "arr.path",
                  preservedFields,
                  doc,
-                 ValueComparator()),
+                 ValueComparator(),
+                 &counter),
         AssertionException,
         5733307);
 }
 
 TEST(DensifyGeneratorTest, ErrorsIfPrefixOfFieldExists) {
     Document doc{{"a", 2}};
+    size_t counter = 0;
     ASSERT_THROWS_CODE(
         GenClass(Value(1),
                  RangeStatement(Value(1), NumericBounds(Value(1), Value(1)), boost::none),
                  "a.b",
                  doc,
                  doc,
-                 ValueComparator()),
+                 ValueComparator(),
+                 &counter),
         AssertionException,
         5733308);
 }
 
 TEST(DensifyGeneratorTest, GeneratesNumericDocumentCorrectly) {
     Document doc{{"a", 2}};
+    size_t counter = 0;
     auto generator =
         GenClass(Value(1),
                  RangeStatement(Value(1), NumericBounds(Value(1), Value(1)), boost::none),
                  "a",
                  Document(),
                  doc,
-                 ValueComparator());
+                 ValueComparator(),
+                 &counter);
     ASSERT_FALSE(generator.done());
     Document docOne{{"a", 1}};
     ASSERT_DOCUMENT_EQ(docOne, generator.getNextDocument());
@@ -190,13 +206,15 @@ TEST(DensifyGeneratorTest, GeneratesNumericDocumentCorrectly) {
 }
 
 TEST(DensifyGeneratorTest, GeneratesNumericDocumentCorrectlyWithoutFinalDoc) {
+    size_t counter = 0;
     auto generator =
         GenClass(Value(1),
                  RangeStatement(Value(1), NumericBounds(Value(1), Value(1)), boost::none),
                  "a",
                  Document(),
                  boost::none,
-                 ValueComparator());
+                 ValueComparator(),
+                 &counter);
     ASSERT_FALSE(generator.done());
     Document docOne{{"a", 1}};
     ASSERT_DOCUMENT_EQ(docOne, generator.getNextDocument());
@@ -206,13 +224,15 @@ TEST(DensifyGeneratorTest, GeneratesNumericDocumentCorrectlyWithoutFinalDoc) {
 TEST(DensifyGeneratorTest, PreservesIncludeFields) {
     Document doc{{"a", 2}, {"b", 2}, {"c", 2}};
     Document preserveFields{{"b", 1}, {"c", 1}};
+    size_t counter = 0;
     auto generator =
         GenClass(Value(1),
                  RangeStatement(Value(1), NumericBounds(Value(1), Value(1)), boost::none),
                  "a",
                  preserveFields,
                  doc,
-                 ValueComparator());
+                 ValueComparator(),
+                 &counter);
     ASSERT_FALSE(generator.done());
     Document docOne{{"b", 1}, {"c", 1}, {"a", 1}};
     ASSERT_DOCUMENT_EQ(docOne, generator.getNextDocument());
@@ -223,13 +243,15 @@ TEST(DensifyGeneratorTest, PreservesIncludeFields) {
 
 TEST(DensifyGeneratorTest, GeneratesNumberOfNumericDocumentsCorrectly) {
     Document doc{{"a", 83}};
+    size_t counter = 0;
     auto generator =
         GenClass(Value(0),
                  RangeStatement(Value(2), NumericBounds(Value(0), Value(10)), boost::none),
                  "a",
                  Document(),
                  doc,
-                 ValueComparator());
+                 ValueComparator(),
+                 &counter);
     for (int curVal = 0; curVal <= 10; curVal += 2) {
         ASSERT_FALSE(generator.done());
         Document nextDoc{{"a", curVal}};
@@ -243,13 +265,15 @@ TEST(DensifyGeneratorTest, GeneratesNumberOfNumericDocumentsCorrectly) {
 TEST(DensifyGeneratorTest, WorksWithNonIntegerStepAndPreserveFields) {
     Document doc{{"a", 2}, {"b", 2}, {"c", 2}};
     Document preserveFields{{"b", 1}, {"c", 1}};
+    size_t counter = 0;
     auto generator =
         GenClass(Value(0),
                  RangeStatement(Value(1.3), NumericBounds(Value(0), Value(10)), boost::none),
                  "a",
                  preserveFields,
                  doc,
-                 ValueComparator());
+                 ValueComparator(),
+                 &counter);
     for (double curVal = 0; curVal <= 10; curVal += 1.3) {
         ASSERT_FALSE(generator.done());
         Document nextDoc{{"b", 1}, {"c", 1}, {"a", curVal}};
@@ -262,13 +286,15 @@ TEST(DensifyGeneratorTest, WorksWithNonIntegerStepAndPreserveFields) {
 
 TEST(DensifyGeneratorTest, GeneratesOffsetFromMaxDocsCorrectly) {
     Document doc{{"a", 83}};
+    size_t counter = 0;
     auto generator =
         GenClass(Value(1),
                  RangeStatement(Value(2), NumericBounds(Value(1), Value(11)), boost::none),
                  "a",
                  Document(),
                  doc,
-                 ValueComparator());
+                 ValueComparator(),
+                 &counter);
     for (int curVal = 1; curVal <= 11; curVal += 2) {
         ASSERT_FALSE(generator.done());
         Document nextDoc{{"a", curVal}};
@@ -282,13 +308,15 @@ TEST(DensifyGeneratorTest, GeneratesOffsetFromMaxDocsCorrectly) {
 TEST(DensifyGeneratorTest, GeneratesAtDottedPathCorrectly) {
     Document doc{{"a", 83}};
     Document preservedFields{{"a", Document{{"b", 1}}}};
+    size_t counter = 0;
     auto generator =
         GenClass(Value(1),
                  RangeStatement(Value(2), NumericBounds(Value(1), Value(11)), boost::none),
                  "a.c",
                  preservedFields,
                  doc,
-                 ValueComparator());
+                 ValueComparator(),
+                 &counter);
     for (int curVal = 1; curVal <= 11; curVal += 2) {
         ASSERT_FALSE(generator.done());
         Document nextDoc{{"a", Document{{"b", 1}, {"c", curVal}}}};
@@ -297,6 +325,7 @@ TEST(DensifyGeneratorTest, GeneratesAtDottedPathCorrectly) {
     ASSERT_FALSE(generator.done());
     ASSERT_DOCUMENT_EQ(doc, generator.getNextDocument());
     ASSERT_TRUE(generator.done());
+    counter = 0;
     // Test deeply nested fields.
     Document secondPreservedFields{{"a", Document{{"b", 1}}}};
     generator = GenClass(Value(1),
@@ -304,7 +333,8 @@ TEST(DensifyGeneratorTest, GeneratesAtDottedPathCorrectly) {
                          "a.c.d",
                          secondPreservedFields,
                          doc,
-                         ValueComparator());
+                         ValueComparator(),
+                         &counter);
     for (int curVal = 1; curVal <= 11; curVal += 2) {
         ASSERT_FALSE(generator.done());
         Document nextDoc{{"a", Document{{"b", 1}, {"c", Document{{"d", curVal}}}}}};
@@ -316,6 +346,7 @@ TEST(DensifyGeneratorTest, GeneratesAtDottedPathCorrectly) {
 }
 
 DEATH_TEST(DensifyGeneratorTest, FailsIfDatesAndUnitNotProvided, "date step") {
+    size_t counter = 0;
     ASSERT_THROWS_CODE(GenClass(makeDate("2021-01-01T00:00:00.000Z"),
                                 RangeStatement(Value(1),
                                                DateBounds(makeDate("2021-01-01T00:00:00.000Z"),
@@ -324,24 +355,28 @@ DEATH_TEST(DensifyGeneratorTest, FailsIfDatesAndUnitNotProvided, "date step") {
                                 "a",
                                 Document(),
                                 Document(),
-                                ValueComparator()),
+                                ValueComparator(),
+                                &counter),
                        AssertionException,
                        5733501);
 }
 
 DEATH_TEST(DensifyGeneratorTest, FailsIfNumberAndUnitProvided, "non-date") {
+    size_t counter = 0;
     ASSERT_THROWS_CODE(
         GenClass(Value(1),
                  RangeStatement(Value(1), NumericBounds(Value(1), Value(10)), TimeUnit::second),
                  "a",
                  Document(),
                  Document(),
-                 ValueComparator()),
+                 ValueComparator(),
+                 &counter),
         AssertionException,
         5733506);
 }
 
 DEATH_TEST(DensifyGeneratorTest, DateMinMustBeLessThanMax, "lower or equal to") {
+    size_t counter = 0;
     ASSERT_THROWS_CODE(GenClass(makeDate("2021-01-01T00:00:02.000Z"),
                                 RangeStatement(Value(1),
                                                DateBounds(makeDate("2021-01-01T00:00:02.000Z"),
@@ -350,12 +385,14 @@ DEATH_TEST(DensifyGeneratorTest, DateMinMustBeLessThanMax, "lower or equal to")
                                 "a",
                                 Document(),
                                 Document(),
-                                ValueComparator()),
+                                ValueComparator(),
+                                &counter),
                        AssertionException,
                        5733502);
 }
 
 DEATH_TEST(DensifyGeneratorTest, DateStepMustBeInt, "integer") {
+    size_t counter = 0;
     ASSERT_THROWS_CODE(GenClass(makeDate("2021-01-01T00:00:00.000Z"),
                                 RangeStatement(Value(1.5),
                                                DateBounds(makeDate("2021-01-01T00:00:00.000Z"),
@@ -364,12 +401,14 @@ DEATH_TEST(DensifyGeneratorTest, DateStepMustBeInt, "integer") {
                                 "a",
                                 Document(),
                                 Document(),
-                                ValueComparator()),
+                                ValueComparator(),
+                                &counter),
                        AssertionException,
                        5733505);
 }
 
 TEST(DensifyGeneratorTest, GeneratesDatesBySecondCorrectly) {
+    size_t counter = 0;
     Document doc{{"a", 83}};
     std::string dateBase = "2021-01-01T00:00:";
     auto generator = GenClass(makeDate("2021-01-01T00:00:01.000Z"),
@@ -380,7 +419,8 @@ TEST(DensifyGeneratorTest, GeneratesDatesBySecondCorrectly) {
                               "a",
                               Document(),
                               doc,
-                              ValueComparator());
+                              ValueComparator(),
+                              &counter);
     for (int curVal = 1; curVal <= 11; curVal += 2) {
         auto appendStr = std::to_string(curVal);
         appendStr.insert(appendStr.begin(), 2 - appendStr.length(), '0');
@@ -395,6 +435,7 @@ TEST(DensifyGeneratorTest, GeneratesDatesBySecondCorrectly) {
 
 TEST(DensifyGeneratorTest, GeneratesDatesByHourCorrectly) {
     Document doc{{"a", 83}};
+    size_t counter = 0;
     std::string dateBase = "2021-01-01T";
     auto generator = GenClass(makeDate("2021-01-01T01:00:00.000Z"),
                               RangeStatement(Value(2),
@@ -404,7 +445,8 @@ TEST(DensifyGeneratorTest, GeneratesDatesByHourCorrectly) {
                               "a",
                               Document(),
                               doc,
-                              ValueComparator());
+                              ValueComparator(),
+                              &counter);
     for (int curVal = 1; curVal <= 15; curVal += 2) {
         auto appendStr = std::to_string(curVal);
         appendStr.insert(appendStr.begin(), 2 - appendStr.length(), '0');
@@ -418,6 +460,7 @@ TEST(DensifyGeneratorTest, GeneratesDatesByHourCorrectly) {
 }
 
 TEST(DensifyGeneratorTest, GeneratesDatesByMonthCorrectly) {
+    size_t counter = 0;
     Document doc{{"a", 83}};
     std::string dateBase = "2021-";
     auto generator = GenClass(makeDate("2021-01-01T01:00:00.000Z"),
@@ -428,7 +471,8 @@ TEST(DensifyGeneratorTest, GeneratesDatesByMonthCorrectly) {
                               "a",
                               Document(),
                               doc,
-                              ValueComparator());
+                              ValueComparator(),
+                              &counter);
     for (int curVal = 1; curVal <= 10; curVal += 2) {
         auto appendStr = std::to_string(curVal);
         appendStr.insert(appendStr.begin(), 2 - appendStr.length(), '0');
diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl
index 3e61364acf2..719894967ed 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -539,3 +539,12 @@ server_parameters:
     cpp_varname: "enableTimeoutOfInactiveSessionCursors"
     cpp_vartype: AtomicWord<bool>
     default: false
+
+  internalQueryMaxAllowedDensifyDocs:
+    description: "Limits the number of documents that $densify is allowed to generate."
+    set_at: [ startup, runtime ]
+    cpp_varname: "internalQueryMaxAllowedDensifyDocs"
+    cpp_vartype: AtomicWord<int>
+    default:  500000
+    validator:
+        gt: 0