SERVER-26194 Account for collation specification size in batch_write_op getWriteSizeBytes

1 files changed, 8 insertions, 3 deletions

diff --git a/src/mongo/s/write_ops/batch_write_op.cpp b/src/mongo/s/write_ops/batch_write_op.cpp
index fbd0b945b23..4f7a4efe52c 100644
--- a/src/mongo/s/write_ops/batch_write_op.cpp
+++ b/src/mongo/s/write_ops/batch_write_op.cpp
@@ -156,14 +156,19 @@ static int getWriteSizeBytes(const WriteOp& writeOp) {
         return item.getDocument().objsize();
     } else if (batchType == BatchedCommandRequest::BatchType_Update) {
         // Note: Be conservative here - it's okay if we send slightly too many batches
-        int estSize = item.getUpdate()->getQuery().objsize() +
-            item.getUpdate()->getUpdateExpr().objsize() + kEstUpdateOverheadBytes;
+        auto collationSize =
+            item.getUpdate()->isCollationSet() ? item.getUpdate()->getCollation().objsize() : 0;
+        auto estSize = item.getUpdate()->getQuery().objsize() +
+            item.getUpdate()->getUpdateExpr().objsize() + collationSize + kEstUpdateOverheadBytes;
         dassert(estSize >= item.getUpdate()->toBSON().objsize());
         return estSize;
     } else {
         dassert(batchType == BatchedCommandRequest::BatchType_Delete);
         // Note: Be conservative here - it's okay if we send slightly too many batches
-        int estSize = item.getDelete()->getQuery().objsize() + kEstDeleteOverheadBytes;
+        auto collationSize =
+            item.getDelete()->isCollationSet() ? item.getDelete()->getCollation().objsize() : 0;
+        auto estSize =
+            item.getDelete()->getQuery().objsize() + collationSize + kEstDeleteOverheadBytes;
         dassert(estSize >= item.getDelete()->toBSON().objsize());
         return estSize;
     }