SERVER-36102 Create initial chunks on appropriate shards for zoned sharding

1 files changed, 260 insertions, 51 deletions

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 0c0370a37e5..ee593935a2f 100644
--- a/src/mongo/db/s/config/initial_split_policy_test.cpp
+++ b/src/mongo/db/s/config/initial_split_policy_test.cpp
@@ -31,6 +31,9 @@
 #include "mongo/platform/basic.h"
 
 #include "mongo/db/s/config/initial_split_policy.h"
+#include "mongo/s/catalog/type_shard.h"
+#include "mongo/s/catalog/type_tags.h"
+#include "mongo/s/config_server_test_fixture.h"
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/log.h"
 
@@ -51,6 +54,20 @@ void assertBSONObjVectorsAreEqual(const std::vector<BSONObj>& expected,
 }
 
 /**
+ * Asserts that the given vectors of ChunkType objects are equal
+ */
+void assertChunkVectorsAreEqual(const std::vector<ChunkType>& expected,
+                                const std::vector<ChunkType>& actual) {
+    ASSERT_EQ(expected.size(), actual.size());
+    for (auto expectedIt = expected.begin(), actualIt = actual.begin();
+         expectedIt != expected.end() && actualIt != actual.end();
+         ++expectedIt, ++actualIt) {
+        ASSERT_BSONOBJ_EQ((*expectedIt).toShardBSON().removeField("lastmod"),
+                          (*actualIt).toShardBSON().removeField("lastmod"));
+    }
+}
+
+/**
  * Returns a test hashed shard key pattern if isHashed is true.
  * Otherwise, returns a regular shard key pattern.
  */
