Update upstream source from tag 'upstream/6.0.0'master

Update to upstream version '6.0.0'
with Debian dir 5604a80ec1c96ca76f25f40d78e6ef855abec322

1 files changed, 381 insertions, 1215 deletions

diff --git a/src/mongo/db/s/balancer/balancer_policy_test.cpp b/src/mongo/db/s/balancer/balancer_policy_test.cpp
index 37b6925a082..0a53200149e 100644
--- a/src/mongo/db/s/balancer/balancer_policy_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy_test.cpp
@@ -35,7 +35,6 @@
 #include "mongo/db/s/balancer/balancer_policy.h"
 #include "mongo/platform/random.h"
 #include "mongo/s/catalog/type_chunk.h"
-#include "mongo/s/chunks_test_util.h"
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
@@ -47,12 +46,9 @@ using std::stringstream;
 using std::vector;
 
 using ShardStatistics = ClusterStatistics::ShardStatistics;
-typedef std::map<ShardId, std::vector<ChunkType>> ShardToChunksMap;
-
-PseudoRandom _random{SecureRandom().nextInt64()};
 
 const auto emptyTagSet = std::set<std::string>();
-const auto kConfigId = ShardId("config");
+const std::string emptyShardVersion = "";
 const auto kShardId0 = ShardId("shard0");
 const auto kShardId1 = ShardId("shard1");
 const auto kShardId2 = ShardId("shard2");
@@ -61,51 +57,6 @@ const auto kShardId4 = ShardId("shard4");
 const auto kShardId5 = ShardId("shard5");
 const NamespaceString kNamespace("TestDB", "TestColl");
 const uint64_t kNoMaxSize = 0;
-const KeyPattern kSKeyPattern(BSON("x" << 1));
-const Timestamp kCollTimestamp{1, 1};
-const OID kCollEpoch;
-
-const UUID& collUUID() {
-    static const UUID kCollectionUUID{UUID::gen()};
-    return kCollectionUUID;
-}
-
-std::vector<ChunkType> makeChunks(const std::vector<std::pair<ShardId, ChunkRange>>& specs) {
-    ChunkVersion chunkVersion{1, 0, kCollEpoch, kCollTimestamp};
-    std::vector<ChunkType> chunks;
-    for (const auto& [shardId, range] : specs) {
-        chunks.push_back(ChunkType(collUUID(), range, chunkVersion, shardId));
-        chunkVersion.incMajor();
-    }
-    return chunks;
-}
-
-RoutingTableHistory makeRoutingTable(const std::vector<ChunkType>& chunks) {
-
-    return RoutingTableHistory::makeNew(kNamespace,
-                                        collUUID(),
-                                        kSKeyPattern,
-                                        nullptr,
-                                        false /* unique */,
-                                        kCollEpoch,
-                                        kCollTimestamp,
-                                        boost::none /* timeseriesFields */,
-                                        boost::none /* reshardingFields */,
-                                        boost::none /* maxChunkSizeBytes */,
-                                        true /* allowMigrations */,
-                                        chunks);
-}
-
-ChunkManager makeChunkManager(const std::vector<ChunkType>& chunks) {
-    DatabaseVersion dbVersion;
-    auto rt = std::make_shared<RoutingTableHistory>(makeRoutingTable(chunks));
-
-    return {kConfigId, std::move(dbVersion), {std::move(rt)}, boost::none /* atClusterTime */};
-}
-
-DistributionStatus makeDistStatus(const ChunkManager& cm, ZoneInfo zoneInfo = ZoneInfo()) {
-    return {kNamespace, std::move(zoneInfo), cm};
-}
 
 /**
  * Constructs a shard statistics vector and a consistent mapping of chunks to shards given the
@@ -114,7 +65,7 @@ DistributionStatus makeDistStatus(const ChunkManager& cm, ZoneInfo zoneInfo = Zo
  *
  * [MinKey, 1), [1, 2), [2, 3) ... [N - 1, MaxKey)
  */
-std::pair<std::pair<ShardStatisticsVector, ShardToChunksMap>, ChunkManager> generateCluster(
+std::pair<ShardStatisticsVector, ShardToChunksMap> generateCluster(
     const vector<std::pair<ShardStatistics, size_t>>& shardsAndNumChunks) {
     int64_t totalNumChunks = 0;
     for (const auto& entry : shardsAndNumChunks) {
@@ -126,9 +77,10 @@ std::pair<std::pair<ShardStatisticsVector, ShardToChunksMap>, ChunkManager> gene
 
     int64_t currentChunk = 0;
 
-    ChunkVersion chunkVersion(1, 0, kCollEpoch, kCollTimestamp);
+    ChunkVersion chunkVersion(1, 0, OID::gen(), Timestamp(1, 1));
+    const UUID uuid = UUID::gen();
 
-    std::vector<ChunkType> chunks;
+    const KeyPattern shardKeyPattern(BSON("x" << 1));
 
     for (auto it = shardsAndNumChunks.begin(); it != shardsAndNumChunks.end(); it++) {
         ShardStatistics shard = std::move(it->first);
@@ -140,586 +92,406 @@ std::pair<std::pair<ShardStatisticsVector, ShardToChunksMap>, ChunkManager> gene
         for (size_t i = 0; i < numChunks; i++, currentChunk++) {
             ChunkType chunk;
 
-            chunk.setCollectionUUID(collUUID());
-            chunk.setMin(currentChunk == 0 ? kSKeyPattern.globalMin() : BSON("x" << currentChunk));
-            chunk.setMax(currentChunk == totalNumChunks - 1 ? kSKeyPattern.globalMax()
+            chunk.setCollectionUUID(uuid);
+            chunk.setMin(currentChunk == 0 ? shardKeyPattern.globalMin()
+                                           : BSON("x" << currentChunk));
+            chunk.setMax(currentChunk == totalNumChunks - 1 ? shardKeyPattern.globalMax()
                                                             : BSON("x" << currentChunk + 1));
             chunk.setShard(shard.shardId);
             chunk.setVersion(chunkVersion);
 
             chunkVersion.incMajor();
 
-            chunkMap[shard.shardId].push_back(chunk);
-            chunks.push_back(std::move(chunk));
+            chunkMap[shard.shardId].push_back(std::move(chunk));
         }
 
         shardStats.push_back(std::move(shard));
     }
 
-    return std::make_pair(std::make_pair(std::move(shardStats), std::move(chunkMap)),
-                          makeChunkManager(chunks));
-}
-
-stdx::unordered_set<ShardId> getAllShardIds(const ShardStatisticsVector& shardStats) {
-    stdx::unordered_set<ShardId> shards;
-    std::transform(shardStats.begin(),
-                   shardStats.end(),
-                   std::inserter(shards, shards.end()),
-                   [](const ShardStatistics& shardStatistics) { return shardStatistics.shardId; });
-    return shards;
+    return std::make_pair(std::move(shardStats), std::move(chunkMap));
 }
 
 MigrateInfosWithReason balanceChunks(const ShardStatisticsVector& shardStats,
                                      const DistributionStatus& distribution,
                                      bool shouldAggressivelyBalance,
                                      bool forceJumbo) {
-    auto availableShards = getAllShardIds(shardStats);
-    return BalancerPolicy::balance(
-        shardStats, distribution, boost::none /* collDataSizeInfo */, &availableShards, forceJumbo);
-}
-
-void checkChunksOnShardForTag(const DistributionStatus& dist,
-                              const ShardId& shardId,
-                              const std::string& zoneName,
-                              const std::vector<ChunkType>& expectedChunks) {
-    auto expectedChunkIt = expectedChunks.cbegin();
-    const auto completed =
-        dist.forEachChunkOnShardInZone(shardId, zoneName, [&](const auto& chunk) {
-            ASSERT(expectedChunkIt != expectedChunks.end())
-                << "forEachChunkOnShardInZone loop found more chunks than expected";
-            ChunkInfo expectedChunkInfo{*expectedChunkIt++};
-            ASSERT_EQ(Chunk(expectedChunkInfo, boost::none /* atClusterTime */).toString(),
-                      chunk.toString());
-            return true;  // continue
-        });
-    ASSERT(completed)
-        << "forEachChunkOnShardInZone loop unexpectedly returned false (did not complete)";
-    ASSERT(expectedChunkIt == expectedChunks.cend())
-        << "forEachChunkOnShardInZone loop did not iterate over all the expected chunks";
+    stdx::unordered_set<ShardId> usedShards;
+    return BalancerPolicy::balance(shardStats, distribution, &usedShards, forceJumbo);
 }
 
 TEST(BalancerPolicy, Basic) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 0},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 3, false, emptyTagSet), 3}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId2, kNoMaxSize, 3, false, emptyTagSet, emptyShardVersion), 3}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 }
 
 TEST(BalancerPolicy, SmallClusterShouldBePerfectlyBalanced) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 1, false, emptyTagSet), 1},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, emptyTagSet), 2},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet), 0}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 1, false, emptyTagSet, emptyShardVersion), 1},
