SERVER-82325 Config server could invariant during balancer roundr7.0.3-rc1r7.0.3

10 files changed, 496 insertions, 540 deletions

diff --git a/src/mongo/db/s/balancer/balance_stats.cpp b/src/mongo/db/s/balancer/balance_stats.cpp
index 0cb04449b04..aa644b3f929 100644
--- a/src/mongo/db/s/balancer/balance_stats.cpp
+++ b/src/mongo/db/s/balancer/balance_stats.cpp
@@ -58,7 +58,7 @@ int64_t getMaxChunkImbalanceCount(const ChunkManager& routingInfo,
     }
 
     routingInfo.forEachChunk([&zoneInfo, &chunkDistributionPerZone](auto chunk) {
-        auto zone = zoneInfo.getZoneForRange(chunk.getRange());
+        auto zone = zoneInfo.getZoneForChunk(chunk.getRange());
         chunkDistributionPerZone[zone][chunk.getShardId()] += 1;
         return true;
     });
diff --git a/src/mongo/db/s/balancer/balancer_chunk_selection_policy.h b/src/mongo/db/s/balancer/balancer_chunk_selection_policy.h
index 1560181a433..e2c54767ffb 100644
--- a/src/mongo/db/s/balancer/balancer_chunk_selection_policy.h
+++ b/src/mongo/db/s/balancer/balancer_chunk_selection_policy.h
@@ -94,6 +94,15 @@ public:
     virtual StatusWith<boost::optional<MigrateInfo>> selectSpecificChunkToMove(
         OperationContext* opCtx, const NamespaceString& nss, const ChunkType& chunk) = 0;
 
+    /**
+     * Asks the chunk selection policy to validate that the specified chunk migration is allowed
+     * given the current rules. Returns OK if the migration won't violate any rules or any other
+     * failed status otherwise.
+     */
+    virtual Status checkMoveAllowed(OperationContext* opCtx,
+                                    const ChunkType& chunk,
+                                    const ShardId& newShardId) = 0;
+
 protected:
     BalancerChunkSelectionPolicy();
 };
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 de6fa0aac63..b89302d3673 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
@@ -46,7 +46,6 @@
 #include "mongo/s/catalog/type_chunk.h"
 #include "mongo/s/catalog/type_collection.h"
 #include "mongo/s/catalog_cache.h"
-#include "mongo/s/chunk_manager.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/request_types/get_stats_for_balancing_gen.h"
 #include "mongo/s/sharding_feature_flags_gen.h"
@@ -70,6 +69,27 @@ StatusWith<DistributionStatus> createCollectionDistributionStatus(
     const NamespaceString& nss,
     const ShardStatisticsVector& allShards,
     const ChunkManager& chunkMgr) {
+    ShardToChunksMap shardToChunksMap;
+
+    // Makes sure there is an entry in shardToChunksMap for every shard, so empty shards will also
+    // be accounted for
+    for (const auto& stat : allShards) {
+        shardToChunksMap[stat.shardId];
+    }
+
+    chunkMgr.forEachChunk([&](const auto& chunkEntry) {
+        ChunkType chunk;
+        chunk.setCollectionUUID(chunkMgr.getUUID());
+        chunk.setMin(chunkEntry.getMin());
+        chunk.setMax(chunkEntry.getMax());
+        chunk.setJumbo(chunkEntry.isJumbo());
+        chunk.setShard(chunkEntry.getShardId());
+        chunk.setVersion(chunkEntry.getLastmod());
+
+        shardToChunksMap[chunkEntry.getShardId()].push_back(chunk);
+
+        return true;
+    });
 
     auto swZoneInfo =
         createCollectionZoneInfo(opCtx, nss, chunkMgr.getShardKeyPattern().getKeyPattern());
@@ -77,7 +97,7 @@ StatusWith<DistributionStatus> createCollectionDistributionStatus(
         return swZoneInfo.getStatus();
     }
 
-    return {DistributionStatus{nss, std::move(swZoneInfo.getValue()), chunkMgr}};
+    return {DistributionStatus{nss, std::move(shardToChunksMap), std::move(swZoneInfo.getValue())}};
 }
 
 stdx::unordered_map<NamespaceString, CollectionDataSizeInfoForBalancing>
@@ -562,6 +582,54 @@ BalancerChunkSelectionPolicyImpl::selectSpecificChunkToMove(OperationContext* op
     return BalancerPolicy::balanceSingleChunk(chunk, shardStats, distribution, dataSizeInfo);
 }
 