@@ -128,20 +145,37 @@ TEST(CalculateHashedSplitPointsTest, NotHashedWithInitialSplitsFails) {
 
 class GenerateInitialSplitChunksTest : public unittest::Test {
 public:
-    const std::vector<BSONObj>& hashedChunkBounds() {
-        return _hashedChunkBounds;
-    }
+    /**
+     * Returns a vector of ChunkType objects for the given chunk ranges.
+     * shardIds[i] is the id of shard for the chunk for chunkRanges[i].
+     * Checks that chunkRanges and shardIds have the same length.
+     */
+    const std::vector<ChunkType> makeChunks(const std::vector<ChunkRange> chunkRanges,
+                                            const std::vector<ShardId> shardIds) {
+        ASSERT_EQ(chunkRanges.size(), shardIds.size());
+        std::vector<ChunkType> chunks;
 
-    const std::vector<BSONObj>& hashedSplitPoints() {
-        return _splitPoints;
+        for (unsigned long i = 0; i < chunkRanges.size(); ++i) {
+            ChunkVersion version(1, 0, OID::gen());
+            ChunkType chunk(_nss, chunkRanges[i], version, shardIds[i]);
+            chunk.setHistory({ChunkHistory(_timeStamp, shardIds[i])});
+            chunks.push_back(chunk);
+        }
+        return chunks;
     }
 
-    ChunkType makeChunk(const BSONObj min, const BSONObj max, const ShardId shardId) {
-        ChunkVersion version(1, 0, OID::gen());
-        ChunkType chunk(_nss, ChunkRange(min, max), version, shardId);
-        chunk.setHistory({ChunkHistory(_timeStamp, shardId)});
-        return chunk;
+    /**
+     * Returns a vector of numShards shard ids with shard names
+     * prefixed by _shardName
+     */
+    const std::vector<ShardId> makeShardIds(const int numShards) {
+        std::vector<ShardId> shardIds;
+        for (int i = 0; i < numShards; i++) {
+            shardIds.push_back(shardId(std::to_string(i)));
+        }
+        return shardIds;
     }
+
     const NamespaceString nss() {
         return _nss;
     }
@@ -150,8 +184,12 @@ public:
         return _shardKeyPattern;
     }
 
-    const std::vector<ShardId> shardIds() {
-        return _shardIds;
+    const KeyPattern& keyPattern() {
+        return _shardKeyPattern.getKeyPattern();
+    }
+
+    const ShardId shardId(std::string shardNum) {
+        return ShardId(_shardName + shardNum);
     }
 
     const Timestamp timeStamp() {
@@ -161,62 +199,233 @@ public:
 private:
     const NamespaceString _nss{"test.foo"};
     const ShardKeyPattern _shardKeyPattern = makeShardKeyPattern(true);
-    const std::vector<ShardId> _shardIds = {ShardId("testShard0"), ShardId("testShard1")};
+    const std::string _shardName = "testShard";
     const Timestamp _timeStamp{Date_t::now()};
-    const KeyPattern& keyPattern = shardKeyPattern().getKeyPattern();
-    const std::vector<BSONObj> _hashedChunkBounds = {keyPattern.globalMin(),
-                                                     BSON("x" << -4611686018427387902LL),
-                                                     BSON("x" << 0),
-                                                     BSON("x" << 4611686018427387902LL),
-                                                     keyPattern.globalMax()};
-    const std::vector<BSONObj> _splitPoints{_hashedChunkBounds.begin() + 1,
-                                            _hashedChunkBounds.end() - 1};
 };
 
-TEST_F(GenerateInitialSplitChunksTest, NoSplitPoints) {
+class GenerateInitialHashedSplitChunksTest : public GenerateInitialSplitChunksTest {
+public:
+    const std::vector<BSONObj>& hashedSplitPoints() {
+        return _splitPoints;
+    }
+
+    const std::vector<ChunkRange>& hashedChunkRanges() {
+        return _chunkRanges;
+    }
+
+private:
+    const std::vector<BSONObj> _splitPoints{
+        BSON("x" << -4611686018427387902LL), BSON("x" << 0), BSON("x" << 4611686018427387902LL)};
+    const std::vector<ChunkRange> _chunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON("x" << -4611686018427387902LL)),
+        ChunkRange(BSON("x" << -4611686018427387902LL), BSON("x" << 0)),
+        ChunkRange(BSON("x" << 0), BSON("x" << 4611686018427387902LL)),
+        ChunkRange(BSON("x" << 4611686018427387902LL), keyPattern().globalMax()),
+    };
+};
+
+TEST_F(GenerateInitialHashedSplitChunksTest, NoSplitPoints) {
     const std::vector<BSONObj> splitPoints;
+    const std::vector<ShardId> shardIds = makeShardIds(2);
     const auto shardCollectionConfig = InitialSplitPolicy::generateShardCollectionInitialChunks(
-        nss(), shardKeyPattern(), shardIds()[0], timeStamp(), splitPoints, shardIds());
+        nss(), shardKeyPattern(), shardIds[0], timeStamp(), splitPoints, shardIds);
 
-    const auto& keyPattern = shardKeyPattern().getKeyPattern();
-    const auto expectedChunk =
-        makeChunk(keyPattern.globalMin(), keyPattern.globalMax(), shardIds()[0]);
-    ASSERT_EQ(1U, shardCollectionConfig.chunks.size());
-    ASSERT_BSONOBJ_EQ(expectedChunk.toShardBSON(), shardCollectionConfig.chunks[0].toShardBSON());
+    // there should only be one chunk
+    const auto expectedChunks = makeChunks(
+        {ChunkRange(keyPattern().globalMin(), keyPattern().globalMax())}, {shardId("0")});
+    assertChunkVectorsAreEqual(expectedChunks, shardCollectionConfig.chunks);
 }
 
-TEST_F(GenerateInitialSplitChunksTest, SplitPointsMoreThanAvailableShards) {
+TEST_F(GenerateInitialHashedSplitChunksTest, SplitPointsMoreThanAvailableShards) {
+    const std::vector<ShardId> shardIds = makeShardIds(2);
     const auto shardCollectionConfig = InitialSplitPolicy::generateShardCollectionInitialChunks(
-        nss(), shardKeyPattern(), shardIds()[0], timeStamp(), hashedSplitPoints(), shardIds());
+        nss(), shardKeyPattern(), shardIds[0], timeStamp(), hashedSplitPoints(), shardIds);
 
-    ASSERT_EQ(hashedSplitPoints().size() + 1, shardCollectionConfig.chunks.size());
+    // // chunks should be distributed in a round-robin manner
+    const std::vector<ChunkType> expectedChunks =
+        makeChunks(hashedChunkRanges(), {shardId("0"), shardId("1"), shardId("0"), shardId("1")});
+    assertChunkVectorsAreEqual(expectedChunks, shardCollectionConfig.chunks);
+}
 
-    // chunks should be distributed in a round-robin manner
-    const std::vector<ShardId> expectedShardIds = {
-        ShardId("testShard0"), ShardId("testShard1"), ShardId("testShard0"), ShardId("testShard1")};
-    for (unsigned long i = 0; i < hashedChunkBounds().size() - 1; ++i) {
-        const auto expectedChunk =
-            makeChunk(hashedChunkBounds()[i], hashedChunkBounds()[i + 1], expectedShardIds[i]);
-        ASSERT_BSONOBJ_EQ(expectedChunk.toShardBSON().removeField("lastmod"),
-                          shardCollectionConfig.chunks[i].toShardBSON().removeField("lastmod"));
+TEST_F(GenerateInitialHashedSplitChunksTest,
+       SplitPointsNumContiguousChunksPerShardsGreaterThanOne) {
+    const std::vector<ShardId> shardIds = makeShardIds(2);
+    const auto shardCollectionConfig = InitialSplitPolicy::generateShardCollectionInitialChunks(
+        nss(), shardKeyPattern(), shardIds[0], timeStamp(), hashedSplitPoints(), shardIds, 2);
+
+    // chunks should be distributed in a round-robin manner two chunks at a time
+    const std::vector<ChunkType> expectedChunks =
+        makeChunks(hashedChunkRanges(), {shardId("0"), shardId("0"), shardId("1"), shardId("1")});
+    assertChunkVectorsAreEqual(expectedChunks, shardCollectionConfig.chunks);
+}
+
+class GenerateShardCollectionInitialZonedChunksTest : public GenerateInitialSplitChunksTest {
+public:
+    /**
+     * Calls generateShardCollectionInitialZonedChunks according to the given arguments
+     * and asserts that returned chunks match with the chunks created using expectedChunkRanges
+     * and expectedShardIds.
+     */
+    void checkGeneratedInitialZoneChunks(const std::vector<TagsType>& tags,
+                                         const int numShards,
+                                         const std::vector<ChunkRange>& expectedChunkRanges,
+                                         const std::vector<ShardId>& expectedShardIds) {
+        const auto shardCollectionConfig =
+            InitialSplitPolicy::generateShardCollectionInitialZonedChunks(
+                nss(),
+                shardKeyPattern(),
+                timeStamp(),
+                tags,
+                makeTagToShards(numShards),
+                makeShardIds(numShards));
+        const std::vector<ChunkType> expectedChunks =
+            makeChunks(expectedChunkRanges, expectedShardIds);
+        assertChunkVectorsAreEqual(expectedChunks, shardCollectionConfig.chunks);
+    }
+
+    const std::string shardKey() {
+        return _shardKey;
+    }
+
+    const std::string zoneName(std::string zoneNum) {
+        return _zoneName + zoneNum;
+    }
+
+    TagsType makeTag(const ChunkRange range, std::string zoneName) {
+        BSONObjBuilder tagDocBuilder;
+        tagDocBuilder.append("_id",
+                             BSON(TagsType::ns(nss().ns()) << TagsType::min(range.getMin())));
+        tagDocBuilder.append(TagsType::ns(), nss().ns());
+        tagDocBuilder.append(TagsType::min(), range.getMin());
+        tagDocBuilder.append(TagsType::max(), range.getMax());
+        tagDocBuilder.append(TagsType::tag(), zoneName);
+        const auto parseStatus = TagsType::fromBSON(tagDocBuilder.obj());
+        uassertStatusOK(parseStatus);
+        return parseStatus.getValue();
     }
+
+    /**
+     * Returns a map of size numTags mapping _zoneName\d to _shardName\d
+     */
+    StringMap<std::vector<ShardId>> makeTagToShards(const int numTags) {
+        StringMap<std::vector<ShardId>> tagToShards;
+        for (int i = 0; i < numTags; i++) {
+            tagToShards[zoneName(std::to_string(i))] = {shardId(std::to_string(i))};
+        }
+        return tagToShards;
+    }
+
+private:
+    const ShardKeyPattern _shardKeyPattern = makeShardKeyPattern(true);
+    const std::string _zoneName = "zoneName";
+    const std::string _shardKey = "x";
+};
+
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, PredefinedZonesSpanFromMinToMax) {
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), keyPattern().globalMax()),  // corresponds to a zone
+    };
+    const std::vector<TagsType> tags = {makeTag(expectedChunkRanges[0], zoneName("0"))};
+    const std::vector<ShardId> expectedShardIds = {shardId("0")};
+    checkGeneratedInitialZoneChunks(tags, 1, expectedChunkRanges, expectedShardIds);
 }
 