+         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 2},
+         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId1, migrations[0].from);
     ASSERT_EQ(kShardId2, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 }
 
 TEST(BalancerPolicy, SingleChunkShouldNotMove) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 1, false, emptyTagSet), 1},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 0}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 1, false, emptyTagSet, emptyShardVersion), 1},
+         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
     {
-        auto [migrations, reason] = balanceChunks(cluster.first, makeDistStatus(cm), true, false);
+        auto [migrations, reason] = balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second), true, false);
         ASSERT(migrations.empty());
         ASSERT_EQ(MigrationReason::none, reason);
     }
     {
-        auto [migrations, reason] = balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        auto [migrations, reason] = balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
         ASSERT(migrations.empty());
         ASSERT_EQ(MigrationReason::none, reason);
     }
 }
 
 TEST(BalancerPolicy, BalanceThresholdObeyed) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet), 2},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, emptyTagSet), 2},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 1, false, emptyTagSet), 1},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 1, false, emptyTagSet), 1}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 2},
+         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 2},
+         {ShardStatistics(kShardId2, kNoMaxSize, 1, false, emptyTagSet, emptyShardVersion), 1},
+         {ShardStatistics(kShardId3, kNoMaxSize, 1, false, emptyTagSet, emptyShardVersion), 1}});
 
     {
-        auto [migrations, reason] = balanceChunks(cluster.first, makeDistStatus(cm), true, false);
+        auto [migrations, reason] = balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second), true, false);
         ASSERT(migrations.empty());
         ASSERT_EQ(MigrationReason::none, reason);
     }
     {
-        auto [migrations, reason] = balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        auto [migrations, reason] = balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
         ASSERT(migrations.empty());
         ASSERT_EQ(MigrationReason::none, reason);
     }
 }
 
 TEST(BalancerPolicy, ParallelBalancing) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet), 0},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet), 0}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId2, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 
     ASSERT_EQ(kShardId1, migrations[1].from);
     ASSERT_EQ(kShardId3, migrations[1].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMin(), migrations[1].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), *migrations[1].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), migrations[1].maxKey);
 }
 
 TEST(BalancerPolicy, ParallelBalancingDoesNotPutChunksOnShardsAboveTheOptimal) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 100, false, emptyTagSet), 100},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 90, false, emptyTagSet), 90},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 90, false, emptyTagSet), 90},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 80, false, emptyTagSet), 80},
-                         {ShardStatistics(kShardId4, kNoMaxSize, 0, false, emptyTagSet), 0},
-                         {ShardStatistics(kShardId5, kNoMaxSize, 0, false, emptyTagSet), 0}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 100, false, emptyTagSet, emptyShardVersion), 100},
+         {ShardStatistics(kShardId1, kNoMaxSize, 90, false, emptyTagSet, emptyShardVersion), 90},
+         {ShardStatistics(kShardId2, kNoMaxSize, 90, false, emptyTagSet, emptyShardVersion), 90},
+         {ShardStatistics(kShardId3, kNoMaxSize, 80, false, emptyTagSet, emptyShardVersion), 80},
+         {ShardStatistics(kShardId4, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId5, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId4, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 
     ASSERT_EQ(kShardId1, migrations[1].from);
     ASSERT_EQ(kShardId5, migrations[1].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMin(), migrations[1].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), *migrations[1].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), migrations[1].maxKey);
 }
 
 TEST(BalancerPolicy, ParallelBalancingDoesNotMoveChunksFromShardsBelowOptimal) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 100, false, emptyTagSet), 100},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 30, false, emptyTagSet), 30},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet), 5},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet), 0}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 100, false, emptyTagSet, emptyShardVersion), 100},
+         {ShardStatistics(kShardId1, kNoMaxSize, 30, false, emptyTagSet, emptyShardVersion), 30},
+         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 5},
+         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(1U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId3, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 }
 
 TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseSourceShardsWithMoveNecessary) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 8, false, emptyTagSet), 8},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet), 0},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet), 0}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 8, false, emptyTagSet, emptyShardVersion), 8},
+         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
     // Here kShardId0 would have been selected as a donor
-    auto availableShards = getAllShardIds(cluster.first);
-    availableShards.erase(kShardId0);
-    const auto [migrations, reason] = BalancerPolicy::balance(cluster.first,
-                                                              makeDistStatus(cm),
-                                                              boost::none /* collDataSizeInfo */,
-                                                              &availableShards,
-                                                              false);
+    stdx::unordered_set<ShardId> usedShards{kShardId0};
+    const auto [migrations, reason] = BalancerPolicy::balance(
+        cluster.first, DistributionStatus(kNamespace, cluster.second), &usedShards, false);
     ASSERT_EQ(1U, migrations.size());
 
     ASSERT_EQ(kShardId1, migrations[0].from);
     ASSERT_EQ(kShardId2, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 }
 
 TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseSourceShardsWithMoveNotNecessary) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 12, false, emptyTagSet), 12},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet), 0},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet), 0}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 12, false, emptyTagSet, emptyShardVersion), 12},
+         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
     // Here kShardId0 would have been selected as a donor
-    auto availableShards = getAllShardIds(cluster.first);
-    availableShards.erase(kShardId0);
-    const auto [migrations, reason] = BalancerPolicy::balance(cluster.first,
-                                                              makeDistStatus(cm),
-                                                              boost::none /* collDataSizeInfo */,
-                                                              &availableShards,
-                                                              false);
+    stdx::unordered_set<ShardId> usedShards{kShardId0};
+    const auto [migrations, reason] = BalancerPolicy::balance(
+        cluster.first, DistributionStatus(kNamespace, cluster.second), &usedShards, false);
     ASSERT_EQ(0U, migrations.size());
 }
 
 TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseDestinationShards) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet), 0},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 1, false, emptyTagSet), 1}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId3, kNoMaxSize, 1, false, emptyTagSet, emptyShardVersion), 1}});
 
     // Here kShardId2 would have been selected as a recipient
-    auto availableShards = getAllShardIds(cluster.first);
-    availableShards.erase(kShardId2);
-    const auto [migrations, reason] = BalancerPolicy::balance(cluster.first,
-                                                              makeDistStatus(cm),
-                                                              boost::none /* collDataSizeInfo */,
-                                                              &availableShards,
-                                                              false);
+    stdx::unordered_set<ShardId> usedShards{kShardId2};
+    const auto [migrations, reason] = BalancerPolicy::balance(
+        cluster.first, DistributionStatus(kNamespace, cluster.second), &usedShards, false);
     ASSERT_EQ(1U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId3, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 }
 
