SERVER-75261 Added accounting for array element overhead for "listCollections", "listIndexes", "_shardsvrCheckMetadataConsistencyParticipant" commands

(cherry picked from commit 3cde7fd5d90f1e6bd16d38cc668963a14671d690)

9 files changed, 165 insertions, 28 deletions

diff --git a/jstests/noPassthrough/list_collections_large_number.js b/jstests/noPassthrough/list_collections_large_number.js
new file mode 100644
index 00000000000..379d4ea5dd7
--- /dev/null
+++ b/jstests/noPassthrough/list_collections_large_number.js
@@ -0,0 +1,34 @@
+/**
+ * Tests that "listCollections" command successfully returns results when the database has a very
+ * large number of collections.
+ * @tags: [
+ *  resource_intensive,
+ * ]
+ */
+(function() {
+"use strict";
+
+const conn = MongoRunner.runMongod({});
+assert.neq(null, conn, "mongod was unable to start up");
+const db = conn.getDB(jsTestName());
+
+const validatorObj = {
+    $jsonSchema: {
+        bsonType: "object",
+        properties: {
+            s: {bsonType: "string", description: "x".repeat(4801)},
+
+        }
+    }
+};
+const nCollections = 3300;
+jsTestLog(`Creating ${nCollections} collections....`);
+for (let i = 0; i < nCollections; i++) {
+    assert.commandWorked(db.createCollection("c_" + i.toPrecision(6), {validator: validatorObj}));
+}
+jsTestLog(`Done creating ${nCollections} collections`);
+assert.commandWorked(db.runCommand({"listCollections": 1}));
+
+// Do not validate collections since that is an expensive action.
+MongoRunner.stopMongod(conn, undefined, {skipValidation: true});
+})();
\ No newline at end of file
diff --git a/src/mongo/db/commands/list_collections.cpp b/src/mongo/db/commands/list_collections.cpp
index 6d42598684e..73f27fdad13 100644
--- a/src/mongo/db/commands/list_collections.cpp
+++ b/src/mongo/db/commands/list_collections.cpp
@@ -504,7 +504,7 @@ public:
                     batchSize = *listCollRequest.getCursor()->getBatchSize();
                 }
 
-                size_t bytesBuffered = 0;
+                FindCommon::BSONArrayResponseSizeTracker responseSizeTracker;
                 for (long long objCount = 0; objCount < batchSize; objCount++) {
                     BSONObj nextDoc;
                     PlanExecutor::ExecState state = exec->getNext(&nextDoc, nullptr);
@@ -515,7 +515,7 @@ public:
 
                     // If we can't fit this result inside the current batch, then we stash it for
                     // later.
-                    if (!FindCommon::haveSpaceForNext(nextDoc, objCount, bytesBuffered)) {
+                    if (!responseSizeTracker.haveSpaceForNext(nextDoc)) {
                         exec->stashResult(nextDoc);
                         break;
                     }
@@ -531,7 +531,7 @@ public:
                             "error"_attr = exc);
                         fassertFailed(5254301);
                     }
-                    bytesBuffered += nextDoc.objsize();
+                    responseSizeTracker.add(nextDoc);
                 }
                 if (exec->isEOF()) {
                     return createListCollectionsCursorReply(
diff --git a/src/mongo/db/commands/list_indexes.cpp b/src/mongo/db/commands/list_indexes.cpp
index 02a37da59ee..046e2eba230 100644
--- a/src/mongo/db/commands/list_indexes.cpp
+++ b/src/mongo/db/commands/list_indexes.cpp
@@ -314,7 +314,7 @@ public:
                                             nss));
 
             std::vector<mongo::ListIndexesReplyItem> firstBatch;
-            size_t bytesBuffered = 0;
+            FindCommon::BSONArrayResponseSizeTracker responseSizeTracker;
             for (long long objCount = 0; objCount < batchSize; objCount++) {
                 BSONObj nextDoc;
                 PlanExecutor::ExecState state = exec->getNext(&nextDoc, nullptr);
@@ -326,7 +326,7 @@ public:
 
                 // If we can't fit this result inside the current batch, then we stash it for
                 // later.
-                if (!FindCommon::haveSpaceForNext(nextDoc, objCount, bytesBuffered)) {
+                if (!responseSizeTracker.haveSpaceForNext(nextDoc)) {
                     exec->stashResult(nextDoc);
                     break;
                 }
@@ -345,7 +345,7 @@ public:
                                           nextDoc.toString(),
                                           exc.toString()));
                 }
