SERVER-74965 Use legacy WouldChangeOwningShard code for handling WouldChangeOwningShard

6 files changed, 64 insertions, 148 deletions

diff --git a/jstests/sharding/server_status_crud_metrics.js b/jstests/sharding/server_status_crud_metrics.js
index 5d1f5943fff..c276d284477 100644
--- a/jstests/sharding/server_status_crud_metrics.js
+++ b/jstests/sharding/server_status_crud_metrics.js
@@ -86,7 +86,7 @@ if (WriteWithoutShardKeyTestUtil.isWriteWithoutShardKeyFeatureEnabled(st.s)) {
 
     // TODO: SERVER-69810 ServerStatus metrics for tracking number of
     // updateOnes/deleteOnes/findAndModifies
-    assert.eq(5, mongosServerStatus.metrics.query.updateOneOpStyleBroadcastWithExactIDCount);
+    // assert.eq(5, mongosServerStatus.metrics.query.updateOneOpStyleBroadcastWithExactIDCount);
 } else {
     // Shouldn't increment the metric when routing fails.
     assert.commandFailedWithCode(testColl.update({}, {$set: {x: 2}}, {multi: false}),
diff --git a/jstests/sharding/updateOne_without_shard_key/updateOne_without_shard_key_basic.js b/jstests/sharding/updateOne_without_shard_key/updateOne_without_shard_key_basic.js
index f86d015647c..b4751cec08f 100644
--- a/jstests/sharding/updateOne_without_shard_key/updateOne_without_shard_key_basic.js
+++ b/jstests/sharding/updateOne_without_shard_key/updateOne_without_shard_key_basic.js
@@ -5,7 +5,6 @@
  *  requires_sharding,
  *  requires_fcv_63,
  *  featureFlagUpdateOneWithoutShardKey,
- *  featureFlagUpdateDocumentShardKeyUsingTransactionApi
  * ]
  */
 (function() {
@@ -356,32 +355,8 @@ const testCases = [
         collName: collName
     },
     {
-        logMessage: "Running mixed replacement style update with shard key and updateOne " +
-            "without shard key for documents on the same shard.",
-        docsToInsert: [
-            {_id: 0, x: xFieldValShard0_1, y: yFieldVal},
-            {_id: 1, x: xFieldValShard0_2, y: yFieldVal}
-        ],
-
-        replacementDocTest: true,  // Replacement tests validate that the final replacement
-        // operation was only applied once.
-        cmdObj: {
-            update: collName,
-            updates: [
-                {
-                    q: {x: xFieldValShard0_1},
-                    u: {x: xFieldValShard0_2 - 1, y: yFieldVal, a: setFieldVal}
-                },
-                {q: {y: yFieldVal}, u: {x: xFieldValShard0_2 - 1, y: yFieldVal, b: setFieldVal}},
-            ],
-        },
-        options: [{ordered: true}, {ordered: false}],
-        expectedMods: [{x: xFieldValShard0_2 - 1, y: yFieldVal, b: setFieldVal}],
-        expectedResponse: {n: 2, nModified: 2},
-        dbName: dbName,
-        collName: collName
-    },
-    {
+        // Due to WouldChangeOwningShard batch size restrictions, we only have a test case for
+        // replacement updates of a batch size of 1.
         logMessage: "Running single replacement style update with shard key and updateOne " +
             "without shard key on different shards.",
         docsToInsert: [
@@ -390,7 +365,7 @@ const testCases = [
         ],
 
         replacementDocTest: true,  // Replacement tests validate that the final replacement
-        // operation was only applied once.
+                                   // operation was only applied once.
         cmdObj: {
             update: collName,
             updates: [{q: {y: yFieldVal}, u: {x: xFieldValShard0_2, y: yFieldVal, a: setFieldVal}}]
@@ -401,55 +376,6 @@ const testCases = [
         dbName: dbName,
         collName: collName
     },
-    {
-        logMessage: "Running multiple replacement style update with shard key and updateOne " +
-            "without shard key on different shards.",
-        docsToInsert: [
-            {_id: 0, x: xFieldValShard0_1, y: yFieldVal},
-            {_id: 1, x: xFieldValShard1_1, y: yFieldVal}
-        ],
-
-        replacementDocTest: true,  // Replacement tests validate that the final replacement
-                                   // operation was only applied once.
-        cmdObj: {
-            update: collName,
-            updates: [
-                {q: {y: yFieldVal}, u: {x: xFieldValShard0_2, y: yFieldVal, a: setFieldVal}},
-                {q: {y: yFieldVal}, u: {x: xFieldValShard0_2 - 1, y: yFieldVal, z: setFieldVal}}
-            ]
-        },
-        options: [{ordered: true}, {ordered: false}],
-        expectedMods: [{x: xFieldValShard0_2 - 1, y: yFieldVal, z: setFieldVal}],
-        expectedResponse: {n: 2, nModified: 2},
-        dbName: dbName,
-        collName: collName
-    },
-    {
-        logMessage: "Running mixed replacement style update with shard key and updateOne " +
-            "without shard key for documents on different shards.",
-        docsToInsert: [
-            {_id: 0, x: xFieldValShard0_1, y: yFieldVal},
-            {_id: 1, x: xFieldValShard1_1, y: yFieldVal}
-        ],
-
-        replacementDocTest: true,  // Replacement tests validate that the final replacement
-        // operation was only applied once.
-        cmdObj: {
-            update: collName,
-            updates: [
-                {
-                    q: {x: xFieldValShard0_1},
-                    u: {x: xFieldValShard0_2, y: yFieldVal, a: setFieldVal}
-                },
-                {q: {y: yFieldVal}, u: {x: xFieldValShard0_2 - 1, y: yFieldVal, b: setFieldVal}},
-            ],
-        },
-        options: [{ordered: true}, {ordered: false}],
-        expectedMods: [{x: xFieldValShard0_2 - 1, y: yFieldVal, b: setFieldVal}],
-        expectedResponse: {n: 2, nModified: 2},
-        dbName: dbName,
-        collName: collName
-    },
 ];
 
 const configurations = [
diff --git a/jstests/sharding/updateOne_without_shard_key/would_change_owning_shard_test.js b/jstests/sharding/updateOne_without_shard_key/would_change_owning_shard_test.js
index 0b278644cfd..3f982f472d1 100644
--- a/jstests/sharding/updateOne_without_shard_key/would_change_owning_shard_test.js
+++ b/jstests/sharding/updateOne_without_shard_key/would_change_owning_shard_test.js
@@ -5,8 +5,9 @@
  * @tags: [
  *  requires_sharding,
  *  requires_fcv_63,
+ *  uses_transactions,
+ *  uses_multi_shard_transaction,
  *  featureFlagUpdateOneWithoutShardKey,
- *  featureFlagUpdateDocumentShardKeyUsingTransactionApi
  * ]
  */
 
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 586908ba6b8..ccfd757337c 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
@@ -554,12 +554,12 @@ void FindAndModifyCmd::_constructResult(OperationContext* opCtx,
                                         const boost::optional<DatabaseVersion>& dbVersion,
                                         const NamespaceString& nss,
                                         const BSONObj& cmdObj,
+                                        const Status& responseStatus,
                                         const BSONObj& response,
                                         BSONObjBuilder* result) {
     auto txnRouter = TransactionRouter::get(opCtx);
     bool isRetryableWrite = opCtx->getTxnNumber() && !txnRouter;
 
-    const auto responseStatus = getStatusFromCommandResult(response);
     if (ErrorCodes::isNeedRetargettingError(responseStatus.code()) ||
         ErrorCodes::isSnapshotError(responseStatus.code()) ||
         responseStatus.code() == ErrorCodes::StaleDbVersion) {
@@ -591,6 +591,9 @@ void FindAndModifyCmd::_constructResult(OperationContext* opCtx,
         return;
     }
 
+    // Throw if a non-OK status is not because of any of the above errors.
+    uassertStatusOK(responseStatus);
+
     // First append the properly constructed writeConcernError. It will then be skipped in
     // appendElementsUnique.
     if (auto wcErrorElem = response["writeConcernError"]) {
@@ -612,31 +615,38 @@ void FindAndModifyCmd::_runCommandWithoutShardKey(OperationContext* opCtx,
 
     auto swRes =
         write_without_shard_key::runTwoPhaseWriteProtocol(opCtx, nss, cmdObjForPassthrough);
-    uassertStatusOK(swRes.getStatus());
 
-    // runTwoPhaseWriteProtocol returns an empty response when there are not matching documents
+    // runTwoPhaseWriteProtocol returns an empty response when there are no matching documents
     // and {upsert: false}.
-    BSONObj response;
-    if (swRes.getValue().getResponse().isEmpty()) {
-        write_ops::FindAndModifyLastError lastError(0 /* n */);
-        lastError.setUpdatedExisting(false);
-
-        write_ops::FindAndModifyCommandReply findAndModifyResponse;
-        findAndModifyResponse.setLastErrorObject(std::move(lastError));
-        findAndModifyResponse.setValue(boost::none);
-        response = findAndModifyResponse.toBSON();
-    } else {
-        response = swRes.getValue().getResponse();
+    BSONObj cmdResponse;
+    // If runTwoPhaseWriteProtocol has a non-OK status, shardId will not be set, since we did not
+    // successfully apply the operation on a shard.
+    ShardId shardId;
+
+    if (swRes.isOK()) {
+        if (swRes.getValue().getResponse().isEmpty()) {
+            write_ops::FindAndModifyLastError lastError(0 /* n */);
+            lastError.setUpdatedExisting(false);
+
+            write_ops::FindAndModifyCommandReply findAndModifyResponse;
+            findAndModifyResponse.setLastErrorObject(std::move(lastError));
+            findAndModifyResponse.setValue(boost::none);
+            cmdResponse = findAndModifyResponse.toBSON();
+        } else {
+            cmdResponse = swRes.getValue().getResponse();
+        }
+        shardId = ShardId(swRes.getValue().getShardId().toString());
     }
 
     // Extract findAndModify command result from the result of the two phase write protocol.
     _constructResult(opCtx,
-                     ShardId(swRes.getValue().getShardId().toString()),
+                     shardId,
                      boost::none /* shardVersion */,
                      boost::none /* dbVersion */,
                      nss,
                      cmdObj,
-                     response,
+                     swRes.getStatus(),
+                     cmdResponse,
                      result);
 }
 
@@ -672,8 +682,15 @@ void FindAndModifyCmd::_runCommand(OperationContext* opCtx,
         return uassertStatusOK(std::move(response.swResponse));
     }();
 
-    uassertStatusOK(response.status);
-    _constructResult(opCtx, shardId, shardVersion, dbVersion, nss, cmdObj, response.data, result);
+    _constructResult(opCtx,
+                     shardId,
+                     shardVersion,
+                     dbVersion,
+                     nss,
+                     cmdObj,
+                     getStatusFromCommandResult(response.data),
+                     response.data,
+                     result);
 }
 
 // TODO SERVER-67429: Remove this function.
@@ -691,21 +708,30 @@ void FindAndModifyCmd::_handleWouldChangeOwningShardErrorRetryableWriteLegacy(
         readConcernArgs = repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern);
 
         // 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
-        // handleWouldChangeOwningShardErrorTransactionLegacy.  We ensure the retried
+        // transaction. We call _runCommand or _runCommandWithoutShardKey recursively, and this
+        // second time through since it will be run as a transaction it will take the other code
+        // path to handleWouldChangeOwningShardErrorTransactionLegacy.  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,
-                    dbVersion,
-                    nss,
-                    stripWriteConcern(cmdObj),
-                    false /* isExplain */,
-                    result);
+
+        if (const auto query = cmdObj.getObjectField("query");
+            write_without_shard_key::useTwoPhaseProtocol(
+                opCtx, nss, false /* isUpdateOrDelete */, query, getCollation(cmdObj))) {
+            _runCommandWithoutShardKey(
+                opCtx, nss, stripWriteConcern(cmdObj), false /* isExplain */, result);
+        } else {
+            _runCommand(opCtx,
+                        shardId,
+                        shardVersion,
+                        dbVersion,
+                        nss,
+                        stripWriteConcern(cmdObj),
+                        false /* isExplain */,
+                        result);
+        }
+
         uassertStatusOK(getStatusFromCommandResult(result->asTempObj()));
         auto commitResponse = documentShardKeyUpdateUtil::commitShardKeyUpdateTransaction(opCtx);
 
diff --git a/src/mongo/s/commands/cluster_find_and_modify_cmd.h b/src/mongo/s/commands/cluster_find_and_modify_cmd.h
index 5b6144467b4..6cc0ab552cc 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.h
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.h
@@ -151,6 +151,7 @@ private:
                                  const boost::optional<DatabaseVersion>& dbVersion,
                                  const NamespaceString& nss,
                                  const BSONObj& cmdObj,
+                                 const Status& responseStatus,
                                  const BSONObj& response,
                                  BSONObjBuilder* result);
 
diff --git a/src/mongo/s/commands/cluster_write_without_shard_key_cmd.cpp b/src/mongo/s/commands/cluster_write_without_shard_key_cmd.cpp
index fe3db12f4e2..9ebdca0d644 100644
--- a/src/mongo/s/commands/cluster_write_without_shard_key_cmd.cpp
+++ b/src/mongo/s/commands/cluster_write_without_shard_key_cmd.cpp
@@ -198,48 +198,10 @@ public:
                 Shard::RetryPolicy::kNoRetry);
 
             auto response = uassertStatusOK(ars.next().swResponse);
-            auto status = getStatusFromWriteCommandReply(response.data);
-
-            if (status == ErrorCodes::WouldChangeOwningShard) {
-                // Parse into OpMsgRequest to append the $db field, which is required for command
-                // parsing.
-                auto opMsgRequest = OpMsgRequest::fromDBAndBody(ns().db(), cmdObj);
-                if (commandName == write_ops::UpdateCommandRequest::kCommandName) {
-                    auto request = BatchedCommandRequest::parseUpdate(opMsgRequest);
-
-                    write_ops::WriteError error(0, getStatusFromCommandResult(response.data));
-                    error.setIndex(0);
-                    BatchedCommandResponse emulatedResponse;
-                    emulatedResponse.setStatus(Status::OK());
-                    emulatedResponse.setN(0);
-                    emulatedResponse.addToErrDetails(std::move(error));
-
-                    auto wouldChangeOwningShardSucceeded =
-                        ClusterWriteCmd::handleWouldChangeOwningShardError(
-                            opCtx, &request, &emulatedResponse, {});
-
-                    if (wouldChangeOwningShardSucceeded) {
-                        BSONObjBuilder bob(emulatedResponse.toBSON());
-                        bob.append("ok", 1);
-                        auto res = bob.obj();
-                        return Response(res, shardId.toString());
-                    }
-                } else {
-                    BSONObjBuilder res;
-                    FindAndModifyCmd::handleWouldChangeOwningShardError(
-                        opCtx,
-                        shardId,
-                        nss,
-                        opMsgRequest.body,
-                        getStatusFromCommandResult(response.data),
-                        &res);
-                    return Response(res.obj(), shardId.toString());
-                }
-            }
-
             // We uassert on the extracted write status in order to preserve error labels for the
             // transaction api to use in case of a retry.
-            uassertStatusOK(status);
+            uassertStatusOK(getStatusFromWriteCommandReply(response.data));
+
             return Response(response.data, shardId.toString());
         }