-TEST(BalancerPolicy, JumboChunksNotMovedWhileEnforcingZones) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 3, false, emptyTagSet), 3},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 3, false, {"a"}), 3}});
-
-    // construct a new chunk map where all the chunks are jumbo except this one
-    const auto& jumboChunk = cluster.second[kShardId0][1];
-
-    std::vector<ChunkType> chunks;
-    cm.forEachChunk([&](const auto& chunk) {
-        ChunkType ct{collUUID(), chunk.getRange(), chunk.getLastmod(), chunk.getShardId()};
-        if (chunk.getLastmod() == jumboChunk.getVersion())
-            ct.setJumbo(false);
-        else
-            ct.setJumbo(true);
-        chunks.emplace_back(std::move(ct));
-        return true;
-    });
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), kSKeyPattern.globalMax(), "a")));
-    const auto distribution = makeDistStatus(makeChunkManager(chunks), std::move(zoneInfo));
-
-    const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
-    ASSERT_EQ(1U, migrations.size());
-    ASSERT_EQ(kShardId0, migrations[0].from);
-    ASSERT_EQ(kShardId1, migrations[0].to);
-    ASSERT_BSONOBJ_EQ(jumboChunk.getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(jumboChunk.getMax(), *migrations[0].maxKey);
-    ASSERT_EQ(MigrationReason::zoneViolation, reason);
-}
-
-TEST(BalancerPolicy, JumboChunksNotMovedWhileEnforcingZonesRandom) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 3, false, emptyTagSet), 3},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 3, false, {"a"}), 3}});
-
-    // construct a new chunk map where all the chunks are jumbo except this one
-    const auto jumboChunkIdx = _random.nextInt64(cluster.second[kShardId0].size());
-    const auto& jumboChunk = cluster.second[kShardId0][jumboChunkIdx];
-
-    std::vector<ChunkType> chunks;
-    cm.forEachChunk([&](const auto& chunk) {
-        ChunkType ct{collUUID(), chunk.getRange(), chunk.getLastmod(), chunk.getShardId()};
-        if (chunk.getLastmod() == jumboChunk.getVersion())
-            ct.setJumbo(false);
-        else
-            ct.setJumbo(true);
-        chunks.emplace_back(std::move(ct));
-        return true;
-    });
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), kSKeyPattern.globalMax(), "a")));
-    const auto distribution = makeDistStatus(makeChunkManager(chunks), std::move(zoneInfo));
-
-    const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
-    ASSERT_EQ(1U, migrations.size());
-    ASSERT_EQ(kShardId0, migrations[0].from);
-    ASSERT_EQ(kShardId1, migrations[0].to);
-    ASSERT_BSONOBJ_EQ(jumboChunk.getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(jumboChunk.getMax(), *migrations[0].maxKey);
-    ASSERT_EQ(MigrationReason::zoneViolation, reason);
-}
-
 TEST(BalancerPolicy, JumboChunksNotMoved) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 0}});
-
-    // construct a new chunk map where all the chunks are jumbo except this one
-    const auto& jumboChunk = cluster.second[kShardId0][1];
-
-    std::vector<ChunkType> chunks;
-    cm.forEachChunk([&](const auto& chunk) {
-        ChunkType ct{collUUID(), chunk.getRange(), chunk.getLastmod(), chunk.getShardId()};
-        if (chunk.getLastmod() == jumboChunk.getVersion())
-            ct.setJumbo(false);
-        else
-            ct.setJumbo(true);
-        chunks.emplace_back(std::move(ct));
-        return true;
-    });
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(makeChunkManager(chunks)), false, false);
-    ASSERT_EQ(1U, migrations.size());
-    ASSERT_EQ(kShardId0, migrations[0].from);
-    ASSERT_EQ(kShardId1, migrations[0].to);
-    ASSERT_BSONOBJ_EQ(jumboChunk.getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(jumboChunk.getMax(), *migrations[0].maxKey);
-    ASSERT_EQ(MigrationReason::chunksImbalance, reason);
-}
-
-TEST(BalancerPolicy, JumboChunksNotMovedRandom) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 0}});
-
-    // construct a new chunk map where all the chunks are jumbo except this one
-    const auto jumboChunkIdx = _random.nextInt64(cluster.second[kShardId0].size());
-    const auto& jumboChunk = cluster.second[kShardId0][jumboChunkIdx];
-
-    std::vector<ChunkType> chunks;
-    cm.forEachChunk([&](const auto& chunk) {
-        ChunkType ct{collUUID(), chunk.getRange(), chunk.getLastmod(), chunk.getShardId()};
-        if (chunk.getLastmod() == jumboChunk.getVersion())
-            ct.setJumbo(false);
-        else
-            ct.setJumbo(true);
-        chunks.emplace_back(std::move(ct));
-        return true;
-    });
+    cluster.second[kShardId0][0].setJumbo(true);
+    cluster.second[kShardId0][1].setJumbo(false);  // Only chunk 1 is not jumbo
+    cluster.second[kShardId0][2].setJumbo(true);
+    cluster.second[kShardId0][3].setJumbo(true);
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(makeChunkManager(chunks)), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
-    ASSERT_BSONOBJ_EQ(jumboChunk.getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(jumboChunk.getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][1].getMin(), migrations[0].minKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][1].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 }
 
 TEST(BalancerPolicy, JumboChunksNotMovedParallel) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 0},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 2, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet), 0}});
-
-    // construct a new chunk map where all the chunks are jumbo except the ones listed below
-    const auto& jumboChunk0 = cluster.second[kShardId0][1];
-    const auto& jumboChunk1 = cluster.second[kShardId2][2];
-
-    std::vector<ChunkType> chunks;
-    cm.forEachChunk([&](const auto& chunk) {
-        ChunkType ct{collUUID(), chunk.getRange(), chunk.getLastmod(), chunk.getShardId()};
-        if (chunk.getLastmod() == jumboChunk0.getVersion() ||
-            chunk.getLastmod() == jumboChunk1.getVersion())
-            ct.setJumbo(false);
-        else
-            ct.setJumbo(true);
-        chunks.emplace_back(std::move(ct));
-        return true;
-    });
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId2, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
+
+    cluster.second[kShardId0][0].setJumbo(true);
+    cluster.second[kShardId0][1].setJumbo(false);  // Only chunk 1 is not jumbo
+    cluster.second[kShardId0][2].setJumbo(true);
+    cluster.second[kShardId0][3].setJumbo(true);
+
+    cluster.second[kShardId2][0].setJumbo(true);
+    cluster.second[kShardId2][1].setJumbo(true);
+    cluster.second[kShardId2][2].setJumbo(false);  // Only chunk 1 is not jumbo
+    cluster.second[kShardId2][3].setJumbo(true);
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(makeChunkManager(chunks)), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
-    ASSERT_BSONOBJ_EQ(jumboChunk0.getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(jumboChunk0.getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][1].getMin(), migrations[0].minKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][1].getMax(), migrations[0].maxKey);
+    ASSERT_EQ(MigrationReason::chunksImbalance, reason);
 
     ASSERT_EQ(kShardId2, migrations[1].from);
     ASSERT_EQ(kShardId3, migrations[1].to);
-    ASSERT_BSONOBJ_EQ(jumboChunk1.getMin(), migrations[1].minKey);
-    ASSERT_BSONOBJ_EQ(jumboChunk1.getMax(), *migrations[1].maxKey);
-
-    ASSERT_EQ(MigrationReason::chunksImbalance, reason);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][2].getMin(), migrations[1].minKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][2].getMax(), migrations[1].maxKey);
 }
 
 TEST(BalancerPolicy, DrainingSingleChunk) {
     // shard0 is draining and chunks will go to shard1, even though it has a lot more chunks
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet), 1},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 5}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet, emptyShardVersion), 1},
+         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 5}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::drain, reason);
 }
 
 TEST(BalancerPolicy, DrainingSingleChunkPerShard) {
     // shard0 and shard2 are draining and chunks will go to shard1 and shard3 in parallel
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet), 1},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 5},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 2, true, emptyTagSet), 1},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet), 5}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet, emptyShardVersion), 1},
+         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 5},
+         {ShardStatistics(kShardId2, kNoMaxSize, 2, true, emptyTagSet, emptyShardVersion), 1},
+         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 5}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::drain, reason);
 
     ASSERT_EQ(kShardId2, migrations[1].from);
     ASSERT_EQ(kShardId3, migrations[1].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMin(), migrations[1].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), *migrations[1].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), migrations[1].maxKey);
 }
 
 TEST(BalancerPolicy, DrainingWithTwoChunksFirstOneSelected) {
     // shard0 is draining and chunks will go to shard1, even though it has a lot more chunks
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet), 2},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 5}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet, emptyShardVersion), 2},
+         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 5}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::drain, reason);
 }
 
 TEST(BalancerPolicy, DrainingMultipleShardsFirstOneSelected) {
     // shard0 and shard1 are both draining with very little chunks in them and chunks will go to
     // shard2, even though it has a lot more chunks that the other two
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, true, emptyTagSet), 1},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, true, emptyTagSet), 2},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet), 16}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, true, emptyTagSet, emptyShardVersion), 1},
+         {ShardStatistics(kShardId1, kNoMaxSize, 5, true, emptyTagSet, emptyShardVersion), 2},
+         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 16}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
 
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId2, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::drain, reason);
 }
 
 TEST(BalancerPolicy, DrainingMultipleShardsWontAcceptChunks) {
     // shard0 has many chunks, but can't move them to shard1 or shard2 because they are draining
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, true, emptyTagSet), 0},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 0, true, emptyTagSet), 0}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 0, true, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId2, kNoMaxSize, 0, true, emptyTagSet, emptyShardVersion), 0}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT(migrations.empty());
 }
 
 TEST(BalancerPolicy, DrainingSingleAppropriateShardFoundDueToTag) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, false, {"NYC"}), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"LAX"}), 4},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 1, true, {"LAX"}), 1}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, false, {"NYC"}, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"LAX"}, emptyShardVersion), 4},
