SERVER-31755 Create intermediate $lookup stage document size limit knob

4 files changed, 133 insertions, 4 deletions

diff --git a/jstests/noPassthrough/lookup_max_intermediate_size.js b/jstests/noPassthrough/lookup_max_intermediate_size.js
new file mode 100644
index 00000000000..33f9976c058
--- /dev/null
+++ b/jstests/noPassthrough/lookup_max_intermediate_size.js
@@ -0,0 +1,111 @@
+// This test verifies that an intermediate $lookup stage can grow larger than 16MB,
+// but no larger than internalLookupStageIntermediateDocumentMaxSizeBytes.
+// @tags: [requires_sharding]
+
+load("jstests/aggregation/extras/utils.js");  // For assertErrorCode.
+
+(function() {
+    "use strict";
+
+    // Used by testPipeline to sort result documents. All _ids must be primitives.
+    function compareId(a, b) {
+        if (a._id < b._id) {
+            return -1;
+        }
+        if (a._id > b._id) {
+            return 1;
+        }
+        return 0;
+    }
+
+    // Helper for testing that pipeline returns correct set of results.
+    function testPipeline(pipeline, expectedResult, collection) {
+        assert.eq(collection.aggregate(pipeline).toArray().sort(compareId),
+                  expectedResult.sort(compareId));
+    }
+
+    function runTest(coll, from) {
+        const db = null;  // Using the db variable is banned in this function.
+
+        //
+        // Confirm aggregation will not fail if intermediate $lookup stage exceeds 16 MB.
+        //
+        assert.commandWorked(coll.insert([
+            {"_id": 3, "same": 1},
+        ]));
+
+        const bigString = new Array(1025).toString();
+        const doc = {_id: new ObjectId(), x: bigString, same: 1};
+        const docSize = Object.bsonsize(doc);
+
+        // Number of documents in lookup to exceed maximum BSON document size.
+        // Using 20 MB instead to be safe.
+        let numDocs = Math.floor(20 * 1024 * 1024 / docSize);
+
+        let bulk = from.initializeUnorderedBulkOp();
+        for (let i = 0; i < numDocs; ++i) {
+            bulk.insert({x: bigString, same: 1});
+        }
+        assert.commandWorked(bulk.execute());
+
+        let pipeline = [
+            {$lookup: {from: "from", localField: "same", foreignField: "same", as: "arr20mb"}},
+            {$project: {_id: 1}}
+        ];
+
+        let expectedResults = [{_id: 3}];
+
+        testPipeline(pipeline, expectedResults, coll);
+
+        //
+        // Confirm aggregation will fail if intermediate $lookup stage exceeds
+        // internalLookupStageIntermediateDocumentMaxSizeBytes, set to 30 MB.
+        //
+
+        // Number of documents to exceed maximum intermediate $lookup stage document size.
+        // Using 35 MB total to be safe (20 MB from previous test + 15 MB).
+        numDocs = Math.floor(15 * 1024 * 1024 / docSize);
+
+        bulk = from.initializeUnorderedBulkOp();
+        for (let i = 0; i < numDocs; ++i) {
+            bulk.insert({x: bigString, same: 1});
+        }
+        assert.commandWorked(bulk.execute());
+
+        pipeline = [
+            {$lookup: {from: "from", localField: "same", foreignField: "same", as: "arr35mb"}},
+            {$project: {_id: 1}}
+        ];
+
+        assertErrorCode(coll, pipeline, 4568);
+    }
+
+    // Run tests on single node.
+    const standalone = MongoRunner.runMongod();
+    const db = standalone.getDB("test");
+
+    assert.commandWorked(db.adminCommand(
+        {setParameter: 1, internalLookupStageIntermediateDocumentMaxSizeBytes: 30 * 1024 * 1024}));
+
+    runTest(db.lookUp, db.from);
+
+    MongoRunner.stopMongod(standalone);
+
+    // Run tests in a sharded environment.
+    const sharded = new ShardingTest({
+        mongos: 1,
+        shards: 2,
+        rs: {
+            nodes: 1,
+            setParameter:
+                {internalLookupStageIntermediateDocumentMaxSizeBytes: 30 * 1024 * 1024}
+        }
+    });
+
+    assert(sharded.adminCommand({enableSharding: "test"}));
+
+    assert(sharded.adminCommand({shardCollection: "test.lookUp", key: {_id: 'hashed'}}));
+    runTest(sharded.getDB('test').lookUp, sharded.getDB('test').from);
+
+    sharded.stop();
+}());
diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index d0fa81f32a3..0ff8f5fb3a6 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -235,14 +235,16 @@ DocumentSource::GetNextResult DocumentSourceLookUp::getNext() {
     std::vector<Value> results;
     int objsize = 0;
 
+    const auto maxBytes = internalLookupStageIntermediateDocumentMaxSizeBytes.load();
     while (auto result = pipeline->getNext()) {
         objsize += result->getApproximateSize();
         uassert(4568,
                 str::stream() << "Total size of documents in " << _fromNs.coll()
-                              << " matching pipeline "
-                              << getUserPipelineDefinition()
-                              << " exceeds maximum document size",
-                objsize <= BSONObjMaxInternalSize);
+                              << " matching pipeline's $lookup stage exceeds "
+                              << maxBytes
+                              << " bytes",
+
+                objsize <= maxBytes);
         results.emplace_back(std::move(*result));
     }
 
