SERVER-26770 Do not move chunks out of shards with below-average utilization

3 files changed, 53 insertions, 29 deletions

diff --git a/src/mongo/db/s/balancer/balancer_policy.cpp b/src/mongo/db/s/balancer/balancer_policy.cpp
index ea8365f6524..87bc5efd223 100644
--- a/src/mongo/db/s/balancer/balancer_policy.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy.cpp
@@ -402,8 +402,29 @@ vector<MigrateInfo> BalancerPolicy::balance(const ShardStatisticsVector& shardSt
     tagsPlusEmpty.push_back("");
 
     for (const auto& tag : tagsPlusEmpty) {
-        while (_singleZoneBalance(
-            shardStats, distribution, tag, imbalanceThreshold, &migrations, &usedShards))
+        const size_t totalNumberOfChunksWithTag =
+            (tag.empty() ? distribution.totalChunks() : distribution.totalChunksWithTag(tag));
+
+        size_t totalNumberOfShardsWithTag = 0;
+
+        for (const auto& stat : shardStats) {
+            if (tag.empty() || stat.shardTags.count(tag)) {
+                totalNumberOfShardsWithTag++;
+            }
+        }
+
+        // Calculate the ceiling of the optimal number of chunks per shard
+        const size_t idealNumberOfChunksPerShardForTag =
+            (totalNumberOfChunksWithTag / totalNumberOfShardsWithTag) +
+            (totalNumberOfChunksWithTag % totalNumberOfShardsWithTag ? 1 : 0);
+
+        while (_singleZoneBalance(shardStats,
+                                  distribution,
+                                  tag,
+                                  idealNumberOfChunksPerShardForTag,
+                                  imbalanceThreshold,
+                                  &migrations,
+                                  &usedShards))
             ;
     }
 
@@ -428,6 +449,7 @@ boost::optional<MigrateInfo> BalancerPolicy::balanceSingleChunk(
 bool BalancerPolicy::_singleZoneBalance(const ShardStatisticsVector& shardStats,
                                         const DistributionStatus& distribution,
                                         const string& tag,
+                                        size_t idealNumberOfChunksPerShardForTag,
                                         size_t imbalanceThreshold,
                                         vector<MigrateInfo>* migrations,
                                         set<ShardId>* usedShards) {
@@ -436,7 +458,9 @@ bool BalancerPolicy::_singleZoneBalance(const ShardStatisticsVector& shardStats,
         return false;
 
     const size_t max = distribution.numberOfChunksInShardWithTag(from, tag);
-    if (max == 0)
+
+    // Do not use a shard if it already has less entries than the optimal per-shard chunk count
+    if (max <= idealNumberOfChunksPerShardForTag)
         return false;
 
     const ShardId to = _getLeastLoadedReceiverShard(shardStats, distribution, tag, *usedShards);
@@ -448,40 +472,18 @@ bool BalancerPolicy::_singleZoneBalance(const ShardStatisticsVector& shardStats,
     }
 
     const size_t min = distribution.numberOfChunksInShardWithTag(to, tag);
-    if (min >= max)
-        return false;
-
-    const size_t totalNumberOfChunksWithTag =
-        (tag.empty() ? distribution.totalChunks() : distribution.totalChunksWithTag(tag));
-
-    size_t totalNumberOfShardsWithTag = 0;
-
-    for (const auto& stat : shardStats) {
-        if (tag.empty() || stat.shardTags.count(tag)) {
-            totalNumberOfShardsWithTag++;
-        }
-    }
-
-    // totalNumberOfShardsWithTag cannot be zero if the to shard is valid
-    invariant(totalNumberOfShardsWithTag);
-    invariant(totalNumberOfChunksWithTag >= max);
-
-    // Calculate the ceiling of the optimal number of chunks per shard
-    const size_t idealNumberOfChunksPerShardWithTag =
-        (totalNumberOfChunksWithTag / totalNumberOfShardsWithTag) +
-        (totalNumberOfChunksWithTag % totalNumberOfShardsWithTag ? 1 : 0);
 
     // Do not use a shard if it already has more entries than the optimal per-shard chunk count
-    if (min >= idealNumberOfChunksPerShardWithTag)
+    if (min >= idealNumberOfChunksPerShardForTag)
         return false;
 
-    const size_t imbalance = max - idealNumberOfChunksPerShardWithTag;
+    const size_t imbalance = max - idealNumberOfChunksPerShardForTag;
 
     LOG(1) << "collection : " << distribution.nss().ns();
     LOG(1) << "zone       : " << tag;
     LOG(1) << "donor      : " << from << " chunks on " << max;
     LOG(1) << "receiver   : " << to << " chunks on " << min;
-    LOG(1) << "ideal      : " << idealNumberOfChunksPerShardWithTag;
+    LOG(1) << "ideal      : " << idealNumberOfChunksPerShardForTag;
     LOG(1) << "threshold  : " << imbalanceThreshold;
 
     // Check whether it is necessary to balance within this zone
diff --git a/src/mongo/db/s/balancer/balancer_policy.h b/src/mongo/db/s/balancer/balancer_policy.h
index c08aee77ceb..63f44c3ca68 100644
--- a/src/mongo/db/s/balancer/balancer_policy.h
+++ b/src/mongo/db/s/balancer/balancer_policy.h
@@ -214,7 +214,11 @@ private:
     /**
      * Selects one chunk for the specified zone (if appropriate) to be moved in order to bring the
      * deviation of the shards chunk contents closer to even across all shards in the specified
-     * zone. Takes into account the shards, which have already been used for migrations.
+     * zone. Takes into account and updates the shards, which have already been used for migrations.
+     *
+     * The 'idealNumberOfChunksPerShardForTag' indicates what is the ideal number of chunks which
+     * each shard must have and is used to determine the imbalance and also to prevent chunks from
+     * moving when not necessary.
      *
      * Returns true if a migration was suggested, false otherwise. This method is intented to be
      * called multiple times until all posible migrations for a zone have been selected.
@@ -222,6 +226,7 @@ private:
     static bool _singleZoneBalance(const ShardStatisticsVector& shardStats,
                                    const DistributionStatus& distribution,
                                    const std::string& tag,
+                                   size_t idealNumberOfChunksPerShardForTag,
                                    size_t imbalanceThreshold,
                                    std::vector<MigrateInfo>* migrations,
                                    std::set<ShardId>* usedShards);
diff --git a/src/mongo/db/s/balancer/balancer_policy_test.cpp b/src/mongo/db/s/balancer/balancer_policy_test.cpp
index 86258c723d8..ea1024a00fb 100644
--- a/src/mongo/db/s/balancer/balancer_policy_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy_test.cpp
@@ -214,6 +214,23 @@ TEST(BalancerPolicy, ParallelBalancingDoesNotPutChunksOnShardsAboveTheOptimal) {
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), migrations[1].maxKey);
 }
 
+TEST(BalancerPolicy, ParallelBalancingDoesNotMoveChunksFromShardsBelowOptimal) {
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 100, false, emptyTagSet, emptyShardVersion), 100},
+         {ShardStatistics(kShardId1, kNoMaxSize, 30, false, emptyTagSet, emptyShardVersion), 30},
+         {ShardStatistics(kShardId2, kNoMaxSize, 5, false, emptyTagSet, emptyShardVersion), 5},
+         {ShardStatistics(kShardId3, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
+
+    const auto migrations(BalancerPolicy::balance(
+        cluster.first, DistributionStatus(kNamespace, cluster.second), false));
+    ASSERT_EQ(1U, migrations.size());
+
+    ASSERT_EQ(kShardId0, migrations[0].from);
+    ASSERT_EQ(kShardId3, migrations[0].to);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
+}
+
 TEST(BalancerPolicy, JumboChunksNotMoved) {
     auto cluster = generateCluster(
         {{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 4},