+         {ShardStatistics(kShardId2, kNoMaxSize, 1, true, {"LAX"}, emptyShardVersion), 1}});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(
         cluster.second[kShardId2][0].getMin(), cluster.second[kShardId2][0].getMax(), "LAX")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId2, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::drain, reason);
 }
 
 TEST(BalancerPolicy, DrainingNoAppropriateShardsFoundDueToTag) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, false, {"NYC"}), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"LAX"}), 4},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 1, true, {"SEA"}), 1}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, false, {"NYC"}, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"LAX"}, emptyShardVersion), 4},
+         {ShardStatistics(kShardId2, kNoMaxSize, 1, true, {"SEA"}, emptyShardVersion), 1}});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(
         cluster.second[kShardId2][0].getMin(), cluster.second[kShardId2][0].getMax(), "SEA")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT(migrations.empty());
 }
 
-TEST(BalancerPolicy, DrainingSingleAppropriateShardFoundMultipleTags) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, true, {}), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"Zone1", "Zone2"}), 2},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 2, false, {"Zone1", "Zone2"}), 2},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 2, false, {}), 2}});
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
-        cluster.second[kShardId0][0].getMin(), cluster.second[kShardId0][0].getMax(), "Zone1")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
-        cluster.second[kShardId0][1].getMin(), cluster.second[kShardId0][1].getMax(), "Zone2")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
-        cluster.second[kShardId0][2].getMin(), cluster.second[kShardId0][2].getMax(), "Zone1")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
-        cluster.second[kShardId0][3].getMin(), cluster.second[kShardId0][3].getMax(), "Zone2")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
-
-    const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
-    ASSERT_EQ(1U, migrations.size());
-    ASSERT_EQ(kShardId0, migrations[0].from);
-    const auto& recipientShard = migrations[0].to;
-    ASSERT(recipientShard == kShardId1 || recipientShard == kShardId2);
-    ASSERT_EQ(MigrationReason::drain, reason);
-}
-
 TEST(BalancerPolicy, NoBalancingDueToAllNodesEitherDrainingOrMaxedOut) {
     // shard0 and shard2 are draining, shard1 is maxed out
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet), 1},
-                         {ShardStatistics(kShardId1, 1, 1, false, emptyTagSet), 6},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 1, true, emptyTagSet), 1}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet, emptyShardVersion), 1},
+         {ShardStatistics(kShardId1, 1, 1, false, emptyTagSet, emptyShardVersion), 6},
+         {ShardStatistics(kShardId2, kNoMaxSize, 1, true, emptyTagSet, emptyShardVersion), 1}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT(migrations.empty());
 }
 
@@ -727,188 +499,137 @@ TEST(BalancerPolicy, BalancerRespectsMaxShardSizeOnlyBalanceToNonMaxed) {
     // Note that maxSize of shard0 is 1, and it is therefore overloaded with currSize = 3. Other
     // shards have maxSize = 0 = unset. Even though the overloaded shard has the least number of
     // less chunks, we shouldn't move chunks to that shard.
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, 1, 3, false, emptyTagSet), 2},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, emptyTagSet), 5},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 10, false, emptyTagSet), 10}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, 1, 3, false, emptyTagSet, emptyShardVersion), 2},
+         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 5},
+         {ShardStatistics(kShardId2, kNoMaxSize, 10, false, emptyTagSet, emptyShardVersion), 10}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId2, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), migrations[0].maxKey);
 }
 
 TEST(BalancerPolicy, BalancerRespectsMaxShardSizeWhenAllBalanced) {
     // Note that maxSize of shard0 is 1, and it is therefore overloaded with currSize = 4. Other
     // shards have maxSize = 0 = unset. We check that being over the maxSize is NOT equivalent to
     // draining, we don't want to empty shards for no other reason than they are over this limit.
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, 1, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet), 4},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 4, false, emptyTagSet), 4}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, 1, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
+         {ShardStatistics(kShardId2, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4}});
 
     const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false);
     ASSERT(migrations.empty());
 }
 
 TEST(BalancerPolicy, BalancerRespectsTagsWhenDraining) {
     // shard1 drains the proper chunk to shard0, even though it is more loaded than shard2
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}), 6},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, true, {"a", "b"}), 1},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, {"b"}), 2}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 6},
+         {ShardStatistics(kShardId1, kNoMaxSize, 5, true, {"a", "b"}, emptyShardVersion), 2},
+         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, {"b"}, emptyShardVersion), 2}});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 7), "a")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 8), kSKeyPattern.globalMax(), "b")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 7), "a")));
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(BSON("x" << 8), kMaxBSONKey, "b")));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId1, migrations[0].from);
     ASSERT_EQ(kShardId0, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::drain, reason);
 }
 
 TEST(BalancerPolicy, BalancerRespectsTagPolicyBeforeImbalance) {
     // There is a large imbalance between shard0 and shard1, but the balancer must first fix the
     // chunks, which are on a wrong shard due to tag policy
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}), 2},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}), 6},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet), 2}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 2},
+         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 6},
+         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 2}});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 100), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 100), "a")));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId2, migrations[0].from);
     ASSERT_EQ(kShardId0, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::zoneViolation, reason);
 }
 
-TEST(BalancerPolicy, OverloadedShardCannotDonateDueToTags) {
-    // There is a large imbalance between shard0 and the other shard, but shard0 can't donate chunks
-    // because it would violate tags.
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}), 5},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet), 0},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet), 0}});
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), kSKeyPattern.globalMax(), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
-
-    const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
-    ASSERT_EQ(0U, migrations.size());
-}
-
 TEST(BalancerPolicy, BalancerFixesIncorrectTagsWithCrossShardViolationOfTags) {
     // The zone policy dictates that the same shard must donate and also receive chunks. The test
     // validates that the same shard is not used as a donor and recipient as part of the same round.
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}), 3},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}), 3},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, {"b"}), 3}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 3},
+         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 3},
+         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, {"b"}, emptyShardVersion), 3}});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 1), "b")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 8), kSKeyPattern.globalMax(), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 1), "b")));
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(BSON("x" << 8), kMaxBSONKey, "a")));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId2, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::zoneViolation, reason);
 }
 
-TEST(BalancerPolicy, BalancerFixesIncorrectTagInOtherwiseBalancedCluster) {
+TEST(BalancerPolicy, BalancerFixesIncorrectTagsInOtherwiseBalancedCluster) {
     // Chunks are balanced across shards, but there are wrong tags, which need to be fixed
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}), 3},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}), 3},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet), 3}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 3},
+         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 3},
+         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 3}});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 10), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 10), "a")));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId2, migrations[0].from);
     ASSERT_EQ(kShardId0, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::zoneViolation, reason);
 }
 
