1 files changed, 64 insertions, 14 deletions

diff --git a/src/mongo/s/write_ops/batch_write_op.cpp b/src/mongo/s/write_ops/batch_write_op.cpp
index ed67f6c0088..a61ee3dd4bf 100644
--- a/src/mongo/s/write_ops/batch_write_op.cpp
+++ b/src/mongo/s/write_ops/batch_write_op.cpp
@@ -59,6 +59,7 @@ struct WriteErrorComp {
 // batches before serializing.
 //
 // TODO: Revisit when we revisit command limits in general
+const int kEstUpdateOverheadBytes = (BSONObjMaxInternalSize - BSONObjMaxUserSize) / 100;
 const int kEstDeleteOverheadBytes = (BSONObjMaxInternalSize - BSONObjMaxUserSize) / 100;
 
 /**
@@ -158,17 +159,51 @@ 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.
-        const auto& update = item.getUpdate();
-        auto estSize = write_ops::getUpdateSizeEstimate(update.getQ(),
-                                                        update.getU(),
-                                                        update.getC(),
-                                                        update.getUpsertSupplied().has_value(),
-                                                        update.getCollation(),
-                                                        update.getArrayFilters(),
-                                                        update.getHint());
+        auto estSize = static_cast<int>(BSONObj::kMinBSONLength);
+        static const auto boolSize = 1;
+
+        // Add the size of the 'collation' field, if present.
+        estSize += !item.getUpdate().getCollation() ? 0
+                                                    : (UpdateOpEntry::kCollationFieldName.size() +
+                                                       item.getUpdate().getCollation()->objsize());
+
+        // Add the size of the 'arrayFilters' field, if present.
+        estSize += !item.getUpdate().getArrayFilters() ? 0 : ([&item]() {
+            auto size = BSONObj::kMinBSONLength + UpdateOpEntry::kArrayFiltersFieldName.size();
+            for (auto&& filter : *item.getUpdate().getArrayFilters()) {
+                size += filter.objsize();
+            }
+            return size;
+        })();
+
+        // Add the sizes of the 'multi' and 'upsert' fields.
+        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());
+
+        // Add the size of the 'c' field if present.
+        if (auto constants = item.getUpdate().getC()) {
+            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;
 
         // When running a debug build, verify that estSize is at least the BSON serialization size.
-        dassert(estSize >= update.toBSON().objsize());
+        dassert(estSize >= item.getUpdate().toBSON().objsize());
         return estSize;
     } else if (batchType == BatchedCommandRequest::BatchType_Delete) {
         // Note: Be conservative here - it's okay if we send slightly too many batches.
@@ -568,9 +603,6 @@ BatchedCommandRequest BatchWriteOp::buildBatchRequest(const TargetedWriteBatch&
             wcb.setStmtIds(std::move(stmtIdsForOp));
         }
 
-        wcb.setBypassEmptyTsReplacement(
-            _clientRequest.getWriteCommandRequestBase().getBypassEmptyTsReplacement());
-
         return wcb;
     }());
 
@@ -583,6 +615,19 @@ BatchedCommandRequest BatchWriteOp::buildBatchRequest(const TargetedWriteBatch&
     if (dbVersion)
         request.setDbVersion(*dbVersion);
 
+    if (_clientRequest.hasWriteConcern()) {
+        if (_clientRequest.isVerboseWC()) {
+            request.setWriteConcern(_clientRequest.getWriteConcern());
+        } else {
+            // Mongos needs to send to the shard with w > 0 so it will be able to see the
+            // writeErrors
+            request.setWriteConcern(upgradeWriteConcern(_clientRequest.getWriteConcern()));
+        }
+    } else if (!TransactionRouter::get(_opCtx)) {
+        // Apply the WC from the opCtx (except if in a transaction).
+        request.setWriteConcern(_opCtx->getWriteConcern().toBSON());
+    }
+
     return request;
 }
 
@@ -765,7 +810,7 @@ void BatchWriteOp::buildClientResponse(BatchedCommandResponse* batchResp) {
     batchResp->setStatus(Status::OK());
 
     // For non-verbose, it's all we need.
-    if (!_opCtx->getWriteConcern().requiresWriteAcknowledgement()) {
+    if (!_clientRequest.isVerboseWC()) {
         return;
     }
 
@@ -816,7 +861,12 @@ void BatchWriteOp::buildClientResponse(BatchedCommandResponse* batchResp) {
         }
     }
 
-    if (!_wcErrors.empty()) {
+    // Only return a write concern error if everything succeeded (unordered or ordered)
+    // OR if something succeeded and we're unordered
+    const bool orderedOps = _clientRequest.getWriteCommandRequestBase().getOrdered();
+    const bool reportWCError =
+        errOps.empty() || (!orderedOps && errOps.size() < _clientRequest.sizeWriteOps());
+    if (!_wcErrors.empty() && reportWCError) {
         WriteConcernErrorDetail* error = new WriteConcernErrorDetail;
 
         // Generate the multi-error message below