+Status BalancerChunkSelectionPolicyImpl::checkMoveAllowed(OperationContext* opCtx,
+                                                          const ChunkType& chunk,
+                                                          const ShardId& newShardId) {
+    auto shardStatsStatus = _clusterStats->getStats(opCtx);
+    if (!shardStatsStatus.isOK()) {
+        return shardStatsStatus.getStatus();
+    }
+
+    const auto catalogClient = ShardingCatalogManager::get(opCtx)->localCatalogClient();
+    const CollectionType collection = catalogClient->getCollection(
+        opCtx, chunk.getCollectionUUID(), repl::ReadConcernLevel::kLocalReadConcern);
+    const auto& nss = collection.getNss();
+
+
+    auto shardStats = std::move(shardStatsStatus.getValue());
+
+    auto routingInfoStatus =
+        Grid::get(opCtx)->catalogCache()->getShardedCollectionRoutingInfoWithPlacementRefresh(opCtx,
+                                                                                              nss);
+    if (!routingInfoStatus.isOK()) {
+        return routingInfoStatus.getStatus();
+    }
+
+    const auto& [cm, _] = routingInfoStatus.getValue();
+
+    const auto collInfoStatus = createCollectionDistributionStatus(opCtx, nss, shardStats, cm);
+    if (!collInfoStatus.isOK()) {
+        return collInfoStatus.getStatus();
+    }
+
+    const DistributionStatus& distribution = collInfoStatus.getValue();
+
+    auto newShardIterator =
+        std::find_if(shardStats.begin(),
+                     shardStats.end(),
+                     [&newShardId](const ClusterStatistics::ShardStatistics& stat) {
+                         return stat.shardId == newShardId;
+                     });
+    if (newShardIterator == shardStats.end()) {
+        return {ErrorCodes::ShardNotFound,
+                str::stream() << "Unable to find constraints information for shard " << newShardId
+                              << ". Move to this shard will be disallowed."};
+    }
+
+    return BalancerPolicy::isShardSuitableReceiver(*newShardIterator,
+                                                   distribution.getZoneForChunk(chunk));
+}
+
 StatusWith<SplitInfoVector> BalancerChunkSelectionPolicyImpl::_getSplitCandidatesForCollection(
     OperationContext* opCtx, const NamespaceString& nss, const ShardStatisticsVector& shardStats) {
     auto routingInfoStatus =
@@ -620,7 +688,7 @@ BalancerChunkSelectionPolicyImpl::_getMigrateCandidatesForCollection(
 
     const DistributionStatus& distribution = collInfoStatus.getValue();
 
-    for (const auto& zoneRangeEntry : distribution.getZoneInfo().zoneRanges()) {
+    for (const auto& zoneRangeEntry : distribution.zoneRanges()) {
         const auto& zoneRange = zoneRangeEntry.second;
 
         const auto chunkAtZoneMin = cm.findIntersectingChunkWithSimpleCollation(zoneRange.min);
diff --git a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.h b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.h
index 60c3a2ca211..ec3d4972418 100644
--- a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.h
+++ b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.h
@@ -57,6 +57,10 @@ public:
     StatusWith<boost::optional<MigrateInfo>> selectSpecificChunkToMove(
         OperationContext* opCtx, const NamespaceString& nss, const ChunkType& chunk) override;
 
+    Status checkMoveAllowed(OperationContext* opCtx,
+                            const ChunkType& chunk,
+                            const ShardId& newShardId) override;
+
 private:
     /**
      * Synchronous method, which iterates the collection's chunks and uses the zones information to
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
index a66a76e9500..8e996c95da7 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
@@ -157,8 +157,8 @@ bool areMergeable(const ChunkType& firstChunk,
                   const ChunkType& secondChunk,
                   const ZoneInfo& collectionZones) {
     return firstChunk.getShard() == secondChunk.getShard() &&
-        collectionZones.getZoneForRange(firstChunk.getRange()) ==
-        collectionZones.getZoneForRange(secondChunk.getRange()) &&
+        collectionZones.getZoneForChunk(firstChunk.getRange()) ==
+        collectionZones.getZoneForChunk(secondChunk.getRange()) &&
         SimpleBSONObjComparator::kInstance.evaluate(firstChunk.getMax() == secondChunk.getMin());
 }
 
@@ -886,8 +886,8 @@ private:
         std::list<ChunkRangeInfoIterator> siblings;
         auto canBeMoveAndMerged = [this](const ChunkRangeInfoIterator& chunkIt,
                                          const ChunkRangeInfoIterator& siblingIt) {
-            auto onSameZone = _zoneInfo.getZoneForRange(chunkIt->range) ==
-                _zoneInfo.getZoneForRange(siblingIt->range);
+            auto onSameZone = _zoneInfo.getZoneForChunk(chunkIt->range) ==
+                _zoneInfo.getZoneForChunk(siblingIt->range);
             auto destinationAvailable = chunkIt->shard == siblingIt->shard ||
                 !_shardInfos.at(siblingIt->shard).isDraining();
             return (onSameZone && destinationAvailable);
diff --git a/src/mongo/db/s/balancer/balancer_policy.cpp b/src/mongo/db/s/balancer/balancer_policy.cpp
index 3d8d56338d8..3f246116b1e 100644
--- a/src/mongo/db/s/balancer/balancer_policy.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy.cpp
@@ -57,106 +57,28 @@ using std::string;
 using std::vector;
 using namespace fmt::literals;
 
-namespace {
-
-ChunkType makeChunkType(const UUID& collUUID, const Chunk& chunk) {
-    ChunkType ct{collUUID, chunk.getRange(), chunk.getLastmod(), chunk.getShardId()};
-    ct.setJumbo(chunk.isJumbo());
-    return ct;
-}
-
-/**
- * Return a vector of zones after they have been normalized according to the given chunk
- * configuration.
- *
- * If a zone covers only partially a chunk, boundaries of that zone will be shrank so that the
- * normalized zone won't overlap with that chunk. The boundaries of a normalized zone will never
- * fall in the middle of a chunk.
- *
- * Additionally the vector will contain also zones for the "NoZone",
- */
-std::vector<ZoneRange> normalizeZones(const ChunkManager& cm, const ZoneInfo& zoneInfo) {
-    std::vector<ZoneRange> normalizedRanges;
-
-    auto lastMax = cm.getShardKeyPattern().getKeyPattern().globalMin();
-
-    for (const auto& [max, zoneRange] : zoneInfo.zoneRanges()) {
-        const auto& minChunk = cm.findIntersectingChunkWithSimpleCollation(zoneRange.min);
-        const auto gtMin =
-            SimpleBSONObjComparator::kInstance.evaluate(zoneRange.min > minChunk.getMin());
-        const auto& normalizedMin = gtMin ? minChunk.getMax() : zoneRange.min;
-
-
-        const auto& maxChunk = cm.findIntersectingChunkWithSimpleCollation(zoneRange.max);
-        const auto gtMax =
-            SimpleBSONObjComparator::kInstance.evaluate(zoneRange.max > maxChunk.getMin()) &&
-            SimpleBSONObjComparator::kInstance.evaluate(
-                zoneRange.max != cm.getShardKeyPattern().getKeyPattern().globalMax());
-        const auto& normalizedMax = gtMax ? maxChunk.getMin() : zoneRange.max;
-
-
-        if (SimpleBSONObjComparator::kInstance.evaluate(normalizedMin == normalizedMax)) {
-            // This zone does not fully contain any chunk thus we can ignore it
-            continue;
-        }
-
-        if (SimpleBSONObjComparator::kInstance.evaluate(normalizedMin != lastMax)) {
-            // The zone is not contiguous with the previous one so we add a kNoZoneRange
-            // does not fully contain any chunk so we will ignore it
-            normalizedRanges.emplace_back(lastMax, normalizedMin, ZoneInfo::kNoZoneName);
-        }
-
-        normalizedRanges.emplace_back(normalizedMin, normalizedMax, zoneRange.zone);
-        lastMax = normalizedMax;
-    }
-
-    const auto& globalMaxKey = cm.getShardKeyPattern().getKeyPattern().globalMax();
-    if (SimpleBSONObjComparator::kInstance.evaluate(lastMax != globalMaxKey)) {
-        normalizedRanges.emplace_back(lastMax, globalMaxKey, ZoneInfo::kNoZoneName);
-    }
-    return normalizedRanges;
-}
-
-}  // namespace
-
 DistributionStatus::DistributionStatus(NamespaceString nss,
-                                       ZoneInfo zoneInfo,
-                                       const ChunkManager& chunkMngr)
-    : _nss(std::move(nss)), _zoneInfo(std::move(zoneInfo)), _chunkMngr(chunkMngr) {
-
-    _normalizedZones = normalizeZones(_chunkMngr, _zoneInfo);
-
-    for (const auto& zoneRange : _normalizedZones) {
-        chunkMngr.forEachOverlappingChunk(
-            zoneRange.min, zoneRange.max, false /* isMaxInclusive */, [&](const auto& chunkInfo) {
-                _shardToZoneSizeMap[chunkInfo.getShardId()][zoneRange.zone]++;
-                return true;
-            });
-    }
-}
+                                       ShardToChunksMap shardToChunksMap,
+                                       ZoneInfo zoneInfo)
+    : _nss(std::move(nss)),
+      _shardChunks(std::move(shardToChunksMap)),
+      _zoneInfo(std::move(zoneInfo)) {}
 
 size_t DistributionStatus::numberOfChunksInShard(const ShardId& shardId) const {
-    const auto shardZonesIt = _shardToZoneSizeMap.find(shardId);
-    if (shardZonesIt == _shardToZoneSizeMap.end()) {
-        return 0;
-    }
-    size_t total = 0;
-    for (const auto& [_, numChunks] : shardZonesIt->second) {
-        total += numChunks;
-    }
-    return total;
+    const auto& shardChunks = getChunks(shardId);
+    return shardChunks.size();
 }
 
-const StringMap<size_t>& DistributionStatus::getChunksPerZoneMap(const ShardId& shardId) const {
-    static const StringMap<size_t> emptyMap;
-    const auto shardZonesIt = _shardToZoneSizeMap.find(shardId);
-    if (shardZonesIt == _shardToZoneSizeMap.end()) {
-        return emptyMap;
-    }
-    return shardZonesIt->second;
+const vector<ChunkType>& DistributionStatus::getChunks(const ShardId& shardId) const {
+    ShardToChunksMap::const_iterator i = _shardChunks.find(shardId);
+    invariant(i != _shardChunks.end());
+
+    return i->second;
 }
 
-const string ZoneInfo::kNoZoneName = "";
+string DistributionStatus::getZoneForChunk(const ChunkType& chunk) const {
+    return _zoneInfo.getZoneForChunk(chunk.getRange());
+}
 
 ZoneInfo::ZoneInfo()
     : _zoneRanges(SimpleBSONObjComparator::kInstance.makeBSONObjIndexedMap<ZoneRange>()) {}
@@ -201,18 +123,18 @@ Status ZoneInfo::addRangeToZone(const ZoneRange& range) {
     return Status::OK();
 }
 
-string ZoneInfo::getZoneForRange(const ChunkRange& chunk) const {
+string ZoneInfo::getZoneForChunk(const ChunkRange& chunk) const {
     const auto minIntersect = _zoneRanges.upper_bound(chunk.getMin());
     const auto maxIntersect = _zoneRanges.lower_bound(chunk.getMax());
 
     // We should never have a partial overlap with a chunk range. If it happens, treat it as if this
     // chunk doesn't belong to a zone
     if (minIntersect != maxIntersect) {
-        return ZoneInfo::kNoZoneName;
+        return "";
     }
 
     if (minIntersect == _zoneRanges.end()) {
-        return ZoneInfo::kNoZoneName;
+        return "";
     }
 
     const ZoneRange& intersectRange = minIntersect->second;
@@ -223,7 +145,7 @@ string ZoneInfo::getZoneForRange(const ChunkRange& chunk) const {
         return intersectRange.zone;
     }
 
-    return ZoneInfo::kNoZoneName;
+    return "";
 }
 
 /**
@@ -262,7 +184,7 @@ Status BalancerPolicy::isShardSuitableReceiver(const ClusterStatistics::ShardSta
                 str::stream() << stat.shardId << " is currently draining."};
     }
 
-    if (chunkZone != ZoneInfo::kNoZoneName && !stat.shardZones.count(chunkZone)) {
+    if (!chunkZone.empty() && !stat.shardZones.count(chunkZone)) {
         return {ErrorCodes::IllegalOperation,
                 str::stream() << stat.shardId << " is not in the correct zone " << chunkZone};
     }
@@ -272,6 +194,7 @@ Status BalancerPolicy::isShardSuitableReceiver(const ClusterStatistics::ShardSta
 
 std::tuple<ShardId, int64_t> BalancerPolicy::_getLeastLoadedReceiverShard(
     const ShardStatisticsVector& shardStats,
+    const DistributionStatus& distribution,
     const CollectionDataSizeInfoForBalancing& collDataSizeInfo,
     const string& zone,
     const stdx::unordered_set<ShardId>& availableShards) {
@@ -305,6 +228,7 @@ std::tuple<ShardId, int64_t> BalancerPolicy::_getLeastLoadedReceiverShard(
 
 std::tuple<ShardId, int64_t> BalancerPolicy::_getMostOverloadedShard(
     const ShardStatisticsVector& shardStats,
+    const DistributionStatus& distribution,
     const CollectionDataSizeInfoForBalancing& collDataSizeInfo,
     const string& chunkZone,
     const stdx::unordered_set<ShardId>& availableShards) {
@@ -380,27 +304,11 @@ boost::optional<MigrateInfo> chooseRandomMigration(stdx::unordered_set<ShardId>*
                 "fromShardId"_attr = donorShard.get(),
                 "toShardId"_attr = recipientShard.get());
 
-    const auto& randomChunk = [&] {
-        const auto numChunksOnDonorShard = distribution.numberOfChunksInShard(donorShard.get());
-        const auto rndChunkIdx = getRandomIndex(numChunksOnDonorShard);
-        ChunkType rndChunk;
-
-        int idx{0};
-        distribution.getChunkManager().forEachChunk([&](const auto& chunk) {
-            if (chunk.getShardId() == donorShard.get() && idx++ == rndChunkIdx) {
-                rndChunk = makeChunkType(distribution.getChunkManager().getUUID(), chunk);
-                return false;
-            }
-            return true;
-        });
-
-        invariant(rndChunk.getShard().isValid());
-        return rndChunk;
-    }();
-
-
-    return MigrateInfo{
-        recipientShard.get(), distribution.nss(), randomChunk, ForceJumbo::kDoNotForce};
+    const auto& chunks = distribution.getChunks(donorShard.get());
+    return MigrateInfo{recipientShard.get(),
+                       distribution.nss(),
+                       chunks[getRandomIndex(chunks.size())],
+                       ForceJumbo::kDoNotForce};
 }
 
 MigrateInfosWithReason BalancerPolicy::balance(
@@ -438,83 +346,72 @@ MigrateInfosWithReason BalancerPolicy::balance(
             if (!availableShards->count(stat.shardId))
                 continue;
 
+            const vector<ChunkType>& chunks = distribution.getChunks(stat.shardId);
+
+            if (chunks.empty())
+                continue;
+
             // Now we know we need to move to chunks off this shard, but only if permitted by the
             // zones policy
             unsigned numJumboChunks = 0;
 
-            const auto& chunksPerZoneMap = distribution.getChunksPerZoneMap(stat.shardId);
-            for (const auto& zoneIt : chunksPerZoneMap) {
-                const auto& zoneName = zoneIt.first;
-                for (const auto& zoneRange : distribution.getNormalizedZones()) {
-                    if (zoneRange.zone != zoneName) {
-                        continue;
-                    }
+            // Since we have to move all chunks, lets just do in order
+            for (const auto& chunk : chunks) {
+                if (chunk.getJumbo()) {
+                    numJumboChunks++;
+                    continue;
+                }
 
-                    distribution.getChunkManager().forEachOverlappingChunk(
-                        zoneRange.min,
-                        zoneRange.max,
-                        false /* isMaxInclusive */,
-                        [&](const auto& chunk) {
-                            if (chunk.getShardId() != stat.shardId) {
-                                return true;  // continue
-                            }
-                            if (chunk.isJumbo()) {
-                                numJumboChunks++;
-                                return true;  // continue
-                            }
-
-                            const auto [to, _] = _getLeastLoadedReceiverShard(
-                                shardStats, collDataSizeInfo, zoneName, *availableShards);
-                            if (!to.isValid()) {
-                                if (migrations.empty()) {
-                                    LOGV2_WARNING(
-                                        21889,
-                                        "Chunk {chunk} is on a draining shard, but no appropriate "
-                                        "recipient found",
-                                        "Chunk is on a draining shard, but no appropriate "
-                                        "recipient found",
-                                        "chunk"_attr = redact(
-                                            makeChunkType(distribution.getChunkManager().getUUID(),
-                                                          chunk)
-                                                .toString()));
-                                }
-                                return true;  // continue
-                            }
-                            invariant(to != stat.shardId);
-
-                            migrations.emplace_back(
-                                to,
-                                chunk.getShardId(),
-                                distribution.nss(),
-                                distribution.getChunkManager().getUUID(),
-                                chunk.getMin(),
-                                boost::none /* max */,
-                                chunk.getLastmod(),
-                                // Always force jumbo chunks to be migrated off draining shards
-                                ForceJumbo::kForceBalancer,
-                                collDataSizeInfo.maxChunkSizeBytes);
-
-                            if (firstReason == MigrationReason::none) {
-                                firstReason = MigrationReason::drain;
-                            }
-                            invariant(availableShards->erase(stat.shardId));
-                            invariant(availableShards->erase(to));
-                            return false;  // break
-                        });
+                const auto zone = distribution.getZoneForChunk(chunk);
 
+                const auto [to, _] = _getLeastLoadedReceiverShard(
+                    shardStats, distribution, collDataSizeInfo, zone, *availableShards);
+                if (!to.isValid()) {
                     if (migrations.empty()) {
-                        LOGV2_WARNING(21890,
-                                      "Unable to find any chunk to move from draining shard "
-                                      "{shardId}. numJumboChunks: {numJumboChunks}",
-                                      "Unable to find any chunk to move from draining shard",
-                                      "shardId"_attr = stat.shardId,
-                                      "numJumboChunks"_attr = numJumboChunks);
+                        LOGV2_WARNING(21889,
+                                      "Chunk {chunk} is on a draining shard, but no appropriate "
+                                      "recipient found",
+                                      "Chunk is on a draining shard, but no appropriate "
+                                      "recipient found",
+                                      "chunk"_attr = redact(chunk.toString()));
                     }
+                    continue;
+                }
 
-                    if (availableShards->size() < 2) {
-                        return std::make_pair(std::move(migrations), firstReason);
-                    }
+                invariant(to != stat.shardId);
+
+                migrations.emplace_back(
+                    to,
+                    chunk.getShard(),
+                    distribution.nss(),
+                    chunk.getCollectionUUID(),
+                    chunk.getMin(),
+                    boost::none /* max */,
+                    chunk.getVersion(),
+                    // Always force jumbo chunks to be migrated off draining shards
+                    ForceJumbo::kForceBalancer,
+                    collDataSizeInfo.maxChunkSizeBytes);
+
+                if (firstReason == MigrationReason::none) {
+                    firstReason = MigrationReason::drain;
                 }
+                invariant(availableShards->erase(stat.shardId));
+                invariant(availableShards->erase(to));
+                break;
+            }
+
+            if (migrations.empty()) {
+                availableShards->erase(stat.shardId);
+                LOGV2_WARNING(21890,
+                              "Unable to find any chunk to move from draining shard "
+                              "{shardId}. numJumboChunks: {numJumboChunks}",
+                              "Unable to find any chunk to move from draining shard",
+                              "shardId"_attr = stat.shardId,
+                              "numJumboChunks"_attr = numJumboChunks);
+            }
+
+            if (availableShards->size() < 2) {
+                return std::make_pair(std::move(migrations), firstReason);
             }
         }
     }