-TEST(BalancerPolicy, BalancerFixesIncorrectTagsInOtherwiseBalancedCluster) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 2, false, {}), 4},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"Zone1", "Zone2"}), 2},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 2, false, {"Zone1", "Zone2"}), 2},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 2, false, {}), 2}});
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
-        cluster.second[kShardId0][0].getMin(), cluster.second[kShardId0][0].getMax(), "Zone1")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
-        cluster.second[kShardId0][1].getMin(), cluster.second[kShardId0][1].getMax(), "Zone2")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
-        cluster.second[kShardId0][2].getMin(), cluster.second[kShardId0][2].getMax(), "Zone1")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
-        cluster.second[kShardId0][3].getMin(), cluster.second[kShardId0][3].getMax(), "Zone2")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
-
-    const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
-    ASSERT_EQ(1U, migrations.size());
-    ASSERT_EQ(kShardId0, migrations[0].from);
-    const auto& recipientShard = migrations[0].to;
-    ASSERT(recipientShard == kShardId1 || recipientShard == kShardId2);
-    ASSERT_EQ(MigrationReason::zoneViolation, reason);
-}
-
-
 TEST(BalancerPolicy, BalancerTagAlreadyBalanced) {
     // Chunks are balanced across shards for the tag.
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 3, false, {"a"}), 2},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"a"}), 2}});
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), kSKeyPattern.globalMax(), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 3, false, {"a"}, emptyShardVersion), 2},
+         {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"a"}, emptyShardVersion), 2}});
+
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, kMaxBSONKey, "a")));
     ASSERT(balanceChunks(cluster.first, distribution, false, false).first.empty());
 }
 
 TEST(BalancerPolicy, BalancerMostOverLoadShardHasMultipleTags) {
     // shard0 has chunks [MinKey, 1), [1, 2), [2, 3), [3, 4), [4, 5), so two chunks each
     // for tag "b" and "c". So [1, 2) is expected to be moved to shard1 in round 1.
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a", "b", "c"}), 5},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 1, false, {"b"}), 1},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 1, false, {"c"}), 1}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a", "b", "c"}, emptyShardVersion), 5},
+         {ShardStatistics(kShardId1, kNoMaxSize, 1, false, {"b"}, emptyShardVersion), 1},
+         {ShardStatistics(kShardId2, kNoMaxSize, 1, false, {"c"}, emptyShardVersion), 1}});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 1), "a")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 3), "b")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 3), BSON("x" << 5), "c")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 1), "a")));
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 3), "b")));
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(BSON("x" << 3), BSON("x" << 5), "c")));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(1U, migrations.size());
@@ -916,47 +637,44 @@ TEST(BalancerPolicy, BalancerMostOverLoadShardHasMultipleTags) {
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][1].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][1].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][1].getMax(), migrations[0].maxKey);
 }
 
 TEST(BalancerPolicy, BalancerMostOverLoadShardHasMultipleTagsSkipTagWithShardInUse) {
     // shard0 has chunks [MinKey, 1), [1, 2), [2, 3), [3, 4), [4, 5), so two chunks each
     // for tag "b" and "c". So [3, 4) is expected to be moved to shard2 because shard1 is
     // in use.
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a", "b", "c"}), 5},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 1, false, {"b"}), 1},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 1, false, {"c"}), 1}});
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 1), "a")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 3), "b")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 3), BSON("x" << 5), "c")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
-
-    auto availableShards = getAllShardIds(cluster.first);
-    availableShards.erase(kShardId1);
-    const auto [migrations, reason] = BalancerPolicy::balance(
-        cluster.first, distribution, boost::none /* collDataSizeInfo */, &availableShards, false);
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a", "b", "c"}, emptyShardVersion), 5},
+         {ShardStatistics(kShardId1, kNoMaxSize, 1, false, {"b"}, emptyShardVersion), 1},
+         {ShardStatistics(kShardId2, kNoMaxSize, 1, false, {"c"}, emptyShardVersion), 1}});
+
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 1), "a")));
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 3), "b")));
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(BSON("x" << 3), BSON("x" << 5), "c")));
+
+    stdx::unordered_set<ShardId> usedShards{kShardId1};
+    const auto [migrations, reason] =
+        BalancerPolicy::balance(cluster.first, distribution, &usedShards, false);
     ASSERT_EQ(1U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId2, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][3].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][3].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][3].getMax(), migrations[0].maxKey);
 }
 
 TEST(BalancerPolicy, BalancerFixesIncorrectTagsInOtherwiseBalancedClusterParallel) {
     // Chunks are balanced across shards, but there are wrong tags, which need to be fixed
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}), 3},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}), 3},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet), 3},
-                         {ShardStatistics(kShardId3, kNoMaxSize, 5, false, emptyTagSet), 3}});
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 3},
+         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 3},
+         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 3},
+         {ShardStatistics(kShardId3, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 3}});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 20), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    DistributionStatus distribution(kNamespace, cluster.second);
+    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 20), "a")));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(2U, migrations.size());
@@ -964,729 +682,177 @@ TEST(BalancerPolicy, BalancerFixesIncorrectTagsInOtherwiseBalancedClusterParalle
     ASSERT_EQ(kShardId2, migrations[0].from);
     ASSERT_EQ(kShardId0, migrations[0].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), *migrations[0].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), migrations[0].maxKey);
     ASSERT_EQ(MigrationReason::zoneViolation, reason);
 
     ASSERT_EQ(kShardId3, migrations[1].from);
     ASSERT_EQ(kShardId1, migrations[1].to);
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId3][0].getMin(), migrations[1].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId3][0].getMax(), *migrations[1].maxKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId3][0].getMax(), migrations[1].maxKey);
     ASSERT_EQ(MigrationReason::zoneViolation, reason);
 }
 
