SERVER-64069 Relax restriction that changing a document's shard key value must have a transaction number

2 files changed, 25 insertions, 16 deletions

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 bcd2d5a746f..657e1e9086d 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
@@ -144,7 +144,7 @@ BSONObj getShardKey(OperationContext* opCtx,
     return shardKey;
 }
 
-void handleWouldChangeOwningShardErrorRetryableWrite(
+void handleWouldChangeOwningShardErrorNonTransaction(
     OperationContext* opCtx,
     const ShardId& shardId,
     const NamespaceString& nss,
@@ -519,7 +519,8 @@ private:
                             const NamespaceString& nss,
                             const BSONObj& cmdObj,
                             BSONObjBuilder* result) {
-        bool isRetryableWrite = opCtx->getTxnNumber() && !TransactionRouter::get(opCtx);
+        auto txnRouter = TransactionRouter::get(opCtx);
+        bool isRetryableWrite = opCtx->getTxnNumber() && !txnRouter;
 
         const auto response = [&] {
             std::vector<AsyncRequestsSender::Request> requests;
@@ -574,13 +575,15 @@ private:
                 // Strip runtime constants because they will be added again when this command is
                 // recursively sent through the service entry point.
                 parsedRequest.setLegacyRuntimeConstants(boost::none);
-                if (isRetryableWrite) {
-                    parsedRequest.setStmtId(0);
-                    handleWouldChangeOwningShardErrorRetryableWrite(
-                        opCtx, shardId, nss, parsedRequest, result);
-                } else {
+                if (txnRouter) {
                     handleWouldChangeOwningShardErrorTransaction(
                         opCtx, nss, responseStatus, parsedRequest, result);
+                } else {
+                    if (isRetryableWrite) {
+                        parsedRequest.setStmtId(0);
+                    }
+                    handleWouldChangeOwningShardErrorNonTransaction(
+                        opCtx, shardId, nss, parsedRequest, result);
                 }
             } else {
                 // TODO SERVER-67429: Remove this branch.
diff --git a/src/mongo/s/commands/cluster_write_cmd.cpp b/src/mongo/s/commands/cluster_write_cmd.cpp
index a43e919b217..358615c1993 100644
--- a/src/mongo/s/commands/cluster_write_cmd.cpp
+++ b/src/mongo/s/commands/cluster_write_cmd.cpp
@@ -144,7 +144,7 @@ boost::optional<WouldChangeOwningShardInfo> getWouldChangeOwningShardErrorInfo(
     }
 }
 
-void handleWouldChangeOwningShardErrorRetryableWrite(OperationContext* opCtx,
+void handleWouldChangeOwningShardErrorNonTransaction(OperationContext* opCtx,
                                                      BatchedCommandRequest* request,
                                                      BatchedCommandResponse* response) {
     // Strip write concern because this command will be sent as part of a
@@ -294,18 +294,24 @@ bool handleWouldChangeOwningShardError(OperationContext* opCtx,
 
     if (feature_flags::gFeatureFlagUpdateDocumentShardKeyUsingTransactionApi.isEnabled(
             serverGlobalParams.featureCompatibility)) {
-        if (isRetryableWrite) {
-            if (MONGO_unlikely(hangAfterThrowWouldChangeOwningShardRetryableWrite.shouldFail())) {
-                LOGV2(5918603, "Hit hangAfterThrowWouldChangeOwningShardRetryableWrite failpoint");
-                hangAfterThrowWouldChangeOwningShardRetryableWrite.pauseWhileSet(opCtx);
-            }
-
-            handleWouldChangeOwningShardErrorRetryableWrite(opCtx, request, response);
-        } else {
+        if (txnRouter) {
             auto updateResult = handleWouldChangeOwningShardErrorTransaction(
                 opCtx, request, response, *wouldChangeOwningShardErrorInfo);
             updatedShardKey = updateResult.updatedShardKey;
             upsertedId = std::move(updateResult.upsertedId);
+        } else {
+            // Updating a document's shard key such that its owning shard changes must be run in a
+            // transaction. If this update is not already in a transaction, complete the update
+            // using an internal transaction.
+            if (isRetryableWrite) {
+                if (MONGO_unlikely(
+                        hangAfterThrowWouldChangeOwningShardRetryableWrite.shouldFail())) {
+                    LOGV2(5918603,
+                          "Hit hangAfterThrowWouldChangeOwningShardRetryableWrite failpoint");
+                    hangAfterThrowWouldChangeOwningShardRetryableWrite.pauseWhileSet(opCtx);
+                }
+            }
+            handleWouldChangeOwningShardErrorNonTransaction(opCtx, request, response);
         }
     } else {
         // TODO SERVER-67429: Delete this branch.