1 files changed, 33 insertions, 3 deletions

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