-TEST(BalancerPolicy, ChunksInNoZoneSpanOnAllShardsWithEmptyZones) {
-    // Balanacer is able to move chunks in the noZone to shards with tags
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, emptyTagSet), 3},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}), 0}});
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 100), kSKeyPattern.globalMax(), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
-
-    const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
-    ASSERT_EQ(1U, migrations.size());
-
-    ASSERT_EQ(kShardId0, migrations[0].from);
-    ASSERT_EQ(kShardId1, migrations[0].to);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), *migrations[0].maxKey);
-    ASSERT_EQ(MigrationReason::chunksImbalance, reason);
-}
-
-TEST(BalancerPolicy, BalancingNoZoneIgnoreTotalShardSize) {
-    // Shard1 is overloaded and contains:
-    // [min, 1) [1, 2) [2, 3] -> zone("a")
-    // [3, 4) [4, 5) [5, 6)   -> NoZone
-    //
-    // But it won't donate any chunk since the
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, {"a"}), 3},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}), 6},
-                         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet), 3}});
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 6), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
-
-    const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
-    ASSERT_EQ(0U, migrations.size());
-}
-
 TEST(BalancerPolicy, BalancerHandlesNoShardsWithTag) {
-    auto [cluster, cm] =
-        generateCluster({{ShardStatistics(kShardId0, kNoMaxSize, 5, false, emptyTagSet), 2},
-                         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, emptyTagSet), 2}});
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 7), "NonExistentZone")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT(balanceChunks(cluster.first, distribution, false, false).first.empty());
-}
-
-TEST(DistributionStatus, OneChunkNoZone) {
-    const auto chunks =
-        makeChunks({{kShardId0, {kSKeyPattern.globalMin(), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-    const auto distStatus = makeDistStatus(cm);
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(1, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShard(kShardId1));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-
-    ASSERT_EQ(1, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(0, distStatus.getZoneInfoForShard(kShardId1).size());
-
-    const auto& noZoneInfo = distStatus.getZoneInfoForShard(kShardId0).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(1, noZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), noZoneInfo.firstChunkMinKey);
-
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, chunks);
-}
-
-TEST(DistributionStatus, OneChunkOneZone) {
-    const auto chunks =
-        makeChunks({{kShardId0, {kSKeyPattern.globalMin(), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), kSKeyPattern.globalMax(), "ZoneA")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(0, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-    ASSERT_EQ(1, distStatus.totalChunksWithTag("ZoneA"));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShard(kShardId1));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneA"));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneA"));
-
-    ASSERT_EQ(1, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(0, distStatus.getZoneInfoForShard(kShardId1).size());
-
-    const auto& shardZoneInfo = distStatus.getZoneInfoForShard(kShardId0).at("ZoneA");
-    ASSERT_EQ(1, shardZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), shardZoneInfo.firstChunkMinKey);
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 2},
+         {ShardStatistics(kShardId1, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 2}});
 
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneA", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneA", chunks);
-}
-
-TEST(DistributionStatus, OneChunkMultipleContiguosZones) {
-    const auto chunks =
-        makeChunks({{kShardId0, {kSKeyPattern.globalMin(), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 0), "ZoneA")));
-    ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 0), kSKeyPattern.globalMax(), "ZoneB")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(1, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("ZoneA"));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("ZoneB"));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShard(kShardId1));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneA"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneB"));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneA"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneB"));
-
-    ASSERT_EQ(1, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(0, distStatus.getZoneInfoForShard(kShardId1).size());
-
-    const auto& shardZoneInfo = distStatus.getZoneInfoForShard(kShardId0).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(1, shardZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), shardZoneInfo.firstChunkMinKey);
-
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneA", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, chunks);
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneA", {});
-}
-
-TEST(DistributionStatus, OneChunkMultipleSparseZones) {
-    const auto chunks =
-        makeChunks({{kShardId0, {kSKeyPattern.globalMin(), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-
-    ZoneInfo zoneInfo;
+    DistributionStatus distribution(kNamespace, cluster.second);
     ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 0), "ZoneA")));
-    ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 10), kSKeyPattern.globalMax(), "ZoneB")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(1, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("ZoneA"));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("ZoneB"));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShard(kShardId1));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneA"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneB"));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneA"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneB"));
-
-    ASSERT_EQ(1, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(0, distStatus.getZoneInfoForShard(kShardId1).size());
-
-    const auto& shardZoneInfo = distStatus.getZoneInfoForShard(kShardId0).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(1, shardZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), shardZoneInfo.firstChunkMinKey);
-
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneA", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneB", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, chunks);
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneA", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneB", {});
-}
-
-TEST(DistributionStatus, MultipleChunksNoZone) {
-    const auto chunks = makeChunks({{kShardId0, {kSKeyPattern.globalMin(), BSON("x" << 0)}},
-                                    {kShardId0, {BSON("x" << 0), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-    const auto distStatus = makeDistStatus(cm);
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(2, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-
-    ASSERT_EQ(2, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShard(kShardId1));
-
-    ASSERT_EQ(2, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-
-    ASSERT_EQ(1, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(0, distStatus.getZoneInfoForShard(kShardId1).size());
-
-    const auto& noZoneInfo = distStatus.getZoneInfoForShard(kShardId0).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(2, noZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), noZoneInfo.firstChunkMinKey);
-
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, chunks);
-}
-
-TEST(DistributionStatus, MultipleChunksDistributedNoZone) {
-    const auto chunks = makeChunks({{kShardId1, {kSKeyPattern.globalMin(), BSON("x" << 0)}},
-                                    {kShardId0, {BSON("x" << 0), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-    const auto distStatus = makeDistStatus(cm);
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(2, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShard(kShardId1));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-
-    ASSERT_EQ(1, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(1, distStatus.getZoneInfoForShard(kShardId1).size());
-
-    const auto& noZoneInfoS0 = distStatus.getZoneInfoForShard(kShardId0).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(1, noZoneInfoS0.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 0), noZoneInfoS0.firstChunkMinKey);
-
-    const auto& noZoneInfoS1 = distStatus.getZoneInfoForShard(kShardId1).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(1, noZoneInfoS1.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), noZoneInfoS1.firstChunkMinKey);
+        distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 7), "NonExistentZone")));
 
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, {chunks[1]});
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {chunks[0]});
-}
-
-TEST(DistributionStatus, MultipleChunksTwoShardsOneZone) {
-    const auto chunks = makeChunks({{kShardId0, {kSKeyPattern.globalMin(), BSON("x" << 0)}},
-                                    {kShardId0, {BSON("x" << 0), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), kSKeyPattern.globalMax(), "ZoneA")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(0, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-    ASSERT_EQ(2, distStatus.totalChunksWithTag("ZoneA"));
-
-    ASSERT_EQ(2, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShard(kShardId1));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-    ASSERT_EQ(2, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneA"));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneA"));
-
-    ASSERT_EQ(1, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(0, distStatus.getZoneInfoForShard(kShardId1).size());
-
-    const auto& noZoneInfo = distStatus.getZoneInfoForShard(kShardId0).at("ZoneA");
-    ASSERT_EQ(2, noZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), noZoneInfo.firstChunkMinKey);
-
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneA", chunks);
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneA", {});
+    ASSERT(balanceChunks(cluster.first, distribution, false, false).first.empty());
 }
 
-TEST(DistributionStatus, MultipleChunksTwoContiguosZones) {
-    const auto chunks = makeChunks({{kShardId0, {kSKeyPattern.globalMin(), BSON("x" << 0)}},
-                                    {kShardId0, {BSON("x" << 0), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
+TEST(DistributionStatus, AddTagRangeOverlap) {
+    DistributionStatus d(kNamespace, ShardToChunksMap{});
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 0), "ZoneA")));
-    ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 0), kSKeyPattern.globalMax(), "ZoneB")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(0, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-    ASSERT_EQ(1, distStatus.totalChunksWithTag("ZoneA"));
-    ASSERT_EQ(1, distStatus.totalChunksWithTag("ZoneB"));
-
-    ASSERT_EQ(2, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShard(kShardId1));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneA"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneB"));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneA"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneB"));
-
-    ASSERT_EQ(2, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(0, distStatus.getZoneInfoForShard(kShardId1).size());
-
-    auto& shardZoneAInfo = distStatus.getZoneInfoForShard(kShardId0).at("ZoneA");
-    ASSERT_EQ(1, shardZoneAInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), shardZoneAInfo.firstChunkMinKey);
-
-    auto& shardZoneBInfo = distStatus.getZoneInfoForShard(kShardId0).at("ZoneB");
-    ASSERT_EQ(1, shardZoneBInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 0), shardZoneBInfo.firstChunkMinKey);
-
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneA", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneB", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneA", {chunks[0]});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneB", {chunks[1]});
-}
+    // Note that there is gap between 10 and 20 for which there is no tag
+    ASSERT_OK(d.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "a")));
+    ASSERT_OK(d.addRangeToZone(ZoneRange(BSON("x" << 20), BSON("x" << 30), "b")));
 
-TEST(DistributionStatus, MultipleChunksTwoZones) {
-    const auto chunks = makeChunks({{kShardId0, {kSKeyPattern.globalMin(), BSON("x" << 0)}},
-                                    {kShardId1, {BSON("x" << 0), BSON("x" << 10)}},
-                                    {kShardId2, {BSON("x" << 10), BSON("x" << 20)}},
-                                    {kShardId2, {BSON("x" << 20), BSON("x" << 30)}},
-                                    {kShardId0, {BSON("x" << 30), BSON("x" << 40)}},
-                                    {kShardId1, {BSON("x" << 40), BSON("x" << 50)}},
-                                    {kShardId2, {BSON("x" << 50), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 20), "ZoneA")));
-    ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 30), kSKeyPattern.globalMax(), "ZoneB")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(1, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(3, distStatus.totalChunksWithTag("ZoneA"));
-    ASSERT_EQ(3, distStatus.totalChunksWithTag("ZoneB"));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-
-    ASSERT_EQ(2, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(2, distStatus.numberOfChunksInShard(kShardId1));
-    ASSERT_EQ(3, distStatus.numberOfChunksInShard(kShardId2));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneA"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneB"));
-
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneA"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneB"));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId2, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId2, "NotExistingZone"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId2, "ZoneA"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId2, "ZoneB"));
-
-    ASSERT_EQ(2, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(2, distStatus.getZoneInfoForShard(kShardId1).size());
-    ASSERT_EQ(3, distStatus.getZoneInfoForShard(kShardId2).size());
-
-    auto& shard0ZoneAInfo = distStatus.getZoneInfoForShard(kShardId0).at("ZoneA");
-    ASSERT_EQ(1, shard0ZoneAInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), shard0ZoneAInfo.firstChunkMinKey);
-    auto& shard0ZoneBInfo = distStatus.getZoneInfoForShard(kShardId0).at("ZoneB");
-    ASSERT_EQ(1, shard0ZoneBInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 30), shard0ZoneBInfo.firstChunkMinKey);
-
-    auto& shard1ZoneAInfo = distStatus.getZoneInfoForShard(kShardId1).at("ZoneA");
-    ASSERT_EQ(1, shard1ZoneAInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 0), shard1ZoneAInfo.firstChunkMinKey);
-    auto& shard1ZoneBInfo = distStatus.getZoneInfoForShard(kShardId1).at("ZoneB");
-    ASSERT_EQ(1, shard1ZoneBInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 40), shard1ZoneBInfo.firstChunkMinKey);
-
-    auto& shard2ZoneAInfo = distStatus.getZoneInfoForShard(kShardId2).at("ZoneA");
-    ASSERT_EQ(1, shard2ZoneAInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 10), shard2ZoneAInfo.firstChunkMinKey);
-    auto& shard2NoZoneInfo = distStatus.getZoneInfoForShard(kShardId2).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(1, shard2NoZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 20), shard2NoZoneInfo.firstChunkMinKey);
-    auto& shard2ZoneBInfo = distStatus.getZoneInfoForShard(kShardId2).at("ZoneB");
-    ASSERT_EQ(1, shard2ZoneBInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 50), shard2ZoneBInfo.firstChunkMinKey);
-
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneA", {chunks[0]});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneB", {chunks[4]});
-
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneA", {chunks[1]});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneB", {chunks[5]});
-
-    checkChunksOnShardForTag(distStatus, kShardId2, ZoneInfo::kNoZoneName, {chunks[3]});
-    checkChunksOnShardForTag(distStatus, kShardId2, "NotExistingZone", {});
-    checkChunksOnShardForTag(distStatus, kShardId2, "ZoneA", {chunks[2]});
-    checkChunksOnShardForTag(distStatus, kShardId2, "ZoneB", {chunks[6]});
+    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
+              d.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 2), "d")));
+    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
+              d.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 5), "d")));
+    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
+              d.addRangeToZone(ZoneRange(BSON("x" << 5), BSON("x" << 9), "d")));
+    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
+              d.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "d")));
+    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
+              d.addRangeToZone(ZoneRange(BSON("x" << 5), BSON("x" << 25), "d")));
+    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
+              d.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 32), "d")));
+    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
+              d.addRangeToZone(ZoneRange(BSON("x" << 25), kMaxBSONKey, "d")));
 }
 