-                bytesBuffered += nextDoc.objsize();
+                responseSizeTracker.add(nextDoc);
             }
 
             if (exec->isEOF()) {
diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript
index a2c337ce8bd..9cab657361e 100644
--- a/src/mongo/db/query/SConscript
+++ b/src/mongo/db/query/SConscript
@@ -410,6 +410,7 @@ env.CppUnitTest(
         "classic_stage_builder_test.cpp",
         "count_command_test.cpp",
         "cursor_response_test.cpp",
+        "find_common_test.cpp",
         "get_executor_test.cpp",
         "getmore_request_test.cpp",
         "hint_parser_test.cpp",
diff --git a/src/mongo/db/query/find_common.cpp b/src/mongo/db/query/find_common.cpp
index 0be74f536ec..4597d9a5233 100644
--- a/src/mongo/db/query/find_common.cpp
+++ b/src/mongo/db/query/find_common.cpp
@@ -132,5 +132,17 @@ std::size_t FindCommon::getBytesToReserveForGetMoreReply(bool isTailable,
     // command metadata to the reply.
     return kMaxBytesToReturnToClientAtOnce;
 }
+bool FindCommon::BSONArrayResponseSizeTracker::haveSpaceForNext(const BSONObj& document) {
+    return FindCommon::haveSpaceForNext(document, _numberOfDocuments, _bsonArraySizeInBytes);
+}
+void FindCommon::BSONArrayResponseSizeTracker::add(const BSONObj& document) {
+    dassert(haveSpaceForNext(document));
+    ++_numberOfDocuments;
+    _bsonArraySizeInBytes += (document.objsize() + kPerDocumentOverheadBytesUpperBound);
+}
 
+// Upper bound of BSON array element overhead. The overhead is 1 byte/doc for the type + 1 byte/doc
+// for the field name's null terminator + 1 byte per digit of the maximum array index value.
+const size_t FindCommon::BSONArrayResponseSizeTracker::kPerDocumentOverheadBytesUpperBound{
+    2 + std::to_string(BSONObjMaxUserSize / BSONObj::kMinBSONLength).length()};
 }  // namespace mongo
diff --git a/src/mongo/db/query/find_common.h b/src/mongo/db/query/find_common.h
index 45f60d2fd51..d38d580bdd2 100644
--- a/src/mongo/db/query/find_common.h
+++ b/src/mongo/db/query/find_common.h
@@ -89,7 +89,7 @@ public:
     // This max may be exceeded by epsilon for output documents that approach the maximum user
     // document size. That is, if we must return a BSONObjMaxUserSize document, then the total
     // response size will be BSONObjMaxUserSize plus the amount of size required for the message
-    // header and the cursor response "envelope". (The envolope contains namespace and cursor id
+    // header and the cursor response "envelope". (The envelope contains namespace and cursor id
     // info.)
     static const size_t kMaxBytesToReturnToClientAtOnce;
 
@@ -148,6 +148,32 @@ public:
     static std::size_t getBytesToReserveForGetMoreReply(bool isTailable,
                                                         size_t firstResultSize,
                                                         size_t batchSize);
+
+    /**
+     * Tracker of a size of a server response presented as a BSON array. Facilitates limiting the
+     * server response size to 16MB + certain epsilon. Accounts for array element and it's overhead
+     * size. Does not account for response "envelope" size.
+     */
+    class BSONArrayResponseSizeTracker {
+        // Upper bound of BSON array element overhead.
+        static const size_t kPerDocumentOverheadBytesUpperBound;
+
+    public:
+        /**
+         * Returns true only if 'document' can be added to the BSON array without violating the
+         * overall response size limit or if it is the first document.
+         */
+        bool haveSpaceForNext(const BSONObj& document);
+
+        /**
+         * Records that 'document' was added to the response.
+         */
+        void add(const BSONObj& document);
+
+    private:
+        std::size_t _numberOfDocuments{0};
+        std::size_t _bsonArraySizeInBytes{0};
+    };
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/find_common_test.cpp b/src/mongo/db/query/find_common_test.cpp
new file mode 100644
index 00000000000..d7dfc10d950
--- /dev/null
+++ b/src/mongo/db/query/find_common_test.cpp
@@ -0,0 +1,74 @@
+/**
+ *    Copyright (C) 2023-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include <string>
+
+#include "mongo/bson/bsonobj.h"
+#include "mongo/bson/bsonobjbuilder.h"
+#include "mongo/db/query/find_common.h"
+
+#include "mongo/unittest/unittest.h"
+
+namespace {
+
+using namespace mongo;
+
+TEST(BSONArrayResponseSizeTrackerTest, AddLargeNumberOfElements) {
+    BSONObjBuilder bsonObjBuilder;
+    {
+        FindCommon::BSONArrayResponseSizeTracker sizeTracker;
+        BSONArrayBuilder arrayBuilder{bsonObjBuilder.subarrayStart("a")};
+        BSONObj emptyObject;
+        while (sizeTracker.haveSpaceForNext(emptyObject)) {
+            sizeTracker.add(emptyObject);
+            arrayBuilder.append(emptyObject);
+        }
+    }
+    // If the BSON object is successfully constructed, then space accounting was correct.
+    bsonObjBuilder.obj();
+}
+TEST(BSONArrayResponseSizeTrackerTest, CanAddAtLeastOneDocument) {
+    auto largeObject = BSON("a" << std::string(16 * 1024 * 1024, 'A'));
+    BSONObj emptyObject;
+    BSONObjBuilder bsonObjBuilder;
+    {
+        FindCommon::BSONArrayResponseSizeTracker sizeTracker;
+        BSONArrayBuilder arrayBuilder{bsonObjBuilder.subarrayStart("a")};
+        // Add an object that is larger than 16MB.
+        ASSERT(sizeTracker.haveSpaceForNext(largeObject));
+        sizeTracker.add(largeObject);
+        arrayBuilder.append(largeObject);
+        ASSERT(!sizeTracker.haveSpaceForNext(emptyObject));
+    }
+    // If the BSON object is successfully constructed, then space accounting was correct.
+    bsonObjBuilder.obj();
+}
+}  // namespace
diff --git a/src/mongo/db/s/metadata_consistency_util.cpp b/src/mongo/db/s/metadata_consistency_util.cpp
index 0d0b878944b..6e74aca45d6 100644
--- a/src/mongo/db/s/metadata_consistency_util.cpp
+++ b/src/mongo/db/s/metadata_consistency_util.cpp
@@ -115,7 +115,7 @@ CheckMetadataConsistencyResponseCursor _makeCursor(
                                                     nss));
 
     std::vector<MetadataInconsistencyItem> firstBatch;
-    size_t bytesBuffered = 0;
+    FindCommon::BSONArrayResponseSizeTracker responseSizeTracker;
     for (long long objCount = 0; objCount < batchSize; objCount++) {
         BSONObj nextDoc;
         PlanExecutor::ExecState state = exec->getNext(&nextDoc, nullptr);
@@ -126,15 +126,14 @@ CheckMetadataConsistencyResponseCursor _makeCursor(
 
         // If we can't fit this result inside the current batch, then we stash it for
         // later.
-        if (!FindCommon::haveSpaceForNext(nextDoc, objCount, bytesBuffered)) {
+        if (!responseSizeTracker.haveSpaceForNext(nextDoc)) {
             exec->stashResult(nextDoc);
             break;
         }
 
-        const auto objsize = nextDoc.objsize();
+        responseSizeTracker.add(nextDoc);
         firstBatch.push_back(MetadataInconsistencyItem::parseOwned(
             IDLParserContext("MetadataInconsistencyItem"), std::move(nextDoc)));
-        bytesBuffered += objsize;
     }
 
     if (exec->isEOF()) {
diff --git a/src/mongo/s/query/cluster_find.cpp b/src/mongo/s/query/cluster_find.cpp
index c9f22408463..543cfd257b6 100644
--- a/src/mongo/s/query/cluster_find.cpp
+++ b/src/mongo/s/query/cluster_find.cpp
@@ -87,11 +87,6 @@ static const BSONObj kSortKeyMetaProjection = BSON("$meta"
                                                    << "sortKey");
 static const BSONObj kGeoNearDistanceMetaProjection = BSON("$meta"
                                                            << "geoNearDistance");
-// We must allow some amount of overhead per result document, since when we make a cursor response
-// the documents are elements of a BSONArray. The overhead is 1 byte/doc for the type + 1 byte/doc
-// for the field name's null terminator + 1 byte per digit in the array index. The index can be no
-// more than 8 decimal digits since the response is at most 16MB, and 16 * 1024 * 1024 < 1 * 10^8.
-static const int kPerDocumentOverheadBytesUpperBound = 10;
 
 const char kFindCmdName[] = "find";
 
@@ -392,7 +387,7 @@ CursorId runQueryWithoutRetrying(OperationContext* opCtx,
     }
 
     auto cursorState = ClusterCursorManager::CursorState::NotExhausted;
-    size_t bytesBuffered = 0;
+    FindCommon::BSONArrayResponseSizeTracker responseSizeTracker;
 
     // This loop will load enough results from the shards for a full first batch.  At first, these
     // results come from the initial batches that were obtained when establishing cursors, but
@@ -419,14 +414,13 @@ CursorId runQueryWithoutRetrying(OperationContext* opCtx,
 
         // If adding this object will cause us to exceed the message size limit, then we stash it
         // for later.
-        if (!FindCommon::haveSpaceForNext(nextObj, results->size(), bytesBuffered)) {
+        if (!responseSizeTracker.haveSpaceForNext(nextObj)) {
             ccc->queueResult(nextObj);
             break;
         }
 
-        // Add doc to the batch. Account for the space overhead associated with returning this doc
-        // inside a BSON array.
-        bytesBuffered += (nextObj.objsize() + kPerDocumentOverheadBytesUpperBound);
+        // Add doc to the batch.
+        responseSizeTracker.add(nextObj);
         results->push_back(std::move(nextObj));
     }
 
@@ -830,7 +824,7 @@ StatusWith<CursorResponse> ClusterFind::runGetMore(OperationContext* opCtx,
     }
 
     std::vector<BSONObj> batch;
-    size_t bytesBuffered = 0;
+    FindCommon::BSONArrayResponseSizeTracker responseSizeTracker;
     long long batchSize = cmd.getBatchSize().value_or(0);
     auto cursorState = ClusterCursorManager::CursorState::NotExhausted;
     BSONObj postBatchResumeToken;
@@ -895,8 +889,7 @@ StatusWith<CursorResponse> ClusterFind::runGetMore(OperationContext* opCtx,
             break;
         }
 
-        if (!FindCommon::haveSpaceForNext(
-                *next.getValue().getResult(), batch.size(), bytesBuffered)) {
+        if (!responseSizeTracker.haveSpaceForNext(*next.getValue().getResult())) {
             pinnedCursor.getValue()->queueResult(*next.getValue().getResult());
             stashedResult = true;
             break;
@@ -905,10 +898,8 @@ StatusWith<CursorResponse> ClusterFind::runGetMore(OperationContext* opCtx,
         // As soon as we get a result, this operation no longer waits.
         awaitDataState(opCtx).shouldWaitForInserts = false;
 
-        // Add doc to the batch. Account for the space overhead associated with returning this doc
-        // inside a BSON array.
-        bytesBuffered +=
-            (next.getValue().getResult()->objsize() + kPerDocumentOverheadBytesUpperBound);
+        // Add doc to the batch.
+        responseSizeTracker.add(*next.getValue().getResult());
         batch.push_back(std::move(*next.getValue().getResult()));
 
         // Update the postBatchResumeToken. For non-$changeStream aggregations, this will be empty.