1 files changed, 297 insertions, 0 deletions
diff --git a/src/mongo/db/ops/write_ops.cpp b/src/mongo/db/ops/write_ops.cpp
index 54cef4d3d2a..de7df85e564 100644
--- a/src/mongo/db/ops/write_ops.cpp
+++ b/src/mongo/db/ops/write_ops.cpp
@@ -54,6 +54,18 @@ using write_ops::WriteCommandRequestBase;
 
 namespace {
 
+// This constant accounts for the null terminator in each field name and the BSONType byte for
+// each element.
+static constexpr int kPerElementOverhead = 2;
+
+// This constant accounts for the size of a bool.
+static constexpr int kBoolSize = 1;
+
+// This constant tracks the overhead for serializing UUIDs. It includes 1 byte for the
+// 'BinDataType', 4 bytes for serializing the integer size of the UUID, and finally, 16 bytes
+// for the UUID itself.
+static const int kUUIDSize = 21;
+
 template <class T>
 void checkOpCountForCommand(const T& op, size_t numOps) {
     uassert(ErrorCodes::InvalidLength,
@@ -78,6 +90,50 @@ void checkOpCountForCommand(const T& op, size_t numOps) {
     }
 }
 
+// Utility which estimates the size of 'WriteCommandRequestBase' when serialized.
+int getWriteCommandRequestBaseSize(const WriteCommandRequestBase& base) {
+    static const int kSizeOfOrderedField =
+        write_ops::WriteCommandRequestBase::kOrderedFieldName.size() + kBoolSize +
+        kPerElementOverhead;
+    static const int kSizeOfBypassDocumentValidationField =
+        write_ops::WriteCommandRequestBase::kBypassDocumentValidationFieldName.size() + kBoolSize +
+        kPerElementOverhead;
+
+    auto estSize = static_cast<int>(BSONObj::kMinBSONLength) + kSizeOfOrderedField +
+        kSizeOfBypassDocumentValidationField;
+
+    if (auto stmtId = base.getStmtId(); stmtId) {
+        estSize += write_ops::WriteCommandRequestBase::kStmtIdFieldName.size() +
+            sizeof(std::int32_t) + kPerElementOverhead;
+    }
+
+    if (auto stmtIds = base.getStmtIds(); stmtIds) {
+        estSize += write_ops::WriteCommandRequestBase::kStmtIdsFieldName.size();
+        estSize += static_cast<int>(BSONObj::kMinBSONLength);
+        estSize +=
+            ((sizeof(std::int32_t) + write_ops::kWriteCommandBSONArrayPerElementOverheadBytes) *
+             stmtIds->size());
+        estSize += kPerElementOverhead;
+    }
+
+    if (auto isTimeseries = base.getIsTimeseriesNamespace(); isTimeseries.has_value()) {
+        estSize += write_ops::WriteCommandRequestBase::kIsTimeseriesNamespaceFieldName.size() +
+            kBoolSize + kPerElementOverhead;
+    }
+
+    if (auto collUUID = base.getCollectionUUID(); collUUID) {
+        estSize += write_ops::WriteCommandRequestBase::kCollectionUUIDFieldName.size() + kUUIDSize +
+            kPerElementOverhead;
+    }
+
+    if (auto encryptionInfo = base.getEncryptionInformation(); encryptionInfo) {
+        estSize += write_ops::WriteCommandRequestBase::kEncryptionInformationFieldName.size() +
+            encryptionInfo->toBSON().objsize() + kPerElementOverhead;
+    }
+
+    return estSize;
+}
+
 }  // namespace
 
 namespace write_ops {
@@ -137,6 +193,247 @@ int32_t getStmtIdForWriteAt(const WriteCommandRequestBase& writeCommandBase, siz
     return kFirstStmtId + writePos;
 }
 
+int estimateRuntimeConstantsSize(const mongo::LegacyRuntimeConstants& constants) {
+    int size = write_ops::UpdateCommandRequest::kLegacyRuntimeConstantsFieldName.size() +
+        static_cast<int>(BSONObj::kMinBSONLength) + kPerElementOverhead;
+
+    // $$NOW
+    size +=
+        LegacyRuntimeConstants::kLocalNowFieldName.size() + sizeof(Date_t) + kPerElementOverhead;
+
+    // $$CLUSTER_TIME
+    size += LegacyRuntimeConstants::kClusterTimeFieldName.size() + sizeof(Timestamp) +
+        kPerElementOverhead;
+
+    // $$JS_SCOPE
+    if (const auto& scope = constants.getJsScope(); scope.has_value()) {
+        size += LegacyRuntimeConstants::kJsScopeFieldName.size() + scope->objsize() +
+            kPerElementOverhead;
+    }
+
+    // $$IS_MR
+    if (const auto& isMR = constants.getIsMapReduce(); isMR.has_value()) {
+        size +=
+            LegacyRuntimeConstants::kIsMapReduceFieldName.size() + kBoolSize + kPerElementOverhead;
+    }
+
+    return size;
+}
+
+int getUpdateSizeEstimate(const BSONObj& q,
+                          const write_ops::UpdateModification& u,
+                          const boost::optional<mongo::BSONObj>& c,
+                          const bool includeUpsertSupplied,
+                          const boost::optional<mongo::BSONObj>& collation,
+                          const boost::optional<std::vector<mongo::BSONObj>>& arrayFilters,
+                          const mongo::BSONObj& hint) {
+    using UpdateOpEntry = write_ops::UpdateOpEntry;
+    int estSize = static_cast<int>(BSONObj::kMinBSONLength);
+
+    // Add the sizes of the 'multi' and 'upsert' fields.
+    estSize += UpdateOpEntry::kUpsertFieldName.size() + kBoolSize + kPerElementOverhead;
+    estSize += UpdateOpEntry::kMultiFieldName.size() + kBoolSize + kPerElementOverhead;
+
+    // Add the size of 'upsertSupplied' field if present.
+    if (includeUpsertSupplied) {
+        estSize += UpdateOpEntry::kUpsertSuppliedFieldName.size() + kBoolSize + kPerElementOverhead;
+    }
+
+    // Add the sizes of the 'q' and 'u' fields.
+    estSize += (UpdateOpEntry::kQFieldName.size() + q.objsize() + kPerElementOverhead +
+                UpdateOpEntry::kUFieldName.size() + u.objsize() + kPerElementOverhead);
+
+    // Add the size of the 'c' field, if present.
+    if (c) {
+        estSize += (UpdateOpEntry::kCFieldName.size() + c->objsize() + kPerElementOverhead);
+    }
+
+    // Add the size of the 'collation' field, if present.
+    if (collation) {
+        estSize += (UpdateOpEntry::kCollationFieldName.size() + collation->objsize() +
+                    kPerElementOverhead);
+    }
+
+    // Add the size of the 'arrayFilters' field, if present.
+    if (arrayFilters) {
+        estSize += ([&]() {
+            auto size = BSONObj::kMinBSONLength + UpdateOpEntry::kArrayFiltersFieldName.size() +
+                kPerElementOverhead;
+            for (auto&& filter : *arrayFilters) {
+                // For each filter, we not only need to account for the size of the filter itself,
+                // but also for the per array element overhead.
+                size += filter.objsize();
+                size += write_ops::kWriteCommandBSONArrayPerElementOverheadBytes;
+            }
+            return size;
+        })();
+    }
+
+    // Add the size of 'hint' field if present.
+    if (!hint.isEmpty()) {
+        estSize += UpdateOpEntry::kHintFieldName.size() + hint.objsize() + kPerElementOverhead;
+    }
+
+    return estSize;
+}
+
+int getDeleteSizeEstimate(const BSONObj& q,
+                          const boost::optional<mongo::BSONObj>& collation,
+                          const mongo::BSONObj& hint) {
+    using DeleteOpEntry = write_ops::DeleteOpEntry;
+
+    static const int kIntSize = 4;
+    int estSize = static_cast<int>(BSONObj::kMinBSONLength);
+
+    // Add the size of the 'q' field.
+    estSize += DeleteOpEntry::kQFieldName.size() + q.objsize() + kPerElementOverhead;
+
+    // Add the size of the 'collation' field, if present.
+    if (collation) {
+        estSize +=
+            DeleteOpEntry::kCollationFieldName.size() + collation->objsize() + kPerElementOverhead;
+    }
+
+    // Add the size of the 'limit' field.
+    estSize += DeleteOpEntry::kMultiFieldName.size() + kIntSize + kPerElementOverhead;
+
+    // Add the size of the 'hint' field, if present.
+    if (!hint.isEmpty()) {
+        estSize += DeleteOpEntry::kHintFieldName.size() + hint.objsize() + kPerElementOverhead;
+    }
+
+    return estSize;
+}
+
+bool verifySizeEstimate(const write_ops::UpdateOpEntry& update) {
+    return write_ops::getUpdateSizeEstimate(update.getQ(),
+                                            update.getU(),
+                                            update.getC(),
+                                            update.getUpsertSupplied().has_value(),
+                                            update.getCollation(),
+                                            update.getArrayFilters(),
+                                            update.getHint()) >= update.toBSON().objsize();
+}
+
+bool verifySizeEstimate(const InsertCommandRequest& insertReq,
+                        const OpMsgRequest* unparsedRequest) {
+    int size = getInsertHeaderSizeEstimate(insertReq);
+    for (auto&& docToInsert : insertReq.getDocuments()) {
+        size += docToInsert.objsize() + kWriteCommandBSONArrayPerElementOverheadBytes;
+    }
+
+    // Return true if 'insertReq' originated from a document sequence and our size estimate exceeds
+    // the size limit.
+    if (unparsedRequest && !unparsedRequest->sequences.empty() && size > BSONObjMaxUserSize) {
+        return true;
+    }
+    return size >= insertReq.toBSON({} /* commandPassthroughFields */).objsize();
+}
+
+bool verifySizeEstimate(const UpdateCommandRequest& updateReq,
+                        const OpMsgRequest* unparsedRequest) {
+    int size = getUpdateHeaderSizeEstimate(updateReq);
+
+    for (auto&& update : updateReq.getUpdates()) {
+        size += getUpdateSizeEstimate(update.getQ(),
+                                      update.getU(),
+                                      update.getC(),
+                                      update.getUpsertSupplied().has_value(),
+                                      update.getCollation(),
+                                      update.getArrayFilters(),
+                                      update.getHint()) +
+            kWriteCommandBSONArrayPerElementOverheadBytes;
+    }
+
+    // Return true if 'updateReq' originated from a document sequence and our size estimate exceeds
+    // the size limit.
+    if (unparsedRequest && !unparsedRequest->sequences.empty() && size > BSONObjMaxUserSize) {
+        return true;
+    }
+    return size >= updateReq.toBSON({} /* commandPassthroughFields */).objsize();
+}
+
+bool verifySizeEstimate(const DeleteCommandRequest& deleteReq,
+                        const OpMsgRequest* unparsedRequest) {
+    int size = getDeleteHeaderSizeEstimate(deleteReq);
+
+    for (auto&& deleteOp : deleteReq.getDeletes()) {
+        size += write_ops::getDeleteSizeEstimate(
+                    deleteOp.getQ(), deleteOp.getCollation(), deleteOp.getHint()) +
+            kWriteCommandBSONArrayPerElementOverheadBytes;
+    }
+
+    // Return true if 'deleteReq' originated from a document sequence and our size estimate exceeds
+    // the size limit.
+    if (unparsedRequest && !unparsedRequest->sequences.empty() && size > BSONObjMaxUserSize) {
+        return true;
+    }
+    return size >= deleteReq.toBSON({} /* commandPassthroughFields */).objsize();
+}
+
+int getInsertHeaderSizeEstimate(const InsertCommandRequest& insertReq) {
+    int size = getWriteCommandRequestBaseSize(insertReq.getWriteCommandRequestBase()) +
+        write_ops::InsertCommandRequest::kDocumentsFieldName.size() + kPerElementOverhead +
+        static_cast<int>(BSONObj::kMinBSONLength);
+
+    size += InsertCommandRequest::kCommandName.size() + kPerElementOverhead +
+        insertReq.getNamespace().size() + 1 /* ns string null terminator */;
+
+    return size;
+}
+
+int getUpdateHeaderSizeEstimate(const UpdateCommandRequest& updateReq) {
+    int size = getWriteCommandRequestBaseSize(updateReq.getWriteCommandRequestBase());
+
+    size += UpdateCommandRequest::kCommandName.size() + kPerElementOverhead +
+        updateReq.getNamespace().size() + 1 /* ns string null terminator */;
+
+    size += write_ops::UpdateCommandRequest::kUpdatesFieldName.size() + kPerElementOverhead +
+        static_cast<int>(BSONObj::kMinBSONLength);
+
+    // Handle legacy runtime constants.
+    if (auto runtimeConstants = updateReq.getLegacyRuntimeConstants();
+        runtimeConstants.has_value()) {
+        size += estimateRuntimeConstantsSize(*runtimeConstants);
+    }
+
+    // Handle let parameters.
+    if (auto let = updateReq.getLet(); let.has_value()) {
+        size += write_ops::UpdateCommandRequest::kLetFieldName.size() + let->objsize() +
+            kPerElementOverhead;
+    }
+    return size;
+}
+
+int getDeleteHeaderSizeEstimate(const DeleteCommandRequest& deleteReq) {
+    int size = getWriteCommandRequestBaseSize(deleteReq.getWriteCommandRequestBase());
+
+    size += DeleteCommandRequest::kCommandName.size() + kPerElementOverhead +
+        deleteReq.getNamespace().size() + 1 /* ns string null terminator */;
+
+    size += write_ops::DeleteCommandRequest::kDeletesFieldName.size() + kPerElementOverhead +
+        static_cast<int>(BSONObj::kMinBSONLength);
+
+    // Handle legacy runtime constants.
+    if (auto runtimeConstants = deleteReq.getLegacyRuntimeConstants();
+        runtimeConstants.has_value()) {
+        size += estimateRuntimeConstantsSize(*runtimeConstants);
+    }
+
+    // Handle let parameters.
+    if (auto let = deleteReq.getLet(); let.has_value()) {
+        size += write_ops::UpdateCommandRequest::kLetFieldName.size() + let->objsize() +
+            kPerElementOverhead;
+    }
+    return size;
+}
+
+bool verifySizeEstimate(const write_ops::DeleteOpEntry& deleteOp) {
+    return write_ops::getDeleteSizeEstimate(deleteOp.getQ(),
+                                            deleteOp.getCollation(),
+                                            deleteOp.getHint()) >= deleteOp.toBSON().objsize();
+}
+
 bool isClassicalUpdateReplacement(const BSONObj& update) {
     // An empty update object will be treated as replacement as firstElementFieldName() returns "".
     return update.firstElementFieldName()[0] != '$';