diff --git a/src/mongo/db/query/query_knobs.cpp b/src/mongo/db/query/query_knobs.cpp
index b8607664b64..741b5293865 100644
--- a/src/mongo/db/query/query_knobs.cpp
+++ b/src/mongo/db/query/query_knobs.cpp
@@ -29,6 +29,8 @@
  */
 
 #include "mongo/db/query/query_knobs.h"
+
+#include "mongo/bson/util/builder.h"
 #include "mongo/db/server_options.h"
 #include "mongo/db/server_parameters.h"
 
@@ -70,6 +72,18 @@ MONGO_EXPORT_SERVER_PARAMETER(internalQueryExecYieldPeriodMS, int, 10);
 
 MONGO_EXPORT_SERVER_PARAMETER(internalQueryFacetBufferSizeBytes, int, 100 * 1024 * 1024);
 
+MONGO_EXPORT_SERVER_PARAMETER(internalLookupStageIntermediateDocumentMaxSizeBytes,
+                              long long,
+                              100 * 1024 * 1024)
+    ->withValidator([](const long long& newVal) {
+        if (newVal < BSONObjMaxInternalSize) {
+            return Status(ErrorCodes::BadValue,
+                          "internalLookupStageIntermediateDocumentMaxSizeBytes must be >= " +
+                              std::to_string(BSONObjMaxInternalSize));
+        }
+        return Status::OK();
+    });
+
 MONGO_EXPORT_SERVER_PARAMETER(internalInsertMaxBatchSize,
                               int,
                               internalQueryExecYieldIterations.load() / 2);
diff --git a/src/mongo/db/query/query_knobs.h b/src/mongo/db/query/query_knobs.h
index cb94cded8f7..f2fe5046296 100644
--- a/src/mongo/db/query/query_knobs.h
+++ b/src/mongo/db/query/query_knobs.h
@@ -118,6 +118,8 @@ const int64_t insertVectorMaxBytes = 256 * 1024;
 // The number of bytes to buffer at once during a $facet stage.
 extern AtomicInt32 internalQueryFacetBufferSizeBytes;
 
+extern AtomicInt64 internalLookupStageIntermediateDocumentMaxSizeBytes;
+
 extern AtomicInt32 internalInsertMaxBatchSize;
 
 extern AtomicInt32 internalDocumentSourceCursorBatchSizeBytes;