SERVER-84336 Fix synchronization of timeseries unordered insertMany when exception is encountered and the thread has multiple write batches (#18099)r7.2.1-rc1

SERVER-84336 Fix exception handling of time-series unordered insert many
with multiple write batches

Threads now attempt to abort any write batches and wait until no other
thread is trying to commit the data before proceeding with tearing down
the command when an exception is encountered.

(cherry picked from commit bc1b48cc91e7ea05c96f7fb8f4ae2a905baeddab)

3 files changed, 359 insertions, 101 deletions

diff --git a/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_insert_many.js b/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_insert_many.js
new file mode 100644
index 00000000000..96a414c3d21
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_insert_many.js
@@ -0,0 +1,152 @@
+/**
+ * Tests insertMany into a time-series collection during a chunk migration. This test is not
+ * checking results but is meant to run in sanitizers to ensure the exception handling in the
+ * time-series insert many code is correct.
+ * @tags: [
+ *  requires_sharding,
+ *  assumes_balancer_off,
+ *  requires_non_retryable_writes,
+ *  does_not_support_transactions,
+ * ]
+ */
+import {extendWorkload} from "jstests/concurrency/fsm_libs/extend_workload.js";
+import {ChunkHelper} from "jstests/concurrency/fsm_workload_helpers/chunks.js";
+import {
+    $config as $baseConfig
+} from 'jstests/concurrency/fsm_workloads/sharded_moveChunk_partitioned.js';
+import {findChunksUtil} from "jstests/sharding/libs/find_chunks_util.js";
+
+export const $config = extendWorkload($baseConfig, function($config, $super) {
+    // A random non-round start value was chosen so that we can verify the rounding behavior that
+    // occurs while routing on mongos.
+    $config.data.startTime = 1021;
+
+    // One minute.
+    $config.data.increment = 1000 * 60;
+
+    // This should generate documents for a span of one month.
+    $config.data.numInitialDocs = 60 * 24 * 30;
+
+    $config.data.bucketPrefix = "system.buckets.";
+
+    $config.data.metaField = 'm';
+    $config.data.timeField = 't';
+
+    $config.threadCount = 10;
+    $config.iterations = 40;
+    $config.startState = "init";
+
+    /**
+     * Perform insertMany with ordered:false that may target multiple buckets across multiple chunks
+     */
+    $config.states.insert = function insert(db, collName, connCache) {
+        var docs = [];
+        for (let i = 0; i < 10; i++) {
+            // Generate a random timestamp between 'startTime' and largest timestamp we inserted.
+            const timer =
+                this.startTime + Math.floor(Random.rand() * this.numInitialDocs * this.increment);
+            const doc = {
+                _id: new ObjectId(),
+                [this.metaField]: 0,
+                [this.timeField]: new Date(timer),
+            };
+            docs.push(doc);
+        }
+
+        // Perform unordered insertMany. When this is done concurrent with chunk migrations we may
+        // get an exception in mongod when only a subset of documents have been (partially)
+        // processed. This will be retried by mongos and no error is observed for the user. This
+        // test is meant to be run with sanitizers to ensure correctness.
+        assert.commandWorked(db[collName].insertMany(docs, {ordered: false}));
+    };
+
+    /**
+     * Moves a random chunk in the target collection.
+     */
+    $config.states.moveChunk = function moveChunk(db, collName, connCache) {
+        const configDB = db.getSiblingDB('config');
+        const ns = db[this.bucketPrefix + collName].getFullName();
+        const chunks = findChunksUtil.findChunksByNs(configDB, ns).toArray();
+        const chunkToMove = chunks[this.tid];
+        const fromShard = chunkToMove.shard;
+
+        // Choose a random shard to move the chunk to.
+        const shardNames = Object.keys(connCache.shards);
+        const destinationShards = shardNames.filter(function(shard) {
+            if (shard !== fromShard) {
+                return shard;
+            }
+        });
+        const toShard = destinationShards[Random.randInt(destinationShards.length)];
+        const waitForDelete = false;
+        ChunkHelper.moveChunk(db,
+                              this.bucketPrefix + collName,
+                              [chunkToMove.min, chunkToMove.max],
+                              toShard,
+                              waitForDelete);
+    };
+
+    $config.states.init = function init(db, collName, connCache) {};
+
+    $config.transitions = {
+        init: {insert: 1},
+        insert: {insert: 7, moveChunk: 1},
+        moveChunk: {insert: 1, moveChunk: 0}
+    };
+
+    $config.setup = function setup(db, collName, cluster) {
+        db[collName].drop();
+
+        assert.commandWorked(db.createCollection(
+            collName, {timeseries: {metaField: this.metaField, timeField: this.timeField}}));
+        cluster.shardCollection(db[collName], {t: 1}, false);
+
+        const bulk = db[collName].initializeUnorderedBulkOp();
+
+        let currentTimeStamp = this.startTime;
+        for (let i = 0; i < this.numInitialDocs; ++i) {
+            currentTimeStamp += this.increment;
+
+            const metaVal = 0;
+            const doc = {
+                _id: new ObjectId(),
+                [this.metaField]: metaVal,
+                [this.timeField]: new Date(currentTimeStamp),
+                // use an invalid tid to not count these documents when we validate at teardown
+                tid: -1,
+            };
+            bulk.insert(doc);
+        }
+
+        let res = bulk.execute();
+        assert.commandWorked(res);
+        assert.eq(this.numInitialDocs, res.nInserted);
+
+        // Pick 'this.threadCount - 1' split points so that we have 'this.threadCount' chunks.
+        const chunkRange = (currentTimeStamp - this.startTime) / this.threadCount;
+        currentTimeStamp = this.startTime;
+        for (let i = 0; i < (this.threadCount - 1); ++i) {
+            currentTimeStamp += chunkRange;
+            assert.commandWorked(ChunkHelper.splitChunkAt(
+                db, this.bucketPrefix + collName, {'control.min.t': new Date(currentTimeStamp)}));
+        }
+
+        // Create an extra chunk on each shard to make sure multi:true operations return correct
+        // metrics in write results.
+        const destinationShards = Object.keys(cluster.getSerializedCluster().shards);
+        for (const destinationShard of destinationShards) {
+            currentTimeStamp += chunkRange;
+            assert.commandWorked(ChunkHelper.splitChunkAt(
+                db, this.bucketPrefix + collName, {'control.min.t': new Date(currentTimeStamp)}));
+
+            ChunkHelper.moveChunk(
+                db,
+                this.bucketPrefix + collName,
+                [{'control.min.t': new Date(currentTimeStamp)}, {'control.min.t': MaxKey}],
+                destinationShard,
+                /*waitForDelete=*/ false);
+        }
+    };
+
+    return $config;
+});
diff --git a/src/mongo/db/ops/write_ops_exec.cpp b/src/mongo/db/ops/write_ops_exec.cpp
index 2a0b8d6f03a..30ccbc32775 100644
--- a/src/mongo/db/ops/write_ops_exec.cpp
+++ b/src/mongo/db/ops/write_ops_exec.cpp
@@ -1000,18 +1000,16 @@ boost::optional<write_ops::WriteError> generateError(OperationContext* opCtx,
         return boost::none;
     }
 
-    boost::optional<Status> overwrittenStatus;
-
     if (status == ErrorCodes::TenantMigrationConflict) {
         hangWriteBeforeWaitingForMigrationDecision.pauseWhileSet(opCtx);
 
-        overwrittenStatus.emplace(
-            tenant_migration_access_blocker::handleTenantMigrationConflict(opCtx, status));
+        Status overwrittenStatus =
+            tenant_migration_access_blocker::handleTenantMigrationConflict(opCtx, status);
 
         // Interruption errors encountered during batch execution fail the entire batch, so throw on
         // such errors here for consistency.
-        if (ErrorCodes::isInterruption(*overwrittenStatus)) {
-            uassertStatusOK(*overwrittenStatus);
+        if (ErrorCodes::isInterruption(overwrittenStatus)) {
+            uassertStatusOK(overwrittenStatus);
         }
 
         // Tenant migration errors, similarly to migration errors consume too much space in the
@@ -1019,25 +1017,34 @@ boost::optional<write_ops::WriteError> generateError(OperationContext* opCtx,
         // 'handleTenantMigrationConflict' above replaces the original status, we need to manually
         // truncate the new reason if the original 'status' was also truncated.
         if (status.reason().empty()) {
-            overwrittenStatus = overwrittenStatus->withReason("");
+            overwrittenStatus = overwrittenStatus.withReason("");
         }
+
+        return generateErrorNoTenantMigration(opCtx, overwrittenStatus, index, numErrors);
+    }
+
+    return generateErrorNoTenantMigration(opCtx, status, index, numErrors);
+}
+
+boost::optional<write_ops::WriteError> generateErrorNoTenantMigration(OperationContext* opCtx,
+                                                                      const Status& status,
+                                                                      int index,
+                                                                      size_t numErrors) noexcept {
+    if (status.isOK()) {
+        return boost::none;
     }
 
     constexpr size_t kMaxErrorReasonsToReport = 1;
     constexpr size_t kMaxErrorSizeToReportAfterMaxReasonsReached = 1024 * 1024;
 
     if (numErrors > kMaxErrorReasonsToReport) {
-        size_t errorSize =
-            overwrittenStatus ? overwrittenStatus->reason().size() : status.reason().size();
-        if (errorSize > kMaxErrorSizeToReportAfterMaxReasonsReached)
-            overwrittenStatus =
-                overwrittenStatus ? overwrittenStatus->withReason("") : status.withReason("");
+        size_t errorSize = status.reason().size();
+        if (errorSize > kMaxErrorSizeToReportAfterMaxReasonsReached) {
+            return write_ops::WriteError(index, status.withReason(""));
+        }
     }
 
-    if (overwrittenStatus)
-        return write_ops::WriteError(index, std::move(*overwrittenStatus));
-    else
-        return write_ops::WriteError(index, status);
+    return write_ops::WriteError(index, status);
 }
 
 void updateRetryStats(OperationContext* opCtx, bool containsRetry) {
@@ -2556,6 +2563,129 @@ void rebuildOptionsWithGranularityFromConfigServer(OperationContext* opCtx,
     }
 }
 
+
+// Waits for all batches to either commit or abort. This function will attempt to acquire commit
+// rights to any batch to mark it as aborted. If another thread alreay have commit rights we will
+// instead wait for the promise when the commit is either successful or failed. Marked as 'noexcept'
+// as we need to safely be able to call this function during exception handling.
+template <typename ErrorGenerator>
+void getTimeseriesBatchResultsBase(OperationContext* opCtx,
+                                   const TimeseriesBatches& batches,
+                                   int64_t start,
+                                   int64_t indexOfLastProcessedBatch,
+                                   bool canContinue,
+                                   std::vector<write_ops::WriteError>* errors,
+                                   boost::optional<repl::OpTime>* opTime,
+                                   boost::optional<OID>* electionId,
+                                   std::vector<size_t>* docsToRetry) {
+    boost::optional<write_ops::WriteError> lastError;
+    if (!errors->empty()) {
+        lastError = errors->back();
+    }
+    invariant(indexOfLastProcessedBatch == (int64_t)batches.size() || lastError);
+
+    for (int64_t itr = 0, size = batches.size(); itr < size; ++itr) {
+        const auto& [batch, index] = batches[itr];
+        if (!batch) {
+            continue;
+        }
+
+        // If there are any unprocessed batches, we mark them as error with the last known
+        // error.
+        if (itr > indexOfLastProcessedBatch &&
+            timeseries::bucket_catalog::claimWriteBatchCommitRights(*batch)) {
+            auto& bucketCatalog = timeseries::bucket_catalog::BucketCatalog::get(opCtx);
+            abort(bucketCatalog, batch, lastError->getStatus());
+            errors->emplace_back(start + index, lastError->getStatus());
+            continue;
+        }
+
+        auto swCommitInfo = timeseries::bucket_catalog::getWriteBatchResult(*batch);
+        if (swCommitInfo.getStatus() == ErrorCodes::TimeseriesBucketCleared) {
+            invariant(docsToRetry, "the 'docsToRetry' cannot be null");
+            docsToRetry->push_back(index);
+            continue;
+        }
+        if (swCommitInfo.getStatus() == ErrorCodes::WriteConflict ||
+            swCommitInfo.getStatus() == ErrorCodes::TemporarilyUnavailable) {
+            docsToRetry->push_back(index);
+            opCtx->recoveryUnit()->abandonSnapshot();
+            continue;
+        }
+        if (auto error =
+                ErrorGenerator{}(opCtx, swCommitInfo.getStatus(), start + index, errors->size())) {
+            errors->emplace_back(std::move(*error));
+            continue;
+        }
+
+        const auto& commitInfo = swCommitInfo.getValue();
+        if (commitInfo.opTime) {
+            *opTime = std::max(opTime->value_or(repl::OpTime()), *commitInfo.opTime);
+        }
+        if (commitInfo.electionId) {
+            *electionId = std::max(electionId->value_or(OID()), *commitInfo.electionId);
+        }
+    }
+
+    // If we cannot continue the request, we should convert all the 'docsToRetry' into an
+    // error.
+    if (!canContinue && docsToRetry) {
+        for (auto&& index : *docsToRetry) {
+            errors->emplace_back(start + index, lastError->getStatus());
+        }
+        docsToRetry->clear();
+    }
+}
+
+void getTimeseriesBatchResults(OperationContext* opCtx,
+                               const TimeseriesBatches& batches,
+                               int64_t start,
+                               int64_t indexOfLastProcessedBatch,
+                               bool canContinue,
+                               std::vector<write_ops::WriteError>* errors,
+                               boost::optional<repl::OpTime>* opTime,
+                               boost::optional<OID>* electionId,
+                               std::vector<size_t>* docsToRetry = nullptr) {
+    auto errorGenerator =
+        [](OperationContext* opCtx, const Status& status, int index, size_t numErrors) {
+            return write_ops_exec::generateError(opCtx, status, index, numErrors);
+        };
+    getTimeseriesBatchResultsBase<decltype(errorGenerator)>(opCtx,
+                                                            batches,
+                                                            start,
+                                                            indexOfLastProcessedBatch,
+                                                            canContinue,
+                                                            errors,
+                                                            opTime,
+                                                            electionId,
+                                                            docsToRetry);
+}
+
+void getTimeseriesBatchResultsNoTenantMigration(
+    OperationContext* opCtx,
+    const TimeseriesBatches& batches,
+    int64_t start,
+    int64_t indexOfLastProcessedBatch,
+    bool canContinue,
+    std::vector<write_ops::WriteError>* errors,
+    boost::optional<repl::OpTime>* opTime,
+    boost::optional<OID>* electionId,
+    std::vector<size_t>* docsToRetry = nullptr) noexcept {
+    auto errorGenerator =
+        [](OperationContext* opCtx, const Status& status, int index, size_t numErrors) {
+            return write_ops_exec::generateErrorNoTenantMigration(opCtx, status, index, numErrors);
+        };
+    getTimeseriesBatchResultsBase<decltype(errorGenerator)>(opCtx,
+                                                            batches,
+                                                            start,
+                                                            indexOfLastProcessedBatch,
+                                                            canContinue,
+                                                            errors,
+                                                            opTime,
+                                                            electionId,
+                                                            docsToRetry);
+}
+
 std::tuple<TimeseriesBatches, TimeseriesStmtIds, size_t /* numInserted */> insertIntoBucketCatalog(
     OperationContext* opCtx,
     size_t start,
@@ -2653,86 +2783,37 @@ std::tuple<TimeseriesBatches, TimeseriesStmtIds, size_t /* numInserted */> inser
         return true;
     };
 
-    if (!indices.empty()) {
-        std::for_each(indices.begin(), indices.end(), insert);
-    } else {
-        for (size_t i = 0; i < numDocs; i++) {
-            if (!insert(i) && request.getOrdered()) {
-                return {std::move(batches), std::move(stmtIds), i};
+    try {
+        if (!indices.empty()) {
+            std::for_each(indices.begin(), indices.end(), insert);
+        } else {
+            for (size_t i = 0; i < numDocs; i++) {
+                if (!insert(i) && request.getOrdered()) {
+                    return {std::move(batches), std::move(stmtIds), i};
+                }
             }
         }
+    } catch (const DBException& ex) {
+        // Exception insert into bucket catalog, append error and wait for all batches that we've
+        // already managed to write into to commit or abort. We need to wait here as pointers to
+        // memory owned by this command is stored in the WriteBatch(es). This ensures that no other
+        // thread may try to access this memory after this command has been torn down due to the
+        // exception.
+
+        boost::optional<repl::OpTime> opTime;
+        boost::optional<OID> electionId;
+        std::vector<size_t> docsToRetry;
+        errors->emplace_back(*write_ops_exec::generateErrorNoTenantMigration(
+            opCtx, ex.toStatus(), 0, errors->size()));
+
+        getTimeseriesBatchResultsNoTenantMigration(
+            opCtx, batches, 0, -1, false, errors, &opTime, &electionId, &docsToRetry);
+        throw;
     }
 
     return {std::move(batches), std::move(stmtIds), request.getDocuments().size()};
 }
 
-void getTimeseriesBatchResults(OperationContext* opCtx,
-                               const TimeseriesBatches& batches,
-                               size_t start,
-                               size_t indexOfLastProcessedBatch,
-                               bool canContinue,
-                               std::vector<write_ops::WriteError>* errors,
-                               boost::optional<repl::OpTime>* opTime,
-                               boost::optional<OID>* electionId,
-                               std::vector<size_t>* docsToRetry = nullptr) {
-    boost::optional<write_ops::WriteError> lastError;
-    if (!errors->empty()) {
-        lastError = errors->back();
-    }
-
-    for (size_t itr = 0; itr < batches.size(); ++itr) {
-        const auto& [batch, index] = batches[itr];
-        if (!batch) {
-            continue;
-        }
-
-        // If there are any unprocessed batches, we mark them as error with the last known
-        // error.
-        if (itr > indexOfLastProcessedBatch &&
-            timeseries::bucket_catalog::claimWriteBatchCommitRights(*batch)) {
-            auto& bucketCatalog = timeseries::bucket_catalog::BucketCatalog::get(opCtx);
-            abort(bucketCatalog, batch, lastError->getStatus());
-            errors->emplace_back(start + index, lastError->getStatus());
-            continue;
-        }
-
-        auto swCommitInfo = timeseries::bucket_catalog::getWriteBatchResult(*batch);
-        if (swCommitInfo.getStatus() == ErrorCodes::TimeseriesBucketCleared) {
-            tassert(6023102, "the 'docsToRetry' cannot be null", docsToRetry);
-            docsToRetry->push_back(index);
-            continue;
-        }
-        if (swCommitInfo.getStatus() == ErrorCodes::WriteConflict ||
-            swCommitInfo.getStatus() == ErrorCodes::TemporarilyUnavailable) {
-            docsToRetry->push_back(index);
-            opCtx->recoveryUnit()->abandonSnapshot();
-            continue;
-        }
-        if (auto error = write_ops_exec::generateError(
-                opCtx, swCommitInfo.getStatus(), start + index, errors->size())) {
-            errors->emplace_back(std::move(*error));
-            continue;
-        }
-
-        const auto& commitInfo = swCommitInfo.getValue();
-        if (commitInfo.opTime) {
-            *opTime = std::max(opTime->value_or(repl::OpTime()), *commitInfo.opTime);
-        }
-        if (commitInfo.electionId) {
-            *electionId = std::max(electionId->value_or(OID()), *commitInfo.electionId);
-        }
-    }
-
-    // If we cannot continue the request, we should convert all the 'docsToRetry' into an
-    // error.
-    if (!canContinue && docsToRetry) {
-        for (auto&& index : *docsToRetry) {
-            errors->emplace_back(start + index, lastError->getStatus());
-        }
-        docsToRetry->clear();
-    }
-}
-
 bool performOrderedTimeseriesWritesAtomically(OperationContext* opCtx,
                                               std::vector<write_ops::WriteError>* errors,
                                               boost::optional<repl::OpTime>* opTime,
@@ -2783,7 +2864,7 @@ std::vector<size_t> performUnorderedTimeseriesWrites(
 
     stdx::unordered_set<timeseries::bucket_catalog::WriteBatch*> handledHere;
     int64_t handledElsewhere = 0;
-    auto guard = ScopeGuard([&handledElsewhere, &request, opCtx]() {
+    auto reportMeasurementsGuard = ScopeGuard([&handledElsewhere, &request, opCtx]() {
         if (handledElsewhere > 0) {
             auto& bucketCatalog = timeseries::bucket_catalog::BucketCatalog::get(opCtx);
             timeseries::bucket_catalog::reportMeasurementsGroupCommitted(
@@ -2799,17 +2880,29 @@ std::vector<size_t> performUnorderedTimeseriesWrites(
             auto stmtIds = isTimeseriesWriteRetryable(opCtx)
                 ? std::move(bucketStmtIds[batch->bucketHandle.bucketId.oid])
                 : std::vector<StmtId>{};
-            canContinue = commitTimeseriesBucket(opCtx,
-                                                 batch,
-                                                 start,
-                                                 index,
-                                                 std::move(stmtIds),
-                                                 errors,
-                                                 opTime,
-                                                 electionId,
-                                                 &docsToRetry,
-                                                 retryAttemptsForDup,
-                                                 request);
+            try {
+                canContinue = commitTimeseriesBucket(opCtx,
+                                                     batch,
+                                                     start,
+                                                     index,
+                                                     std::move(stmtIds),
+                                                     errors,
+                                                     opTime,
+                                                     electionId,
+                                                     &docsToRetry,
+                                                     retryAttemptsForDup,
+                                                     request);
+            } catch (const DBException& ex) {
+                // Exception during commit, append error and wait for all our batches to commit or
+                // abort. We need to wait here as pointers to memory owned by this command is stored
+                // in the WriteBatch(es). This ensures that no other thread may try to access this
+                // memory after this command has been torn down due to the exception.
+                errors->emplace_back(*write_ops_exec::generateErrorNoTenantMigration(
+                    opCtx, ex.toStatus(), start + index, errors->size()));
+                getTimeseriesBatchResultsNoTenantMigration(
+                    opCtx, batches, 0, itr, canContinue, errors, opTime, electionId, &docsToRetry);
+                throw;
+            }
 
             batch.reset();
             if (!canContinue) {
diff --git a/src/mongo/db/ops/write_ops_exec.h b/src/mongo/db/ops/write_ops_exec.h
index e52997b8475..20237385cd1 100644
--- a/src/mongo/db/ops/write_ops_exec.h
+++ b/src/mongo/db/ops/write_ops_exec.h
@@ -150,6 +150,8 @@ long long performDelete(OperationContext* opCtx,
 
 /**
  * Generates a WriteError for a given Status.
+ *
+ * This function may throw.
  */
 boost::optional<write_ops::WriteError> generateError(OperationContext* opCtx,
                                                      const Status& status,
@@ -157,6 +159,17 @@ boost::optional<write_ops::WriteError> generateError(OperationContext* opCtx,
                                                      size_t numErrors);
 
 /**
+ * Generates a WriteError for a given Status. Does not handle tenant migration errors.
+ *
+ * Marked as 'noexcept' as we need to safely be able to call this function during exception
+ * handling.
+ */
+boost::optional<write_ops::WriteError> generateErrorNoTenantMigration(OperationContext* opCtx,
+                                                                      const Status& status,
+                                                                      int index,
+                                                                      size_t numErrors) noexcept;
+
+/**
  * Updates the retryable write stats if the write op contains retry.
  */
 void updateRetryStats(OperationContext* opCtx, bool containsRetry);