SERVER-60464 Enable balancer for sharded time-series collections

3 files changed, 99 insertions, 17 deletions

diff --git a/jstests/sharding/timeseries_balancer.js b/jstests/sharding/timeseries_balancer.js
new file mode 100644
index 00000000000..25480503738
--- /dev/null
+++ b/jstests/sharding/timeseries_balancer.js
@@ -0,0 +1,94 @@
+/**
+ * Test inserts into sharded timeseries collection with the balancer on.
+ *
+ * @tags: [
+ *   requires_fcv_51,
+ * ]
+ */
+
+(function() {
+"use strict";
+
+load("jstests/core/timeseries/libs/timeseries.js");  // For 'TimeseriesTest' helpers.
+
+Random.setRandomSeed();
+
+const dbName = 'testDB';
+const collName = 'testColl';
+const timeField = 'time';
+const metaField = 'hostid';
+
+// Connections.
+const st = new ShardingTest({shards: 2, rs: {nodes: 2}});
+const mongos = st.s0;
+
+// Sanity checks.
+if (!TimeseriesTest.shardedtimeseriesCollectionsEnabled(st.shard0)) {
+    jsTestLog("Skipping test because the sharded time-series collection feature flag is disabled");
+    st.stop();
+    return;
+}
+
+// Databases and collections.
+assert.commandWorked(mongos.adminCommand({enableSharding: dbName}));
+const mainDB = mongos.getDB(dbName);
+
+// Helpers.
+let currentId = 0;
+function generateId() {
+    return currentId++;
+}
+
+const largeStr = "a".repeat(10000);
+
+function generateBatch(size) {
+    return TimeseriesTest.generateHosts(size).map((host, index) => Object.assign(host, {
+        _id: generateId(),
+        [metaField]: index,
+        [timeField]: new Date(currentId),
+        largeField: largeStr,
+    }));
+}
+
+st.startBalancer();
+function runTest(shardKey) {
+    assert.commandWorked(mainDB.createCollection(
+        collName, {timeseries: {timeField: timeField, metaField: metaField}}));
+    const coll = mainDB.getCollection(collName);
+
+    // Shard timeseries collection.
+    assert.commandWorked(coll.createIndex(shardKey));
+    assert.commandWorked(mongos.adminCommand({
+        shardCollection: `${dbName}.${collName}`,
+        key: shardKey,
+    }));
+
+    // Insert a large dataset so that the balancer is guranteed to split the chunks.
+    let bulk = coll.initializeUnorderedBulkOp();
+    const numDocs = 100000;
+    const firstBatch = generateBatch(numDocs);
+    for (let doc of firstBatch) {
+        bulk.insert(doc);
+    }
+    assert.commandWorked(bulk.execute());
+
+    // Ensure that each shard has at least one chunk after the split.
+    const primaryShard = st.getPrimaryShard(dbName);
+    const otherShard = st.getOther(primaryShard);
+    assert.soon(
+        () => {
+            const counts = st.chunkCounts(`system.buckets.${collName}`, dbName);
+            return counts[primaryShard.shardName] >= 1 && counts[otherShard.shardName] >= 1;
+        },
+        () => {
+            return tojson(mongos.getDB("config").getCollection("chunks").find().toArray());
+        });
+
+    // Verify that all the documents still exist in the collection.
+    assert.eq(coll.find().itcount(), numDocs);
+    assert(coll.drop());
+}
+
+runTest({time: 1});
+st.stop();
+})();
diff --git a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
index 1c8d241778b..67ce067c200 100644
--- a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
@@ -296,16 +296,6 @@ StatusWith<SplitInfoVector> BalancerChunkSelectionPolicyImpl::selectChunksToSpli
     for (const auto& coll : collections) {
         const NamespaceString& nss(coll.getNss());
 
-        if (coll.getTimeseriesFields()) {
-            LOGV2_DEBUG(5559200,
-                        1,
-                        "Not splitting collection {namespace}; explicitly disabled.",
-                        "Not splitting explicitly disabled collection",
-                        "namespace"_attr = nss,
-                        "timeseriesFields"_attr = coll.getTimeseriesFields());
-            continue;
-        }
-
         auto candidatesStatus = _getSplitCandidatesForCollection(opCtx, nss, shardStats);
         if (candidatesStatus == ErrorCodes::NamespaceNotFound) {
             // Namespace got dropped before we managed to get to it, so just skip it
@@ -375,7 +365,7 @@ StatusWith<MigrateInfoVector> BalancerChunkSelectionPolicyImpl::selectChunksToMo
     for (const auto& coll : collections) {
         const NamespaceString& nss(coll.getNss());
 
-        if (!coll.getAllowBalance() || !coll.getAllowMigrations() || coll.getTimeseriesFields()) {
+        if (!coll.getAllowBalance() || !coll.getAllowMigrations()) {
             LOGV2_DEBUG(21851,
                         1,
                         "Not balancing collection {namespace}; explicitly disabled.",
diff --git a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp
index 1b2e9f33dcc..4e1f232fc51 100644
--- a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp
@@ -231,7 +231,7 @@ TEST_F(BalancerChunkSelectionTest, TagRangeMaxNotAlignedWithChunkMax) {
                               {BSON(kPattern << -15), kKeyPattern.globalMax()}});
 }
 
-TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCannotBeAutoSplitted) {
+TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCanBeAutoSplitted) {
     // Set up two shards in the metadata, each one with its own tag
     ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(),
                                                     ShardType::ConfigNS,
@@ -274,15 +274,14 @@ TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCannotBeAutoSplit
         auto candidateChunksStatus = _chunkSelectionPolicy.get()->selectChunksToSplit(opCtx.get());
         ASSERT_OK(candidateChunksStatus.getStatus());
 
-        // No chunks to split since the coll is a sharded time-series collection
-        ASSERT_EQUALS(0U, candidateChunksStatus.getValue().size());
+        ASSERT_EQUALS(1U, candidateChunksStatus.getValue().size());
     });
 
     expectGetStatsCommands(2);
     future.default_timed_get();
 }
 
-TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCannotBeBalanced) {
+TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCanBeBalanced) {
     // Set up two shards in the metadata.
     ASSERT_OK(catalogClient()->insertConfigDocument(
         operationContext(), ShardType::ConfigNS, kShard0, kMajorityWriteConcern));
@@ -321,8 +320,7 @@ TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCannotBeBalanced)
         auto candidateChunksStatus = _chunkSelectionPolicy.get()->selectChunksToMove(opCtx.get());
         ASSERT_OK(candidateChunksStatus.getStatus());
 
-        // No chunks to move since the coll is a sharded time-series collection
-        ASSERT_EQUALS(0, candidateChunksStatus.getValue().size());
+        ASSERT_EQUALS(1, candidateChunksStatus.getValue().size());
     });
 
     expectGetStatsCommands(2);