SERVER-47233 Consider WriteOp complete when no shards left to target

(cherry picked from commit 2f6cac5d4b98969621b48decd2cadaed9e8faf14)
(cherry picked from commit 03a97185d0900fe46e86d404cb172f63644ef4c3)

2 files changed, 103 insertions, 1 deletions

diff --git a/src/mongo/s/write_ops/batch_write_exec_test.cpp b/src/mongo/s/write_ops/batch_write_exec_test.cpp
index f41cc6245d4..cbc00a2934e 100644
--- a/src/mongo/s/write_ops/batch_write_exec_test.cpp
+++ b/src/mongo/s/write_ops/batch_write_exec_test.cpp
@@ -799,6 +799,106 @@ TEST_F(BatchWriteExecTest, RetryableErrorReturnedFromMultiWriteWithShard1FirstOK
     ASSERT_EQ(2, response.getNModified());
 }
 
+TEST_F(BatchWriteExecTest, RetryableErrorReturnedFromWriteWithShard1SSVShard2OK) {
+    BatchedCommandRequest request([&] {
+        write_ops::Update updateOp(nss);
+        updateOp.setWriteCommandBase([] {
+            write_ops::WriteCommandBase writeCommandBase;
+            writeCommandBase.setOrdered(false);
+            return writeCommandBase;
+        }());
+        updateOp.setUpdates(
+            std::vector{write_ops::UpdateOpEntry(BSON("_id" << 150), {BSON("x" << 1)})});
+        return updateOp;
+    }());
+    request.setWriteConcern(BSONObj());
+
+    static const auto epoch = OID::gen();
+
+    // This allows the batch to target each write operation to perform this test
+    class MultiShardTargeter : public MockNSTargeter {
+    public:
+        using MockNSTargeter::MockNSTargeter;
+
+        StatusWith<std::vector<ShardEndpoint>> targetUpdate(
+            OperationContext* opCtx, const write_ops::UpdateOpEntry& updateDoc) const override {
+            if (targetAll) {
+                return std::vector<ShardEndpoint>{
+                    ShardEndpoint(kShardName1, ChunkVersion(100, 200, epoch)),
+                    ShardEndpoint(kShardName2, ChunkVersion(101, 200, epoch))};
+            } else {
+                return std::vector<ShardEndpoint>{
+                    ShardEndpoint(kShardName2, ChunkVersion(101, 200, epoch))};
+            }
+        }
+
+        bool targetAll = true;
+    };
+
+    MultiShardTargeter multiShardNSTargeter(
+        nss,
+        {MockRange(ShardEndpoint(kShardName1, ChunkVersion(100, 200, epoch)),
+                   BSON("sk" << MINKEY),
+                   BSON("sk" << 10)),
+         MockRange(ShardEndpoint(kShardName2, ChunkVersion(101, 200, epoch)),
+                   BSON("sk" << 10),
+                   BSON("sk" << MAXKEY))});
+
+    auto future = launchAsync([&] {
+        BatchedCommandResponse response;
+        BatchWriteExecStats stats;
+        BatchWriteExec::executeBatch(
+            operationContext(), multiShardNSTargeter, request, &response, &stats);
+
+        return response;
+    });
+
+    onCommandForPoolExecutor([&](const RemoteCommandRequest& request) {
+        ASSERT_EQ(kTestShardHost1, request.target);
+
+        BatchedCommandResponse response;
+        response.setStatus(Status::OK());
+        response.setNModified(0);
+        response.setN(0);
+        response.addToErrDetails([&] {
+            WriteErrorDetail* errDetail = new WriteErrorDetail();
+            errDetail->setIndex(0);
+            errDetail->setStatus({ErrorCodes::StaleShardVersion, "Stale shard version"});
+            errDetail->setErrInfo([&] {
+                Status ssvStatus(StaleConfigInfo(nss,
+                                                 ChunkVersion(101, 200, epoch),
+                                                 ChunkVersion(105, 200, epoch)),
+                                 "Migration happened");
+                BSONObjBuilder builder;
+                ssvStatus.serializeErrorToBSON(&builder);
+                return builder.obj();
+            }());
+            return errDetail;
+        }());
+
+        // This simulates a migration of the last chunk on shard 1 to shard 2, which means that
+        // future rounds on the batchExecutor should only target shard 2
+        multiShardNSTargeter.targetAll = false;
+        return response.toBSON();
+    });
+
+    onCommandForPoolExecutor([&](const RemoteCommandRequest& request) {
+        ASSERT_EQ(kTestShardHost2, request.target);
+
+        BatchedCommandResponse response;
+        response.setStatus(Status::OK());
+        response.setNModified(1);
+        response.setN(1);
+        return response.toBSON();
+    });
+
+    auto response = future.default_timed_get();
+    ASSERT_OK(response.getTopLevelStatus());
+    ASSERT_EQ(1, response.getNModified());
+    ASSERT_EQ(1, response.getN());
+    ASSERT_FALSE(response.isErrDetailsSet());
+}
+
 //
 // Test retryable errors
 //
diff --git a/src/mongo/s/write_ops/write_op.cpp b/src/mongo/s/write_ops/write_op.cpp
index c85870d07bc..1607029682d 100644
--- a/src/mongo/s/write_ops/write_op.cpp
+++ b/src/mongo/s/write_ops/write_op.cpp
@@ -109,7 +109,9 @@ Status WriteOp::targetWrites(OperationContext* opCtx,
         }
     }
 
-    _state = WriteOpState_Pending;
+    // If all operations currently targeted were successful on a previous round we might have 0
+    // childOps, that would mean that the operation is finished.
+    _state = _childOps.size() ? WriteOpState_Pending : WriteOpState_Completed;
     return Status::OK();
 }