SERVER-44289 permit writeConcern on retryable writes that move doc onto another shard

(cherry picked from commit 9861e230bcb9bb37ed1cbb504ed4ff177a0300da)
(partial cherry pick from commit 81238fa87afbe52a9658547f63c79fac126862f1)

6 files changed, 61 insertions, 14 deletions

diff --git a/jstests/sharding/update_shard_key_doc_moves_shards.js b/jstests/sharding/update_shard_key_doc_moves_shards.js
index 9567b807b1e..cb1e2c9af32 100644
--- a/jstests/sharding/update_shard_key_doc_moves_shards.js
+++ b/jstests/sharding/update_shard_key_doc_moves_shards.js
@@ -294,6 +294,24 @@ assert.commandFailedWithCode(session.abortTransaction_forTesting(), ErrorCodes.N
 
 mongos.getDB(kDbName).foo.drop();
 
+// ----Assert that specifying writeConcern succeeds----
+
+shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+cleanupOrphanedDocs(st, ns);
+
+assert.commandWorked(sessionDB.foo.update({x: 4}, {$set: {x: 1000}}, {writeConcern: {w: 1}}));
+
+assert.commandWorked(sessionDB.runCommand({
+    findAndModify: 'foo',
+    query: {x: 78},
+    update: {$set: {x: 250}},
+    lsid: {id: UUID()},
+    txnNumber: NumberLong(1),
+    writeConcern: {w: 1},
+}));
+
+mongos.getDB(kDbName).foo.drop();
+
 // ----Assert retryable write result has WCE when the internal commitTransaction fails----
 
 shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
@@ -314,14 +332,13 @@ res = sessionDB.foo.update({x: 4}, {$set: {x: 1000}});
 assert.commandWorkedIgnoringWriteConcernErrors(res);
 assert.eq(12345, res.getWriteConcernError().code);
 
-let findAndModCmd = {
+res = sessionDB.runCommand({
     findAndModify: 'foo',
     query: {x: 78},
     update: {$set: {x: 250}},
     lsid: {id: UUID()},
     txnNumber: NumberLong(1),
-};
-res = sessionDB.runCommand(findAndModCmd);
+});
 assert.commandWorkedIgnoringWriteConcernErrors(res);
 assert.eq(res.writeConcernError.code, 12345);
 assert(res.writeConcernError.errmsg.includes("dummy error"));
diff --git a/src/mongo/s/cluster_commands_helpers.cpp b/src/mongo/s/cluster_commands_helpers.cpp
index 616433389c5..ba4ae319197 100644
--- a/src/mongo/s/cluster_commands_helpers.cpp
+++ b/src/mongo/s/cluster_commands_helpers.cpp
@@ -264,6 +264,18 @@ BSONObj appendAllowImplicitCreate(BSONObj cmdObj, bool allow) {
     return newCmdBuilder.obj();
 }
 
+BSONObj stripWriteConcern(const BSONObj& cmdObj) {
+    BSONObjBuilder output;
+    for (const auto& elem : cmdObj) {
+        const auto name = elem.fieldNameStringData();
+        if (name == WriteConcernOptions::kWriteConcernField) {
+            continue;
+        }
+        output.append(elem);
+    }
+    return output.obj();
+}
+
 std::vector<AsyncRequestsSender::Response> scatterGatherUnversionedTargetAllShards(
     OperationContext* opCtx,
     StringData dbName,
diff --git a/src/mongo/s/cluster_commands_helpers.h b/src/mongo/s/cluster_commands_helpers.h
index 433013c76b6..48bdd58e2d8 100644
--- a/src/mongo/s/cluster_commands_helpers.h
+++ b/src/mongo/s/cluster_commands_helpers.h
@@ -81,6 +81,11 @@ BSONObj appendShardVersion(BSONObj cmdObj, ChunkVersion version);
 BSONObj appendAllowImplicitCreate(BSONObj cmdObj, bool allow);
 
 /**
+ * Returns a copy of 'cmdObj' with the writeConcern removed.
+ */
+BSONObj stripWriteConcern(const BSONObj& cmdObj);
+
+/**
  * Utility for dispatching unversioned commands to all shards in a cluster.
  *
  * Returns a non-OK status if a failure occurs on *this* node during execution. Otherwise, returns
diff --git a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
index 1f70c6080a0..6d561285d94 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
@@ -310,9 +310,13 @@ private:
                     // Re-run the findAndModify command that will change the shard key value in a
                     // transaction. We call _runCommand recursively, and this second time through
                     // since it will be run as a transaction it will take the other code path to
-                    // updateShardKeyValueOnWouldChangeOwningShardError.
+                    // updateShardKeyValueOnWouldChangeOwningShardError.  We ensure the retried
+                    // operation does not include WC inside the transaction by stripping it from the
+                    // cmdObj.  The transaction commit will still use the WC, because it uses the WC
+                    // from the opCtx (which has been set previously in Strategy).
                     documentShardKeyUpdateUtil::startTransactionForShardKeyUpdate(opCtx);
-                    _runCommand(opCtx, shardId, shardVersion, nss, cmdObj, result);
+                    _runCommand(
+                        opCtx, shardId, shardVersion, nss, stripWriteConcern(cmdObj), result);
                     auto commitResponse =
                         documentShardKeyUpdateUtil::commitShardKeyUpdateTransaction(opCtx);
 
diff --git a/src/mongo/s/commands/cluster_write_cmd.cpp b/src/mongo/s/commands/cluster_write_cmd.cpp
index 917589675de..ef49fa8cfa5 100644
--- a/src/mongo/s/commands/cluster_write_cmd.cpp
+++ b/src/mongo/s/commands/cluster_write_cmd.cpp
@@ -190,14 +190,14 @@ boost::optional<WouldChangeOwningShardInfo> getWouldChangeOwningShardErrorInfo(
  * inserts the new one. Returns whether or not we actually complete the delete and insert.
  */
 bool handleWouldChangeOwningShardError(OperationContext* opCtx,
-                                       const BatchedCommandRequest& request,
+                                       BatchedCommandRequest* request,
                                        BatchedCommandResponse* response,
                                        BatchWriteExecStats stats) {
     auto txnRouter = TransactionRouter::get(opCtx);
     bool isRetryableWrite = opCtx->getTxnNumber() && !txnRouter;
 
     auto wouldChangeOwningShardErrorInfo =
-        getWouldChangeOwningShardErrorInfo(opCtx, request, response, !isRetryableWrite);
+        getWouldChangeOwningShardErrorInfo(opCtx, *request, response, !isRetryableWrite);
     if (!wouldChangeOwningShardErrorInfo)
         return false;
 
@@ -215,12 +215,17 @@ bool handleWouldChangeOwningShardError(OperationContext* opCtx,
             auto& readConcernArgs = repl::ReadConcernArgs::get(opCtx);
             readConcernArgs = repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern);
 
+            // Ensure the retried operation does not include WC inside the transaction.  The
+            // transaction commit will still use the WC, because it uses the WC from the opCtx
+            // (which has been set previously in Strategy).
+            request->unsetWriteConcern();
+
             documentShardKeyUpdateUtil::startTransactionForShardKeyUpdate(opCtx);
             // Clear the error details from the response object before sending the write again
             response->unsetErrDetails();
-            ClusterWriter::write(opCtx, request, &stats, response);
+            ClusterWriter::write(opCtx, *request, &stats, response);
             wouldChangeOwningShardErrorInfo =
-                getWouldChangeOwningShardErrorInfo(opCtx, request, response, !isRetryableWrite);
+                getWouldChangeOwningShardErrorInfo(opCtx, *request, response, !isRetryableWrite);
 
             // If we do not get WouldChangeOwningShard when re-running the update, the document has
             // been modified or deleted concurrently and we do not need to delete it and insert a
@@ -229,9 +234,9 @@ bool handleWouldChangeOwningShardError(OperationContext* opCtx,
                 wouldChangeOwningShardErrorInfo &&
                 documentShardKeyUpdateUtil::updateShardKeyForDocument(
                     opCtx,
-                    request.getNS(),
+                    request->getNS(),
                     wouldChangeOwningShardErrorInfo.get(),
-                    boost::get_optional_value_or(request.getWriteCommandBase().getStmtId(), 0));
+                    boost::get_optional_value_or(request->getWriteCommandBase().getStmtId(), 0));
 
             // If the operation was an upsert, record the _id of the new document.
             if (updatedShardKey && wouldChangeOwningShardErrorInfo->getShouldUpsert()) {
@@ -276,9 +281,9 @@ bool handleWouldChangeOwningShardError(OperationContext* opCtx,
             // Delete the original document and insert the new one
             updatedShardKey = documentShardKeyUpdateUtil::updateShardKeyForDocument(
                 opCtx,
-                request.getNS(),
+                request->getNS(),
                 wouldChangeOwningShardErrorInfo.get(),
-                boost::get_optional_value_or(request.getWriteCommandBase().getStmtId(), 0));
+                boost::get_optional_value_or(request->getWriteCommandBase().getStmtId(), 0));
 
             // If the operation was an upsert, record the _id of the new document.
             if (updatedShardKey && wouldChangeOwningShardErrorInfo->getShouldUpsert()) {
@@ -446,7 +451,7 @@ private:
         bool updatedShardKey = false;
         if (_batchedRequest.getBatchType() == BatchedCommandRequest::BatchType_Update) {
             updatedShardKey =
-                handleWouldChangeOwningShardError(opCtx, batchedRequest, &response, stats);
+                handleWouldChangeOwningShardError(opCtx, &batchedRequest, &response, stats);
         }
 
         // Populate the lastError object based on the write response
diff --git a/src/mongo/s/write_ops/batched_command_request.h b/src/mongo/s/write_ops/batched_command_request.h
index b49e893e5aa..be207399320 100644
--- a/src/mongo/s/write_ops/batched_command_request.h
+++ b/src/mongo/s/write_ops/batched_command_request.h
@@ -91,6 +91,10 @@ public:
         _writeConcern = writeConcern.getOwned();
     }
 
+    void unsetWriteConcern() {
+        _writeConcern = boost::none;
+    }
+
     bool hasWriteConcern() const {
         return _writeConcern.is_initialized();
     }