SERVER-54824 Make mongos write size estimation respect hint field

2 files changed, 51 insertions, 0 deletions

diff --git a/jstests/core/update_with_large_hint.js b/jstests/core/update_with_large_hint.js
new file mode 100644
index 00000000000..0b2521337ce
--- /dev/null
+++ b/jstests/core/update_with_large_hint.js
@@ -0,0 +1,36 @@
+// Test that write size estimation in mongos respects 'hint' field.
+// @tags: [
+//   requires_sharding,
+// ]
+(function() {
+"use strict";
+
+const coll = db.update_with_large_hint;
+coll.drop();
+
+const longHint = "x".repeat(1000);
+assert.commandWorked(coll.createIndex({[longHint]: 1}));
+assert.commandWorked(coll.insert({_id: 0}));
+
+assert.commandWorked(coll.runCommand("update", {
+    updates: [{q: {_id: 0}, u: {$set: {x: 1}}, hint: {[longHint]: 1}}],
+}));
+
+assert.commandWorked(coll.runCommand("delete", {
+    deletes: [{q: {_id: 0}, limit: 1, hint: {[longHint]: 1}}],
+}));
+
+assert.commandWorked(coll.dropIndexes());
+assert.commandWorked(coll.insert({_id: 0}));
+
+// Both commands should fail because hint does not correspond to the existing index.
+assert.commandFailedWithCode(coll.runCommand("update", {
+    updates: [{q: {_id: 0}, u: {$set: {x: 1}}, hint: {[longHint]: 1}}],
+}),
+                             ErrorCodes.BadValue);
+
+assert.commandFailedWithCode(coll.runCommand("delete", {
+    deletes: [{q: {_id: 0}, limit: 1, hint: {[longHint]: 1}}],
+}),
+                             ErrorCodes.BadValue);
+}());
\ No newline at end of file
diff --git a/src/mongo/s/write_ops/batch_write_op.cpp b/src/mongo/s/write_ops/batch_write_op.cpp
index 03fd3ae99f1..3337d8f5d90 100644
--- a/src/mongo/s/write_ops/batch_write_op.cpp
+++ b/src/mongo/s/write_ops/batch_write_op.cpp
@@ -184,6 +184,11 @@ int getWriteSizeBytes(const WriteOp& writeOp) {
         estSize += UpdateOpEntry::kUpsertFieldName.size() + boolSize;
         estSize += UpdateOpEntry::kMultiFieldName.size() + boolSize;
 
+        // Add the size of 'upsertSupplied' field if present.
+        if (auto upsertSupplied = item.getUpdate().getUpsertSupplied()) {
+            estSize += UpdateOpEntry::kUpsertSuppliedFieldName.size() + boolSize;
+        }
+
         // Add the sizes of the 'q' and 'u' fields.
         estSize += (UpdateOpEntry::kQFieldName.size() + item.getUpdate().getQ().objsize() +
                     UpdateOpEntry::kUFieldName.size() + item.getUpdate().getU().objsize());
@@ -193,6 +198,11 @@ int getWriteSizeBytes(const WriteOp& writeOp) {
             estSize += UpdateOpEntry::kCFieldName.size() + item.getUpdate().getC()->objsize();
         }
 
+        // Add the size of 'hint' field if present.
+        if (auto hint = item.getUpdate().getHint(); !hint.isEmpty()) {
+            estSize += UpdateOpEntry::kHintFieldName.size() + hint.objsize();
+        }
+
         // Finally, add the constant updateOp overhead size.
         estSize += kEstUpdateOverheadBytes;
 
@@ -212,6 +222,11 @@ int getWriteSizeBytes(const WriteOp& writeOp) {
         // Add the size of the 'limit' field.
         estSize += DeleteOpEntry::kMultiFieldName.size() + intSize;
 
+        // Add the size of 'hint' field if present.
+        if (auto hint = item.getDelete().getHint(); !hint.isEmpty()) {
+            estSize += DeleteOpEntry::kHintFieldName.size() + hint.objsize();
+        }
+
         // Add the size of the 'q' field, plus the constant deleteOp overhead size.
         estSize += kEstDeleteOverheadBytes +
             (DeleteOpEntry::kQFieldName.size() + item.getDelete().getQ().objsize());