-TEST_F(GenerateInitialSplitChunksTest, SplitPointsNumContiguousChunksPerShardsGreaterThanOne) {
-    const auto shardCollectionConfig = InitialSplitPolicy::generateShardCollectionInitialChunks(
-        nss(), shardKeyPattern(), shardIds()[0], timeStamp(), hashedSplitPoints(), shardIds(), 2);
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, PredefinedZonesSpanDoNotSpanFromMinToMax) {
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),
+        ChunkRange(BSON(shardKey() << 0), BSON(shardKey() << 10)),  // corresponds to a zone
+        ChunkRange(BSON(shardKey() << 10), keyPattern().globalMax()),
+    };
+    const std::vector<TagsType> tags = {makeTag(expectedChunkRanges[1], zoneName("0"))};
+    const std::vector<ShardId> expectedShardIds = {shardId("0"), shardId("0"), shardId("1")};
+    checkGeneratedInitialZoneChunks(tags, 2, expectedChunkRanges, expectedShardIds);
+}
 
-    ASSERT_EQ(hashedSplitPoints().size() + 1, shardCollectionConfig.chunks.size());
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, PredefinedZonesContainGlobalMin) {
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),  // corresponds to a zone
+        ChunkRange(BSON(shardKey() << 0), keyPattern().globalMax()),
+    };
+    const std::vector<TagsType> tags = {makeTag(expectedChunkRanges[0], zoneName("0"))};
+    const std::vector<ShardId> expectedShardIds = {shardId("0"), shardId("0")};
+    checkGeneratedInitialZoneChunks(tags, 2, expectedChunkRanges, expectedShardIds);
+}
 