@@ -522,88 +419,68 @@ MigrateInfosWithReason BalancerPolicy::balance(
     // 2) Check for chunks, which are on the wrong shard and must be moved off of it
     if (!distribution.zones().empty()) {
         for (const auto& stat : shardStats) {
-
             if (!availableShards->count(stat.shardId))
                 continue;
 
-            const auto& chunksPerZoneMap = distribution.getChunksPerZoneMap(stat.shardId);
-            for (const auto& zoneIt : chunksPerZoneMap) {
-                const auto& zoneName = zoneIt.first;
+            const vector<ChunkType>& chunks = distribution.getChunks(stat.shardId);
+
+            for (const auto& chunk : chunks) {
+                const string zone = distribution.getZoneForChunk(chunk);
 
-                if (zoneName == ZoneInfo::kNoZoneName)
+                if (zone.empty())
                     continue;
 
-                if (stat.shardZones.count(zoneName))
+                if (stat.shardZones.count(zone))
                     continue;
 
-                for (const auto& zoneRange : distribution.getNormalizedZones()) {
-                    if (zoneRange.zone != zoneName) {
-                        continue;
+                if (chunk.getJumbo()) {
+                    LOGV2_WARNING(
+                        21891,
+                        "Chunk {chunk} violates zone {zone}, but it is jumbo and cannot be moved",
+                        "Chunk violates zone, but it is jumbo and cannot be moved",
+                        "chunk"_attr = redact(chunk.toString()),
+                        "zone"_attr = redact(zone));
+
+                    continue;
+                }
+
+                const auto [to, _] = _getLeastLoadedReceiverShard(
+                    shardStats, distribution, collDataSizeInfo, zone, *availableShards);
+                if (!to.isValid()) {
+                    if (migrations.empty()) {
+                        LOGV2_WARNING(21892,
+                                      "Chunk {chunk} violates zone {zone}, but no appropriate "
+                                      "recipient found",
+                                      "Chunk violates zone, but no appropriate recipient found",
+                                      "chunk"_attr = redact(chunk.toString()),
+                                      "zone"_attr = redact(zone));
                     }
-                    distribution.getChunkManager().forEachOverlappingChunk(
-                        zoneRange.min,
-                        zoneRange.max,
-                        false /* isMaxInclusive */,
-                        [&](const auto& chunk) {
-                            if (chunk.getShardId() != stat.shardId) {
-                                return true;  // continue
-                            }
-                            if (chunk.isJumbo()) {
-                                LOGV2_WARNING(
-                                    21891,
-                                    "Chunk {chunk} violates zone {zone}, but it is jumbo and "
-                                    "cannot be moved",
-                                    "Chunk violates zone, but it is jumbo and cannot be moved",
-                                    "chunk"_attr =
-                                        redact(makeChunkType(
-                                                   distribution.getChunkManager().getUUID(), chunk)
-                                                   .toString()),
-                                    "zone"_attr = redact(zoneName));
-                                return true;  // continue
-                            }
-
-                            const auto [to, _] = _getLeastLoadedReceiverShard(
-                                shardStats, collDataSizeInfo, zoneName, *availableShards);
-                            if (!to.isValid()) {
-                                if (migrations.empty()) {
-                                    LOGV2_WARNING(
-                                        21892,
-                                        "Chunk {chunk} violates zone {zone}, but no appropriate "
-                                        "recipient found",
-                                        "Chunk violates zone, but no appropriate recipient found",
-                                        "chunk"_attr = redact(
-                                            makeChunkType(distribution.getChunkManager().getUUID(),
-                                                          chunk)
-                                                .toString()),
-                                        "zone"_attr = redact(zoneName));
-                                }
-                                return true;  // continue
-                            }
-                            invariant(to != stat.shardId);
-
-                            migrations.emplace_back(to,
-                                                    chunk.getShardId(),
-                                                    distribution.nss(),
-                                                    distribution.getChunkManager().getUUID(),
-                                                    chunk.getMin(),
-                                                    boost::none /* max */,
-                                                    chunk.getLastmod(),
-                                                    forceJumbo ? ForceJumbo::kForceBalancer
-                                                               : ForceJumbo::kDoNotForce,
-                                                    collDataSizeInfo.maxChunkSizeBytes);
-
-                            if (firstReason == MigrationReason::none) {
-                                firstReason = MigrationReason::zoneViolation;
-                            }
-                            invariant(availableShards->erase(stat.shardId));
-                            invariant(availableShards->erase(to));
-                            return false;  // break
-                        });
+                    continue;
                 }
 
-                if (availableShards->size() < 2) {
-                    return std::make_pair(std::move(migrations), firstReason);
+                invariant(to != stat.shardId);
+
+                migrations.emplace_back(to,
+                                        chunk.getShard(),
+                                        distribution.nss(),
+                                        chunk.getCollectionUUID(),
+                                        chunk.getMin(),
+                                        boost::none /* max */,
+                                        chunk.getVersion(),
+                                        forceJumbo ? ForceJumbo::kForceBalancer
+                                                   : ForceJumbo::kDoNotForce,
+                                        collDataSizeInfo.maxChunkSizeBytes);
+
+                if (firstReason == MigrationReason::none) {
+                    firstReason = MigrationReason::zoneViolation;
                 }
+                invariant(availableShards->erase(stat.shardId));
+                invariant(availableShards->erase(to));
+                break;
+            }
+
+            if (availableShards->size() < 2) {
+                return std::make_pair(std::move(migrations), firstReason);
             }
         }
     }
@@ -611,14 +488,14 @@ MigrateInfosWithReason BalancerPolicy::balance(
     // 3) for each zone balance
 
     vector<string> zonesPlusEmpty(distribution.zones().begin(), distribution.zones().end());
-    zonesPlusEmpty.push_back(ZoneInfo::kNoZoneName);
+    zonesPlusEmpty.push_back("");
 
     for (const auto& zone : zonesPlusEmpty) {
         size_t numShardsInZone = 0;
         int64_t totalDataSizeOfShardsWithZone = 0;
 
         for (const auto& stat : shardStats) {
-            if (zone == ZoneInfo::kNoZoneName || stat.shardZones.count(zone)) {
+            if (zone.empty() || stat.shardZones.count(zone)) {
                 const auto& shardSizeIt = collDataSizeInfo.shardToDataSizeMap.find(stat.shardId);
                 if (shardSizeIt == collDataSizeInfo.shardToDataSizeMap.end()) {
                     // Skip if stats not available (may happen if add|remove shard during a round)
@@ -632,16 +509,17 @@ MigrateInfosWithReason BalancerPolicy::balance(
         // Skip zones which have no shards assigned to them. This situation is not harmful, but
         // should not be possible so warn the operator to correct it.
         if (numShardsInZone == 0) {
-            if (zone != ZoneInfo::kNoZoneName) {
-                LOGV2_WARNING(21893,
-                              "Zone {zone} in collection {namespace} has no assigned shards and "
-                              "chunks which fall into it cannot be balanced. This should be "
-                              "corrected by either assigning shards to the zone or by deleting it.",
-                              "Zone in collection has no assigned shards and chunks which fall "
-                              "into it cannot be balanced. This should be corrected by either "
-                              "assigning shards to the zone or by deleting it.",
-                              "zone"_attr = redact(zone),
-                              logAttrs(distribution.nss()));
+            if (!zone.empty()) {
+                LOGV2_WARNING(
+                    21893,
+                    "Zone {zone} in collection {namespace} has no assigned shards and chunks "
+                    "which fall into it cannot be balanced. This should be corrected by either "
+                    "assigning shards to the zone or by deleting it.",
+                    "Zone in collection has no assigned shards and chunks which fall into it "
+                    "cannot be balanced. This should be corrected by either assigning shards "
+                    "to the zone or by deleting it.",
+                    "zone"_attr = redact(zone),
+                    logAttrs(distribution.nss()));
             }
             continue;
         }
@@ -683,18 +561,10 @@ boost::optional<MigrateInfo> BalancerPolicy::balanceSingleChunk(
     const ShardStatisticsVector& shardStats,
     const DistributionStatus& distribution,
     const CollectionDataSizeInfoForBalancing& collDataSizeInfo) {
-    const auto& zone = distribution.getZoneInfo().getZoneForRange(chunk.getRange());
-
-    stdx::unordered_set<ShardId> availableShards;
-    std::transform(shardStats.begin(),
-                   shardStats.end(),
-                   std::inserter(availableShards, availableShards.end()),
-                   [](const ClusterStatistics::ShardStatistics& shardStatistics) -> ShardId {
-                       return shardStatistics.shardId;
-                   });
+    const string zone = distribution.getZoneForChunk(chunk);
 
     const auto [newShardId, _] = _getLeastLoadedReceiverShard(
-        shardStats, collDataSizeInfo, zone, stdx::unordered_set<ShardId>());
+        shardStats, distribution, collDataSizeInfo, zone, stdx::unordered_set<ShardId>());
     if (!newShardId.isValid() || newShardId == chunk.getShard()) {
         return boost::optional<MigrateInfo>();
     }
@@ -712,12 +582,12 @@ bool BalancerPolicy::_singleZoneBalanceBasedOnDataSize(
     stdx::unordered_set<ShardId>* availableShards,
     ForceJumbo forceJumbo) {
     const auto [from, fromSize] =
-        _getMostOverloadedShard(shardStats, collDataSizeInfo, zone, *availableShards);
+        _getMostOverloadedShard(shardStats, distribution, collDataSizeInfo, zone, *availableShards);
     if (!from.isValid())
         return false;
 
-    const auto [to, toSize] =
-        _getLeastLoadedReceiverShard(shardStats, collDataSizeInfo, zone, *availableShards);
+    const auto [to, toSize] = _getLeastLoadedReceiverShard(
+        shardStats, distribution, collDataSizeInfo, zone, *availableShards);
     if (!to.isValid()) {
         if (migrations->empty()) {
             LOGV2(6581600, "No available shards to take chunks for zone", "zone"_attr = zone);
@@ -750,50 +620,33 @@ bool BalancerPolicy::_singleZoneBalanceBasedOnDataSize(
         return false;
     }
 
+    const vector<ChunkType>& chunks = distribution.getChunks(from);
+
     unsigned numJumboChunks = 0;
-    bool chunkFound = false;
 
-    const auto& fromShardId = from;
-    const auto& toShardId = to;
+    for (const auto& chunk : chunks) {
+        if (distribution.getZoneForChunk(chunk) != zone)
+            continue;
 
-    for (const auto& zoneRange : distribution.getNormalizedZones()) {
-        if (zoneRange.zone != zone) {
+        if (chunk.getJumbo()) {
+            numJumboChunks++;
             continue;
         }
 
-        distribution.getChunkManager().forEachOverlappingChunk(
-            zoneRange.min, zoneRange.max, false /* isMaxInclusive */, [&](const auto& chunk) {
-                if (chunk.getShardId() != fromShardId) {
-                    return true;  // continue
-                }
-
-                if (chunk.isJumbo()) {
-                    numJumboChunks++;
-                    return true;  // continue
-                }
-
-                migrations->emplace_back(toShardId,
-                                         chunk.getShardId(),
-                                         distribution.nss(),
-                                         distribution.getChunkManager().getUUID(),
-                                         chunk.getMin(),
-                                         boost::none /* max */,
-                                         chunk.getLastmod(),
-                                         forceJumbo,
-                                         collDataSizeInfo.maxChunkSizeBytes);
-                invariant(availableShards->erase(chunk.getShardId()));
-                invariant(availableShards->erase(toShardId));
-                chunkFound = true;
-                return false;  // break
-            });
-
-        if (chunkFound) {
-            return chunkFound;
-        }
+        migrations->emplace_back(to,
+                                 chunk.getShard(),
+                                 distribution.nss(),
+                                 chunk.getCollectionUUID(),
+                                 chunk.getMin(),
+                                 boost::none /* max */,
+                                 chunk.getVersion(),
+                                 forceJumbo,
+                                 collDataSizeInfo.maxChunkSizeBytes);
+        invariant(availableShards->erase(chunk.getShard()));
+        invariant(availableShards->erase(to));
+        return true;
     }
 
-    invariant(!chunkFound);
-
     if (numJumboChunks) {
         LOGV2_WARNING(6581602,
                       "Shard has only jumbo chunks for this collection and cannot be balanced",
diff --git a/src/mongo/db/s/balancer/balancer_policy.h b/src/mongo/db/s/balancer/balancer_policy.h
index 7a5f1e49a79..18ae3828212 100644
--- a/src/mongo/db/s/balancer/balancer_policy.h
+++ b/src/mongo/db/s/balancer/balancer_policy.h
@@ -39,12 +39,13 @@
 #include "mongo/db/s/balancer/cluster_statistics.h"
 #include "mongo/db/shard_id.h"
 #include "mongo/s/catalog/type_chunk.h"
-#include "mongo/s/chunk_manager.h"
 #include "mongo/s/request_types/move_range_request_gen.h"
 #include "mongo/s/shard_version.h"
 #include "mongo/util/concurrency/with_lock.h"
 
 namespace mongo {
+
+
 struct ZoneRange {
     ZoneRange(const BSONObj& a_min, const BSONObj& a_max, const std::string& _zone);
 
@@ -193,7 +194,7 @@ typedef stdx::variant<Status, StatusWith<DataSizeResponse>, StatusWith<NumMerged
     BalancerStreamActionResponse;
 
 typedef std::vector<ClusterStatistics::ShardStatistics> ShardStatisticsVector;
-typedef std::map<ShardId, StringMap<size_t>> ShardToZoneSizeMap;
+typedef std::map<ShardId, std::vector<ChunkType>> ShardToChunksMap;
 
 /*
  * Keeps track of info needed for data size aware balancing.
@@ -212,8 +213,6 @@ struct CollectionDataSizeInfoForBalancing {
  */
 class ZoneInfo {
 public:
-    static const std::string kNoZoneName;
-
     ZoneInfo();
     ZoneInfo(ZoneInfo&&) = default;
 
@@ -231,10 +230,10 @@ public:
     }
 
     /**
-     * Using the set of zones added so far, returns what zone corresponds to the specified range.
+     * Using the set of zones added so far, returns what zone corresponds to the specified chunk.
      * Returns an empty string if the chunk doesn't fall into any zone.
      */
-    std::string getZoneForRange(const ChunkRange& chunkRange) const;
+    std::string getZoneForChunk(const ChunkRange& chunkRange) const;
 
     /**
      * Returns all zone ranges defined.
@@ -243,14 +242,6 @@ public:
         return _zoneRanges;
     }
 
-    const ZoneRange& getZoneRange(const std::string& zoneName) const {
-        for (const auto& [_, zoneRange] : _zoneRanges) {
-            if (zoneRange.zone == zoneName)
-                return zoneRange;
-        }
-        MONGO_UNREACHABLE;
-    }
-
 private:
     // Map of zone max key to the zone description
     BSONObjIndexedMap<ZoneRange> _zoneRanges;
@@ -277,7 +268,7 @@ class DistributionStatus final {
     DistributionStatus& operator=(const DistributionStatus&) = delete;
 
 public:
-    DistributionStatus(NamespaceString nss, ZoneInfo zoneInfo, const ChunkManager& chunkMngr);
+    DistributionStatus(NamespaceString nss, ShardToChunksMap shardToChunksMap, ZoneInfo zoneInfo);
     DistributionStatus(DistributionStatus&&) = default;
     ~DistributionStatus() {}
 
@@ -289,45 +280,57 @@ 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 number of chunks in the specified shard.
      */
     size_t numberOfChunksInShard(const ShardId& shardId) const;
 
     /**
-     * Returns all zones defined for the collection.
+     * Returns all chunks for the specified shard.
      */
-    const std::set<std::string>& zones() const {
-        return _zoneInfo.allZones();
-    }
+    const std::vector<ChunkType>& getChunks(const ShardId& shardId) const;
 
-    const ChunkManager& getChunkManager() const {
-        return _chunkMngr;
+    /**
+     * Returns all zone ranges defined for the collection.
+     */
+    const BSONObjIndexedMap<ZoneRange>& zoneRanges() const {
+        return _zoneInfo.zoneRanges();
     }
 
-    const std::vector<ZoneRange>& getNormalizedZones() const {
-        return _normalizedZones;
+    /**
+     * Returns all zones defined for the collection.
+     */
+    const std::set<std::string>& zones() const {
+        return _zoneInfo.allZones();
     }
 
-    const ZoneInfo& getZoneInfo() const {
+    /**
+     * Direct access to zone info
+     */
+    ZoneInfo& zoneInfo() {
         return _zoneInfo;
     }
 
-    const StringMap<size_t>& getChunksPerZoneMap(const ShardId& shardId) const;
+    /**
+     * Using the set of zones defined for the collection, returns what zone corresponds to the
+     * specified chunk. If the chunk doesn't fall into any zone returns the empty string.
+     */
+    std::string getZoneForChunk(const ChunkType& chunk) const;
 
 private:
     // Namespace for which this distribution applies
     NamespaceString _nss;
 
-    // Map that tracks how many chunks every shard is owning in each zone
-    // shardId -> zoneName -> numChunks
-    ShardToZoneSizeMap _shardToZoneSizeMap;
+    // Map of what chunks are owned by each shard
+    ShardToChunksMap _shardChunks;
 
     // Info for zones.
     ZoneInfo _zoneInfo;
-
-    std::vector<ZoneRange> _normalizedZones;
-
-    ChunkManager _chunkMngr;
 };
 
 class BalancerPolicy {
@@ -379,6 +382,7 @@ private:
      */
     static std::tuple<ShardId, int64_t> _getLeastLoadedReceiverShard(
         const ShardStatisticsVector& shardStats,
+        const DistributionStatus& distribution,
         const CollectionDataSizeInfoForBalancing& collDataSizeInfo,
         const std::string& zone,
         const stdx::unordered_set<ShardId>& availableShards);
@@ -389,6 +393,7 @@ private:
      */
     static std::tuple<ShardId, int64_t> _getMostOverloadedShard(
         const ShardStatisticsVector& shardStats,
+        const DistributionStatus& distribution,
         const CollectionDataSizeInfoForBalancing& collDataSizeInfo,
         const std::string& zone,
         const stdx::unordered_set<ShardId>& availableShards);
diff --git a/src/mongo/db/s/balancer/balancer_policy_test.cpp b/src/mongo/db/s/balancer/balancer_policy_test.cpp
index c3dea8cf9d6..58705275401 100644
--- a/src/mongo/db/s/balancer/balancer_policy_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy_test.cpp
@@ -45,11 +45,9 @@ using std::stringstream;
 using std::vector;
 
 using ShardStatistics = ClusterStatistics::ShardStatistics;
-typedef std::map<ShardId, std::vector<ChunkType>> ShardToChunksMap;
 
 const auto emptyZoneSet = std::set<std::string>();
 const std::string emptyShardVersion = "";
-const auto kConfigId = ShardId("config");
 const auto kShardId0 = ShardId("shard0");
 const auto kShardId1 = ShardId("shard1");
 const auto kShardId2 = ShardId("shard2");
@@ -57,41 +55,6 @@ const auto kShardId3 = ShardId("shard3");
 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;
-}
-
-RoutingTableHistory makeRoutingTable(const std::vector<ChunkType>& chunks) {
-
-    return RoutingTableHistory::makeNew(kNamespace,
-                                        collUUID(),
-                                        kSKeyPattern,
-                                        nullptr /* defaultCollator */,
-                                        false /* unique */,
-                                        kCollEpoch,
-                                        kCollTimestamp,
-                                        boost::none /* timeseriesFields */,
-                                        boost::none /* reshardingFields */,
-                                        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)}, kCollTimestamp};
-}
-
-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
@@ -100,7 +63,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<ShardStatistics>& statsVector) {
 
     // Distribute one chunk per shard, no matter the owned data size.
@@ -111,9 +74,10 @@ std::pair<std::pair<ShardStatisticsVector, ShardToChunksMap>, ChunkManager> gene
 
     int64_t currentChunk = 0;
 
-    ChunkVersion chunkVersion({kCollEpoch, kCollTimestamp}, {1, 0});
+    ChunkVersion chunkVersion({OID::gen(), Timestamp(1, 1)}, {1, 0});
+    const UUID uuid = UUID::gen();
 
-    std::vector<ChunkType> chunks;
+    const KeyPattern shardKeyPattern(BSON("x" << 1));
 
     for (const auto& shard : statsVector) {
         // Ensure that an entry is created
@@ -121,23 +85,20 @@ std::pair<std::pair<ShardStatisticsVector, ShardToChunksMap>, ChunkManager> gene
 
         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));
         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));
+    return std::make_pair(std::move(shardStats), std::move(chunkMap));
 }
 
 stdx::unordered_set<ShardId> getAllShardIds(const ShardStatisticsVector& shardStats) {
@@ -174,7 +135,7 @@ MigrateInfosWithReason balanceChunks(const ShardStatisticsVector& shardStats,
 }
 
 TEST(BalancerPolicy, Basic) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          4 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -189,8 +150,8 @@ TEST(BalancerPolicy, Basic) {
                                          emptyShardVersion,
                                          ShardStatistics::use_bytes_t())});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = 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);
@@ -199,7 +160,7 @@ TEST(BalancerPolicy, Basic) {
 }
 
 TEST(BalancerPolicy, SmallSingleChunkShouldNotMove) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -209,19 +170,24 @@ TEST(BalancerPolicy, SmallSingleChunkShouldNotMove) {
                          ShardStatistics(kShardId1, 0, false, emptyZoneSet, emptyShardVersion)});
 
     {
-        auto [migrations, reason] = balanceChunks(cluster.first, makeDistStatus(cm), true, false);
+        auto [migrations, reason] = balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), 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, ZoneInfo()),
+                          false,
+                          false);
         ASSERT(migrations.empty());
         ASSERT_EQ(MigrationReason::none, reason);
     }
 }
 
 TEST(BalancerPolicy, BalanceThresholdObeyed) {
-    auto [cluster, cm] = generateCluster({
+    auto cluster = generateCluster({
         ShardStatistics(kShardId0,
                         2 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                         false,
@@ -250,19 +216,24 @@ TEST(BalancerPolicy, BalanceThresholdObeyed) {
     });
 
     {
-        auto [migrations, reason] = balanceChunks(cluster.first, makeDistStatus(cm), true, false);
+        auto [migrations, reason] = balanceChunks(
+            cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), 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, ZoneInfo()),
+                          false,
+                          false);
         ASSERT(migrations.empty());
         ASSERT_EQ(MigrationReason::none, reason);
     }
 }
 
 TEST(BalancerPolicy, ParallelBalancing) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          4 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -278,8 +249,8 @@ TEST(BalancerPolicy, ParallelBalancing) {
                          ShardStatistics(kShardId2, 0, false, emptyZoneSet, emptyShardVersion),
                          ShardStatistics(kShardId3, 0, false, emptyZoneSet, emptyShardVersion)});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -293,7 +264,7 @@ TEST(BalancerPolicy, ParallelBalancing) {
 }
 
 TEST(BalancerPolicy, ParallelBalancingDoesNotScheduleMigrationsOnShardsAboveTheThreshold) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          100 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -321,8 +292,8 @@ TEST(BalancerPolicy, ParallelBalancingDoesNotScheduleMigrationsOnShardsAboveTheT
                          ShardStatistics(kShardId4, 0, false, emptyZoneSet, emptyShardVersion),
                          ShardStatistics(kShardId5, 0, false, emptyZoneSet, emptyShardVersion)});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -336,7 +307,7 @@ TEST(BalancerPolicy, ParallelBalancingDoesNotScheduleMigrationsOnShardsAboveTheT
 }
 
 TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseSourceShardsWithMoveNecessary) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          8 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -359,7 +330,7 @@ TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseSourceShardsWithMoveNe
     availableShards.erase(kShardId0);
     const auto [migrations, reason] =
         BalancerPolicy::balance(shardStats,
-                                makeDistStatus(cm),
+                                DistributionStatus(kNamespace, cluster.second, ZoneInfo()),
                                 buildDataSizeInfoForBalancingFromShardStats(shardStats),
                                 &availableShards,
                                 false);
@@ -372,7 +343,7 @@ TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseSourceShardsWithMoveNe
 }
 
 TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseDestinationShards) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          4 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -400,7 +371,7 @@ TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseDestinationShards) {
     availableShards.erase(kShardId2);
     const auto [migrations, reason] =
         BalancerPolicy::balance(shardStats,
-                                makeDistStatus(cm),
+                                DistributionStatus(kNamespace, cluster.second, ZoneInfo()),
                                 buildDataSizeInfoForBalancingFromShardStats(shardStats),
                                 &availableShards,
                                 false);
@@ -413,7 +384,7 @@ TEST(BalancerPolicy, ParallelBalancingNotSchedulingOnInUseDestinationShards) {
 }
 
 TEST(BalancerPolicy, ParallelBalancingDoesNotMoveDataFromShardsBelowIdealZoneSize) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          100 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -434,8 +405,8 @@ TEST(BalancerPolicy, ParallelBalancingDoesNotMoveDataFromShardsBelowIdealZoneSiz
                                          ShardStatistics::use_bytes_t()),
                          ShardStatistics(kShardId3, 0, false, emptyZoneSet, emptyShardVersion)});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT_EQ(1U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -445,7 +416,7 @@ TEST(BalancerPolicy, ParallelBalancingDoesNotMoveDataFromShardsBelowIdealZoneSiz
 }
 
 TEST(BalancerPolicy, JumboChunksNotMoved) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          4 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -454,30 +425,15 @@ TEST(BalancerPolicy, JumboChunksNotMoved) {
                                          ShardStatistics::use_bytes_t()),
                          ShardStatistics(kShardId1, 0, false, emptyZoneSet, emptyShardVersion)});
 
+    cluster.second[kShardId0][0].setJumbo(true);
 
-    // construct a new chunk map where the following chunk is jumbo
-    const auto& jumboChunk = cluster.second[kShardId0][0];
-
-    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(true);
-        else
-            ct.setJumbo(false);
-        chunks.emplace_back(std::move(ct));
-        return true;
-    });
-
-    auto newCm = makeChunkManager(chunks);
-
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(newCm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT(migrations.empty());
 }
 
 TEST(BalancerPolicy, JumboChunksNotMovedParallel) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          4 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -493,32 +449,18 @@ TEST(BalancerPolicy, JumboChunksNotMovedParallel) {
                                          ShardStatistics::use_bytes_t()),
                          ShardStatistics(kShardId3, 0, false, emptyZoneSet, emptyShardVersion)});
 
-    // construct a new chunk map where all the following chunks are jumbo
-    const auto& jumboChunk0 = cluster.second[kShardId0][0];
-    const auto& jumboChunk1 = cluster.second[kShardId2][0];
-
-    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(true);
-        else
-            ct.setJumbo(false);
-        chunks.emplace_back(std::move(ct));
-        return true;
-    });
+    cluster.second[kShardId0][0].setJumbo(true);
 
