SERVER-80488 Avoid traversing routing table in balancer split chunk policyr5.0.21-rc0r5.0.21

5 files changed, 166 insertions, 139 deletions

diff --git a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
index 6af2597ffd2..8371a27b950 100644
--- a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
@@ -57,6 +57,30 @@ using SplitInfoVector = BalancerChunkSelectionPolicy::SplitInfoVector;
 
 namespace {
 
+StatusWith<ZoneInfo> createCollectionZoneInfo(OperationContext* opCtx,
+                                              const NamespaceString& nss,
+                                              const KeyPattern& keyPattern) {
+    ZoneInfo zoneInfo;
+    const auto swCollectionTags =
+        Grid::get(opCtx)->catalogClient()->getTagsForCollection(opCtx, nss);
+    if (!swCollectionTags.isOK()) {
+        return swCollectionTags.getStatus().withContext(
+            str::stream() << "Unable to load tags for collection " << nss);
+    }
+    const auto& collectionTags = swCollectionTags.getValue();
+
+    for (const auto& tag : collectionTags) {
+        auto status =
+            zoneInfo.addRangeToZone(ZoneRange{keyPattern.extendRangeBound(tag.getMinKey(), false),
+                                              keyPattern.extendRangeBound(tag.getMaxKey(), false),
+                                              tag.getTag()});
+        if (!status.isOK()) {
+            return status;
+        }
+    }
+    return {std::move(zoneInfo)};
+}
+
 /**
  * Does a linear pass over the information cached in the specified chunk manager and extracts chunk
  * distribution and chunk placement information which is needed by the balancer policy.
@@ -88,31 +112,13 @@ StatusWith<DistributionStatus> createCollectionDistributionStatus(
         return true;
     });
 
-    const auto swCollectionTags =
-        Grid::get(opCtx)->catalogClient()->getTagsForCollection(opCtx, nss);
-    if (!swCollectionTags.isOK()) {
-        return swCollectionTags.getStatus().withContext(
-            str::stream() << "Unable to load tags for collection " << nss);
-    }
-    const auto& collectionTags = swCollectionTags.getValue();
-
-    DistributionStatus distribution(nss, std::move(shardToChunksMap));
-
-    // Cache the collection tags
-    const auto& keyPattern = chunkMgr.getShardKeyPattern().getKeyPattern();
-
-    for (const auto& tag : collectionTags) {
-        auto status = distribution.addRangeToZone(
-            ZoneRange(keyPattern.extendRangeBound(tag.getMinKey(), false),
-                      keyPattern.extendRangeBound(tag.getMaxKey(), false),
-                      tag.getTag()));
-
-        if (!status.isOK()) {
-            return status;
-        }
+    auto swZoneInfo =
+        createCollectionZoneInfo(opCtx, nss, chunkMgr.getShardKeyPattern().getKeyPattern());
+    if (!swZoneInfo.isOK()) {
+        return swZoneInfo.getStatus();
     }
 
-    return {std::move(distribution)};
+    return {DistributionStatus{nss, std::move(shardToChunksMap), std::move(swZoneInfo.getValue())}};
 }
 
 /**
@@ -183,12 +189,12 @@ private:
  * range boundaries.
  */
 void getSplitCandidatesToEnforceTagRanges(const ChunkManager& cm,
-                                          const DistributionStatus& distribution,
+                                          const ZoneInfo& zoneInfo,
                                           SplitCandidatesBuffer* splitCandidates) {
     const auto& globalMax = cm.getShardKeyPattern().getKeyPattern().globalMax();
 
     // For each tag range, find chunks that need to be split.
-    for (const auto& tagRangeEntry : distribution.tagRanges()) {
+    for (const auto& tagRangeEntry : zoneInfo.zoneRanges()) {
         const auto& tagRange = tagRangeEntry.second;
 
         const auto chunkAtZoneMin = cm.findIntersectingChunkWithSimpleCollation(tagRange.min);
@@ -533,26 +539,26 @@ StatusWith<SplitInfoVector> BalancerChunkSelectionPolicyImpl::_getSplitCandidate
 
     const auto& cm = routingInfoStatus.getValue();
 
-    const auto collInfoStatus = createCollectionDistributionStatus(opCtx, nss, shardStats, cm);
-    if (!collInfoStatus.isOK()) {
-        return collInfoStatus.getStatus();
+    const auto swCollZoneInfo =
+        createCollectionZoneInfo(opCtx, nss, cm.getShardKeyPattern().getKeyPattern());
+    if (!swCollZoneInfo.isOK()) {
+        return swCollZoneInfo.getStatus();
     }
-
-    const DistributionStatus& distribution = collInfoStatus.getValue();
+    const auto& collZoneInfo = swCollZoneInfo.getValue();
 
     // Accumulate split points for the same chunk together
     SplitCandidatesBuffer splitCandidates(nss, cm.getVersion());
 
     if (nss == NamespaceString::kLogicalSessionsNamespace) {
-        if (!distribution.tags().empty()) {
+        if (!collZoneInfo.allZones().empty()) {
             LOGV2_WARNING(4562401,
                           "Ignoring zones for the sessions collection",
-                          "tags"_attr = distribution.tags());
+                          "tags"_attr = collZoneInfo.allZones());
         }
 
         getSplitCandidatesForSessionsCollection(opCtx, cm, &splitCandidates);
     } else {
-        getSplitCandidatesToEnforceTagRanges(cm, distribution, &splitCandidates);
+        getSplitCandidatesToEnforceTagRanges(cm, collZoneInfo, &splitCandidates);
     }
 
     return splitCandidates.done();
diff --git a/src/mongo/db/s/balancer/balancer_policy.cpp b/src/mongo/db/s/balancer/balancer_policy.cpp
index b632f1d79d8..ec4041b9667 100644
--- a/src/mongo/db/s/balancer/balancer_policy.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy.cpp
@@ -60,8 +60,12 @@ const size_t kDefaultImbalanceThreshold = 1;
 
 }  // namespace
 
-DistributionStatus::DistributionStatus(NamespaceString nss, ShardToChunksMap shardToChunksMap)
-    : _nss(std::move(nss)), _shardChunks(std::move(shardToChunksMap)) {}
+DistributionStatus::DistributionStatus(NamespaceString nss,
+                                       ShardToChunksMap shardToChunksMap,
+                                       ZoneInfo zoneInfo)
+    : _nss(std::move(nss)),
+      _shardChunks(std::move(shardToChunksMap)),
+      _zoneInfo(std::move(zoneInfo)) {}
 
 size_t DistributionStatus::totalChunks() const {
     size_t total = 0;
@@ -110,10 +114,6 @@ const vector<ChunkType>& DistributionStatus::getChunks(const ShardId& shardId) c
     return i->second;
 }
 
-Status DistributionStatus::addRangeToZone(const ZoneRange& range) {
-    return _zoneInfo.addRangeToZone(range);
-}
-
 string DistributionStatus::getTagForChunk(const ChunkType& chunk) const {
     return _zoneInfo.getZoneForChunk(chunk.getRange());
 }
diff --git a/src/mongo/db/s/balancer/balancer_policy.h b/src/mongo/db/s/balancer/balancer_policy.h
index 2b685b79bbb..b8403c3881f 100644
--- a/src/mongo/db/s/balancer/balancer_policy.h
+++ b/src/mongo/db/s/balancer/balancer_policy.h
@@ -86,6 +86,7 @@ typedef std::map<ShardId, std::vector<ChunkType>> ShardToChunksMap;
 class ZoneInfo {
 public:
     ZoneInfo();
+    ZoneInfo(ZoneInfo&&) = default;
 
     /**
      * Appends the specified range to the set of ranges tracked for this collection and checks if
@@ -131,7 +132,7 @@ class DistributionStatus {
     DistributionStatus& operator=(const DistributionStatus&) = delete;
 
 public:
-    DistributionStatus(NamespaceString nss, ShardToChunksMap shardToChunksMap);
+    DistributionStatus(NamespaceString nss, ShardToChunksMap shardToChunksMap, ZoneInfo zoneInfo);
     DistributionStatus(DistributionStatus&&) = default;
 
     /**
@@ -142,12 +143,6 @@ public:
     }
 
     /**
-     * Appends the specified range to the set of ranges tracked for this collection and checks if
-     * it overlaps with existing ranges.
-     */
-    Status addRangeToZone(const ZoneRange& range);
-
-    /**
      * Returns total number of chunks across all shards.
      */
     size_t totalChunks() const;
diff --git a/src/mongo/db/s/balancer/balancer_policy_test.cpp b/src/mongo/db/s/balancer/balancer_policy_test.cpp
index 7d0c56e72ae..8f3b340c82d 100644
--- a/src/mongo/db/s/balancer/balancer_policy_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy_test.cpp
@@ -130,8 +130,8 @@ TEST(BalancerPolicy, Basic) {
          {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
          {ShardStatistics(kShardId2, kNoMaxSize, 3, false, emptyTagSet, emptyShardVersion), 3}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
@@ -146,8 +146,8 @@ TEST(BalancerPolicy, SmallClusterShouldBePerfectlyBalanced) {
          {ShardStatistics(kShardId1, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 2},
          {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId1, migrations[0].from);
     ASSERT_EQ(kShardId2, migrations[0].to);
@@ -161,10 +161,13 @@ TEST(BalancerPolicy, SingleChunkShouldNotMove) {
         {{ShardStatistics(kShardId0, kNoMaxSize, 1, false, emptyTagSet, emptyShardVersion), 1},
          {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
-    ASSERT(balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), true, false)
-               .empty());
     ASSERT(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false)
+        balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), true, false)
+            .empty());
+    ASSERT(
+        balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false)
             .empty());
 }
 
@@ -175,10 +178,13 @@ TEST(BalancerPolicy, BalanceThresholdObeyed) {
          {ShardStatistics(kShardId2, kNoMaxSize, 1, false, emptyTagSet, emptyShardVersion), 1},
          {ShardStatistics(kShardId3, kNoMaxSize, 1, false, emptyTagSet, emptyShardVersion), 1}});
 
-    ASSERT(balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), true, false)
-               .empty());
     ASSERT(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false)
+        balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), true, false)
+            .empty());
+    ASSERT(
+        balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false)
             .empty());
 }
 
