1 files changed, 25 insertions, 4 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 07427b21879..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
@@ -47,9 +47,9 @@
 #include "mongo/s/catalog/type_collection.h"
 #include "mongo/s/catalog/type_tags.h"
 #include "mongo/s/catalog_cache.h"
-#include "mongo/s/chunk_manager.h"
 #include "mongo/s/grid.h"
 #include "mongo/util/str.h"
+
 namespace mongo {
 
 using MigrateInfoVector = BalancerChunkSelectionPolicy::MigrateInfoVector;
@@ -90,6 +90,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.setNS(nss);
+        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());
@@ -97,7 +118,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())}};
 }
 
 /**
@@ -505,7 +526,7 @@ Status BalancerChunkSelectionPolicyImpl::checkMoveAllowed(OperationContext* opCt
     }
 
     return BalancerPolicy::isShardSuitableReceiver(*newShardIterator,
-                                                   distribution.getTagForRange(chunk.getRange()));
+                                                   distribution.getTagForChunk(chunk));
 }
 
 StatusWith<SplitInfoVector> BalancerChunkSelectionPolicyImpl::_getSplitCandidatesForCollection(
@@ -565,7 +586,7 @@ StatusWith<MigrateInfoVector> BalancerChunkSelectionPolicyImpl::_getMigrateCandi
 
     const DistributionStatus& distribution = collInfoStatus.getValue();
 
-    for (const auto& tagRangeEntry : distribution.getZoneInfo().zoneRanges()) {
+    for (const auto& tagRangeEntry : distribution.tagRanges()) {
         const auto& tagRange = tagRangeEntry.second;
 
         const auto chunkAtZoneMin = cm.findIntersectingChunkWithSimpleCollation(tagRange.min);