-    auto newCm = makeChunkManager(chunks);
+    cluster.second[kShardId2][0].setJumbo(true);
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(newCm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT(migrations.empty());
 }
 
 TEST(BalancerPolicy, DrainingFromShardWithFewData) {
     // shard1 is draining and chunks will go to shard0, even though it has a lot more data
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          20 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false /* draining */,
@@ -532,8 +474,8 @@ TEST(BalancerPolicy, DrainingFromShardWithFewData) {
                                          emptyShardVersion,
                                          ShardStatistics::use_bytes_t())});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId1, migrations[0].from);
     ASSERT_EQ(kShardId0, migrations[0].to);
@@ -543,7 +485,7 @@ TEST(BalancerPolicy, DrainingFromShardWithFewData) {
 
 TEST(BalancerPolicy, DrainingSingleChunkPerShard) {
     // shard0 and shard2 are draining and chunks will go to shard1 and shard3 in parallel
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          2 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          true,
@@ -559,8 +501,8 @@ TEST(BalancerPolicy, DrainingSingleChunkPerShard) {
                                          ShardStatistics::use_bytes_t()),
                          ShardStatistics(kShardId3, 0, false, emptyZoneSet, emptyShardVersion)});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT_EQ(2U, migrations.size());
 
     ASSERT_EQ(kShardId0, migrations[0].from);
