SERVER-26579 Do not move chunks to shards with above optimal chunk count

4 files changed, 56 insertions, 28 deletions

diff --git a/src/mongo/db/s/balancer/balancer_policy.cpp b/src/mongo/db/s/balancer/balancer_policy.cpp
index 912d0fdb15b..5b53cfc0320 100644
--- a/src/mongo/db/s/balancer/balancer_policy.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy.cpp
@@ -298,10 +298,11 @@ vector<MigrateInfo> BalancerPolicy::balance(const ShardStatisticsVector& shardSt
     // migrations for the same shard.
     set<ShardId> usedShards;
 
-    // 1) Check for shards, which are in draining mode and must have chunks moved off of them
+    // 1) Check for shards, which are in draining mode or are above the size limit and must have
+    // chunks moved off of them
     {
         for (const auto& stat : shardStats) {
-            if (!stat.isDraining)
+            if (!stat.isDraining && !stat.isSizeExceeded())
                 continue;
 
             if (usedShards.count(stat.shardId))
@@ -469,6 +470,10 @@ bool BalancerPolicy::_singleZoneBalance(const ShardStatisticsVector& shardStats,
         (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)
+        return false;
+
     const size_t imbalance = max - idealNumberOfChunksPerShardWithTag;
 
     LOG(1) << "collection : " << distribution.nss().ns();
diff --git a/src/mongo/db/s/balancer/balancer_policy_test.cpp b/src/mongo/db/s/balancer/balancer_policy_test.cpp
index ebfe146843c..35da7c7bf7f 100644
--- a/src/mongo/db/s/balancer/balancer_policy_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy_test.cpp
@@ -52,6 +52,8 @@ const auto kShardId0 = ShardId("shard0");
 const auto kShardId1 = ShardId("shard1");
 const auto kShardId2 = ShardId("shard2");
 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;
 
@@ -177,6 +179,30 @@ TEST(BalancerPolicy, ParallelBalancing) {
     ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), migrations[1].maxKey);
 }
 
+TEST(BalancerPolicy, ParallelBalancingDoesNotPutChunksOnShardsAboveTheOptimal) {
+    auto cluster = generateCluster(
+        {{ShardStatistics(kShardId0, kNoMaxSize, 100, false, emptyTagSet, emptyShardVersion), 100},
+         {ShardStatistics(kShardId1, kNoMaxSize, 90, false, emptyTagSet, emptyShardVersion), 90},
+         {ShardStatistics(kShardId2, kNoMaxSize, 90, false, emptyTagSet, emptyShardVersion), 90},
+         {ShardStatistics(kShardId3, kNoMaxSize, 80, false, emptyTagSet, emptyShardVersion), 80},
+         {ShardStatistics(kShardId4, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0},
+         {ShardStatistics(kShardId5, kNoMaxSize, 0, false, emptyTagSet, emptyShardVersion), 0}});
+
+    const auto migrations(BalancerPolicy::balance(
+        cluster.first, DistributionStatus(kNamespace, cluster.second), false));
+    ASSERT_EQ(2U, migrations.size());
+
+    ASSERT_EQ(kShardId0, migrations[0].from);
+    ASSERT_EQ(kShardId4, migrations[0].to);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
+
+    ASSERT_EQ(kShardId1, migrations[1].from);
+    ASSERT_EQ(kShardId5, migrations[1].to);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMin(), migrations[1].minKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId1][0].getMax(), migrations[1].maxKey);
+}
+
 TEST(BalancerPolicy, JumboChunksNotMoved) {
     auto cluster = generateCluster(
         {{ShardStatistics(kShardId0, kNoMaxSize, 2, false, emptyTagSet, emptyShardVersion), 4},
@@ -354,7 +380,7 @@ TEST(BalancerPolicy, NoBalancingDueToAllNodesEitherDrainingOrMaxedOut) {
     ASSERT(migrations.empty());
 }
 
-TEST(BalancerPolicy, BalancerRespectsMaxShardSizeOnlyBalanceToNonMaxed) {
+TEST(BalancerPolicy, BalancerMovesChunksOffSizeMaxedShards) {
     // Note that maxSize of shard0 is 1, and it is therefore overloaded with currSize = 3. Other
     // shards have maxSize = 0 = unset. Even though the overloaded shard has the least number of
     // less chunks, we shouldn't move chunks to that shard.
@@ -366,24 +392,10 @@ TEST(BalancerPolicy, BalancerRespectsMaxShardSizeOnlyBalanceToNonMaxed) {
     const auto migrations(BalancerPolicy::balance(
         cluster.first, DistributionStatus(kNamespace, cluster.second), false));
     ASSERT_EQ(1U, migrations.size());
-    ASSERT_EQ(kShardId2, migrations[0].from);
+    ASSERT_EQ(kShardId0, migrations[0].from);
     ASSERT_EQ(kShardId1, migrations[0].to);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMin(), migrations[0].minKey);
-    ASSERT_BSONOBJ_EQ(cluster.second[kShardId2][0].getMax(), migrations[0].maxKey);
-}
-
-TEST(BalancerPolicy, BalancerRespectsMaxShardSizeWhenAllBalanced) {
-    // Note that maxSize of shard0 is 1, and it is therefore overloaded with currSize = 4. Other
-    // shards have maxSize = 0 = unset. We check that being over the maxSize is NOT equivalent to
-    // draining, we don't want to empty shards for no other reason than they are over this limit.
-    auto cluster = generateCluster(
-        {{ShardStatistics(kShardId0, 1, 4, false, emptyTagSet, emptyShardVersion), 4},
-         {ShardStatistics(kShardId1, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4},
-         {ShardStatistics(kShardId2, kNoMaxSize, 4, false, emptyTagSet, emptyShardVersion), 4}});
-
-    const auto migrations(BalancerPolicy::balance(
-        cluster.first, DistributionStatus(kNamespace, cluster.second), false));
-    ASSERT(migrations.empty());
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMin(), migrations[0].minKey);
+    ASSERT_BSONOBJ_EQ(cluster.second[kShardId0][0].getMax(), migrations[0].maxKey);
 }
 
 TEST(BalancerPolicy, BalancerRespectsTagsWhenDraining) {
diff --git a/src/mongo/db/s/balancer/cluster_statistics.cpp b/src/mongo/db/s/balancer/cluster_statistics.cpp
index 495a54b0493..125f11d3551 100644
--- a/src/mongo/db/s/balancer/cluster_statistics.cpp
+++ b/src/mongo/db/s/balancer/cluster_statistics.cpp
@@ -39,8 +39,6 @@ ClusterStatistics::ClusterStatistics() = default;
 
 ClusterStatistics::~ClusterStatistics() = default;
 
-ClusterStatistics::ShardStatistics::ShardStatistics() = default;
-
 ClusterStatistics::ShardStatistics::ShardStatistics(ShardId inShardId,
                                                     uint64_t inMaxSizeMB,
                                                     uint64_t inCurrSizeMB,
@@ -48,9 +46,9 @@ ClusterStatistics::ShardStatistics::ShardStatistics(ShardId inShardId,
                                                     std::set<std::string> inShardTags,
                                                     std::string inMongoVersion)
     : shardId(std::move(inShardId)),
-      maxSizeMB(std::move(inMaxSizeMB)),
-      currSizeMB(std::move(inCurrSizeMB)),
-      isDraining(std::move(inIsDraining)),
+      maxSizeMB(inMaxSizeMB),
+      currSizeMB(inCurrSizeMB),
+      isDraining(inIsDraining),
       shardTags(std::move(inShardTags)),
       mongoVersion(std::move(inMongoVersion)) {}
 
@@ -62,6 +60,14 @@ bool ClusterStatistics::ShardStatistics::isSizeMaxed() const {
     return currSizeMB >= maxSizeMB;
 }
 
+bool ClusterStatistics::ShardStatistics::isSizeExceeded() const {
+    if (!maxSizeMB || !currSizeMB) {
+        return false;
+    }
+
+    return currSizeMB > maxSizeMB;
+}
+
 BSONObj ClusterStatistics::ShardStatistics::toBSON() const {
     BSONObjBuilder builder;
     builder.append("id", shardId.toString());
diff --git a/src/mongo/db/s/balancer/cluster_statistics.h b/src/mongo/db/s/balancer/cluster_statistics.h
index 8963720ee6f..2717c42b7ee 100644
--- a/src/mongo/db/s/balancer/cluster_statistics.h
+++ b/src/mongo/db/s/balancer/cluster_statistics.h
@@ -57,7 +57,6 @@ public:
      */
     struct ShardStatistics {
     public:
-        ShardStatistics();
         ShardStatistics(ShardId shardId,
                         uint64_t maxSizeMB,
                         uint64_t currSizeMB,
@@ -66,12 +65,18 @@ public:
                         std::string mongoVersion);
 
         /**
-         * Returns if a shard cannot receive any new chunks because it has reached the per-shard
-         * data size limit.
+         * Returns true if a shard is not allowed to receive any new chunks because it has reached
+         * the per-shard data size limit.
          */
         bool isSizeMaxed() const;
 
         /**
+         * Returns true if a shard must be relieved (if possible) of some of the chunks it hosts
+         * because it has exceeded its per-shard data size limit.
+         */
+        bool isSizeExceeded() const;
+
+        /**
          * Returns BSON representation of this shard's statistics, for reporting purposes.
          */
         BSONObj toBSON() const;