-TEST(DistributionStatus, MultipleChunksMulitpleZoneRanges) {
-    const auto chunks = makeChunks({{kShardId0, {kSKeyPattern.globalMin(), BSON("x" << 0)}},
-
-                                    {kShardId1, {BSON("x" << 0), BSON("x" << 10)}},   // ZoneA
-                                    {kShardId0, {BSON("x" << 10), BSON("x" << 20)}},  // ZoneA
-
-                                    {kShardId1, {BSON("x" << 20), BSON("x" << 30)}},
-                                    {kShardId0, {BSON("x" << 30), BSON("x" << 40)}},
-
-                                    {kShardId1, {BSON("x" << 40), BSON("x" << 50)}},  // ZoneA
-                                    {kShardId0, {BSON("x" << 50), BSON("x" << 60)}},  // ZoneA
+TEST(DistributionStatus, ChunkTagsSelectorWithRegularKeys) {
+    DistributionStatus d(kNamespace, ShardToChunksMap{});
 
-                                    {kShardId1, {BSON("x" << 60), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
+    ASSERT_OK(d.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "a")));
+    ASSERT_OK(d.addRangeToZone(ZoneRange(BSON("x" << 10), BSON("x" << 20), "b")));
+    ASSERT_OK(d.addRangeToZone(ZoneRange(BSON("x" << 20), BSON("x" << 30), "c")));
 
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 0), BSON("x" << 20), "ZoneA")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 40), BSON("x" << 60), "ZoneA")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(4, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(4, distStatus.totalChunksWithTag("ZoneA"));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-
-    ASSERT_EQ(4, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(4, distStatus.numberOfChunksInShard(kShardId1));
+    {
+        ChunkType chunk;
+        chunk.setMin(kMinBSONKey);
+        chunk.setMax(BSON("x" << 1));
+        ASSERT_EQUALS("", d.getTagForChunk(chunk));
+    }
 
-    ASSERT_EQ(2, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-    ASSERT_EQ(2, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneA"));
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 0));
+        chunk.setMax(BSON("x" << 1));
+        ASSERT_EQUALS("", d.getTagForChunk(chunk));
+    }
 
-    ASSERT_EQ(2, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-    ASSERT_EQ(2, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneA"));
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 1));
+        chunk.setMax(BSON("x" << 5));
+        ASSERT_EQUALS("a", d.getTagForChunk(chunk));
+    }
 
-    ASSERT_EQ(2, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(2, distStatus.getZoneInfoForShard(kShardId1).size());
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 10));
+        chunk.setMax(BSON("x" << 20));
+        ASSERT_EQUALS("b", d.getTagForChunk(chunk));
+    }
 
-    auto& shard0ZoneAInfo = distStatus.getZoneInfoForShard(kShardId0).at("ZoneA");
-    ASSERT_EQ(2, shard0ZoneAInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 10), shard0ZoneAInfo.firstChunkMinKey);
-    auto& shard0NoZoneInfo = distStatus.getZoneInfoForShard(kShardId0).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(2, shard0NoZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(kSKeyPattern.globalMin(), shard0NoZoneInfo.firstChunkMinKey);
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 15));
+        chunk.setMax(BSON("x" << 20));
+        ASSERT_EQUALS("b", d.getTagForChunk(chunk));
+    }
 
-    auto& shard1ZoneAInfo = distStatus.getZoneInfoForShard(kShardId1).at("ZoneA");
-    ASSERT_EQ(2, shard1ZoneAInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 0), shard1ZoneAInfo.firstChunkMinKey);
-    auto& shard1NoZoneInfo = distStatus.getZoneInfoForShard(kShardId1).at(ZoneInfo::kNoZoneName);
-    ASSERT_EQ(2, shard1NoZoneInfo.numChunks);
-    ASSERT_BSONOBJ_EQ(BSON("x" << 20), shard1NoZoneInfo.firstChunkMinKey);
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 25));
+        chunk.setMax(BSON("x" << 30));
+        ASSERT_EQUALS("c", d.getTagForChunk(chunk));
+    }
 
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, {chunks[0], chunks[4]});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneA", {chunks[2], chunks[6]});
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 35));
+        chunk.setMax(BSON("x" << 40));
+        ASSERT_EQUALS("", d.getTagForChunk(chunk));
+    }
 
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {chunks[3], chunks[7]});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneA", {chunks[1], chunks[5]});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-}
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 30));
+        chunk.setMax(kMaxBSONKey);
+        ASSERT_EQUALS("", d.getTagForChunk(chunk));
+    }
 
-TEST(DistributionStatus, MultipleChunksMulitpleZoneRangesNotAligned) {
-    const auto chunks =
-        makeChunks({{kShardId0, {kSKeyPattern.globalMin(), BSON("x" << 0)}},     // Zone A
-                    {kShardId0, {BSON("x" << 0), BSON("x" << 10)}},              // Zone A
-                    {kShardId2, {BSON("x" << 10), BSON("x" << 20)}},             // No Zone
-                    {kShardId2, {BSON("x" << 20), BSON("x" << 30)}},             // Zone A
-                    {kShardId1, {BSON("x" << 30), BSON("x" << 40)}},             // No Zone
-                    {kShardId1, {BSON("x" << 40), BSON("x" << 50)}},             // Zone B
-                    {kShardId0, {BSON("x" << 50), BSON("x" << 60)}},             // No Zone
-                    {kShardId0, {BSON("x" << 60), BSON("x" << 70)}},             // Zone B
-                    {kShardId2, {BSON("x" << 70), kSKeyPattern.globalMax()}}});  // No zone
-    const auto cm = makeChunkManager(chunks);
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(
-        zoneInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 15), "ZoneA")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 15), BSON("x" << 35), "ZoneA")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 40), BSON("x" << 55), "ZoneB")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 55), BSON("x" << 75), "ZoneB")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    ASSERT_EQ(kNamespace, distStatus.nss());
-
-    ASSERT_EQ(4, distStatus.totalChunksWithTag(ZoneInfo::kNoZoneName));
-    ASSERT_EQ(3, distStatus.totalChunksWithTag("ZoneA"));
-    ASSERT_EQ(2, distStatus.totalChunksWithTag("ZoneB"));
-    ASSERT_EQ(0, distStatus.totalChunksWithTag("NotExistingZone"));
-
-    ASSERT_EQ(4, distStatus.numberOfChunksInShard(kShardId0));
-    ASSERT_EQ(2, distStatus.numberOfChunksInShard(kShardId1));
-    ASSERT_EQ(3, distStatus.numberOfChunksInShard(kShardId2));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(2, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneA"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId0, "ZoneB"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId0, "NotExistingZone"));
-
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId1, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneA"));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId1, "ZoneB"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId1, "NotExistingZone"));
-
-    ASSERT_EQ(2, distStatus.numberOfChunksInShardWithTag(kShardId2, ZoneInfo::kNoZoneName));
-    ASSERT_EQ(1, distStatus.numberOfChunksInShardWithTag(kShardId2, "ZoneA"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId2, "ZoneB"));
-    ASSERT_EQ(0, distStatus.numberOfChunksInShardWithTag(kShardId2, "NotExistingZone"));
-
-    ASSERT_EQ(3, distStatus.getZoneInfoForShard(kShardId0).size());
-    ASSERT_EQ(2, distStatus.getZoneInfoForShard(kShardId1).size());
-    ASSERT_EQ(2, distStatus.getZoneInfoForShard(kShardId2).size());
-
-    checkChunksOnShardForTag(distStatus, kShardId0, ZoneInfo::kNoZoneName, {chunks[6]});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneA", {chunks[0], chunks[1]});
-    checkChunksOnShardForTag(distStatus, kShardId0, "ZoneB", {chunks[7]});
-    checkChunksOnShardForTag(distStatus, kShardId0, "NotExistingZone", {});
-
-    checkChunksOnShardForTag(distStatus, kShardId1, ZoneInfo::kNoZoneName, {chunks[4]});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneA", {});
-    checkChunksOnShardForTag(distStatus, kShardId1, "ZoneB", {chunks[5]});
-    checkChunksOnShardForTag(distStatus, kShardId1, "NotExistingZone", {});
-
-    checkChunksOnShardForTag(distStatus, kShardId2, ZoneInfo::kNoZoneName, {chunks[2], chunks[8]});
-    checkChunksOnShardForTag(distStatus, kShardId2, "ZoneA", {chunks[3]});
-    checkChunksOnShardForTag(distStatus, kShardId2, "ZoneB", {});
-    checkChunksOnShardForTag(distStatus, kShardId2, "NotExistingZone", {});
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 40));
+        chunk.setMax(kMaxBSONKey);
+        ASSERT_EQUALS("", d.getTagForChunk(chunk));
+    }
 }
 
