SERVER-70768 Balancer use wrong chunk size for jumbo chunks

(cherry picked from commit b8c329e86ae43522ff0fa52cdecbc23f1b823c00)

8 files changed, 250 insertions, 67 deletions

diff --git a/jstests/sharding/jumbo1.js b/jstests/sharding/jumbo1.js
deleted file mode 100644
index fa5d13b9f2a..00000000000
--- a/jstests/sharding/jumbo1.js
+++ /dev/null
@@ -1,46 +0,0 @@
-(function() {
-'use strict';
-
-load("jstests/sharding/libs/find_chunks_util.js");
-
-var s = new ShardingTest({shards: 2, other: {chunkSize: 1}});
-
-assert.commandWorked(s.s.adminCommand({enablesharding: "test", primaryShard: s.shard1.shardName}));
-assert.commandWorked(
-    s.s.adminCommand({addShardToZone: s.shard0.shardName, zone: 'finalDestination'}));
-
-// Set the chunk range with a zone that will cause the chunk to be in the wrong place so the
-// balancer will be forced to attempt to move it out.
-assert.commandWorked(s.s.adminCommand({shardcollection: "test.foo", key: {x: 1}}));
-assert.commandWorked(s.s.adminCommand(
-    {updateZoneKeyRange: 'test.foo', min: {x: 0}, max: {x: MaxKey}, zone: 'finalDestination'}));
-
-var db = s.getDB("test");
-
-const big = 'X'.repeat(1024 * 1024);  // 1MB
-
-// Insert 3MB of documents to create a jumbo chunk, and use the same shard key in all of
-// them so that the chunk cannot be split.
-var bulk = db.foo.initializeUnorderedBulkOp();
-for (var i = 0; i < 3; i++) {
-    bulk.insert({x: 0, big: big});
-}
-
-assert.commandWorked(bulk.execute());
-
-s.startBalancer();
-
-// Wait for the balancer to try to move the chunk and mark it as jumbo.
-assert.soon(() => {
-    let chunk = findChunksUtil.findOneChunkByNs(s.getDB('config'), 'test.foo', {min: {x: 0}});
-    if (chunk == null) {
-        // Balancer hasn't run and enforce the zone boundaries yet.
-        return false;
-    }
-
-    assert.eq(s.shard1.shardName, chunk.shard, `${tojson(chunk)} was moved by the balancer`);
-    return chunk.jumbo;
-});
-
-s.stop();
-})();
diff --git a/jstests/sharding/jumbo_chunks.js b/jstests/sharding/jumbo_chunks.js
new file mode 100644
index 00000000000..4e03817b68e
--- /dev/null
+++ b/jstests/sharding/jumbo_chunks.js
@@ -0,0 +1,178 @@
+/**
+ * Test jumbo chunks
+ *
+ * @tags: [
+ *    # Command `configureCollectionBalancing` was added in v5.3
+ *    requires_fcv_53,
+ * ]
+ */
+
+(function() {
+'use strict';
+
+load("jstests/sharding/libs/find_chunks_util.js");
+
+function bulkInsert(coll, keyValue, sizeMBytes) {
+    const big = 'X'.repeat(1024 * 1024);  // 1MB
+    var bulk = coll.initializeUnorderedBulkOp();
+    for (var i = 0; i < sizeMBytes; i++) {
+        bulk.insert({x: keyValue, big: big});
+    }
+    assert.commandWorked(bulk.execute());
+}
+
+function assertNumJumboChunks(configDB, ns, expectedNumJumboChunks) {
+    assert.eq(findChunksUtil.countChunksForNs(configDB, ns, {jumbo: true}), expectedNumJumboChunks);
+}
+
+function setGlobalChunkSize(st, chunkSizeMBytes) {
+    // Set global chunk size
+    assert.commandWorked(
+        st.s.getDB("config").settings.update({_id: 'chunksize'},
+                                             {$set: {value: chunkSizeMBytes}},
+                                             {upsert: true, writeConcern: {w: 'majority'}}));
+}
+
+function setCollectionChunkSize(st, ns, chunkSizeMBytes) {
+    assert.commandWorked(
+        st.s.adminCommand({configureCollectionBalancing: ns, chunkSize: chunkSizeMBytes}));
+}
+
+// Test setup
+var st = new ShardingTest({shards: 2, other: {chunkSize: 1}});
+
+assert.commandWorked(
+    st.s.adminCommand({enablesharding: "test", primaryShard: st.shard1.shardName}));
+assert.commandWorked(st.s.adminCommand({addShardToZone: st.shard0.shardName, zone: 'zoneShard0'}));
+
+// Try to move unsuccessfully a 3MB chunk and check it gets marked as jumbo
+{
+    // Set the chunk range with a zone that will cause the chunk to be in the wrong place so the
+    // balancer will be forced to attempt to move it out.
+    assert.commandWorked(st.s.adminCommand({shardcollection: "test.foo", key: {x: 1}}));
+    assert.commandWorked(st.s.adminCommand(
+        {updateZoneKeyRange: 'test.foo', min: {x: 0}, max: {x: MaxKey}, zone: 'zoneShard0'}));
+
+    var db = st.getDB("test");
+
+    const big = 'X'.repeat(1024 * 1024);  // 1MB
+
+    // Insert 3MB of documents to create a jumbo chunk, and use the same shard key in all of
+    // them so that the chunk cannot be split.
+    var bulk = db.foo.initializeUnorderedBulkOp();
+    for (var i = 0; i < 3; i++) {
+        bulk.insert({x: 0, big: big});
+    }
+
+    assert.commandWorked(bulk.execute());
+
+    st.startBalancer();
+
+    // Wait for the balancer to try to move the chunk and check it gets marked as jumbo.
+    assert.soon(() => {
+        let chunk = findChunksUtil.findOneChunkByNs(st.getDB('config'), 'test.foo', {min: {x: 0}});
+        if (chunk == null) {
+            // Balancer hasn't run and enforce the zone boundaries yet.
+            return false;
+        }
+
+        assert.eq(st.shard1.shardName, chunk.shard, `${tojson(chunk)} was moved by the balancer`);
+        return chunk.jumbo;
+    });
+
+    st.stopBalancer();
+}
+
+// Move successfully a 3MB chunk
+// Collection chunkSize must prevail over global chunkSize setting
+//     global chunkSize     -> 1MB
+//     collection chunkSize -> 5MB
+{
+    const collName = "collA";
+    const coll = st.s.getDB("test").getCollection(collName);
+    const configDB = st.s.getDB("config");
+    const splitPoint = 0;
+
+    assert.commandWorked(st.s.adminCommand({shardcollection: coll.getFullName(), key: {x: 1}}));
+    assert.commandWorked(st.s.adminCommand({
+        updateZoneKeyRange: coll.getFullName(),
+        min: {x: splitPoint},
+        max: {x: MaxKey},
+        zone: 'zoneShard0'
+    }));
+
+    bulkInsert(coll, splitPoint, 3);
+
+    setCollectionChunkSize(st, coll.getFullName(), 5);
+    setGlobalChunkSize(st, 1);
+
+    // Move the 3MB chunk to shard0
+    st.startBalancer();
+    st.awaitCollectionBalance(coll);
+    st.stopBalancer();
+
+    const chunk =
+        findChunksUtil.findOneChunkByNs(configDB, coll.getFullName(), {min: {x: splitPoint}});
+
+    // Verify chunk has been moved to shard0
+    assert.eq(st.shard0.shardName,
+              chunk.shard,
+              `${tojson(chunk)} was not moved to ${tojson(st.shard0.shardName)}`);
+    assertNumJumboChunks(configDB, coll.getFullName(), 0);
+
+    coll.drop();
+}
+
+// Try to move unsuccessfully a 3MB chunk and mark it as jumbo
+// Collection chunkSize must prevail over global chunkSize setting
+//     global chunkSize     -> 5MB
+//     collection chunkSize -> 1MB
+{
+    const collName = "collB";
+    const coll = st.s.getDB("test").getCollection(collName);
+    const configDB = st.s.getDB("config");
+    const splitPoint = 0;
+
+    assert.commandWorked(st.s.adminCommand({shardcollection: coll.getFullName(), key: {x: 1}}));
+    assert.commandWorked(st.s.adminCommand({
+        updateZoneKeyRange: coll.getFullName(),
+        min: {x: splitPoint},
+        max: {x: MaxKey},
+        zone: 'zoneShard0'
+    }));
+
+    bulkInsert(coll, splitPoint, 3);
+
+    setCollectionChunkSize(st, coll.getFullName(), 1);
+    setGlobalChunkSize(st, 5);
+
+    // Try to move the 3MB chunk and mark it as jumbo
+    st.startBalancer();
+
+    assert.soon(() => {
+        const chunk =
+            findChunksUtil.findOneChunkByNs(configDB, coll.getFullName(), {min: {x: splitPoint}});
+        if (chunk == null) {
+            // Balancer hasn't run and enforce the zone boundaries yet.
+            return false;
+        }
+
+        return chunk.jumbo;
+    });
+
+    st.stopBalancer();
+
+    const chunk =
+        findChunksUtil.findOneChunkByNs(configDB, coll.getFullName(), {min: {x: splitPoint}});
+
+    // Verify chunk hasn't been moved to shard0 and it's jumbo
+    assert.eq(st.shard1.shardName,
+              chunk.shard,
+              `${tojson(chunk)} was moved to ${tojson(st.shard0.shardName)}`);
+    assertNumJumboChunks(configDB, coll.getFullName(), 1);
+
+    coll.drop();
+}
+
+st.stop();
+})();
diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index 7432d038afc..f4752f3d12d 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -1097,7 +1097,7 @@ int Balancer::_moveChunks(OperationContext* opCtx,
                 opCtx, migrateInfo.uuid, repl::ReadConcernLevel::kMajorityReadConcern);
 
             ShardingCatalogManager::get(opCtx)->splitOrMarkJumbo(
-                opCtx, collection.getNss(), migrateInfo.minKey);
+                opCtx, collection.getNss(), migrateInfo.minKey, migrateInfo.getMaxChunkSizeBytes());
             continue;
         }
 
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 df80d254e12..d52421edc9e 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
@@ -402,7 +402,8 @@ public:
                                         std::move(collectionChunks),
                                         std::move(shardInfos),
                                         std::move(collectionZones),