@@ -189,8 +195,8 @@ TEST(BalancerPolicy, ParallelBalancing) {
          {ShardStatistics(kShardId2, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
          {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -215,8 +221,8 @@ TEST(BalancerPolicy, ParallelBalancingDoesNotPutChunksOnShardsAboveTheOptimal) {
          {ShardStatistics(kShardId4, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
          {ShardStatistics(kShardId5, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -239,8 +245,8 @@ TEST(BalancerPolicy, ParallelBalancingDoesNotMoveChunksFromShardsBelowOptimal) {
          {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 5},
          {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(1U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -259,8 +265,11 @@ TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseSourceShardsWithMoveNe
 
     // Here kShardId0 would have been selected as a donor
     stdx::unordered_set<ShardId> availableShards{kShardId1, kShardId2, kShardId3};
-    const auto migrations(BalancerPolicy::balance(
-        cluster.first, DistributionStatus(kNamespace, cluster.second), &availableShards, false));
+    const auto migrations(
+        BalancerPolicy::balance(cluster.first,
+                                DistributionStatus(kNamespace, cluster.second, ZoneInfo()),
+                                &availableShards,
+                                false));
     ASSERT_EQ(1U, migrations.size());
 
     ASSERT_EQ(kShardId1, migrations[0].from);
@@ -279,8 +288,11 @@ TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseSourceShardsWithMoveNo
 
     // Here kShardId0 would have been selected as a donor
     stdx::unordered_set<ShardId> availableShards{kShardId1, kShardId2, kShardId3};
-    const auto migrations(BalancerPolicy::balance(
-        cluster.first, DistributionStatus(kNamespace, cluster.second), &availableShards, false));
+    const auto migrations(
+        BalancerPolicy::balance(cluster.first,
+                                DistributionStatus(kNamespace, cluster.second, ZoneInfo()),
+                                &availableShards,
+                                false));
     ASSERT_EQ(0U, migrations.size());
 }
 
@@ -293,8 +305,11 @@ TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseDestinationShards) {
 
     // Here kShardId2 would have been selected as a recipient
     stdx::unordered_set<ShardId> availableShards{kShardId0, kShardId1, kShardId3};
-    const auto migrations(BalancerPolicy::balance(
-        cluster.first, DistributionStatus(kNamespace, cluster.second), &availableShards, false));
+    const auto migrations(
+        BalancerPolicy::balance(cluster.first,
+                                DistributionStatus(kNamespace, cluster.second, ZoneInfo()),
+                                &availableShards,
+                                false));
     ASSERT_EQ(1U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -314,8 +329,8 @@ TEST(BalancerPolicy, JumboChunksNotMoved) {
     cluster.second[kShardId0][2].setJumbo(true);
     cluster.second[kShardId0][3].setJumbo(true);
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
@@ -341,8 +356,8 @@ TEST(BalancerPolicy, JumboChunksNotMovedParallel) {
     cluster.second[kShardId2][2].setJumbo(false);  // Only chunk 1 is not jumbo
     cluster.second[kShardId2][3].setJumbo(true);
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -364,8 +379,8 @@ TEST(BalancerPolicy, DrainingSingleChunk) {
         {{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet, emptyShardVersion), 1},
          {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 5}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
@@ -382,8 +397,8 @@ TEST(BalancerPolicy, DrainingSingleChunkPerShard) {
          {ShardStatistics(kShardId2, kNoMaxSize, 2, true, emptyTagSet, emptyShardVersion), 1},
          {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 5}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -405,8 +420,8 @@ TEST(BalancerPolicy, DrainingWithTwoChunksFirstOneSelected) {
         {{ShardStatistics(kShardId0, kNoMaxSize, 2, true, emptyTagSet, emptyShardVersion), 2},
          {ShardStatistics(kShardId1, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 5}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
@@ -423,8 +438,8 @@ TEST(BalancerPolicy, DrainingMultipleShardsFirstOneSelected) {
          {ShardStatistics(kShardId1, kNoMaxSize, 5, true, emptyTagSet, emptyShardVersion), 2},
          {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 16}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId2, migrations[0].to);
@@ -440,8 +455,8 @@ TEST(BalancerPolicy, DrainingMultipleShardsWontAcceptChunks) {
          {ShardStatistics(kShardId1, kNoMaxSize, 0, true, emptyTagSet, emptyShardVersion), 0},
          {ShardStatistics(kShardId2, kNoMaxSize, 0, true, emptyTagSet, emptyShardVersion), 0}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT(migrations.empty());
 }
 
@@ -451,9 +466,10 @@ TEST(BalancerPolicy, DrainingSingleAppropriateShardFoundDueToTag) {
          {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"LAX"}, emptyShardVersion), 4},
          {ShardStatistics(kShardId2, kNoMaxSize, 1, true, {"LAX"}, emptyShardVersion), 1}});
 
-    DistributionStatus distribution(kNamespace, cluster.second);
-    ASSERT_OK(distribution.addRangeToZone(ZoneRange(
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
         cluster.second[kShardId2][0].getMin(), cluster.second[kShardId2][0].getMax(), "LAX")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto migrations(balanceChunks(cluster.first, distribution, false, false));
     ASSERT_EQ(1U, migrations.size());
@@ -470,9 +486,10 @@ TEST(BalancerPolicy, DrainingNoAppropriateShardsFoundDueToTag) {
          {ShardStatistics(kShardId1, kNoMaxSize, 2, false, {"LAX"}, emptyShardVersion), 4},
          {ShardStatistics(kShardId2, kNoMaxSize, 1, true, {"SEA"}, emptyShardVersion), 1}});
 
-    DistributionStatus distribution(kNamespace, cluster.second);
-    ASSERT_OK(distribution.addRangeToZone(ZoneRange(
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
         cluster.second[kShardId2][0].getMin(), cluster.second[kShardId2][0].getMax(), "SEA")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto migrations(balanceChunks(cluster.first, distribution, false, false));
     ASSERT(migrations.empty());
@@ -485,8 +502,8 @@ TEST(BalancerPolicy, NoBalancingDueToAllNodesEitherDrainingOrMaxedOut) {
          {ShardStatistics(kShardId1, 1, 1, false, emptyTagSet, emptyShardVersion), 6},
          {ShardStatistics(kShardId2, kNoMaxSize, 1, true, emptyTagSet, emptyShardVersion), 1}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT(migrations.empty());
 }
 
@@ -499,8 +516,8 @@ TEST(BalancerPolicy, BalancerRespectsMaxShardSizeOnlyBalanceToNonMaxed) {
          {ShardStatistics(kShardId1, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 5},
          {ShardStatistics(kShardId2, kNoMaxSize, 10, false, emptyTagSet, emptyShardVersion), 10}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId2, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
@@ -517,8 +534,8 @@ TEST(BalancerPolicy, BalancerRespectsMaxShardSizeWhenAllBalanced) {
          {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
          {ShardStatistics(kShardId2, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4}});
 
-    const auto migrations(
-        balanceChunks(cluster.first, DistributionStatus(kNamespace, cluster.second), false, false));
+    const auto migrations(balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false));
     ASSERT(migrations.empty());
 }
 
@@ -529,9 +546,10 @@ TEST(BalancerPolicy, BalancerRespectsTagsWhenDraining) {
          {ShardStatistics(kShardId1, kNoMaxSize, 5, true, {"a", "b"}, emptyShardVersion), 2},
          {ShardStatistics(kShardId2, kNoMaxSize, 5, false, {"b"}, emptyShardVersion), 2}});
 
-    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")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 7), "a")));
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 8), kMaxBSONKey, "b")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto migrations(balanceChunks(cluster.first, distribution, false, false));
     ASSERT_EQ(1U, migrations.size());
@@ -550,8 +568,9 @@ TEST(BalancerPolicy, BalancerRespectsTagPolicyBeforeImbalance) {
          {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 6},
          {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 2}});
 
-    DistributionStatus distribution(kNamespace, cluster.second);
-    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 100), "a")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 100), "a")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto migrations(balanceChunks(cluster.first, distribution, false, false));
     ASSERT_EQ(1U, migrations.size());
@@ -570,9 +589,10 @@ TEST(BalancerPolicy, BalancerFixesIncorrectTagsWithCrossShardViolationOfTags) {
          {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 3},
          {ShardStatistics(kShardId2, kNoMaxSize, 5, false, {"b"}, emptyShardVersion), 3}});
 
-    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")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 1), "b")));
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(BSON("x" << 8), kMaxBSONKey, "a")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto migrations(balanceChunks(cluster.first, distribution, false, false));
     ASSERT_EQ(1U, migrations.size());
@@ -590,8 +610,9 @@ TEST(BalancerPolicy, BalancerFixesIncorrectTagsInOtherwiseBalancedCluster) {
          {ShardStatistics(kShardId1, kNoMaxSize, 5, false, {"a"}, emptyShardVersion), 3},
          {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 3}});
 
-    DistributionStatus distribution(kNamespace, cluster.second);
-    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 10), "a")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 10), "a")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto migrations(balanceChunks(cluster.first, distribution, false, false));
     ASSERT_EQ(1U, migrations.size());
@@ -608,8 +629,9 @@ TEST(BalancerPolicy, BalancerTagAlreadyBalanced) {
         {{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")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, kMaxBSONKey, "a")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
     ASSERT(balanceChunks(cluster.first, distribution, false, false).empty());
 }
 
@@ -621,10 +643,11 @@ TEST(BalancerPolicy, BalancerMostOverLoadShardHasMultipleTags) {
          {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")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, 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")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto migrations(balanceChunks(cluster.first, distribution, false, false));
     ASSERT_EQ(1U, migrations.size());
@@ -644,10 +667,11 @@ TEST(BalancerPolicy, BalancerMostOverLoadShardHasMultipleTagsSkipTagWithShardInU
          {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")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, 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")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     stdx::unordered_set<ShardId> availableShards{kShardId0, kShardId2, kShardId3};
     const auto migrations(
@@ -668,8 +692,9 @@ TEST(BalancerPolicy, BalancerFixesIncorrectTagsInOtherwiseBalancedClusterParalle
          {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 3},
          {ShardStatistics(kShardId3, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 3}});
 
-    DistributionStatus distribution(kNamespace, cluster.second);
-    ASSERT_OK(distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 20), "a")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 20), "a")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto migrations(balanceChunks(cluster.first, distribution, false, false));
     ASSERT_EQ(2U, migrations.size());
@@ -692,42 +717,42 @@ TEST(BalancerPolicy, BalancerHandlesNoShardsWithTag) {
         {{ShardStatistics(kShardId0, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 2},
          {ShardStatistics(kShardId1, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 2}});
 
-    DistributionStatus distribution(kNamespace, cluster.second);
-    ASSERT_OK(
-        distribution.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 7), "NonExistentZone")));
+    ZoneInfo zoneInfo;
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 7), "NonExistentZone")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     ASSERT(balanceChunks(cluster.first, distribution, false, false).empty());
 }
 
 TEST(DistributionStatus, AddTagRangeOverlap) {
-    DistributionStatus d(kNamespace, ShardToChunksMap{});
+    ZoneInfo zInfo;
 
     // 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")));
+    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "a")));
+    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 20), BSON("x" << 30), "b")));
 
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              d.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 2), "d")));
+              zInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 2), "d")));
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              d.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 5), "d")));
+              zInfo.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 5), "d")));
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              d.addRangeToZone(ZoneRange(BSON("x" << 5), BSON("x" << 9), "d")));
+              zInfo.addRangeToZone(ZoneRange(BSON("x" << 5), BSON("x" << 9), "d")));
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              d.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "d")));
+              zInfo.addRangeToZone(ZoneRange(BSON("x" << 1), BSON("x" << 10), "d")));
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              d.addRangeToZone(ZoneRange(BSON("x" << 5), BSON("x" << 25), "d")));
+              zInfo.addRangeToZone(ZoneRange(BSON("x" << 5), BSON("x" << 25), "d")));
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              d.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 32), "d")));
+              zInfo.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 32), "d")));
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              d.addRangeToZone(ZoneRange(BSON("x" << 25), kMaxBSONKey, "d")));
+              zInfo.addRangeToZone(ZoneRange(BSON("x" << 25), kMaxBSONKey, "d")));
 }
 
 TEST(DistributionStatus, ChunkTagsSelectorWithRegularKeys) {
-    DistributionStatus d(kNamespace, ShardToChunksMap{});
-
-    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 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")));
+    DistributionStatus d(kNamespace, ShardToChunksMap{}, std::move(zInfo));
 
     {
         ChunkType chunk;
@@ -794,11 +819,12 @@ TEST(DistributionStatus, ChunkTagsSelectorWithRegularKeys) {
 }
 
 TEST(DistributionStatus, ChunkTagsSelectorWithMinMaxKeys) {
-    DistributionStatus d(kNamespace, ShardToChunksMap{});
 
-    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")));
+    ZoneInfo zInfo;
+    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << -100), "a")));
+    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << -10), BSON("x" << 10), "b")));
+    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 100), kMaxBSONKey, "c")));
+    DistributionStatus d(kNamespace, ShardToChunksMap{}, std::move(zInfo));
 
     {
         ChunkType chunk;
diff --git a/src/mongo/db/s/config/sharding_catalog_manager_zone_operations.cpp b/src/mongo/db/s/config/sharding_catalog_manager_zone_operations.cpp
index a9a19dc7a41..c0c4b36eeb0 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager_zone_operations.cpp
+++ b/src/mongo/db/s/config/sharding_catalog_manager_zone_operations.cpp
@@ -66,7 +66,7 @@ Status checkForOverlappingZonedKeyRange(OperationContext* opCtx,
                                         const ChunkRange& range,
                                         const std::string& zoneName,
                                         const KeyPattern& shardKeyPattern) {
-    DistributionStatus chunkDist(nss, ShardToChunksMap{});
+    ZoneInfo zoneInfo;
 
     auto tagStatus = configServer->exhaustiveFindOnConfig(opCtx,
                                                           kConfigPrimarySelector,
@@ -89,7 +89,7 @@ Status checkForOverlappingZonedKeyRange(OperationContext* opCtx,
         // Always extend ranges to full shard key to be compatible with tags created before
         // the zone commands were implemented.
         const auto& parsedTagDoc = tagParseStatus.getValue();
-        auto overlapStatus = chunkDist.addRangeToZone(
+        auto overlapStatus = zoneInfo.addRangeToZone(
             ZoneRange(shardKeyPattern.extendRangeBound(parsedTagDoc.getMinKey(), false),
                       shardKeyPattern.extendRangeBound(parsedTagDoc.getMaxKey(), false),
                       parsedTagDoc.getTag()));
@@ -99,7 +99,7 @@ Status checkForOverlappingZonedKeyRange(OperationContext* opCtx,
     }
 
     auto overlapStatus =
-        chunkDist.addRangeToZone(ZoneRange(range.getMin(), range.getMax(), zoneName));
+        zoneInfo.addRangeToZone(ZoneRange(range.getMin(), range.getMax(), zoneName));
     if (!overlapStatus.isOK()) {
         return overlapStatus;
     }