SERVER-76615 implement shardDistribution functionality in reshardCollection cmd

8 files changed, 490 insertions, 37 deletions

diff --git a/src/mongo/db/s/config/configsvr_reshard_collection_cmd.cpp b/src/mongo/db/s/config/configsvr_reshard_collection_cmd.cpp
index 0e9e1f82e9e..e23acdfbcfb 100644
--- a/src/mongo/db/s/config/configsvr_reshard_collection_cmd.cpp
+++ b/src/mongo/db/s/config/configsvr_reshard_collection_cmd.cpp
@@ -127,6 +127,19 @@ public:
                     *presetChunks, opCtx, ShardKeyPattern(request().getKey()).getKeyPattern());
             }
 
+            if (!resharding::gFeatureFlagReshardingImprovements.isEnabled(
+                    serverGlobalParams.featureCompatibility)) {
+                uassert(
+                    ErrorCodes::InvalidOptions,
+                    "Resharding improvements is not enabled, reject shardDistribution parameter",
+                    !request().getShardDistribution().has_value());
+            }
+
+            if (const auto& shardDistribution = request().getShardDistribution()) {
+                resharding::validateShardDistribution(
+                    *shardDistribution, opCtx, ShardKeyPattern(request().getKey()));
+            }
+
             // Returns boost::none if there isn't any work to be done by the resharding operation.
             auto instance =
                 ([&]() -> boost::optional<std::shared_ptr<const ReshardingCoordinator>> {
@@ -178,6 +191,7 @@ public:
                     coordinatorDoc.setZones(request().getZones());
                     coordinatorDoc.setPresetReshardedChunks(request().get_presetReshardedChunks());
                     coordinatorDoc.setNumInitialChunks(request().getNumInitialChunks());
+                    coordinatorDoc.setShardDistribution(request().getShardDistribution());
 
                     opCtx->setAlwaysInterruptAtStepDownOrUp_UNSAFE();
                     auto instance = getOrCreateReshardingCoordinator(opCtx, coordinatorDoc);
diff --git a/src/mongo/db/s/config/initial_split_policy.cpp b/src/mongo/db/s/config/initial_split_policy.cpp
index 8f07990319c..29fd01e4e6a 100644
--- a/src/mongo/db/s/config/initial_split_policy.cpp
+++ b/src/mongo/db/s/config/initial_split_policy.cpp
@@ -29,6 +29,9 @@
 
 #include "mongo/db/s/config/initial_split_policy.h"
 
+#include "mongo/base/error_codes.h"
+#include "mongo/bson/bsonobj.h"
+#include "mongo/bson/simple_bsonobj_comparator.h"
 #include "mongo/client/read_preference.h"
 #include "mongo/db/bson/dotted_path_support.h"
 #include "mongo/db/curop.h"
@@ -149,6 +152,67 @@ StringMap<std::vector<ShardId>> buildTagsToShardIdsMap(OperationContext* opCtx,
     return tagToShardIds;
 }
 
+/**
+ * Returns a set of split points to ensure that chunk boundaries will align with the zone
+ * ranges.
+ */
+BSONObjSet extractSplitPointsFromZones(const ShardKeyPattern& shardKey,
+                                       const boost::optional<std::vector<TagsType>>& zones) {
+    auto splitPoints = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+
+    if (!zones) {
+        return splitPoints;
+    }
+
+    for (const auto& zone : *zones) {
+        splitPoints.insert(zone.getMinKey());
+        splitPoints.insert(zone.getMaxKey());
+    }
+
+    const auto keyPattern = shardKey.getKeyPattern();
+    splitPoints.erase(keyPattern.globalMin());
+    splitPoints.erase(keyPattern.globalMax());
+
+    return splitPoints;
+}
+
+/*
+ * Returns a map mapping shard id to a set of zone tags.
+ */
+stdx::unordered_map<ShardId, stdx::unordered_set<std::string>> buildShardIdToTagsMap(
+    OperationContext* opCtx, const std::vector<ShardKeyRange>& shards) {
+    stdx::unordered_map<ShardId, stdx::unordered_set<std::string>> shardIdToTags;
+    if (shards.empty()) {
+        return shardIdToTags;
+    }
+
+    // Get all docs in config.shards through a query instead of going through the shard registry
+    // because we need the zones as well
+    const auto configServer = Grid::get(opCtx)->shardRegistry()->getConfigShard();
+    const auto shardDocs = uassertStatusOK(
+        configServer->exhaustiveFindOnConfig(opCtx,
+                                             ReadPreferenceSetting(ReadPreference::Nearest),
+                                             repl::ReadConcernLevel::kMajorityReadConcern,
+                                             NamespaceString::kConfigsvrShardsNamespace,
+                                             BSONObj(),
+                                             BSONObj(),
+                                             boost::none));
+    uassert(
+        7661502, str::stream() << "Could not find any shard documents", !shardDocs.docs.empty());
+
+    for (const auto& shard : shards) {
+        shardIdToTags[shard.getShard()] = {};
+    }
+
+    for (const auto& shardDoc : shardDocs.docs) {
+        auto parsedShard = uassertStatusOK(ShardType::fromBSON(shardDoc));
+        for (const auto& tag : parsedShard.getTags()) {
+            shardIdToTags[ShardId(parsedShard.getName())].insert(tag);
+        }
+    }
+
+    return shardIdToTags;
+}
 }  // namespace
 
 std::vector<BSONObj> InitialSplitPolicy::calculateHashedSplitPoints(
@@ -647,7 +711,7 @@ BSONObjSet SamplingBasedSplitPolicy::createFirstSplitPoints(OperationContext* op
         }
     }
 
-    auto splitPoints = _extractSplitPointsFromZones(shardKey);
+    auto splitPoints = extractSplitPointsFromZones(shardKey, _zones);
     if (splitPoints.size() < static_cast<size_t>(_numInitialChunks - 1)) {
         // The BlockingResultsMerger underlying the $mergeCursors stage records how long was
         // spent waiting for samples from the donor shards. It doing so requires the CurOp
@@ -719,25 +783,6 @@ InitialSplitPolicy::ShardCollectionConfig SamplingBasedSplitPolicy::createFirstC
     return {std::move(chunks)};
 }
 
-BSONObjSet SamplingBasedSplitPolicy::_extractSplitPointsFromZones(const ShardKeyPattern& shardKey) {
-    auto splitPoints = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-
-    if (!_zones) {
-        return splitPoints;
-    }
-
-    for (const auto& zone : *_zones) {
-        splitPoints.insert(zone.getMinKey());
-        splitPoints.insert(zone.getMaxKey());
-    }
-
-    const auto keyPattern = shardKey.getKeyPattern();
-    splitPoints.erase(keyPattern.globalMin());
-    splitPoints.erase(keyPattern.globalMax());
-
-    return splitPoints;
-}
-
 void SamplingBasedSplitPolicy::_appendSplitPointsFromSample(BSONObjSet* splitPoints,
                                                             const ShardKeyPattern& shardKey,
                                                             int nToAppend) {
@@ -837,4 +882,126 @@ boost::optional<BSONObj> SamplingBasedSplitPolicy::PipelineDocumentSource::getNe
     return val->toBson();
 }
 
+ShardDistributionSplitPolicy ShardDistributionSplitPolicy::make(
+    OperationContext* opCtx,
+    const ShardKeyPattern& shardKey,
+    std::vector<ShardKeyRange> shardDistribution,
+    boost::optional<std::vector<TagsType>> zones) {
+    uassert(7661501, "ShardDistribution should not be empty", shardDistribution.size() > 0);
+    return ShardDistributionSplitPolicy(shardDistribution, zones);
+}
+
+ShardDistributionSplitPolicy::ShardDistributionSplitPolicy(
+    std::vector<ShardKeyRange>& shardDistribution, boost::optional<std::vector<TagsType>> zones)
+    : _shardDistribution(std::move(shardDistribution)), _zones(std::move(zones)) {}
+
+InitialSplitPolicy::ShardCollectionConfig ShardDistributionSplitPolicy::createFirstChunks(
+    OperationContext* opCtx,
+    const ShardKeyPattern& shardKeyPattern,
+    const SplitPolicyParams& params) {
+    const auto& keyPattern = shardKeyPattern.getKeyPattern();
+    if (_zones) {
+        for (auto& zone : *_zones) {
+            zone.setMinKey(keyPattern.extendRangeBound(zone.getMinKey(), false));
+            zone.setMaxKey(keyPattern.extendRangeBound(zone.getMaxKey(), false));
+        }
+    }
+
+    auto splitPoints = extractSplitPointsFromZones(shardKeyPattern, _zones);
+    std::vector<ChunkType> chunks;
+    uassert(ErrorCodes::InvalidOptions,
+            "ShardDistribution without min/max is not supported.",
+            _shardDistribution[0].getMin());
+
+    unsigned long shardDistributionIdx = 0;
+    const auto currentTime = VectorClock::get(opCtx)->getTime();
+    const auto validAfter = currentTime.clusterTime().asTimestamp();
+    ChunkVersion version({OID::gen(), validAfter}, {1, 0});
+    for (const auto& splitPoint : splitPoints) {
+        _appendChunks(params, splitPoint, keyPattern, shardDistributionIdx, version, chunks);
+    }
+    _appendChunks(
+        params, keyPattern.globalMax(), keyPattern, shardDistributionIdx, version, chunks);
+
+    if (_zones) {
+        _checkShardsMatchZones(opCtx, chunks, *_zones);
+    }
+
+    return {std::move(chunks)};
+}
+
+void ShardDistributionSplitPolicy::_appendChunks(const SplitPolicyParams& params,
+                                                 const BSONObj& splitPoint,
+                                                 const KeyPattern& keyPattern,
+                                                 unsigned long& shardDistributionIdx,
+                                                 ChunkVersion& version,
+                                                 std::vector<ChunkType>& chunks) {
+    while (shardDistributionIdx < _shardDistribution.size()) {
+        auto shardMin =
+            keyPattern.extendRangeBound(*_shardDistribution[shardDistributionIdx].getMin(), false);
+        auto shardMax =
+            keyPattern.extendRangeBound(*_shardDistribution[shardDistributionIdx].getMax(), false);
+        auto lastChunkMax =
+            chunks.empty() ? keyPattern.globalMin() : chunks.back().getRange().getMax();
+        /* When we compare a defined shard range with a splitPoint, there are three cases:
+         * 1. The whole shard range is on the left side of the splitPoint -> Add this shard as a
+         * whole chunk and move to next shard.
+         * 2. The splitPoint is in the middle of the shard range. -> Append (shardMin,
+         * splitPoint) as a chunk and move to next split point.
+         * 3. The whole shard range is on the right side of the splitPoint -> Move to the next
+         * splitPoint.
+         * This algorithm relies on the shardDistribution is continuous and complete to be
+         * correct, which is validated in the cmd handler.
+         */
+        if (SimpleBSONObjComparator::kInstance.evaluate(shardMin < splitPoint)) {
+            // The whole shard range is on the left side of the splitPoint.
+            if (SimpleBSONObjComparator::kInstance.evaluate(shardMax <= splitPoint)) {
+                appendChunk(params,
+                            lastChunkMax,
+                            shardMax,
+                            &version,
+                            _shardDistribution[shardDistributionIdx].getShard(),
+                            &chunks);
+                lastChunkMax = shardMax;
+                shardDistributionIdx++;
+            } else {  // The splitPoint is in the middle of the shard range.
+                appendChunk(params,
+                            lastChunkMax,
+                            splitPoint,
+                            &version,
+                            _shardDistribution[shardDistributionIdx].getShard(),
+                            &chunks);
+                lastChunkMax = splitPoint;
+                return;
+            }
+        } else {  // The whole shard range is on the right side of the splitPoint.
+            return;
+        }
+    }
+}
+
+void ShardDistributionSplitPolicy::_checkShardsMatchZones(
+    OperationContext* opCtx,
+    const std::vector<ChunkType>& chunks,
+    const std::vector<mongo::TagsType>& zones) {
+    ZoneInfo zoneInfo;
+    auto shardIdToTags = buildShardIdToTagsMap(opCtx, _shardDistribution);
+    for (const auto& zone : zones) {
+        uassertStatusOK(
+            zoneInfo.addRangeToZone({zone.getMinKey(), zone.getMaxKey(), zone.getTag()}));
+    }
+
+    for (const auto& chunk : chunks) {
+        auto zoneFromCmdParameter = zoneInfo.getZoneForChunk({chunk.getMin(), chunk.getMax()});
+        auto iter = shardIdToTags.find(chunk.getShard());
+        uassert(ErrorCodes::InvalidOptions,
+                str::stream() << "Specified zones and shardDistribution are conflicting with the "
+                                 "existing shard/zone, shard "
+                              << chunk.getShard() << "doesn't belong to zone "
+                              << zoneFromCmdParameter,
+                iter != shardIdToTags.end() &&
+                    iter->second.find(zoneFromCmdParameter) != iter->second.end());
+    }
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/config/initial_split_policy.h b/src/mongo/db/s/config/initial_split_policy.h
index 1d8774a781c..6a5fcfcb996 100644
--- a/src/mongo/db/s/config/initial_split_policy.h
+++ b/src/mongo/db/s/config/initial_split_policy.h
@@ -38,6 +38,7 @@
 #include "mongo/db/shard_id.h"
 #include "mongo/s/catalog/type_chunk.h"
 #include "mongo/s/catalog/type_tags.h"
+#include "mongo/s/resharding/common_types_gen.h"
 #include "mongo/s/shard_key_pattern.h"
 #include "mongo/util/string_map.h"
 namespace mongo {
@@ -331,12 +332,6 @@ private:
         MakePipelineOptions opts = {});
 
     /**
-     * Returns a set of split points to ensure that chunk boundaries will align with the zone
-     * ranges.
-     */
-    BSONObjSet _extractSplitPointsFromZones(const ShardKeyPattern& shardKey);
-
-    /**
      * Append split points based from the samples taken from the collection.
      */
     void _appendSplitPointsFromSample(BSONObjSet* splitPoints,
@@ -348,4 +343,41 @@ private:
     boost::optional<std::vector<TagsType>> _zones;
     std::unique_ptr<SampleDocumentSource> _samples;
 };
+
+class ShardDistributionSplitPolicy : public InitialSplitPolicy {
+public:
+    static ShardDistributionSplitPolicy make(OperationContext* opCtx,
+                                             const ShardKeyPattern& shardKey,
+                                             std::vector<ShardKeyRange> shardDistribution,
+                                             boost::optional<std::vector<TagsType>> zones);
+
+    ShardDistributionSplitPolicy(std::vector<ShardKeyRange>& shardDistribution,
+                                 boost::optional<std::vector<TagsType>> zones);
+
+    ShardCollectionConfig createFirstChunks(OperationContext* opCtx,
+                                            const ShardKeyPattern& shardKeyPattern,
+                                            const SplitPolicyParams& params) override;
+
+private:
+    /**
+     * Given a splitPoint, create chunks from _shardDistribution until passing the splitPoint.
+     */
+    void _appendChunks(const SplitPolicyParams& params,
+                       const BSONObj& splitPoint,
+                       const KeyPattern& keyPattern,
+                       unsigned long& shardDistributionIdx,
+                       ChunkVersion& version,
+                       std::vector<ChunkType>& chunks);
+
+    /**
+     * Check the chunks created from command parameter "zones" and "shardDistribution" are
+     * satisfying the existing zone mapping rules in config.
+     */
+    void _checkShardsMatchZones(OperationContext* opCtx,
+                                const std::vector<ChunkType>& chunks,
+                                const std::vector<mongo::TagsType>& zones);
+    std::vector<ShardKeyRange> _shardDistribution;
+    boost::optional<std::vector<TagsType>> _zones;
+};
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/config/initial_split_policy_test.cpp b/src/mongo/db/s/config/initial_split_policy_test.cpp
index 648cbcb9719..e896bd7d28f 100644
--- a/src/mongo/db/s/config/initial_split_policy_test.cpp
+++ b/src/mongo/db/s/config/initial_split_policy_test.cpp
@@ -2372,5 +2372,150 @@ TEST_F(SamplingBasedInitSplitTest, ZeroInitialChunks) {
     }
 }
 
+
+class ShardDistributionInitSplitTest : public SingleChunkPerTagSplitPolicyTest {
+public:
+    std::unique_ptr<ShardDistributionSplitPolicy> makeInitialSplitPolicy(
+        std::vector<ShardKeyRange>& shardDistribution,
+        boost::optional<std::vector<TagsType>> zones) {
+        return std::make_unique<ShardDistributionSplitPolicy>(shardDistribution, zones);
+    }
+
+    /**
+     * Calls createFirstChunks() according to the given arguments and asserts that returned chunks
+     * match with the chunks created using expectedChunkRanges and expectedShardIds.
+     */
+    void checkGeneratedInitialZoneChunks(std::unique_ptr<ShardDistributionSplitPolicy> splitPolicy,
+                                         const ShardKeyPattern& shardKeyPattern,
+                                         const std::vector<ShardType>& shardList,
+                                         const std::vector<ShardKeyRange>& shardDistribution,
+                                         const std::vector<ChunkRange>& expectedChunkRanges,
+                                         const std::vector<ShardId>& expectedShardIds) {
+        const ShardId primaryShard("doesntMatter");
+
+        const auto shardCollectionConfig = splitPolicy->createFirstChunks(
+            operationContext(), shardKeyPattern, {UUID::gen(), primaryShard});
+
+        ASSERT_EQ(expectedShardIds.size(), expectedChunkRanges.size());
+        ASSERT_EQ(expectedChunkRanges.size(), shardCollectionConfig.chunks.size());
+        for (size_t i = 0; i < shardCollectionConfig.chunks.size(); ++i) {
+            // Check the chunk range matches the expected range.
+            ASSERT_EQ(expectedChunkRanges[i], shardCollectionConfig.chunks[i].getRange());
+
+            // Check that the shardId matches the expected.
+            const auto& actualShardId = shardCollectionConfig.chunks[i].getShard();
+            ASSERT_EQ(expectedShardIds[i], actualShardId);
+        }
+    }
+};
+
+TEST_F(ShardDistributionInitSplitTest, WithoutZones) {
+    const NamespaceString ns = NamespaceString::createNamespaceString_forTest("foo", "bar");
+    const ShardKeyPattern shardKey(BSON("y" << 1));
+
+    std::vector<ShardType> shardList;
+    shardList.emplace_back(
+        ShardType(shardId("0").toString(), "rs0/fakeShard0:123", {std::string("zoneA")}));
+    shardList.emplace_back(
+        ShardType(shardId("1").toString(), "rs1/fakeShard1:123", {std::string("zoneB")}));
+
+    setupShards(shardList);
+    shardRegistry()->reload(operationContext());
+    ShardKeyRange range0(shardId("0"));
+    range0.setMin(BSON("y" << MINKEY));
+    range0.setMax(BSON("y" << 0));
+    ShardKeyRange range1(shardId("1"));
+    range1.setMin(BSON("y" << 0));
+    range1.setMax(BSON("y" << MAXKEY));
+    std::vector<ShardKeyRange> shardDistribution = {range0, range1};
+
+    std::vector<ChunkRange> expectedChunkRanges = {ChunkRange(BSON("y" << MINKEY), BSON("y" << 0)),
+                                                   ChunkRange(BSON("y" << 0), BSON("y" << MAXKEY))};
+    std::vector<ShardId> expectedShardForEachChunk = {shardId("0"), shardId("1")};
+
+    checkGeneratedInitialZoneChunks(
+        makeInitialSplitPolicy(shardDistribution, boost::none /*zones*/),
+        shardKey,
+        shardList,
+        shardDistribution,
+        expectedChunkRanges,
+        expectedShardForEachChunk);
+}
+
+TEST_F(ShardDistributionInitSplitTest, ZonesConflictShardDistribution) {
+    const NamespaceString ns = NamespaceString::createNamespaceString_forTest("foo", "bar");
+    const ShardKeyPattern shardKey(BSON("y" << 1));
+
+    std::vector<ShardType> shardList;
+    shardList.emplace_back(
+        ShardType(shardId("0").toString(), "rs0/fakeShard0:123", {std::string("zoneA")}));
+    shardList.emplace_back(
+        ShardType(shardId("1").toString(), "rs1/fakeShard1:123", {std::string("zoneB")}));
+
+    setupShards(shardList);
+    shardRegistry()->reload(operationContext());
+
+    std::vector<TagsType> zones;
+    zones.emplace_back(nss(), "zoneB", ChunkRange(BSON("y" << MINKEY), BSON("y" << 0)));
+    zones.emplace_back(nss(), "zoneA", ChunkRange(BSON("y" << 0), BSON("y" << MAXKEY)));
+
+    ShardKeyRange range0(shardId("0"));
+    range0.setMin(BSON("y" << MINKEY));
+    range0.setMax(BSON("y" << 0));
+    ShardKeyRange range1(shardId("1"));
+    range1.setMin(BSON("y" << 0));
+    range1.setMax(BSON("y" << MAXKEY));
+    std::vector<ShardKeyRange> shardDistribution = {range0, range1};
+
+    SplitPolicyParams params{UUID::gen(), shardId("0")};
+    auto initSplitPolicy = makeInitialSplitPolicy(shardDistribution, zones);
+    ASSERT_THROWS(initSplitPolicy->createFirstChunks(operationContext(), shardKey, params),
+                  DBException);
+}
+
+TEST_F(ShardDistributionInitSplitTest, InterleaveWithZones) {
+    const NamespaceString ns = NamespaceString::createNamespaceString_forTest("foo", "bar");
+    const ShardKeyPattern shardKey(BSON("y" << 1));
+
+    std::vector<ShardType> shardList;
+    shardList.emplace_back(ShardType(shardId("0").toString(),
+                                     "rs0/fakeShard0:123",
+                                     {std::string("zoneA"), std::string("zoneB")}));
+    shardList.emplace_back(
+        ShardType(shardId("1").toString(), "rs1/fakeShard1:123", {std::string("zoneB")}));
+
+    setupShards(shardList);
+    shardRegistry()->reload(operationContext());
+
+    std::vector<TagsType> zones;
+    zones.emplace_back(nss(), "zoneA", ChunkRange(BSON("y" << MINKEY), BSON("y" << 0)));
+    zones.emplace_back(nss(), "zoneB", ChunkRange(BSON("y" << 0), BSON("y" << MAXKEY)));
+
+    ShardKeyRange range0(shardId("0"));
+    range0.setMin(BSON("y" << MINKEY));
+    range0.setMax(BSON("y" << -1));
+    ShardKeyRange range1(shardId("0"));
+    range1.setMin(BSON("y" << -1));
+    range1.setMax(BSON("y" << 1));
+    ShardKeyRange range2(shardId("1"));
+    range2.setMin(BSON("y" << 1));
+    range2.setMax(BSON("y" << MAXKEY));
+    std::vector<ShardKeyRange> shardDistribution = {range0, range1, range2};
+
+    std::vector<ChunkRange> expectedChunkRanges = {ChunkRange(BSON("y" << MINKEY), BSON("y" << -1)),
+                                                   ChunkRange(BSON("y" << -1), BSON("y" << 0)),
+                                                   ChunkRange(BSON("y" << 0), BSON("y" << 1)),
+                                                   ChunkRange(BSON("y" << 1), BSON("y" << MAXKEY))};
+    std::vector<ShardId> expectedShardForEachChunk = {
+        shardId("0"), shardId("0"), shardId("0"), shardId("1")};
+
+    checkGeneratedInitialZoneChunks(makeInitialSplitPolicy(shardDistribution, zones),
+                                    shardKey,
+                                    shardList,
+                                    shardDistribution,
+                                    expectedChunkRanges,
+                                    expectedShardForEachChunk);
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/s/resharding/coordinator_document.idl b/src/mongo/db/s/resharding/coordinator_document.idl
index 8dbd638d98b..9556515cbce 100644
--- a/src/mongo/db/s/resharding/coordinator_document.idl
+++ b/src/mongo/db/s/resharding/coordinator_document.idl
@@ -133,3 +133,8 @@ structs:
                 type: ReshardingCoordinatorMetrics
                 description: "Metrics related to the coordinator."
                 optional: true
+            shardDistribution:
+                type: array<ShardKeyRange>
+                description: "The key ranges for the new shard key. This should be continuous and complete."
+                optional: true
+                
\ No newline at end of file
diff --git a/src/mongo/db/s/resharding/resharding_coordinator_service.cpp b/src/mongo/db/s/resharding/resharding_coordinator_service.cpp
index 5952f0876f8..5f8801fef19 100644
--- a/src/mongo/db/s/resharding/resharding_coordinator_service.cpp
+++ b/src/mongo/db/s/resharding/resharding_coordinator_service.cpp
@@ -67,6 +67,7 @@
 #include "mongo/s/request_types/flush_resharding_state_change_gen.h"
 #include "mongo/s/request_types/flush_routing_table_cache_updates_gen.h"
 #include "mongo/s/resharding/resharding_coordinator_service_conflicting_op_in_progress_info.h"
+#include "mongo/s/resharding/resharding_feature_flag_gen.h"
 #include "mongo/s/sharding_feature_flags_gen.h"
 #include "mongo/s/write_ops/batched_command_request.h"
 #include "mongo/s/write_ops/batched_command_response.h"
@@ -1100,17 +1101,31 @@ ReshardingCoordinatorExternalStateImpl::calculateParticipantShardsAndChunks(
             }
         }
 
-        auto initialSplitter = SamplingBasedSplitPolicy::make(opCtx,
-                                                              coordinatorDoc.getSourceNss(),
-                                                              shardKey,
-                                                              numInitialChunks,
-                                                              std::move(parsedZones));
-
-        // Note: The resharding initial split policy doesn't care about what is the real primary
-        // shard, so just pass in a random shard.
-        const SplitPolicyParams splitParams{coordinatorDoc.getReshardingUUID(),
-                                            *donorShardIds.begin()};
-        auto splitResult = initialSplitter.createFirstChunks(opCtx, shardKey, splitParams);
+        InitialSplitPolicy::ShardCollectionConfig splitResult;
+        if (const auto& shardDistribution = coordinatorDoc.getShardDistribution()) {
+            uassert(ErrorCodes::InvalidOptions,
+                    "Resharding improvements is not enabled, should not have shardDistribution in "
+                    "coordinatorDoc",
+                    resharding::gFeatureFlagReshardingImprovements.isEnabled(
+                        serverGlobalParams.featureCompatibility));
+            auto initialSplitter = ShardDistributionSplitPolicy::make(
+                opCtx, shardKey, *shardDistribution, std::move(parsedZones));
+            const SplitPolicyParams splitParams{coordinatorDoc.getReshardingUUID(),
+                                                *donorShardIds.begin()};
+            splitResult = initialSplitter.createFirstChunks(opCtx, shardKey, splitParams);
+        } else {
+            auto initialSplitter = SamplingBasedSplitPolicy::make(opCtx,
+                                                                  coordinatorDoc.getSourceNss(),
+                                                                  shardKey,
+                                                                  numInitialChunks,
+                                                                  std::move(parsedZones));
+            // Note: The resharding initial split policy doesn't care about what is the real primary
+            // shard, so just pass in a random shard.
+            const SplitPolicyParams splitParams{coordinatorDoc.getReshardingUUID(),
+                                                *donorShardIds.begin()};
+            splitResult = initialSplitter.createFirstChunks(opCtx, shardKey, splitParams);
+        }
+
         initialChunks = std::move(splitResult.chunks);
 
         for (const auto& chunk : initialChunks) {
diff --git a/src/mongo/db/s/resharding/resharding_util.cpp b/src/mongo/db/s/resharding/resharding_util.cpp
index c2b42b2e1a3..6102c6ed66d 100644
--- a/src/mongo/db/s/resharding/resharding_util.cpp
+++ b/src/mongo/db/s/resharding/resharding_util.cpp
@@ -29,6 +29,8 @@
 
 #include "mongo/db/s/resharding/resharding_util.h"
 
+#include "mongo/bson/simple_bsonobj_comparator.h"
+#include "mongo/s/resharding/common_types_gen.h"
 #include <fmt/format.h>
 
 #include "mongo/bson/bsonobj.h"
@@ -403,5 +405,65 @@ boost::optional<Milliseconds> estimateRemainingRecipientTime(bool applyingBegan,
     return {};
 }
 
+void validateShardDistribution(const std::vector<ShardKeyRange>& shardDistribution,
+                               OperationContext* opCtx,
+                               const ShardKeyPattern& keyPattern) {
+    boost::optional<bool> hasMinMax = boost::none;
+    std::vector<ShardKeyRange> validShards;
+    for (const auto& shard : shardDistribution) {
+        uassertStatusOK(Grid::get(opCtx)->shardRegistry()->getShard(opCtx, shard.getShard()));
+        uassert(ErrorCodes::InvalidOptions,
+                "ShardKeyRange should have a pair of min/max or none of them",
+                !(shard.getMax().has_value() ^ shard.getMin().has_value()));
+        uassert(ErrorCodes::InvalidOptions,
+                "ShardKeyRange min should follow shard key's keyPattern",
+                (!shard.getMin().has_value()) || keyPattern.isShardKey(*shard.getMin()));
+        uassert(ErrorCodes::InvalidOptions,
+                "ShardKeyRange max should follow shard key's keyPattern",
+                (!shard.getMax().has_value()) || keyPattern.isShardKey(*shard.getMax()));
+
+        // Check all shardKeyRanges have min/max or none of them has min/max.
+        if (hasMinMax.has_value()) {
+            uassert(ErrorCodes::InvalidOptions,
+                    "All ShardKeyRanges should have the same min/max pattern",
+                    !(*hasMinMax ^ shard.getMax().has_value()));
+        } else {
+            hasMinMax = shard.getMax().has_value();
+        }
+
+        validShards.push_back(shard);
+    }
+
+    // If the shardDistribution contains min/max, validate whether they are continuous and complete.
+    if (hasMinMax && *hasMinMax) {
+        std::sort(validShards.begin(),
+                  validShards.end(),
+                  [](const ShardKeyRange& a, const ShardKeyRange& b) {
+                      return SimpleBSONObjComparator::kInstance.evaluate(*a.getMin() < *b.getMin());
+                  });
+
+        uassert(
+            ErrorCodes::InvalidOptions,
+            "ShardKeyRange must start at global min for the new shard key",
+            SimpleBSONObjComparator::kInstance.evaluate(validShards.front().getMin().value() ==
+                                                        keyPattern.getKeyPattern().globalMin()));
+        uassert(ErrorCodes::InvalidOptions,
+                "ShardKeyRange must end at global max for the new shard key",
+                SimpleBSONObjComparator::kInstance.evaluate(
+                    validShards.back().getMax().value() == keyPattern.getKeyPattern().globalMax()));
+
+        boost::optional<BSONObj> prevMax = boost::none;
+        for (const auto& shard : validShards) {
+            if (prevMax) {
+                uassert(ErrorCodes::InvalidOptions,
+                        "ShardKeyRanges must be continuous",
+                        SimpleBSONObjComparator::kInstance.evaluate(prevMax.value() ==
+                                                                    *shard.getMin()));
+            }
+            prevMax = *shard.getMax();
+        }
+    }
+}
+
 }  // namespace resharding
 }  // namespace mongo
diff --git a/src/mongo/db/s/resharding/resharding_util.h b/src/mongo/db/s/resharding/resharding_util.h
index c4541ab9835..83a72955cc9 100644
--- a/src/mongo/db/s/resharding/resharding_util.h
+++ b/src/mongo/db/s/resharding/resharding_util.h
@@ -322,5 +322,18 @@ std::vector<std::shared_ptr<Instance>> getReshardingStateMachines(OperationConte
     return result;
 }
 
+/**
+ * Validate the shardDistribution parameter in reshardCollection cmd, which should satisfy the
+ * following properties:
+ * - The shardKeyRanges should be continuous and cover the full data range.
+ * - Every shardKeyRange should be on the same key.
+ * - A shardKeyRange should either have no min/max or have a min/max pair.
+ * - All shardKeyRanges in the array should have the same min/max pattern.
+ * Not satisfying the rules above will cause an uassert failure.
+ */
+void validateShardDistribution(const std::vector<ShardKeyRange>& shardDistribution,
+                               OperationContext* opCtx,
+                               const ShardKeyPattern& keyPattern);
+
 }  // namespace resharding
 }  // namespace mongo