@@ -576,7 +518,7 @@ TEST(BalancerPolicy, DrainingSingleChunkPerShard) {
 TEST(BalancerPolicy, DrainingMultipleShardsFirstOneSelected) {
     // shard0 and shard1 are both draining with very little data in them and chunks will go to
     // shard2, even though it has a lot more data that the other two
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          50 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false /* draining */,
@@ -596,8 +538,8 @@ TEST(BalancerPolicy, DrainingMultipleShardsFirstOneSelected) {
                                          emptyShardVersion,
                                          ShardStatistics::use_bytes_t())});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
 
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].to);
@@ -606,7 +548,7 @@ TEST(BalancerPolicy, DrainingMultipleShardsFirstOneSelected) {
 
 TEST(BalancerPolicy, DrainingMultipleShardsWontAcceptMigrations) {
     // shard0 has many data, but can't move them to shard1 or shard2 because they are draining
-    auto [cluster, cm] = generateCluster(
+    auto cluster = generateCluster(
         {ShardStatistics(kShardId0,
                          20 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                          false /* draining */,
@@ -616,14 +558,14 @@ TEST(BalancerPolicy, DrainingMultipleShardsWontAcceptMigrations) {
          ShardStatistics(kShardId1, 0, true /* draining */, emptyZoneSet, emptyShardVersion),
          ShardStatistics(kShardId2, 0, true /* draining */, emptyZoneSet, emptyShardVersion)});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT_EQ(1U, migrations.size());
     ASSERT_EQ(kShardId0, migrations[0].to);
 }
 
 TEST(BalancerPolicy, DrainingSingleAppropriateShardFoundDueToZone) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          2 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -646,7 +588,7 @@ TEST(BalancerPolicy, DrainingSingleAppropriateShardFoundDueToZone) {
     ZoneInfo zoneInfo;
     ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
         cluster.second[kShardId2][0].getMin(), cluster.second[kShardId2][0].getMax(), "LAX")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(1U, migrations.size());
@@ -657,7 +599,7 @@ TEST(BalancerPolicy, DrainingSingleAppropriateShardFoundDueToZone) {
 }
 
 TEST(BalancerPolicy, DrainingNoAppropriateShardsFoundDueToZone) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          2 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -680,14 +622,14 @@ TEST(BalancerPolicy, DrainingNoAppropriateShardsFoundDueToZone) {
     ZoneInfo zoneInfo;
     ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(
         cluster.second[kShardId2][0].getMin(), cluster.second[kShardId2][0].getMax(), "SEA")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
 
     const auto [migrations, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT(migrations.empty());
 }
 
 TEST(BalancerPolicy, NoBalancingDueToAllNodesDraining) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          2 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          true,
@@ -701,14 +643,14 @@ TEST(BalancerPolicy, NoBalancingDueToAllNodesDraining) {
                                          emptyShardVersion,
                                          ShardStatistics::use_bytes_t())});
 
-    const auto [migrations, reason] =
-        balanceChunks(cluster.first, makeDistStatus(cm), false, false);
+    const auto [migrations, reason] = balanceChunks(
+        cluster.first, DistributionStatus(kNamespace, cluster.second, ZoneInfo()), false, false);
     ASSERT(migrations.empty());
 }
 
 TEST(BalancerPolicy, BalancerRespectsZonesWhenDraining) {
     // shard1 drains the proper chunk to shard0, even though it is more loaded than shard2
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          5 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -729,9 +671,9 @@ TEST(BalancerPolicy, BalancerRespectsZonesWhenDraining) {
                                          ShardStatistics::use_bytes_t())});
 
     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));
