SERVER-66289 Update write size estimation logic in DocumentSourceWriter

2 files changed, 74 insertions, 0 deletions

diff --git a/src/mongo/db/ops/write_ops.cpp b/src/mongo/db/ops/write_ops.cpp
index 0ef8f329558..8e19bfd3117 100644
--- a/src/mongo/db/ops/write_ops.cpp
+++ b/src/mongo/db/ops/write_ops.cpp
@@ -136,6 +136,68 @@ int32_t getStmtIdForWriteAt(const WriteCommandRequestBase& writeCommandBase, siz
     return kFirstStmtId + writePos;
 }
 
+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;
+
+    // This constant accounts for the null terminator in each field name and the BSONType byte for
+    // each element.
+    static const int kPerElementOverhead = 2;
+    static const int kBoolSize = 1;
+    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;
+}
+
 bool isClassicalUpdateReplacement(const BSONObj& update) {
     // An empty update object will be treated as replacement as firstElementFieldName() returns "".
     return update.firstElementFieldName()[0] != '$';
diff --git a/src/mongo/db/ops/write_ops.h b/src/mongo/db/ops/write_ops.h
index d78791fdfdc..a1201412a33 100644
--- a/src/mongo/db/ops/write_ops.h
+++ b/src/mongo/db/ops/write_ops.h
@@ -104,6 +104,18 @@ const std::vector<BSONObj>& arrayFiltersOf(const T& opEntry) {
 }
 
 /**
+ * Utility which estimates the size in bytes of an update statement with the given parameters, when
+ * serialized in the format used for the update command.
+ */
+int getUpdateSizeEstimate(const BSONObj& q,
+                          const write_ops::UpdateModification& u,
+                          const boost::optional<mongo::BSONObj>& c,
+                          bool includeUpsertSupplied,
+                          const boost::optional<mongo::BSONObj>& collation,
+                          const boost::optional<std::vector<mongo::BSONObj>>& arrayFilters,
+                          const mongo::BSONObj& hint);
+
+/**
  * If the response from a write command contains any write errors, it will throw the first one. All
  * the remaining errors will be disregarded.
  *