SERVER-44341 Round-robin policy for shardCollection pre-splitting on zones

Bug:
 when pre-splitting a collection during shardCollection using existing zones,
 we choose the first shard associated with the zone to place a chunk on.
 This can cause a problem if there are many shards associated with the same zone,
 the balancer will still schedule migrations to balance the zones afterward.

Implemented solution:
 For each new chunk, the shard will be selected in round-robin fashion among the ones associated with its zone.
 The new `tagToIndx` StringMap is used to keep the incrementing couters
 for each zone (tag).

2 files changed, 69 insertions, 10 deletions

diff --git a/src/mongo/db/s/config/initial_split_policy.cpp b/src/mongo/db/s/config/initial_split_policy.cpp
index ca164a0264f..680fd4e54e4 100644
--- a/src/mongo/db/s/config/initial_split_policy.cpp
+++ b/src/mongo/db/s/config/initial_split_policy.cpp
@@ -226,6 +226,19 @@ InitialSplitPolicy::generateShardCollectionInitialZonedChunks(
         return shardIdsForGaps[indx++ % shardIdsForGaps.size()];
     };
 
+    auto nextShardIdForTag = [&, tagToIndx = StringMap<size_t>()](const auto& tag) mutable {
+        const auto it = tagToShards.find(tag.getTag());
+        invariant(it != tagToShards.end());
+        const auto& shardIdsForTag = it->second;
+        uassert(50973,
+                str::stream()
+                    << "Cannot shard collection " << nss.ns() << " due to zone " << tag.getTag()
+                    << " which is not assigned to a shard. Please assign this zone to a shard.",
+                !shardIdsForTag.empty());
+        const auto nextShardIndx = tagToIndx[tag.getTag()]++;
+        return shardIdsForTag[nextShardIndx % shardIdsForTag.size()];
+    };
+
     std::vector<ChunkType> chunks;
 
     ChunkVersion version(1, 0, OID::gen());
@@ -244,21 +257,12 @@ InitialSplitPolicy::generateShardCollectionInitialZonedChunks(
         }
 
         // Create chunk for the actual tag - [tag.getMinKey, tag.getMaxKey)
-        const auto it = tagToShards.find(tag.getTag());
-        invariant(it != tagToShards.end());
-        const auto& shardIdsForChunk = it->second;
-        uassert(50973,
-                str::stream()
-                    << "Cannot shard collection " << nss.ns() << " due to zone " << tag.getTag()
-                    << " which is not assigned to a shard. Please assign this zone to a shard.",
-                !shardIdsForChunk.empty());
-
         appendChunk(nss,
                     tag.getMinKey(),
                     tag.getMaxKey(),
                     &version,
                     validAfter,
-                    shardIdsForChunk[0],
+                    nextShardIdForTag(tag),
                     &chunks);
         lastChunkMax = tag.getMaxKey();
     }
diff --git a/src/mongo/db/s/config/initial_split_policy_test.cpp b/src/mongo/db/s/config/initial_split_policy_test.cpp
index 45447ecd729..68fdb44abd8 100644
--- a/src/mongo/db/s/config/initial_split_policy_test.cpp
+++ b/src/mongo/db/s/config/initial_split_policy_test.cpp
@@ -418,6 +418,61 @@ TEST_F(GenerateShardCollectionInitialZonedChunksTest, NumRemainingChunksGreaterT
     checkGeneratedInitialZoneChunks(tags, 2, expectedChunkRanges, expectedShardIds);
 }
 
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, MultipleChunksToOneZoneWithMultipleShards) {
+    const auto zone0 = zoneName("Z0");
+
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),   // zone0
+        ChunkRange(BSON(shardKey() << 0), BSON(shardKey() << 10)),     // gap
+        ChunkRange(BSON(shardKey() << 10), BSON(shardKey() << 20)),    // zone0
+        ChunkRange(BSON(shardKey() << 20), keyPattern().globalMax()),  // zone0
+    };
+    const std::vector<TagsType> tags = {
+        makeTag(expectedChunkRanges.at(0), zone0),
+        makeTag(expectedChunkRanges.at(2), zone0),
+        makeTag(expectedChunkRanges.at(3), zone0),
+    };
+    const StringMap<std::vector<ShardId>> tagToShards = {{zone0, {{"S0"}, {"S1"}}}};
+    const std::vector<ShardId> shardIdsForGaps = {{"G0"}};
+    const std::vector<ShardId> expectedShardIds = {{"S0"}, {"G0"}, {"S1"}, {"S0"}};
+    const auto shardCollectionConfig =
+        InitialSplitPolicy::generateShardCollectionInitialZonedChunks(
+            nss(), shardKeyPattern(), timeStamp(), tags, tagToShards, shardIdsForGaps);
+
+    const std::vector<ChunkType> expectedChunks = makeChunks(expectedChunkRanges, expectedShardIds);
+    assertChunkVectorsAreEqual(expectedChunks, shardCollectionConfig.chunks);
+};
+
+TEST_F(GenerateShardCollectionInitialZonedChunksTest,
+       MultipleChunksToInterleavedZonesWithMultipleShards) {
+    const auto zone0 = zoneName("Z0");
+    const auto zone1 = zoneName("Z1");
+
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),   // zone0
+        ChunkRange(BSON(shardKey() << 0), BSON(shardKey() << 10)),     // zone1
+        ChunkRange(BSON(shardKey() << 10), BSON(shardKey() << 20)),    // gap
+        ChunkRange(BSON(shardKey() << 20), keyPattern().globalMax()),  // zone0
+    };
+    const std::vector<TagsType> tags = {
+        makeTag(expectedChunkRanges.at(0), zone0),
+        makeTag(expectedChunkRanges.at(1), zone1),
+        makeTag(expectedChunkRanges.at(3), zone0),
+    };
+    const StringMap<std::vector<ShardId>> tagToShards = {
+        {zone0, {{"Z0-S0"}, {"Z0-S1"}}},
+        {zone1, {{"Z1-S0"}, {"Z1-S1"}}},
+    };
+    const std::vector<ShardId> shardIdsForGaps = {{"G0"}};
+    const std::vector<ShardId> expectedShardIds = {{"Z0-S0"}, {"Z1-S0"}, {"G0"}, {"Z0-S1"}};
+    const auto shardCollectionConfig =
+        InitialSplitPolicy::generateShardCollectionInitialZonedChunks(
+            nss(), shardKeyPattern(), timeStamp(), tags, tagToShards, shardIdsForGaps);
+
+    const std::vector<ChunkType> expectedChunks = makeChunks(expectedChunkRanges, expectedShardIds);
+    assertChunkVectorsAreEqual(expectedChunks, shardCollectionConfig.chunks);
+};
+
 TEST_F(GenerateShardCollectionInitialZonedChunksTest, ZoneNotAssociatedWithAnyShardShouldFail) {
     const auto zone1 = zoneName("0");
     const auto zone2 = zoneName("1");