+    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, reason] = balanceChunks(cluster.first, distribution, false, false);
     ASSERT_EQ(1U, migrations.size());
@@ -743,7 +685,7 @@ TEST(BalancerPolicy, BalancerRespectsZonesWhenDraining) {
 
 TEST(BalancerPolicy, BalancerZoneAlreadyBalanced) {
     // Chunks are balanced across shards for the zone.
-    auto [cluster, cm] = generateCluster({
+    auto cluster = generateCluster({
         ShardStatistics(kShardId0,
                         3 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                         false,
@@ -759,14 +701,13 @@ TEST(BalancerPolicy, BalancerZoneAlreadyBalanced) {
     });
 
     ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), kSKeyPattern.globalMax(), "a")));
-    const auto distribution = makeDistStatus(cm, std::move(zoneInfo));
+    ASSERT_OK(zoneInfo.addRangeToZone(ZoneRange(kMinBSONKey, kMaxBSONKey, "a")));
+    DistributionStatus distribution(kNamespace, cluster.second, std::move(zoneInfo));
     ASSERT(balanceChunks(cluster.first, distribution, false, false).first.empty());
 }
 
 TEST(BalancerPolicy, BalancerHandlesNoShardsWithZone) {
-    auto [cluster, cm] =
+    auto cluster =
         generateCluster({ShardStatistics(kShardId0,
                                          5 * ChunkSizeSettingsType::kDefaultMaxChunkSizeBytes,
                                          false,
@@ -781,9 +722,8 @@ TEST(BalancerPolicy, BalancerHandlesNoShardsWithZone) {
                                          ShardStatistics::use_bytes_t())});
 
     ZoneInfo zoneInfo;
-    ASSERT_OK(zoneInfo.addRangeToZone(
-        ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 7), "NonExistentZone")));
-    const auto distribution = makeDistStatus(cm, std::move(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).first.empty());
 }