-                                        smallChunkSizeThresholdBytes));
+                                        smallChunkSizeThresholdBytes,
+                                        maxChunkSizeBytes));
     }
 
     DefragmentationPhaseEnum getType() const override {
@@ -465,7 +466,8 @@ public:
             availableShards->erase(targetSibling->shard);
             auto smallChunkVersion = getShardVersion(opCtx, nextSmallChunk->shard, _nss);
             _outstandingMigrations.emplace_back(nextSmallChunk, targetSibling);
-            return _outstandingMigrations.back().asMigrateInfo(_uuid, _nss, smallChunkVersion);
+            return _outstandingMigrations.back().asMigrateInfo(
+                _uuid, _nss, smallChunkVersion, _maxChunkSizeBytes);
         }
 
         return boost::none;
@@ -717,7 +719,8 @@ private:
 
         MigrateInfo asMigrateInfo(const UUID& collUuid,
                                   const NamespaceString& nss,
-                                  const ChunkVersion& version) const {
+                                  const ChunkVersion& version,
+                                  uint64_t maxChunkSizeBytes) const {
             return MigrateInfo(chunkToMergeWith->shard,
                                chunkToMove->shard,
                                nss,
@@ -725,7 +728,8 @@ private:
                                chunkToMove->range.getMin(),
                                chunkToMove->range.getMax(),
                                version,
-                               MoveChunkRequest::ForceJumbo::kDoNotForce);
+                               MoveChunkRequest::ForceJumbo::kDoNotForce,
+                               maxChunkSizeBytes);
         }
 
         ChunkRange asMergedRange() const {
@@ -790,6 +794,8 @@ private:
 
     const int64_t _smallChunkSizeThresholdBytes;
 
+    const uint64_t _maxChunkSizeBytes;
+
     bool _aborted{false};
 
     DefragmentationPhaseEnum _nextPhase{DefragmentationPhaseEnum::kMergeChunks};
@@ -799,7 +805,8 @@ private:
                             std::vector<ChunkType>&& collectionChunks,
                             stdx::unordered_map<ShardId, ShardInfo>&& shardInfos,
                             ZoneInfo&& collectionZones,
-                            uint64_t smallChunkSizeThresholdBytes)
+                            uint64_t smallChunkSizeThresholdBytes,
+                            uint64_t maxChunkSizeBytes)
         : _nss(nss),
           _uuid(uuid),
           _collectionChunks(),
