SERVER-59180 Implement update rewrite and routing for sharded timeseries collections

2 files changed, 101 insertions, 8 deletions

diff --git a/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_inserts.js b/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_inserts.js
index 4b8e496dd05..fa052f3c730 100644
--- a/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_inserts.js
+++ b/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_inserts.js
@@ -34,6 +34,17 @@ var $config = extendWorkload($config, function($config, $super) {
 
     $config.data.bucketPrefix = "system.buckets.";
 
+    $config.data.timeField = 't';
+    $config.data.metaField = 'm';
+
+    $config.data.generateMetaFieldValueForInitialInserts = () => {
+        return Math.floor(Random.rand() * $config.data.numMetaCount);
+    };
+
+    $config.data.generateMetaFieldValueForInsertStage = (i) => {
+        return i % $config.data.numMetaCount;
+    };
+
     $config.threadCount = 10;
     $config.iterations = 40;
     $config.startState = "init";
@@ -49,8 +60,8 @@ var $config = extendWorkload($config, function($config, $super) {
                 this.startTime + Math.floor(Random.rand() * this.numInitialDocs * this.increment);
             const doc = {
                 _id: new ObjectId(),
-                t: new Date(timer),
-                m: Math.floor(Random.rand() * this.numMetaCount)
+                [this.timeField]: new Date(timer),
+                [this.metaField]: this.generateMetaFieldValueForInsertStage(this.tid),
             };
             assertAlways.commandWorked(db[collName].insert(doc));
             assertAlways.commandWorked(db[this.nonShardCollName].insert(doc));
@@ -121,8 +132,11 @@ var $config = extendWorkload($config, function($config, $super) {
         // buckets with all documents.
         const verifyBucketIndex = (bucketIndex) => {
             const unpackStage = {
-                "$_internalUnpackBucket":
-                    {"timeField": "t", "metaField": "m", "bucketMaxSpanSeconds": NumberInt(3600)}
+                "$_internalUnpackBucket": {
+                    "timeField": this.timeField,
+                    "metaField": this.metaField,
+                    "bucketMaxSpanSeconds": NumberInt(3600)
+                }
             };
             const bucketColl = db.getCollection(`system.buckets.${collName}`);
             const numDocsInBuckets =
@@ -151,8 +165,8 @@ var $config = extendWorkload($config, function($config, $super) {
         db[collName].drop();
         db[this.nonShardCollName].drop();
 
-        assertAlways.commandWorked(
-            db.createCollection(collName, {timeseries: {timeField: "t", metaField: "m"}}));
+        assertAlways.commandWorked(db.createCollection(
+            collName, {timeseries: {timeField: this.timeField, metaField: this.metaField}}));
         cluster.shardCollection(db[collName], {t: 1}, false);
 
         // Create indexes to verify index integrity during the teardown state.
@@ -171,8 +185,8 @@ var $config = extendWorkload($config, function($config, $super) {
 
             const doc = {
                 _id: new ObjectId(),
-                t: new Date(currentTimeStamp),
-                m: i % this.numMetaCount
+                [this.timeField]: new Date(currentTimeStamp),
+                [this.metaField]: this.generateMetaFieldValueForInitialInserts(i),
             };
             bulk.insert(doc);
             bulkUnsharded.insert(doc);
diff --git a/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_updates.js b/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_updates.js
new file mode 100644
index 00000000000..e700d90f095
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_updates.js
@@ -0,0 +1,79 @@
+/**
+ * Tests the updates into sharded time-series collection during a chunk migration. To ensure the
+ * correctness, the test does the same writes into an unsharded collection and verifies that the
+ * number of documents remain the same at the end. This test also checks that indexes on the
+ * time-series buckets collection remain consistent after the test run.
+ * @tags: [
+ *  requires_sharding,
+ *  assumes_balancer_off,
+ *  requires_non_retryable_writes,
+ * ]
+ */
+'use strict';
+
+const numValues = 10;
+
+load("jstests/core/timeseries/libs/timeseries.js");  // For 'TimeseriesTest' helpers.
+load('jstests/concurrency/fsm_workloads/random_moveChunk_timeseries_inserts.js');
+
+var $config = extendWorkload($config, function($config, $super) {
+    $config.states.init = function(db, collName, connCache) {
+        if (TimeseriesTest.timeseriesCollectionsEnabled(db) &&
+            TimeseriesTest.shardedtimeseriesCollectionsEnabled(db) &&
+            TimeseriesTest.timeseriesUpdatesAndDeletesEnabled(db) &&
+            TimeseriesTest.shardedTimeseriesUpdatesAndDeletesEnabled(db)) {
+            this.featureFlagDisabled = false;
+        } else {
+            jsTestLog(
+                "Skipping executing this test as the requisite feature flags are not enabled.");
+        }
+
+        $super.states.init(db, collName);
+    };
+
+    $config.data.generateMetaFieldValueForInitialInserts = () => {
+        let meta = {};
+        // Insert a document with a field for every thread to test concurrent updates on the
+        // same document.
+        for (let i = 0; i < $config.threadCount; i++) {
+            meta["tid" + i] = Random.randInt(numValues);
+        }
+        return meta;
+    };
+
+    $config.data.generateMetaFieldValueForInsertStage = (tid) => {
+        return {["tid" + tid]: Random.randInt(numValues)};
+    };
+
+    $config.states.update = function(db, collName, connCache) {
+        if (this.featureFlagDisabled) {
+            return;
+        }
+
+        const shardedColl = db[collName];
+        const unshardedColl = db[this.nonShardCollName];
+        const updateField = "tid" + this.tid;
+        const oldValue = Random.randInt(numValues);
+
+        // Updates some measurements along the field owned by this thread in both sharded and
+        // unsharded ts collections.
+        jsTestLog("Executing update state on: " + collName + " on field " + updateField);
+        assertAlways.commandWorked(
+            shardedColl.update({[this.metaField]: {[updateField]: {$gte: oldValue}}},
+                               {$inc: {[this.metaField + "." + updateField]: 1}},
+                               {multi: true}));
+        assertAlways.commandWorked(
+            unshardedColl.update({[this.metaField]: {[updateField]: {$gte: oldValue}}},
+                                 {$inc: {[this.metaField + "." + updateField]: 1}},
+                                 {multi: true}));
+    };
+
+    $config.transitions = {
+        init: {insert: 1},
+        insert: {insert: 0.4, moveChunk: 0.1, update: 0.5},
+        update: {insert: 0.5, moveChunk: 0.1, update: 0.4},
+        moveChunk: {insert: 0.4, moveChunk: 0.1, update: 0.5},
+    };
+
+    return $config;
+});