@@ -796,7 +736,7 @@ TEST(DistributionStatus, AddZoneRangeOverlap) {
     ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 20), BSON("x" << 30), "b")));
 
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
-              zInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << 2), "d")));
+              zInfo.addRangeToZone(ZoneRange(kMinBSONKey, BSON("x" << 2), "d")));
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
               zInfo.addRangeToZone(ZoneRange(BSON("x" << -1), BSON("x" << 5), "d")));
     ASSERT_EQ(ErrorCodes::RangeOverlapConflict,
@@ -808,7 +748,7 @@ TEST(DistributionStatus, AddZoneRangeOverlap) {
     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")));
+              zInfo.addRangeToZone(ZoneRange(BSON("x" << 25), kMaxBSONKey, "d")));
 }
 
 TEST(DistributionStatus, ChunkZonesSelectorWithRegularKeys) {
@@ -818,35 +758,129 @@ TEST(DistributionStatus, ChunkZonesSelectorWithRegularKeys) {
     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.getZoneForRange({kSKeyPattern.globalMin(), BSON("x" << 1)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName, zInfo.getZoneForRange({BSON("x" << 0), BSON("x" << 1)}));
-    ASSERT_EQUALS("a", zInfo.getZoneForRange({BSON("x" << 1), BSON("x" << 5)}));
-    ASSERT_EQUALS("b", zInfo.getZoneForRange({BSON("x" << 10), BSON("x" << 20)}));
-    ASSERT_EQUALS("b", zInfo.getZoneForRange({BSON("x" << 15), BSON("x" << 20)}));
-    ASSERT_EQUALS("c", zInfo.getZoneForRange({BSON("x" << 25), BSON("x" << 30)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName, zInfo.getZoneForRange({BSON("x" << 35), BSON("x" << 40)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForRange({BSON("x" << 30), kSKeyPattern.globalMax()}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForRange({BSON("x" << 40), kSKeyPattern.globalMax()}));
+    DistributionStatus d(kNamespace, ShardToChunksMap{}, std::move(zInfo));
+
+    {
+        ChunkType chunk;
+        chunk.setMin(kMinBSONKey);
+        chunk.setMax(BSON("x" << 1));
+        ASSERT_EQUALS("", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 0));
+        chunk.setMax(BSON("x" << 1));
+        ASSERT_EQUALS("", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 1));
+        chunk.setMax(BSON("x" << 5));
+        ASSERT_EQUALS("a", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 10));
+        chunk.setMax(BSON("x" << 20));
+        ASSERT_EQUALS("b", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 15));
+        chunk.setMax(BSON("x" << 20));
+        ASSERT_EQUALS("b", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 25));
+        chunk.setMax(BSON("x" << 30));
+        ASSERT_EQUALS("c", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 35));
+        chunk.setMax(BSON("x" << 40));
+        ASSERT_EQUALS("", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 30));
+        chunk.setMax(kMaxBSONKey);
+        ASSERT_EQUALS("", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 40));
+        chunk.setMax(kMaxBSONKey);
+        ASSERT_EQUALS("", d.getZoneForChunk(chunk));
+    }
 }
 
 TEST(DistributionStatus, ChunkZonesSelectorWithMinMaxKeys) {
     ZoneInfo zInfo;
-    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(kSKeyPattern.globalMin(), BSON("x" << -100), "a")));
+
+    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), kSKeyPattern.globalMax(), "c")));
-
-    ASSERT_EQUALS("a", zInfo.getZoneForRange({kSKeyPattern.globalMin(), BSON("x" << -100)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForRange({BSON("x" << -100), BSON("x" << -11)}));
-    ASSERT_EQUALS("b", zInfo.getZoneForRange({BSON("x" << -10), BSON("x" << 0)}));
-    ASSERT_EQUALS("b", zInfo.getZoneForRange({BSON("x" << 0), BSON("x" << 10)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName, zInfo.getZoneForRange({BSON("x" << 10), BSON("x" << 20)}));
-    ASSERT_EQUALS(ZoneInfo::kNoZoneName,
-                  zInfo.getZoneForRange({BSON("x" << 10), BSON("x" << 100)}));
-    ASSERT_EQUALS("c", zInfo.getZoneForRange({BSON("x" << 200), kSKeyPattern.globalMax()}));
+    ASSERT_OK(zInfo.addRangeToZone(ZoneRange(BSON("x" << 100), kMaxBSONKey, "c")));
+
+    DistributionStatus d(kNamespace, ShardToChunksMap{}, std::move(zInfo));
+
+    {
+        ChunkType chunk;
+        chunk.setMin(kMinBSONKey);
+        chunk.setMax(BSON("x" << -100));
+        ASSERT_EQUALS("a", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << -100));
+        chunk.setMax(BSON("x" << -11));
+        ASSERT_EQUALS("", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << -10));
+        chunk.setMax(BSON("x" << 0));
+        ASSERT_EQUALS("b", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 0));
+        chunk.setMax(BSON("x" << 10));
+        ASSERT_EQUALS("b", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 10));
+        chunk.setMax(BSON("x" << 20));
+        ASSERT_EQUALS("", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 10));
+        chunk.setMax(BSON("x" << 100));
+        ASSERT_EQUALS("", d.getZoneForChunk(chunk));
+    }
+
+    {
+        ChunkType chunk;
+        chunk.setMin(BSON("x" << 200));
+        chunk.setMax(kMaxBSONKey);
+        ASSERT_EQUALS("c", d.getZoneForChunk(chunk));
+    }
 }
 
 }  // namespace