-    // chunks should be distributed in a round-robin manner two chunks at a time
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, PredefinedZonesContainGlobalMax) {
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),
+        ChunkRange(BSON(shardKey() << 0), keyPattern().globalMax()),  // corresponds to a zone
+    };
+    const std::vector<TagsType> tags = {makeTag(expectedChunkRanges[1], zoneName("0"))};
+    const std::vector<ShardId> expectedShardIds = {shardId("0"), shardId("0")};
+    checkGeneratedInitialZoneChunks(tags, 2, expectedChunkRanges, expectedShardIds);
+}
+
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, PredefinedZonesContainGlobalMinAndMax) {
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),  // corresponds to a zone
+        ChunkRange(BSON(shardKey() << 0), BSON(shardKey() << 10)),
+        ChunkRange(BSON(shardKey() << 10), keyPattern().globalMax()),  // corresponds to a zone
+    };
+    const std::vector<TagsType> tags = {makeTag(expectedChunkRanges[0], zoneName("0")),
+                                        makeTag(expectedChunkRanges[2], zoneName("1"))};
+    const std::vector<ShardId> expectedShardIds = {shardId("0"), shardId("0"), shardId("1")};
+    checkGeneratedInitialZoneChunks(tags, 2, expectedChunkRanges, expectedShardIds);
+}
+
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, PredefinedZonesContiguous) {
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),
+        ChunkRange(BSON(shardKey() << 0), BSON(shardKey() << 10)),   // corresponds to a zone
+        ChunkRange(BSON(shardKey() << 10), BSON(shardKey() << 20)),  // corresponds to a zone
+        ChunkRange(BSON(shardKey() << 20), keyPattern().globalMax()),
+    };
+    const std::vector<TagsType> tags = {makeTag(expectedChunkRanges[1], zoneName("1")),
+                                        makeTag(expectedChunkRanges[2], zoneName("0"))};
     const std::vector<ShardId> expectedShardIds = {
-        ShardId("testShard0"), ShardId("testShard0"), ShardId("testShard1"), ShardId("testShard1")};
-    for (unsigned long i = 0; i < hashedChunkBounds().size() - 1; ++i) {
-        const auto expectedChunk =
-            makeChunk(hashedChunkBounds()[i], hashedChunkBounds()[i + 1], expectedShardIds[i]);
-        ASSERT_BSONOBJ_EQ(expectedChunk.toShardBSON().removeField("lastmod"),
-                          shardCollectionConfig.chunks[i].toShardBSON().removeField("lastmod"));
-    }
+        shardId("0"), shardId("1"), shardId("0"), shardId("1")};
+    checkGeneratedInitialZoneChunks(tags, 2, expectedChunkRanges, expectedShardIds);
+}
+
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, PredefinedZonesNotContiguous) {
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),
+        ChunkRange(BSON(shardKey() << 0), BSON(shardKey() << 10)),  // corresponds to a zone
+        ChunkRange(BSON(shardKey() << 10), BSON(shardKey() << 20)),
+        ChunkRange(BSON(shardKey() << 20), BSON(shardKey() << 30)),  // corresponds to a zone
+        ChunkRange(BSON(shardKey() << 30), keyPattern().globalMax()),
+    };
+    const std::vector<TagsType> tags = {makeTag(expectedChunkRanges[1], zoneName("0")),
+                                        makeTag(expectedChunkRanges[3], zoneName("1"))};
+    const std::vector<ShardId> expectedShardIds = {
+        shardId("0"), shardId("0"), shardId("1"), shardId("1"), shardId("2")};
+    checkGeneratedInitialZoneChunks(tags, 3, expectedChunkRanges, expectedShardIds);
+}
+
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, NumRemainingChunksGreaterThanNumShards) {
+    const std::vector<ChunkRange> expectedChunkRanges = {
+        ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)),
+        ChunkRange(BSON(shardKey() << 0), BSON(shardKey() << 10)),  // corresponds to a zone
+        ChunkRange(BSON(shardKey() << 10), BSON(shardKey() << 20)),
+        ChunkRange(BSON(shardKey() << 20), BSON(shardKey() << 30)),
+        ChunkRange(BSON(shardKey() << 30), keyPattern().globalMax()),
+    };
+    const std::vector<TagsType> tags = {makeTag(expectedChunkRanges[1], zoneName("0")),
+                                        makeTag(expectedChunkRanges[3], zoneName("1"))};
+    // shard assignment should wrap around to the first shard
+    const std::vector<ShardId> expectedShardIds = {
+        shardId("0"), shardId("0"), shardId("1"), shardId("1"), shardId("0")};
+    checkGeneratedInitialZoneChunks(tags, 2, expectedChunkRanges, expectedShardIds);
+}
+
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, EmptyTagsShouldFail) {
+    const std::vector<ChunkRange> expectedChunkRanges;
+    const std::vector<TagsType> tags;
+    const std::vector<ShardId> expectedShardIds;
+    ASSERT_THROWS_CODE(
+        checkGeneratedInitialZoneChunks(tags, 1, expectedChunkRanges, expectedShardIds),
+        AssertionException,
+        ErrorCodes::InvalidOptions);
 }
 
 }  // namespace