-TEST(DistributionStatus, forEachChunkOnShardInZoneExitCondition) {
-    const auto chunks = makeChunks({{kShardId0, {kSKeyPattern.globalMin(), BSON("x" << 0)}},
-
-                                    {kShardId1, {BSON("x" << 0), BSON("x" << 10)}},   // ZoneA
-                                    {kShardId0, {BSON("x" << 10), BSON("x" << 20)}},  // ZoneA
-
-                                    {kShardId1, {BSON("x" << 20), BSON("x" << 30)}},
-                                    {kShardId0, {BSON("x" << 30), BSON("x" << 40)}},
-
-                                    {kShardId1, {BSON("x" << 40), BSON("x" << 50)}},  // ZoneA
-                                    {kShardId0, {BSON("x" << 50), BSON("x" << 60)}},  // ZoneA
-
-                                    {kShardId1, {BSON("x" << 60), kSKeyPattern.globalMax()}}});
-    const auto cm = makeChunkManager(chunks);
-
-    ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 0), BSON("x" << 20), "ZoneA")));
-    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 40), BSON("x" << 60), "ZoneA")));
-    const auto distStatus = makeDistStatus(cm, std::move(zoneInfo));
-
-    const auto assertLoopStopAt =
-        [&](const ShardId& shardId, const std::string& zoneName, const size_t stopCount) {
-            size_t numChunkIterated{0};
+TEST(DistributionStatus, ChunkTagsSelectorWithMinMaxKeys) {
+    DistributionStatus d(kNamespace, ShardToChunksMap{});
 
-            const auto completed =
-                distStatus.forEachChunkOnShardInZone(shardId, zoneName, [&](const auto& chunk) {
-                    if (++numChunkIterated == stopCount) {
-                        return false;  // break
-                    }
-                    return true;  // continue
-                });
+    ASSERT_OK(d.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << -100), "a")));
+    ASSERT_OK(d.addRangeToZone(ZoneRange(BSON("x" << -10), BSON("x" << 10), "b")));
+    ASSERT_OK(d.addRangeToZone(ZoneRange(BSON("x" << 100), kMaxBSONKey, "c")));
 
-            ASSERT(!completed) << "forEachChunkOnShardInZone loop did not stop";
-            ASSERT_EQ(stopCount, numChunkIterated);
-        };
-
-    assertLoopStopAt(kShardId0, ZoneInfo::kNoZoneName, 1);
-    assertLoopStopAt(kShardId0, ZoneInfo::kNoZoneName, 2);
-
-    assertLoopStopAt(kShardId1, ZoneInfo::kNoZoneName, 1);
-    assertLoopStopAt(kShardId1, ZoneInfo::kNoZoneName, 2);
-
-    assertLoopStopAt(kShardId0, "ZoneA", 1);
-    assertLoopStopAt(kShardId0, "ZoneA", 2);
+    {
+        ChunkType chunk;
+        chunk.setMin(kMinBSONKey);
+        chunk.setMax(BSON("x" << -100));
+        ASSERT_EQUALS("a", d.getTagForChunk(chunk));
+    }
 
-    assertLoopStopAt(kShardId1, "ZoneA", 1);
-    assertLoopStopAt(kShardId1, "ZoneA", 2);
-}
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << -100));
+        chunk.setMax(BSON("x" << -11));
+        ASSERT_EQUALS("", d.getTagForChunk(chunk));
+    }
 
-TEST(ZoneInfo, AddTagRangeOverlap) {
-    ZoneInfo zInfo;
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << -10));
+        chunk.setMax(BSON("x" << 0));
+        ASSERT_EQUALS("b", d.getTagForChunk(chunk));
+    }
 
-    // Note that there is gap between 10 and 20 for which there is no tag
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "a")));
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 20), BSON("x" << 30), "b")));
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 0));
+        chunk.setMax(BSON("x" << 10));
+        ASSERT_EQUALS("b", d.getTagForChunk(chunk));
+    }
 
-    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              zInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 2), "d")));
-    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              zInfo.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 5), "d")));
-    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              zInfo.addRangeToZone(ZoneRange(BSON("x" << 5), BSON("x" << 9), "d")));
-    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              zInfo.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "d")));
-    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              zInfo.addRangeToZone(ZoneRange(BSON("x" << 5), BSON("x" << 25), "d")));
-    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              zInfo.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 32), "d")));
-    ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              zInfo.addRangeToZone(ZoneRange(BSON("x" << 25), kSKeyPattern.globalMax(), "d")));
-}
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 10));
+        chunk.setMax(BSON("x" << 20));
+        ASSERT_EQUALS("", d.getTagForChunk(chunk));
+    }
 
-TEST(ZoneInfo, ChunkTagsSelectorWithRegularKeys) {
-    ZoneInfo zInfo;
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "a")));
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 10), BSON("x" << 20), "b")));
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 20), BSON("x" << 30), "c")));
-
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForChunk({kSKeyPattern.globalMin(), BSON("x" << 1)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName, zInfo.getZoneForChunk({BSON("x" << 0), BSON("x" << 1)}));
-    ASSERT_EQUALS("a", zInfo.getZoneForChunk({BSON("x" << 1), BSON("x" << 5)}));
-    ASSERT_EQUALS("b", zInfo.getZoneForChunk({BSON("x" << 10), BSON("x" << 20)}));
-    ASSERT_EQUALS("b", zInfo.getZoneForChunk({BSON("x" << 15), BSON("x" << 20)}));
-    ASSERT_EQUALS("c", zInfo.getZoneForChunk({BSON("x" << 25), BSON("x" << 30)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName, zInfo.getZoneForChunk({BSON("x" << 35), BSON("x" << 40)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForChunk({BSON("x" << 30), kSKeyPattern.globalMax()}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForChunk({BSON("x" << 40), kSKeyPattern.globalMax()}));
-}
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 10));
+        chunk.setMax(BSON("x" << 100));
+        ASSERT_EQUALS("", d.getTagForChunk(chunk));
+    }
 
-TEST(ZoneInfo, ChunkTagsSelectorWithMinMaxKeys) {
-
-    ZoneInfo zInfo;
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << -100), "a")));
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << -10), BSON("x" << 10), "b")));
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 100), kSKeyPattern.globalMax(), "c")));
-
-    ASSERT_EQUALS("a", zInfo.getZoneForChunk({kSKeyPattern.globalMin(), BSON("x" << -100)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForChunk({BSON("x" << -100), BSON("x" << -11)}));
-    ASSERT_EQUALS("b", zInfo.getZoneForChunk({BSON("x" << -10), BSON("x" << 0)}));
-    ASSERT_EQUALS("b", zInfo.getZoneForChunk({BSON("x" << 0), BSON("x" << 10)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName, zInfo.getZoneForChunk({BSON("x" << 10), BSON("x" << 20)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForChunk({BSON("x" << 10), BSON("x" << 100)}));
-    ASSERT_EQUALS("c", zInfo.getZoneForChunk({BSON("x" << 200), kSKeyPattern.globalMax()}));
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 200));
+        chunk.setMax(kMaxBSONKey);
+        ASSERT_EQUALS("c", d.getTagForChunk(chunk));
+    }
 }
 
 }  // namespace