SERVER-43096 Pre-allocate smaller buffer for "distinct" commands

(cherry picked from commit 22256673a9db8f9cf44bdb149fe20a599ecb7e9a)

2 files changed, 40 insertions, 15 deletions

diff --git a/jstests/noPassthrough/distinct_cmd_response_too_large.js b/jstests/noPassthrough/distinct_cmd_response_too_large.js
new file mode 100644
index 00000000000..79f8d73c82c
--- /dev/null
+++ b/jstests/noPassthrough/distinct_cmd_response_too_large.js
@@ -0,0 +1,23 @@
+/**
+ * Confirms that mongod will return an error when the result generated for a distinct command
+ * exceeds MaxBSONSize.
+ */
+(function() {
+"use strict";
+
+const conn = MongoRunner.runMongod();
+const db = conn.getDB('test');
+const coll = db.test;
+
+const largeString = new Array(1000 * 1000).join('x');
+
+let bulk = coll.initializeUnorderedBulkOp();
+for (let x = 0; x < 17; ++x) {
+    bulk.insert({_id: (largeString + x.toString())});
+}
+assert.commandWorked(bulk.execute());
+
+assert.commandFailedWithCode(db.runCommand({distinct: "test", key: "_id", query: {}}), 17217);
+
+MongoRunner.stopMongod(conn);
+})();
diff --git a/src/mongo/db/commands/distinct.cpp b/src/mongo/db/commands/distinct.cpp
index 507894ffb4a..24b1c2518dc 100644
--- a/src/mongo/db/commands/distinct.cpp
+++ b/src/mongo/db/commands/distinct.cpp
@@ -228,13 +228,11 @@ public:
 
         const auto key = cmdObj[ParsedDistinct::kKeyField].valuestrsafe();
 
-        int bufSize = BSONObjMaxUserSize - 4096;
-        BufBuilder bb(bufSize);
-        char* start = bb.buf();
-
-        BSONArrayBuilder arr(bb);
+        std::vector<BSONObj> distinctValueHolder;
         BSONElementSet values(executor.getValue()->getCanonicalQuery()->getCollator());
 
+        const int kMaxResponseSize = BSONObjMaxUserSize - 4096;
+        size_t listApproxBytes = 0;
         BSONObj obj;
         PlanExecutor::ExecState state;
         while (PlanExecutor::ADVANCED == (state = executor.getValue()->getNext(&obj, NULL))) {
@@ -251,15 +249,16 @@ public:
                 if (values.count(elt)) {
                     continue;
                 }
-                int currentBufPos = bb.len();
 
-                uassert(17217,
-                        "distinct too big, 16mb cap",
-                        (currentBufPos + elt.size() + 1024) < bufSize);
+                // This is an approximate size check which safeguards against use of unbounded
+                // memory by the distinct command. We perform a more precise check at the end of
+                // this method to confirm that the response size is less than 16MB.
+                listApproxBytes += elt.size();
+                uassert(17217, "distinct too big, 16mb cap", listApproxBytes < kMaxResponseSize);
 
-                arr.append(elt);
-                BSONElement x(start + currentBufPos);
-                values.insert(x);
+                auto distinctObj = elt.wrap();
+                values.insert(distinctObj.firstElement());
+                distinctValueHolder.push_back(std::move(distinctObj));
             }
         }
 
@@ -293,10 +292,13 @@ public:
             curOp->debug().execStats = execStatsBob.obj();
         }
 
-        verify(start == bb.buf());
-
-        result.appendArray("values", arr.done());
+        BSONArrayBuilder valueListBuilder(result.subarrayStart("values"));
+        for (const auto& value : values) {
+            valueListBuilder.append(value);
+        }
+        valueListBuilder.doneFast();
 
+        uassert(31299, "distinct too big, 16mb cap", result.len() < kMaxResponseSize);
         return true;
     }