diff --git a/src/mongo/db/s/config/initial_split_policy.cpp b/src/mongo/db/s/config/initial_split_policy.cpp
index cb67583c0b3..f347a00d0ce 100644
--- a/src/mongo/db/s/config/initial_split_policy.cpp
+++ b/src/mongo/db/s/config/initial_split_policy.cpp
@@ -85,7 +85,7 @@ ShardId selectBestShard(const ChunkDistributionMap& chunkMap,
                         const ZoneInfo& zoneInfo,
                         const ZoneShardMap& zoneToShards,
                         const ChunkRange& chunkRange) {
-    auto zone = zoneInfo.getZoneForRange(chunkRange);
+    auto zone = zoneInfo.getZoneForChunk(chunkRange);
     auto iter = zoneToShards.find(zone);
 
     uassert(4952605,
diff --git a/src/mongo/s/chunk_manager.h b/src/mongo/s/chunk_manager.h
index c917a1121b2..31536992a42 100644
--- a/src/mongo/s/chunk_manager.h
+++ b/src/mongo/s/chunk_manager.h
@@ -717,23 +717,6 @@ public:
             shardKey);
     }
 
-    template <typename Callable>
-    void forEachOverlappingChunk(const BSONObj& min,
-                                 const BSONObj& max,
-                                 bool isMaxInclusive,
-                                 Callable&& handler) const {
-        _rt->optRt->forEachOverlappingChunk(
-            min,
-            max,
-            isMaxInclusive,
-            [this, handler = std::forward<Callable>(handler)](const auto& chunkInfo) mutable {
-                if (!handler(Chunk{*chunkInfo, _clusterTime})) {
-                    return false;
-                }
-                return true;
-            });
-    }
-
     /**
      * Returns true if a document with the given "shardKey" is owned by the shard with the given
      * "shardId" in this routing table. If "shardKey" is empty returns false. If "shardKey" is not a