SERVER-54710 Improve checks for overlarge BSON when generating explain output

(cherry picked from commit 4fc3991bf64b33ca5f5237722bc563f8eb1a552a)

3 files changed, 73 insertions, 10 deletions

diff --git a/jstests/noPassthrough/explain_output_truncation.js b/jstests/noPassthrough/explain_output_truncation.js
new file mode 100644
index 00000000000..60fbc5d7d3b
--- /dev/null
+++ b/jstests/noPassthrough/explain_output_truncation.js
@@ -0,0 +1,52 @@
+/**
+ * Test that explain output is correctly truncated when it grows too large.
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/analyze_plan.js");
+
+const dbName = "test";
+const collName = jsTestName();
+const explainSizeParam = "internalQueryExplainSizeThresholdBytes";
+
+const conn = MongoRunner.runMongod({});
+assert.neq(conn, null, "mongod failed to start up");
+
+const testDb = conn.getDB(dbName);
+const coll = testDb[collName];
+coll.drop();
+
+assert.commandWorked(coll.createIndex({a: 1}));
+
+// Explain output should show a simple IXSCAN => FETCH => SORT plan with no truncation.
+let explain = coll.find({a: 1, b: 1}).sort({c: 1}).explain();
+let sortStage = getPlanStage(explain.queryPlanner.winningPlan, "SORT");
+assert.neq(sortStage, null, explain);
+let fetchStage = getPlanStage(sortStage, "FETCH");
+assert.neq(fetchStage, null, explain);
+let ixscanStage = getPlanStage(sortStage, "IXSCAN");
+assert.neq(ixscanStage, null, explain);
+
+// Calculate the size of explain output without the index scan. If the explain size threshold is set
+// near this amount, then the IXSCAN stage will need to be truncated.
+assert.neq(ixscanStage, null, explain);
+const newExplainSize = Object.bsonsize(explain.queryPlanner) - Object.bsonsize(ixscanStage) - 10;
+assert.gt(newExplainSize, 0);
+
+// Reduce the size at which we start truncating explain output. If we explain the same query again,
+// then the FETCH stage should be present, but the IXSCAN stage should be truncated.
+assert.commandWorked(testDb.adminCommand({setParameter: 1, [explainSizeParam]: newExplainSize}));
+
+explain = coll.find({a: 1, b: 1}).sort({c: 1}).explain();
+assert(planHasStage(testDb, explain, "SORT"), explain);
+fetchStage = getPlanStage(explain.queryPlanner.winningPlan, "FETCH");
+assert.neq(fetchStage, null, explain);
+assert(fetchStage.hasOwnProperty("inputStage"), explain);
+assert(fetchStage.inputStage.hasOwnProperty("warning"), explain);
+assert.eq(
+    fetchStage.inputStage.warning, "stats tree exceeded BSON size limit for explain", explain);
+assert(!planHasStage(testDb, explain, "IXSCAN"), explain);
+
+MongoRunner.stopMongod(conn);
+}());
diff --git a/src/mongo/db/query/explain.cpp b/src/mongo/db/query/explain.cpp
index 86ec5d85f9a..4c42804dce4 100644
--- a/src/mongo/db/query/explain.cpp
+++ b/src/mongo/db/query/explain.cpp
@@ -309,10 +309,9 @@ void Explain::statsToBSON(const PlanStageStats& stats,
     invariant(bob);
     invariant(topLevelBob);
 
-    // Stop as soon as the BSON object we're building exceeds 10 MB.
-    static const int kMaxStatsBSONSize = 10 * 1024 * 1024;
-    if (topLevelBob->len() > kMaxStatsBSONSize) {
-        bob->append("warning", "stats tree exceeded 10 MB");
+    // Stop as soon as the BSON object we're building exceeds the limit.
+    if (topLevelBob->len() > internalQueryExplainSizeThresholdBytes.load()) {
+        bob->append("warning", "stats tree exceeded BSON size limit for explain");
         return;
     }
 
@@ -400,12 +399,11 @@ void Explain::statsToBSON(const PlanStageStats& stats,
         bob->appendBool("isPartial", spec->isPartial);
         bob->append("indexVersion", spec->indexVersion);
 
-        BSONObjBuilder indexBoundsBob;
+        BSONObjBuilder indexBoundsBob(bob->subobjStart("indexBounds"));
         indexBoundsBob.append("startKey", spec->startKey);
         indexBoundsBob.append("startKeyInclusive", spec->startKeyInclusive);
         indexBoundsBob.append("endKey", spec->endKey);
         indexBoundsBob.append("endKeyInclusive", spec->endKeyInclusive);
-        bob->append("indexBounds", indexBoundsBob.obj());
     } else if (STAGE_DELETE == stats.stageType) {
         DeleteStats* spec = static_cast<DeleteStats*>(stats.specific.get());
 
@@ -430,7 +428,8 @@ void Explain::statsToBSON(const PlanStageStats& stats,
         bob->append("indexVersion", spec->indexVersion);
         bob->append("direction", spec->direction > 0 ? "forward" : "backward");
 
-        if ((topLevelBob->len() + spec->indexBounds.objsize()) > kMaxStatsBSONSize) {
+        if ((topLevelBob->len() + spec->indexBounds.objsize()) >
+            internalQueryExplainSizeThresholdBytes.load()) {
             bob->append("warning", "index bounds omitted due to BSON size limit");
         } else {
             bob->append("indexBounds", spec->indexBounds);
@@ -496,7 +495,8 @@ void Explain::statsToBSON(const PlanStageStats& stats,
         bob->append("indexVersion", spec->indexVersion);
         bob->append("direction", spec->direction > 0 ? "forward" : "backward");
 
-        if ((topLevelBob->len() + spec->indexBounds.objsize()) > kMaxStatsBSONSize) {
+        if ((topLevelBob->len() + spec->indexBounds.objsize()) >
+            internalQueryExplainSizeThresholdBytes.load()) {
             bob->append("warning", "index bounds omitted due to BSON size limit");
         } else {
             bob->append("indexBounds", spec->indexBounds);
@@ -598,9 +598,8 @@ void Explain::statsToBSON(const PlanStageStats& stats,
     // the output more readable by saving a level of nesting. Name the field 'inputStage'
     // rather than 'inputStages'.
     if (1 == stats.children.size()) {
-        BSONObjBuilder childBob;
+        BSONObjBuilder childBob(bob->subobjStart("inputStage"));
         statsToBSON(*stats.children[0], verbosity, &childBob, topLevelBob);
-        bob->append("inputStage", childBob.obj());
         return;
     }
 
diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl
index f9ae4b34198..71148f35dae 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -403,3 +403,15 @@ server_parameters:
     cpp_varname: "internalQueryDesugarWhereToFunction"
     cpp_vartype: AtomicWord<bool>
     default: false
+
+  internalQueryExplainSizeThresholdBytes:
+    description: "Number of bytes after which explain should start truncating portions of its
+    output."
+    set_at: [ startup, runtime ]
+    cpp_varname: "internalQueryExplainSizeThresholdBytes"
+    cpp_vartype: AtomicWord<int>
+    default:
+      expr: 10 * 1024 * 1024
+    validator:
+        gt: 0
+        lte: { expr: BSONObjMaxInternalSize }