@@ -810,7 +817,8 @@ private:
           _actionableMerges(),
           _outstandingMerges(),
           _zoneInfo(std::move(collectionZones)),
-          _smallChunkSizeThresholdBytes(smallChunkSizeThresholdBytes) {
+          _smallChunkSizeThresholdBytes(smallChunkSizeThresholdBytes),
+          _maxChunkSizeBytes(maxChunkSizeBytes) {
 
         // Load the collection routing table in a std::list to ease later manipulation
         for (auto&& chunk : collectionChunks) {
diff --git a/src/mongo/db/s/balancer/balancer_policy.cpp b/src/mongo/db/s/balancer/balancer_policy.cpp
index 59a1b5b1197..0cf71995d04 100644
--- a/src/mongo/db/s/balancer/balancer_policy.cpp
+++ b/src/mongo/db/s/balancer/balancer_policy.cpp
@@ -37,6 +37,7 @@
 
 #include "mongo/db/s/balancer/type_migration.h"
 #include "mongo/logv2/log.h"
+#include "mongo/s/balancer_configuration.h"
 #include "mongo/s/catalog/type_shard.h"
 #include "mongo/s/catalog/type_tags.h"
 #include "mongo/s/grid.h"
@@ -495,8 +496,19 @@ MigrateInfosWithReason BalancerPolicy::balance(
                 }
 
                 invariant(to != stat.shardId);
-                migrations.emplace_back(
-                    to, distribution.nss(), chunk, MoveChunkRequest::ForceJumbo::kForceBalancer);
+
+                auto maxChunkSizeBytes = [&]() -> boost::optional<int64_t> {
+                    if (collDataSizeInfo.has_value()) {
+                        return collDataSizeInfo->maxChunkSizeBytes;
+                    }
+                    return boost::none;
+                }();
+
+                migrations.emplace_back(to,
+                                        distribution.nss(),
+                                        chunk,
+                                        MoveChunkRequest::ForceJumbo::kForceBalancer,
+                                        maxChunkSizeBytes);
                 if (firstReason == MigrationReason::none) {
                     firstReason = MigrationReason::drain;
                 }
@@ -564,11 +576,20 @@ MigrateInfosWithReason BalancerPolicy::balance(
                 }
 
                 invariant(to != stat.shardId);
+
+                auto maxChunkSizeBytes = [&]() -> boost::optional<int64_t> {
+                    if (collDataSizeInfo.has_value()) {
+                        return collDataSizeInfo->maxChunkSizeBytes;
+                    }
+                    return boost::none;
+                }();
+
                 migrations.emplace_back(to,
                                         distribution.nss(),
                                         chunk,
                                         forceJumbo ? MoveChunkRequest::ForceJumbo::kForceBalancer
-                                                   : MoveChunkRequest::ForceJumbo::kDoNotForce);
+                                                   : MoveChunkRequest::ForceJumbo::kDoNotForce,
+                                        maxChunkSizeBytes);
                 if (firstReason == MigrationReason::none) {
                     firstReason = MigrationReason::zoneViolation;
                 }
@@ -859,7 +880,8 @@ string ZoneRange::toString() const {
 MigrateInfo::MigrateInfo(const ShardId& a_to,
                          const NamespaceString& a_nss,
                          const ChunkType& a_chunk,
-                         const MoveChunkRequest::ForceJumbo a_forceJumbo)
+                         const MoveChunkRequest::ForceJumbo a_forceJumbo,
+                         boost::optional<int64_t> maxChunkSizeBytes)
     : nss(a_nss), uuid(a_chunk.getCollectionUUID()) {
     invariant(a_to.isValid());
 
@@ -870,6 +892,7 @@ MigrateInfo::MigrateInfo(const ShardId& a_to,
     maxKey = a_chunk.getMax();
     version = a_chunk.getVersion();
     forceJumbo = a_forceJumbo;
+    optMaxChunkSizeBytes = maxChunkSizeBytes;
 }
 
 MigrateInfo::MigrateInfo(const ShardId& a_to,
@@ -921,6 +944,10 @@ string MigrateInfo::toString() const {
                          << ", to " << to;
 }
 
+boost::optional<int64_t> MigrateInfo::getMaxChunkSizeBytes() const {
+    return optMaxChunkSizeBytes;
+}
+
 SplitInfo::SplitInfo(const ShardId& inShardId,
                      const NamespaceString& inNss,
                      const ChunkVersion& inCollectionVersion,
diff --git a/src/mongo/db/s/balancer/balancer_policy.h b/src/mongo/db/s/balancer/balancer_policy.h
index af9377999a6..35f0c8ada30 100644
--- a/src/mongo/db/s/balancer/balancer_policy.h
+++ b/src/mongo/db/s/balancer/balancer_policy.h
@@ -59,7 +59,8 @@ struct MigrateInfo {
     MigrateInfo(const ShardId& a_to,
                 const NamespaceString& a_nss,
                 const ChunkType& a_chunk,
-                MoveChunkRequest::ForceJumbo a_forceJumbo);
+                MoveChunkRequest::ForceJumbo a_forceJumbo,
+                boost::optional<int64_t> maxChunkSizeBytes = boost::none);
 
     MigrateInfo(const ShardId& a_to,
                 const ShardId& a_from,
@@ -77,6 +78,8 @@ struct MigrateInfo {
 
     std::string toString() const;
 
+    boost::optional<int64_t> getMaxChunkSizeBytes() const;
+
     NamespaceString nss;
     UUID uuid;
     ShardId to;
diff --git a/src/mongo/db/s/config/sharding_catalog_manager.h b/src/mongo/db/s/config/sharding_catalog_manager.h
index 44fefe25496..f6e791e0589 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager.h
+++ b/src/mongo/db/s/config/sharding_catalog_manager.h
@@ -348,7 +348,8 @@ public:
      */
     void splitOrMarkJumbo(OperationContext* opCtx,
                           const NamespaceString& nss,
-                          const BSONObj& minKey);
+                          const BSONObj& minKey,
+                          boost::optional<int64_t> optMaxChunkSizeBytes);
 
     /**
      * In a transaction, sets the 'allowMigrations' to the requested state and bumps the collection
diff --git a/src/mongo/db/s/config/sharding_catalog_manager_chunk_operations.cpp b/src/mongo/db/s/config/sharding_catalog_manager_chunk_operations.cpp
index 76af2ce840a..3c02f089b4c 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager_chunk_operations.cpp
+++ b/src/mongo/db/s/config/sharding_catalog_manager_chunk_operations.cpp
@@ -1764,19 +1764,31 @@ void ShardingCatalogManager::bumpMultipleCollectionVersionsAndChangeMetadataInTx
 
 void ShardingCatalogManager::splitOrMarkJumbo(OperationContext* opCtx,
                                               const NamespaceString& nss,
-                                              const BSONObj& minKey) {
+                                              const BSONObj& minKey,
+                                              boost::optional<int64_t> optMaxChunkSizeBytes) {
     const auto cm = uassertStatusOK(
         Grid::get(opCtx)->catalogCache()->getShardedCollectionRoutingInfoWithRefresh(opCtx, nss));
     auto chunk = cm.findIntersectingChunkWithSimpleCollation(minKey);
 
     try {
-        const auto splitPoints = uassertStatusOK(shardutil::selectChunkSplitPoints(
-            opCtx,
-            chunk.getShardId(),
-            nss,
-            cm.getShardKeyPattern(),
-            ChunkRange(chunk.getMin(), chunk.getMax()),
-            Grid::get(opCtx)->getBalancerConfiguration()->getMaxChunkSizeBytes()));
+        const auto maxChunkSizeBytes = [&]() -> int64_t {
+            if (optMaxChunkSizeBytes.has_value()) {
+                return *optMaxChunkSizeBytes;
+            }
+
+            auto coll = Grid::get(opCtx)->catalogClient()->getCollection(
+                opCtx, nss, repl::ReadConcernLevel::kMajorityReadConcern);
+            return coll.getMaxChunkSizeBytes().value_or(
+                Grid::get(opCtx)->getBalancerConfiguration()->getMaxChunkSizeBytes());
+        }();
+
+        const auto splitPoints = uassertStatusOK(
+            shardutil::selectChunkSplitPoints(opCtx,
+                                              chunk.getShardId(),
+                                              nss,
+                                              cm.getShardKeyPattern(),
+                                              ChunkRange(chunk.getMin(), chunk.getMax()),
+                                              maxChunkSizeBytes));
 
         if (splitPoints.empty()) {
             LOGV2(21873,