SERVER-74610 Support fromMigrate option in timeseries::performAtomicWrites()

4 files changed, 20 insertions, 13 deletions

diff --git a/src/mongo/db/exec/timeseries_write.cpp b/src/mongo/db/exec/timeseries_write.cpp
index 0493a078cc7..fe304e3b843 100644
--- a/src/mongo/db/exec/timeseries_write.cpp
+++ b/src/mongo/db/exec/timeseries_write.cpp
@@ -77,7 +77,8 @@ void TimeseriesWriteStage::_writeToTimeseriesBuckets() {
         write_ops::DeleteOpEntry deleteEntry(BSON("_id" << bucketId), false);
         write_ops::DeleteCommandRequest op(collection()->ns(), {deleteEntry});
         // TODO (SERVER-73093): Handles the write failures through retry.
-        auto result = timeseries::performAtomicWrites(opCtx(), collection(), _currentBucketRid, op);
+        auto result = timeseries::performAtomicWrites(
+            opCtx(), collection(), _currentBucketRid, op, /*fromMigrate=*/false);
     } else {
         auto timeseriesOptions = collection()->getTimeseriesOptions();
         auto metaFieldName = timeseriesOptions->getMetaField();
@@ -94,7 +95,8 @@ void TimeseriesWriteStage::_writeToTimeseriesBuckets() {
         write_ops::UpdateOpEntry updateEntry(BSON("_id" << bucketId), std::move(u));
         write_ops::UpdateCommandRequest op(collection()->ns(), {updateEntry});
         // TODO (SERVER-73093): Handles the write failures through retry.
-        auto result = timeseries::performAtomicWrites(opCtx(), collection(), _currentBucketRid, op);
+        auto result = timeseries::performAtomicWrites(
+            opCtx(), collection(), _currentBucketRid, op, /*fromMigrate=*/false);
     }
 }
 
diff --git a/src/mongo/db/timeseries/timeseries_write_util.cpp b/src/mongo/db/timeseries/timeseries_write_util.cpp
index 48d808aa67e..6fda7516f4a 100644
--- a/src/mongo/db/timeseries/timeseries_write_util.cpp
+++ b/src/mongo/db/timeseries/timeseries_write_util.cpp
@@ -143,12 +143,12 @@ BSONObj makeNewDocumentForWrite(
     return makeNewDocument(bucketId, metadata, minmax->min(), minmax->max(), dataBuilders);
 }
 
-Status performAtomicWrites(
-    OperationContext* opCtx,
-    const CollectionPtr& coll,
-    const RecordId& recordId,
-    const stdx::variant<write_ops::UpdateCommandRequest, write_ops::DeleteCommandRequest>&
-        modificationOp) try {
+Status performAtomicWrites(OperationContext* opCtx,
+                           const CollectionPtr& coll,
+                           const RecordId& recordId,
+                           const stdx::variant<write_ops::UpdateCommandRequest,
+                                               write_ops::DeleteCommandRequest>& modificationOp,
+                           bool fromMigrate) try {
     invariant(!opCtx->lockState()->inAWriteUnitOfWork());
     invariant(!opCtx->inMultiDocumentTransaction());
 
@@ -181,7 +181,9 @@ Status performAtomicWrites(
                 if (const auto& stmtIds = updateOp.getStmtIds()) {
                     args.stmtIds = *stmtIds;
                 }
-                args.source = OperationSource::kTimeseriesDelete;
+                if (fromMigrate) {
+                    args.source = OperationSource::kFromMigrate;
+                }
 
                 BSONObj diffFromUpdate;
                 const BSONObj* diffOnIndexes =
@@ -208,7 +210,7 @@ Status performAtomicWrites(
                     record_id_helpers::keyForOID(deleteOp.getDeletes().front().getQ()["_id"].OID());
                 invariant(recordId == deleteId);
                 collection_internal::deleteDocument(
-                    opCtx, coll, kUninitializedStmtId, recordId, &curOp->debug());
+                    opCtx, coll, kUninitializedStmtId, recordId, &curOp->debug(), fromMigrate);
             }},
         modificationOp);
 
diff --git a/src/mongo/db/timeseries/timeseries_write_util.h b/src/mongo/db/timeseries/timeseries_write_util.h
index ecc432183c5..e99c9ad9340 100644
--- a/src/mongo/db/timeseries/timeseries_write_util.h
+++ b/src/mongo/db/timeseries/timeseries_write_util.h
@@ -65,5 +65,6 @@ Status performAtomicWrites(OperationContext* opCtx,
                            const CollectionPtr& coll,
                            const RecordId& recordId,
                            const stdx::variant<write_ops::UpdateCommandRequest,
-                                               write_ops::DeleteCommandRequest>& modificationOp);
+                                               write_ops::DeleteCommandRequest>& modificationOp,
+                           bool fromMigrate);
 }  // namespace mongo::timeseries
diff --git a/src/mongo/db/timeseries/timeseries_write_util_test.cpp b/src/mongo/db/timeseries/timeseries_write_util_test.cpp
index 3688ccd5c5b..e797ec3db60 100644
--- a/src/mongo/db/timeseries/timeseries_write_util_test.cpp
+++ b/src/mongo/db/timeseries/timeseries_write_util_test.cpp
@@ -223,7 +223,8 @@ TEST_F(TimeseriesWriteUtilTest, PerformAtomicDelete) {
 
         op.setWriteCommandRequestBase(std::move(base));
 
-        ASSERT_OK(performAtomicWrites(opCtx, bucketsColl.getCollection(), recordId, op));
+        ASSERT_OK(performAtomicWrites(
+            opCtx, bucketsColl.getCollection(), recordId, op, /*fromMigrate=*/false));
     }
 
     // Checks the document is removed.
@@ -285,7 +286,8 @@ TEST_F(TimeseriesWriteUtilTest, PerformAtomicUpdate) {
 
         op.setWriteCommandRequestBase(std::move(base));
 
-        ASSERT_OK(performAtomicWrites(opCtx, bucketsColl.getCollection(), recordId, op));
+        ASSERT_OK(performAtomicWrites(
+            opCtx, bucketsColl.getCollection(), recordId, op, /*fromMigrate=*/false));
     }
 
     // Checks the document is updated.