SERVER-76378 Estimate metadata size

(cherry picked from commit a759429a26c33a176a7b85c2e240934d471f7c6e)

6 files changed, 150 insertions, 2 deletions

diff --git a/jstests/noPassthrough/metadata_size_estimate.js b/jstests/noPassthrough/metadata_size_estimate.js
new file mode 100644
index 00000000000..55e224d2487
--- /dev/null
+++ b/jstests/noPassthrough/metadata_size_estimate.js
@@ -0,0 +1,55 @@
+// Test the impact of having too many roles
+// @tags: [requires_sharding]
+
+(function() {
+'use strict';
+
+// Use a relatively small record size to more reliably hit a tipping point where the write batching
+// logic thinks we have more space available for metadata than we really do.
+const kDataBlockSize = 64 * 1024;
+const kDataBlock = 'x'.repeat(kDataBlockSize);
+const kBSONMaxObjSize = 16 * 1024 * 1024;
+const kNumRows = (kBSONMaxObjSize / kDataBlockSize) + 5;
+
+function runTest(conn) {
+    const admin = conn.getDB('admin');
+    assert.commandWorked(admin.runCommand({createUser: 'admin', pwd: 'pwd', roles: ['root']}));
+    assert(admin.auth('admin', 'pwd'));
+
+    // Create more than 16KB of role data.
+    // These roles are grouped into a meta-role to avoid calls to `usersInfo` unexpectedly
+    // overflowing from duplication of roles/inheritedRoles plus showPrivileges.
+    const userRoles = [];
+    for (let i = 0; i < 10000; ++i) {
+        userRoles.push({db: 'qwertyuiopasdfghjklzxcvbnm_' + i, role: 'read'});
+    }
+    assert.commandWorked(
+        admin.runCommand({createRole: 'bigRole', roles: userRoles, privileges: []}));
+    assert.commandWorked(admin.runCommand({createUser: 'user', pwd: 'pwd', roles: ['bigRole']}));
+    admin.logout();
+
+    assert(admin.auth('user', 'pwd'));
+    const db = conn.getDB(userRoles[0].db);
+
+    // Fill a collection with enough rows to necessitate paging.
+    for (let i = 1; i <= kNumRows; ++i) {
+        assert.commandWorked(db.myColl.insert({_id: i, data: kDataBlock}));
+    }
+    // Verify initial write.
+    assert.eq(kNumRows, db.myColl.count({}));
+
+    // Create an aggregation which will batch up to kMaxWriteBatchSize or 16MB
+    // (not counting metadata)
+    assert.eq(0, db.myColl.aggregate([{"$out": 'yourColl'}]).itcount(), 'Aggregation failed');
+
+    // Verify the $out stage completed.
+    assert.eq(db.myColl.count({}), db.yourColl.count({}));
+    assert.eq(kNumRows, db.yourColl.count({}));
+}
+
+{
+    const st = new ShardingTest({mongos: 1, config: 1, shards: 1});
+    runTest(st.s0);
+    st.stop();
+}
+})();
diff --git a/src/mongo/db/auth/role_name.h b/src/mongo/db/auth/role_name.h
index 81055b5ac20..1cd68bf1a62 100644
--- a/src/mongo/db/auth/role_name.h
+++ b/src/mongo/db/auth/role_name.h
@@ -60,6 +60,18 @@ public:
     BSONObj toBSON() const;
 
     /**
+     * Serialized length (in bytes) of object returned by toBSON().
+     */
+    std::size_t getBSONObjSize() const {
+        return 4UL +                          // BSONObj size
+            1UL + ("role"_sd).size() + 1UL +  // FieldName elem type, FieldName, NULL.
+            4UL + getRole().size() + 1UL +    // Length of name data, name data, NULL.
+            1UL + ("db"_sd).size() + 1UL +    // DB field elem type, "db", NULL.
+            4UL + getDB().size() + 1UL +      // DB value length, DB value, NULL.
+            1UL;                              // EOD marker.
+    }
+
+    /**
      * Gets the name of the role excluding the "@dbname" component.
      */
     StringData getRole() const {
diff --git a/src/mongo/db/auth/user_name.h b/src/mongo/db/auth/user_name.h
index 4654ce4dabb..fb33d021999 100644
--- a/src/mongo/db/auth/user_name.h
+++ b/src/mongo/db/auth/user_name.h
@@ -81,6 +81,18 @@ public:
     BSONObj toBSON() const;
 
     /**
+     * Serialized length (in bytes) of object returned by toBSON().
+     */
+    std::size_t getBSONObjSize() const {
+        return 4UL +                          // BSONObj size
+            1UL + ("user"_sd).size() + 1UL +  // FieldName elem type, FieldName, NULL.
+            4UL + getUser().size() + 1UL +    // Length of name data, name data, NULL.
+            1UL + ("db"_sd).size() + 1UL +    // DB field elem type, "db", NULL.
+            4UL + getDB().size() + 1UL +      // DB value length, DB value, NULL.
+            1UL;                              // EOD marker.
+    }
+
+    /**
      * Gets the user part of a UserName.
      */
     StringData getUser() const {
diff --git a/src/mongo/db/pipeline/document_source_writer.h b/src/mongo/db/pipeline/document_source_writer.h
index 98c8590d650..bcc1ab56535 100644
--- a/src/mongo/db/pipeline/document_source_writer.h
+++ b/src/mongo/db/pipeline/document_source_writer.h
@@ -38,6 +38,7 @@
 #include "mongo/db/pipeline/document_source.h"
 #include "mongo/db/read_concern.h"
 #include "mongo/db/storage/recovery_unit.h"
+#include "mongo/rpc/metadata/impersonated_user_metadata.h"
 
 namespace mongo {
 using namespace fmt::literals;
@@ -203,8 +204,20 @@ DocumentSource::GetNextResult DocumentSourceWriter<B>::doGetNext() {
             _initialized = true;
         }
 
+        // While most metadata attached to a command is limited to less than a KB,
+        // Impersonation metadata may grow to an arbitrary size.
+        // Ask the active Client how much impersonation metadata we'll use for it,
+        // and assume the rest can fit in the 16KB already built into BSONObjMaxUserSize.
+        const auto estimatedMetadataSizeBytes =
+            rpc::estimateImpersonatedUserMetadataSize(pExpCtx->opCtx);
+        uassert(7637800,
+                "Unable to proceed with write while metadata size ({}KB) exceeds {}KB"_format(
+                    estimatedMetadataSizeBytes / 1024, BSONObjMaxUserSize / 1024),
+                estimatedMetadataSizeBytes <= BSONObjMaxUserSize);
+
+        const auto maxBatchSizeBytes = BSONObjMaxUserSize - estimatedMetadataSizeBytes;
         BatchedObjects batch;
-        int bufferedBytes = 0;
+        std::size_t bufferedBytes = 0;
 
         auto nextInput = pSource->getNext();
         for (; nextInput.isAdvanced(); nextInput = pSource->getNext()) {
@@ -215,7 +228,7 @@ DocumentSource::GetNextResult DocumentSourceWriter<B>::doGetNext() {
 
             bufferedBytes += objSize;
             if (!batch.empty() &&
-                (bufferedBytes > BSONObjMaxUserSize ||
+                (bufferedBytes > maxBatchSizeBytes ||
                  batch.size() >= write_ops::kMaxWriteBatchSize)) {
                 spill(std::move(batch));
                 batch.clear();
diff --git a/src/mongo/rpc/metadata/impersonated_user_metadata.cpp b/src/mongo/rpc/metadata/impersonated_user_metadata.cpp
index 64b78931f39..3ee1a8c4320 100644
--- a/src/mongo/rpc/metadata/impersonated_user_metadata.cpp
+++ b/src/mongo/rpc/metadata/impersonated_user_metadata.cpp
@@ -95,5 +95,55 @@ void writeAuthDataToImpersonatedUserMetadata(OperationContext* opCtx, BSONObjBui
     metadata.serialize(&section);
 }
 
+std::size_t estimateImpersonatedUserMetadataSize(OperationContext* opCtx) {
+    if (!opCtx) {
+        return 0;
+    }
+
+    // Otherwise construct a metadata section from the list of authenticated users/roles
+    auto authSession = AuthorizationSession::get(opCtx->getClient());
+    auto userNames = authSession->getImpersonatedUserNames();
+    auto roleNames = authSession->getImpersonatedRoleNames();
+    if (!userNames.more() && !roleNames.more()) {
+        userNames = authSession->getAuthenticatedUserNames();
+        roleNames = authSession->getAuthenticatedRoleNames();
+    }
+
+    // If there are no users/roles being impersonated just exit
+    if (!userNames.more() && !roleNames.more()) {
+        return 0;
+    }
+
+    std::size_t ret = 4 +                                   // BSONObj size
+        1 + kImpersonationMetadataSectionName.size() + 1 +  // "$audit" sub-object key
+        4;                                                  // $audit object length
+
+    // BSONArrayType + "impersonatedUsers" + NULL + BSONArray Length
+    ret += 1 + ImpersonatedUserMetadata::kUsersFieldName.size() + 1 + 4;
+    for (std::size_t i = 0; userNames.more(); userNames.next(), ++i) {
+        // BSONType::Object + strlen(indexId) + NULL byte
+        // to_string(i).size() will be log10(i) plus some rounding and fuzzing.
+        // Increment prior to taking the log so that we never take log10(0) which is NAN.
+        // This estimates one extra byte every time we reach (i % 10) == 9.
+        ret += 1 + static_cast<std::size_t>(1.1 + log10(i + 1)) + 1;
+        ret += userNames.get().getBSONObjSize();
+    }
+    // EOD terminator for impersonatedUsers
+    ++ret;
+
+    // BSONArrayType + "impersonatedRoles" + NULL + BSONArray Length
+    ret += 1 + ImpersonatedUserMetadata::kRolesFieldName.size() + 1 + 4;
+    for (std::size_t i = 0; roleNames.more(); roleNames.next(), ++i) {
+        // Same calculation as for UserNames above.
+        ret += 1 + static_cast<std::size_t>(1.1 + log10(i + 1)) + 1;
+        ret += roleNames.get().getBSONObjSize();
+    }
+
+    // EOD terminators for: impersonatedRoles, $audit, and metadata
+    ret += 1 + 1 + 1;
+
+    return ret;
+}
+
 }  // namespace rpc
 }  // namespace mongo
diff --git a/src/mongo/rpc/metadata/impersonated_user_metadata.h b/src/mongo/rpc/metadata/impersonated_user_metadata.h
index 71c927e80b1..21f010d5447 100644
--- a/src/mongo/rpc/metadata/impersonated_user_metadata.h
+++ b/src/mongo/rpc/metadata/impersonated_user_metadata.h
@@ -70,5 +70,11 @@ void readImpersonatedUserMetadata(const BSONElement& elem, OperationContext* opC
  */
 void writeAuthDataToImpersonatedUserMetadata(OperationContext* opCtx, BSONObjBuilder* out);
 
+/*
+ * Estimates the size of impersonation metadata which will be written by
+ * writeAuthDataToImpersonatedUserMetadata.
+ */
+std::size_t estimateImpersonatedUserMetadataSize(OperationContext* opCtx);
+
 }  // namespace rpc
 }  // namespace mongo