SERVER-40317 Fail query when $facet intermediate output exceeds 100MB

Co-authored-by: Justin Seyster <justin.seyster@mongodb.com>
Co-authored-by: Jacob Evans <jacob.evans@10gen.com>
(cherry picked from commit 6b6e686be20ed63446111445982513ebfb94a8cb)

4 files changed, 45 insertions, 8 deletions

diff --git a/src/mongo/db/pipeline/document_source_facet.cpp b/src/mongo/db/pipeline/document_source_facet.cpp
index 9e2f13bf07d..1ac15267687 100644
--- a/src/mongo/db/pipeline/document_source_facet.cpp
+++ b/src/mongo/db/pipeline/document_source_facet.cpp
@@ -58,10 +58,13 @@ using std::string;
 using std::vector;
 
 DocumentSourceFacet::DocumentSourceFacet(std::vector<FacetPipeline> facetPipelines,
-                                         const intrusive_ptr<ExpressionContext>& expCtx)
+                                         const intrusive_ptr<ExpressionContext>& expCtx,
+                                         size_t bufferSizeBytes,
+                                         size_t maxOutputDocBytes)
     : DocumentSource(expCtx),
-      _teeBuffer(TeeBuffer::create(facetPipelines.size())),
-      _facets(std::move(facetPipelines)) {
+      _teeBuffer(TeeBuffer::create(facetPipelines.size(), bufferSizeBytes)),
+      _facets(std::move(facetPipelines)),
+      _maxOutputDocSizeBytes(maxOutputDocBytes) {
     for (size_t facetId = 0; facetId < _facets.size(); ++facetId) {
         auto& facet = _facets[facetId];
         facet.pipeline->addInitialSource(
@@ -148,8 +151,12 @@ REGISTER_DOCUMENT_SOURCE(facet,
                          DocumentSourceFacet::createFromBson);
 
 intrusive_ptr<DocumentSourceFacet> DocumentSourceFacet::create(
-    std::vector<FacetPipeline> facetPipelines, const intrusive_ptr<ExpressionContext>& expCtx) {
-    return new DocumentSourceFacet(std::move(facetPipelines), expCtx);
+    std::vector<FacetPipeline> facetPipelines,
+    const intrusive_ptr<ExpressionContext>& expCtx,
+    size_t bufferSizeBytes,
+    size_t maxOutputDocBytes) {
+    return new DocumentSourceFacet(
+        std::move(facetPipelines), expCtx, bufferSizeBytes, maxOutputDocBytes);
 }
 
 void DocumentSourceFacet::setSource(DocumentSource* source) {
@@ -170,6 +177,17 @@ DocumentSource::GetNextResult DocumentSourceFacet::getNext() {
         return GetNextResult::makeEOF();
     }
 
+    const size_t maxBytes = _maxOutputDocSizeBytes;
+    auto ensureUnderMemoryLimit = [ usedBytes = 0ul, &maxBytes ](long long additional) mutable {
+        usedBytes += additional;
+        uassert(4031700,
+                str::stream() << "document constructed by $facet is " << usedBytes
+                              << " bytes, which exceeds the limit of "
+                              << maxBytes
+                              << " bytes",
+                usedBytes <= maxBytes);
+    };
+
     vector<vector<Value>> results(_facets.size());
     bool allPipelinesEOF = false;
     while (!allPipelinesEOF) {
@@ -178,6 +196,7 @@ DocumentSource::GetNextResult DocumentSourceFacet::getNext() {
             const auto& pipeline = _facets[facetId].pipeline;
             auto next = pipeline->getSources().back()->getNext();
             for (; next.isAdvanced(); next = pipeline->getSources().back()->getNext()) {
+                ensureUnderMemoryLimit(next.getDocument().getApproximateSize());
                 results[facetId].emplace_back(next.releaseDocument());
             }
             allPipelinesEOF = allPipelinesEOF && next.isEOF();
@@ -323,6 +342,6 @@ intrusive_ptr<DocumentSource> DocumentSourceFacet::createFromBson(
         facetPipelines.emplace_back(facetName, std::move(pipeline));
     }
 
-    return new DocumentSourceFacet(std::move(facetPipelines), expCtx);
+    return DocumentSourceFacet::create(std::move(facetPipelines), expCtx);
 }
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_facet.h b/src/mongo/db/pipeline/document_source_facet.h
index 1d1ab36b163..ecc1b96f9f4 100644
--- a/src/mongo/db/pipeline/document_source_facet.h
+++ b/src/mongo/db/pipeline/document_source_facet.h
@@ -99,7 +99,9 @@ public:
 
     static boost::intrusive_ptr<DocumentSourceFacet> create(
         std::vector<FacetPipeline> facetPipelines,
-        const boost::intrusive_ptr<ExpressionContext>& expCtx);
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        size_t bufferSizeBytes = internalQueryFacetBufferSizeBytes.load(),
+        size_t maxOutputDocBytes = internalQueryFacetMaxOutputDocSizeBytes.load());
 
     /**
      * Blocking call. Will consume all input and produces one output document.
@@ -153,13 +155,17 @@ protected:
 
 private:
     DocumentSourceFacet(std::vector<FacetPipeline> facetPipelines,
-                        const boost::intrusive_ptr<ExpressionContext>& expCtx);
+                        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                        size_t bufferSizeBytes,
+                        size_t maxOutputDocBytes);
 
     Value serialize(boost::optional<ExplainOptions::Verbosity> explain = boost::none) const final;
 
     boost::intrusive_ptr<TeeBuffer> _teeBuffer;
     std::vector<FacetPipeline> _facets;
 
+    const size_t _maxOutputDocSizeBytes;
+
     bool _done = false;
 };
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_knobs.cpp b/src/mongo/db/query/query_knobs.cpp
index bdde620c87f..9e4a9aa3931 100644
--- a/src/mongo/db/query/query_knobs.cpp
+++ b/src/mongo/db/query/query_knobs.cpp
@@ -72,6 +72,15 @@ MONGO_EXPORT_SERVER_PARAMETER(internalQueryExecYieldPeriodMS, int, 10);
 
 MONGO_EXPORT_SERVER_PARAMETER(internalQueryFacetBufferSizeBytes, int, 100 * 1024 * 1024);
 
+MONGO_EXPORT_SERVER_PARAMETER(internalQueryFacetMaxOutputDocSizeBytes, long long, 100 * 1024 * 1024)
+    ->withValidator([](const long long& newVal) {
+        if (newVal <= 0) {
+            return Status(ErrorCodes::BadValue,
+                          "internalQueryFacetMaxOutputDocSizeBytes must be positive");
+        }
+        return Status::OK();
+    });
+
 MONGO_EXPORT_SERVER_PARAMETER(internalLookupStageIntermediateDocumentMaxSizeBytes,
                               long long,
                               100 * 1024 * 1024)
diff --git a/src/mongo/db/query/query_knobs.h b/src/mongo/db/query/query_knobs.h
index 035915267e8..ed4304de4ae 100644
--- a/src/mongo/db/query/query_knobs.h
+++ b/src/mongo/db/query/query_knobs.h
@@ -118,6 +118,9 @@ const int64_t insertVectorMaxBytes = 256 * 1024;
 // The number of bytes to buffer at once during a $facet stage.
 extern AtomicInt32 internalQueryFacetBufferSizeBytes;
 
+// The maximum size in bytes of the $facet stage's output document.
+extern AtomicInt64 internalQueryFacetMaxOutputDocSizeBytes;
+
 extern AtomicInt64 internalLookupStageIntermediateDocumentMaxSizeBytes;
 
 extern AtomicInt32 internalInsertMaxBatchSize;