SERVER-61365 - adding unit test for batch writes resulting in tenantMigrationAborted errors

3 files changed, 271 insertions, 1 deletions

diff --git a/src/mongo/s/write_ops/batch_write_exec.cpp b/src/mongo/s/write_ops/batch_write_exec.cpp
index 4d5bf523d07..8f4b59916ca 100644
--- a/src/mongo/s/write_ops/batch_write_exec.cpp
+++ b/src/mongo/s/write_ops/batch_write_exec.cpp
@@ -300,6 +300,7 @@ void BatchWriteExec::executeBatch(OperationContext* opCtx,
                     TrackedErrors trackedErrors;
                     trackedErrors.startTracking(ErrorCodes::StaleShardVersion);
                     trackedErrors.startTracking(ErrorCodes::StaleDbVersion);
+                    trackedErrors.startTracking(ErrorCodes::TenantMigrationAborted);
 
                     LOGV2_DEBUG(22907,
                                 4,
@@ -338,6 +339,8 @@ void BatchWriteExec::executeBatch(OperationContext* opCtx,
                     const auto& staleShardErrors =
                         trackedErrors.getErrors(ErrorCodes::StaleShardVersion);
                     const auto& staleDbErrors = trackedErrors.getErrors(ErrorCodes::StaleDbVersion);
+                    const auto& tenantMigrationAbortedErrors =
+                        trackedErrors.getErrors(ErrorCodes::TenantMigrationAborted);
 
                     if (!staleShardErrors.empty()) {
                         invariant(staleDbErrors.empty());
@@ -351,6 +354,10 @@ void BatchWriteExec::executeBatch(OperationContext* opCtx,
                         ++stats->numStaleDbBatches;
                     }
 
+                    if (!tenantMigrationAbortedErrors.empty()) {
+                        ++stats->numTenantMigrationAbortedErrors;
+                    }
+
                     if (response.shardHostAndPort) {
                         // Remember that we successfully wrote to this shard
                         // NOTE: This will record lastOps for shards where we actually didn't update
diff --git a/src/mongo/s/write_ops/batch_write_exec.h b/src/mongo/s/write_ops/batch_write_exec.h
index 75fc070cc83..86ce5d7508b 100644
--- a/src/mongo/s/write_ops/batch_write_exec.h
+++ b/src/mongo/s/write_ops/batch_write_exec.h
@@ -85,7 +85,11 @@ typedef std::map<ConnectionString, HostOpTime> HostOpTimeMap;
 
 class BatchWriteExecStats {
 public:
-    BatchWriteExecStats() : numRounds(0), numStaleShardBatches(0), numStaleDbBatches(0) {}
+    BatchWriteExecStats()
+        : numRounds(0),
+          numStaleShardBatches(0),
+          numStaleDbBatches(0),
+          numTenantMigrationAbortedErrors(0) {}
 
     void noteWriteAt(const HostAndPort& host, repl::OpTime opTime, const OID& electionId);
     void noteTargetedShard(const ShardId& shardId);
@@ -103,6 +107,8 @@ public:
     int numStaleShardBatches;
     // Number of stale batches due to StaleDbVersion
     int numStaleDbBatches;
+    // Number of tenant migration aborted errors
+    int numTenantMigrationAbortedErrors;
 
 private:
     std::set<ShardId> _targetedShards;
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 8634bc290fa..e24c95b36e4 100644
--- a/src/mongo/s/write_ops/batch_write_exec_test.cpp
+++ b/src/mongo/s/write_ops/batch_write_exec_test.cpp
@@ -157,6 +157,57 @@ BSONObj expectInsertsReturnStaleDbVersionErrorsBase(const NamespaceString& nss,
 }
 
 /**
+ * Expects to send tenantMigrationAborted error for the numberOfFailedOps given.
+ * If
+ */
+BSONObj expectInsertsReturnTenantMigrationAbortedErrorsBase(
+    const NamespaceString& nss,
+    const std::vector<BSONObj>& expected,
+    const executor::RemoteCommandRequest& request,
+    int numberOfFailedOps) {
+    ASSERT_EQUALS(nss.db(), request.dbname);
+
+    const auto opMsgRequest(OpMsgRequest::fromDBAndBody(request.dbname, request.cmdObj));
+    const auto actualBatchedInsert(BatchedCommandRequest::parseInsert(opMsgRequest));
+    ASSERT_EQUALS(nss.toString(), actualBatchedInsert.getNS().ns());
+
+    const auto& inserted = actualBatchedInsert.getInsertRequest().getDocuments();
+    ASSERT_EQUALS(expected.size(), inserted.size());
+
+    auto itInserted = inserted.begin();
+    auto itExpected = expected.begin();
+
+    for (; itInserted != inserted.end(); itInserted++, itExpected++) {
+        ASSERT_BSONOBJ_EQ(*itExpected, *itInserted);
+    }
+
+    BSONObjBuilder tenantMigrationAbortedResponse;
+    tenantMigrationAbortedResponse.append("ok", 1);
+    tenantMigrationAbortedResponse.append("n", int(inserted.size()));
+
+    int expectedSize = int(expected.size());
+    invariant(numberOfFailedOps >= 0 && numberOfFailedOps <= expectedSize,
+              str::stream() << "Expected numberOfFailedOps value to be between 0 and "
+                            << expectedSize << " but found " << numberOfFailedOps << ".");
+    int i = expectedSize - numberOfFailedOps;
+
+    std::vector<BSONObj> errors;
+    for (; i < expectedSize; i++) {
+        BSONObjBuilder errorBuilder;
+        errorBuilder.append("index", i);
+        errorBuilder.append("code", int(ErrorCodes::TenantMigrationAborted));
+        errorBuilder.append("errmsg", "mock tenantmigrationaborted error");
+        errors.push_back(errorBuilder.obj());
+        // ordered bulk only return one error and stop.
+        if (actualBatchedInsert.getWriteCommandRequestBase().getOrdered())
+            break;
+    }
+    tenantMigrationAbortedResponse.append("writeErrors", errors);
+
+    return tenantMigrationAbortedResponse.obj();
+}
+
+/**
  * Mimics a single shard backend for a particular collection which can be initialized with a
  * set of write command results to return.
  */
@@ -247,6 +298,14 @@ public:
         });
     }
 
+    void expectInsertsReturnTenantMigrationAbortedErrors(const std::vector<BSONObj>& expected,
+                                                         int numberOfFailedOps) {
+        onCommandForPoolExecutor([&](const executor::RemoteCommandRequest& request) {
+            return expectInsertsReturnTenantMigrationAbortedErrorsBase(
+                nss, expected, request, numberOfFailedOps);
+        });
+    }
+
     void expectInsertsReturnError(const std::vector<BSONObj>& expected,
                                   const BatchedCommandResponse& errResponse) {
         onCommandForPoolExecutor([&](const executor::RemoteCommandRequest& request) {
@@ -1625,6 +1684,204 @@ TEST_F(BatchWriteExecTest, StaleEpochIsNotRetryable) {
     future.default_timed_get();
 }
 
+TEST_F(BatchWriteExecTest, TenantMigrationAbortedErrorOrderedOp) {
+    const std::vector<BSONObj> expected{BSON("x" << 1), BSON("x" << 2), BSON("x" << 3)};
+    BatchedCommandRequest request([&] {
+        write_ops::InsertCommandRequest insertOp(nss);
+        insertOp.setWriteCommandRequestBase([] {
+            write_ops::WriteCommandRequestBase writeCommandBase;
+            writeCommandBase.setOrdered(true);
+            return writeCommandBase;
+        }());
+        insertOp.setDocuments(expected);
+        return insertOp;
+    }());
+    request.setWriteConcern(BSONObj());
+
+    // Execute request
+    auto future = launchAsync([&] {
+        BatchedCommandResponse response;
+        BatchWriteExecStats stats;
+        BatchWriteExec::executeBatch(
+            operationContext(), singleShardNSTargeter, request, &response, &stats);
+        ASSERT(response.getOk());
+
+        ASSERT_EQUALS(1, stats.numTenantMigrationAbortedErrors);
+    });
+
+    expectInsertsReturnTenantMigrationAbortedErrors(expected, expected.size());
+    expectInsertsReturnSuccess(expected);
+
+    future.default_timed_get();
+}
+
+TEST_F(BatchWriteExecTest, TenantMigrationAbortedErrorUnorderedOp) {
+    const std::vector<BSONObj> expected{BSON("x" << 1), BSON("x" << 2), BSON("x" << 3)};
+    BatchedCommandRequest request([&] {
+        write_ops::InsertCommandRequest insertOp(nss);
+        insertOp.setWriteCommandRequestBase([] {
+            write_ops::WriteCommandRequestBase writeCommandBase;
+            writeCommandBase.setOrdered(false);
+            return writeCommandBase;
+        }());
+        insertOp.setDocuments(expected);
+        return insertOp;
+    }());
+    request.setWriteConcern(BSONObj());
+
+    // Execute request
+    auto future = launchAsync([&] {
+        BatchedCommandResponse response;
+        BatchWriteExecStats stats;
+        BatchWriteExec::executeBatch(
+            operationContext(), singleShardNSTargeter, request, &response, &stats);
+        ASSERT(response.getOk());
+
+        ASSERT_EQUALS(1, stats.numTenantMigrationAbortedErrors);
+    });
+
+    expectInsertsReturnTenantMigrationAbortedErrors(expected, expected.size());
+    expectInsertsReturnSuccess(expected);
+
+    future.default_timed_get();
+}
+
+TEST_F(BatchWriteExecTest, MultipleTenantMigrationAbortedErrorUnorderedOp) {
+    const std::vector<BSONObj> expected{BSON("x" << 1), BSON("x" << 2), BSON("x" << 3)};
+    BatchedCommandRequest request([&] {
+        write_ops::InsertCommandRequest insertOp(nss);
+        insertOp.setWriteCommandRequestBase([] {
+            write_ops::WriteCommandRequestBase writeCommandBase;
+            writeCommandBase.setOrdered(false);
+            return writeCommandBase;
+        }());
+        insertOp.setDocuments(expected);
+        return insertOp;
+    }());
+    request.setWriteConcern(BSONObj());
+
+    const int numTenantMigrationAbortedErrors = 3;
+
+    // Execute request
+    auto future = launchAsync([&] {
+        BatchedCommandResponse response;
+        BatchWriteExecStats stats;
+        BatchWriteExec::executeBatch(
+            operationContext(), singleShardNSTargeter, request, &response, &stats);
+        ASSERT(response.getOk());
+
+        ASSERT_EQUALS(numTenantMigrationAbortedErrors, stats.numTenantMigrationAbortedErrors);
+    });
+
+    for (int i = 0; i < numTenantMigrationAbortedErrors; i++) {
+        expectInsertsReturnTenantMigrationAbortedErrors(expected, expected.size());
+    }
+    expectInsertsReturnSuccess(expected);
+
+    future.default_timed_get();
+}
+
+TEST_F(BatchWriteExecTest, MultipleTenantMigrationAbortedErrorOrderedOp) {
+    const std::vector<BSONObj> expected{BSON("x" << 1), BSON("x" << 2), BSON("x" << 3)};
+    BatchedCommandRequest request([&] {
+        write_ops::InsertCommandRequest insertOp(nss);
+        insertOp.setWriteCommandRequestBase([] {
+            write_ops::WriteCommandRequestBase writeCommandBase;
+            writeCommandBase.setOrdered(true);
+            return writeCommandBase;
+        }());
+        insertOp.setDocuments(expected);
+        return insertOp;
+    }());
+    request.setWriteConcern(BSONObj());
+
+    const int numTenantMigrationAbortedErrors = 3;
+
+    // Execute request
+    auto future = launchAsync([&] {
+        BatchedCommandResponse response;
+        BatchWriteExecStats stats;
+        BatchWriteExec::executeBatch(
+            operationContext(), singleShardNSTargeter, request, &response, &stats);
+        ASSERT(response.getOk());
+
+        ASSERT_EQUALS(numTenantMigrationAbortedErrors, stats.numTenantMigrationAbortedErrors);
+    });
+
+    for (int i = 0; i < numTenantMigrationAbortedErrors; i++) {
+        expectInsertsReturnTenantMigrationAbortedErrors(expected, expected.size());
+    }
+    expectInsertsReturnSuccess(expected);
+
+    future.default_timed_get();
+}
+
+TEST_F(BatchWriteExecTest, PartialTenantMigrationAbortedErrorOrderedOp) {
+    const std::vector<BSONObj> expected{BSON("x" << 1), BSON("x" << 2), BSON("x" << 3)};
+    BatchedCommandRequest request([&] {
+        write_ops::InsertCommandRequest insertOp(nss);
+        insertOp.setWriteCommandRequestBase([] {
+            write_ops::WriteCommandRequestBase writeCommandBase;
+            writeCommandBase.setOrdered(true);
+            return writeCommandBase;
+        }());
+        insertOp.setDocuments(expected);
+        return insertOp;
+    }());
+    request.setWriteConcern(BSONObj());
+
+    // Execute request
+    auto future = launchAsync([&] {
+        BatchedCommandResponse response;
+        BatchWriteExecStats stats;
+        BatchWriteExec::executeBatch(
+            operationContext(), singleShardNSTargeter, request, &response, &stats);
+        ASSERT(response.getOk());
+
+        ASSERT_EQUALS(1, stats.numTenantMigrationAbortedErrors);
+    });
+
+    const std::vector<BSONObj> expected_retries{BSON("x" << 2), BSON("x" << 3)};
+    int numberOfFailedOps = expected_retries.size();
+    expectInsertsReturnTenantMigrationAbortedErrors(expected, numberOfFailedOps);
+    expectInsertsReturnSuccess(expected_retries);
+
+    future.default_timed_get();
+}
+
+TEST_F(BatchWriteExecTest, PartialTenantMigrationErrorUnorderedOp) {
+    const std::vector<BSONObj> expected{BSON("x" << 1), BSON("x" << 2), BSON("x" << 3)};
+    BatchedCommandRequest request([&] {
+        write_ops::InsertCommandRequest insertOp(nss);
+        insertOp.setWriteCommandRequestBase([] {
+            write_ops::WriteCommandRequestBase writeCommandBase;
+            writeCommandBase.setOrdered(false);
+            return writeCommandBase;
+        }());
+        insertOp.setDocuments(expected);
+        return insertOp;
+    }());
+    request.setWriteConcern(BSONObj());
+
+    // Execute request
+    auto future = launchAsync([&] {
+        BatchedCommandResponse response;
+        BatchWriteExecStats stats;
+        BatchWriteExec::executeBatch(
+            operationContext(), singleShardNSTargeter, request, &response, &stats);
+        ASSERT(response.getOk());
+
+        ASSERT_EQUALS(1, stats.numTenantMigrationAbortedErrors);
+    });
+
+    const std::vector<BSONObj> expected_retries{BSON("x" << 2), BSON("x" << 3)};
+    int numberOfFailedOps = expected_retries.size();
+    expectInsertsReturnTenantMigrationAbortedErrors(expected, numberOfFailedOps);
+    expectInsertsReturnSuccess(expected_retries);
+
+    future.default_timed_get();
+}
+
 /**
  * Mimics a two shard backend with a targeting error on the first shard.
  */