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

(cherry picked from commit 4fc3991bf64b33ca5f5237722bc563f8eb1a552a)
(cherry picked from commit 754b0efe8549fe2d5fa2847676985709c60412a7)
(cherry picked from commit fc232126bd406e6c85ec45b869c88d7a238e60c5)

3 files changed, 25 insertions, 10 deletions

diff --git a/src/mongo/db/query/explain.cpp b/src/mongo/db/query/explain.cpp
index 0627cd78912..0067182a3df 100644
--- a/src/mongo/db/query/explain.cpp
+++ b/src/mongo/db/query/explain.cpp
@@ -307,10 +307,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;
     }
 
@@ -397,12 +396,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());
 
@@ -428,7 +426,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);
@@ -499,7 +498,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);
@@ -595,9 +595,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.cpp b/src/mongo/db/query/query_knobs.cpp
index 4245820cc6b..2cd2c61f022 100644
--- a/src/mongo/db/query/query_knobs.cpp
+++ b/src/mongo/db/query/query_knobs.cpp
@@ -136,4 +136,17 @@ MONGO_EXPORT_SERVER_PARAMETER(internalQueryMaxAddToSetBytes, int, 100 * 1024 * 1
         }
         return Status::OK();
     });
+
+MONGO_EXPORT_SERVER_PARAMETER(internalQueryExplainSizeThresholdBytes, int, 10 * 1024 * 1024)
+    ->withValidator([](const int& newVal) {
+        if (newVal <= 0) {
+            return Status(ErrorCodes::BadValue,
+                          "internalQueryExplainSizeThresholdBytes must be positive");
+        } else if (newVal > BSONObjMaxInternalSize) {
+            return Status(ErrorCodes::BadValue,
+                          "internalQueryExplainSizeThresholdBytes cannot exceed max BSON size");
+        } else {
+            return Status::OK();
+        }
+    });
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_knobs.h b/src/mongo/db/query/query_knobs.h
index f95f0aa82bf..6165be149f8 100644
--- a/src/mongo/db/query/query_knobs.h
+++ b/src/mongo/db/query/query_knobs.h
@@ -156,4 +156,7 @@ extern AtomicBool internalQueryProhibitBlockingMergeOnMongoS;
 extern AtomicInt32 internalQueryMaxPushBytes;
 
 extern AtomicInt32 internalQueryMaxAddToSetBytes;
+
+// The number of bytes after which explain should start truncating portions of its output.
+extern AtomicInt32 internalQueryExplainSizeThresholdBytes;
 }  // namespace mongo