SERVER-48740 Modify ChunkMap to use std::vector internally

8 files changed, 584 insertions, 151 deletions

diff --git a/src/mongo/db/s/collection_metadata.cpp b/src/mongo/db/s/collection_metadata.cpp
index d83793ca0f7..984fa6a1059 100644
--- a/src/mongo/db/s/collection_metadata.cpp
+++ b/src/mongo/db/s/collection_metadata.cpp
@@ -130,8 +130,9 @@ Status CollectionMetadata::checkChunkIsValid(const ChunkType& chunk) const {
         existingChunk.getMax().woCompare(chunk.getMax())) {
         return {ErrorCodes::StaleShardVersion,
                 str::stream() << "Unable to find chunk with the exact bounds "
-                              << ChunkRange(chunk.getMin(), chunk.getMax()).toString()
-                              << " at collection version " << getCollVersion().toString()};
+                              << chunk.getRange().toString() << " at collection version "
+                              << getCollVersion().toString()
+                              << " found existing chunk: " << existingChunk.toString()};
     }
 
     return Status::OK();
diff --git a/src/mongo/s/catalog/type_chunk.cpp b/src/mongo/s/catalog/type_chunk.cpp
index 9e78e767649..390dd8e4e0c 100644
--- a/src/mongo/s/catalog/type_chunk.cpp
+++ b/src/mongo/s/catalog/type_chunk.cpp
@@ -173,6 +173,10 @@ boost::optional<ChunkRange> ChunkRange::overlapWith(ChunkRange const& other) con
                       le(_maxKey, other._maxKey) ? _maxKey : other._maxKey);
 }
 
+bool ChunkRange::overlaps(const ChunkRange& other) const {
+    return _minKey.woCompare(other._maxKey) < 0 && _maxKey.woCompare(other._minKey) > 0;
+}
+
 ChunkRange ChunkRange::unionWith(ChunkRange const& other) const {
     auto le = [](auto const& a, auto const& b) { return a.woCompare(b) <= 0; };
     return ChunkRange(le(_minKey, other._minKey) ? _minKey : other._minKey,
diff --git a/src/mongo/s/catalog/type_chunk.h b/src/mongo/s/catalog/type_chunk.h
index 3e602cb191c..5dd58a4cc5c 100644
--- a/src/mongo/s/catalog/type_chunk.h
+++ b/src/mongo/s/catalog/type_chunk.h
@@ -111,6 +111,11 @@ public:
     boost::optional<ChunkRange> overlapWith(ChunkRange const& other) const;
 
     /**
+     * Returns true if there is any overlap between the two ranges.
+     */
+    bool overlaps(const ChunkRange& other) const;
+
+    /**
      * Returns a range that includes *this and other. If the ranges do not overlap, it includes
      * all the space between, as well.
      */
diff --git a/src/mongo/s/chunk_manager.cpp b/src/mongo/s/chunk_manager.cpp
index 8a9b7263cad..b138a5c14d7 100644
--- a/src/mongo/s/chunk_manager.cpp
+++ b/src/mongo/s/chunk_manager.cpp
@@ -67,23 +67,71 @@ void checkAllElementsAreOfType(BSONType type, const BSONObj& o) {
             allElementsAreOfType(type, o));
 }
 
+void appendChunkTo(std::vector<std::shared_ptr<ChunkInfo>>& chunks,
+                   const std::shared_ptr<ChunkInfo>& chunk) {
+    if (!chunks.empty() && chunk->getRange().overlaps(chunks.back()->getRange())) {
+        if (chunk->getLastmod() > chunks.back()->getLastmod()) {
+            chunks.pop_back();
+            chunks.push_back(chunk);
+        }
+    } else {
+        chunks.push_back(chunk);
+    }
+}
+
+// This function processes the passed in chunks by removing the older versions of any overlapping
+// chunks. The resulting chunks must be ordered by the maximum bound and not have any
+// overlapping chunks. In order to process the original set of chunks correctly which may have
+// chunks from older versions of the map that overlap, this algorithm would need to sort by
+// ascending minimum bounds before processing it. However, since we want to take advantage of the
+// precomputed KeyString representations of the maximum bounds, this function implements the same
+// algorithm by reverse sorting the chunks by the maximum before processing but then must
+// reverse the resulting collection before it is returned.
+std::vector<std::shared_ptr<ChunkInfo>> flatten(const std::vector<ChunkType>& changedChunks) {
+    if (changedChunks.empty())
+        return std::vector<std::shared_ptr<ChunkInfo>>();
+
+    std::vector<std::shared_ptr<ChunkInfo>> changedChunkInfos(changedChunks.size());
+    std::transform(changedChunks.begin(),
+                   changedChunks.end(),
+                   changedChunkInfos.begin(),
+                   [](const auto& c) { return std::make_shared<ChunkInfo>(c); });
+
+    std::sort(changedChunkInfos.begin(), changedChunkInfos.end(), [](const auto& a, const auto& b) {
+        return a->getMaxKeyString() > b->getMaxKeyString();
+    });
+
+    std::vector<std::shared_ptr<ChunkInfo>> flattened;
+    flattened.reserve(changedChunkInfos.size());
+    flattened.push_back(changedChunkInfos[0]);
+
+    for (size_t i = 1; i < changedChunkInfos.size(); ++i) {
+        appendChunkTo(flattened, changedChunkInfos[i]);
+    }
+
+    std::reverse(flattened.begin(), flattened.end());
+
+    return flattened;
+}
+
 }  // namespace
 
-ShardVersionMap ChunkMap::constructShardVersionMap(const OID& epoch) const {
+ShardVersionMap ChunkMap::constructShardVersionMap() const {
     ShardVersionMap shardVersions;
-    ChunkInfoMap::const_iterator current = _chunkMap.cbegin();
+    ChunkVector::const_iterator current = _chunkMap.cbegin();
 
     boost::optional<BSONObj> firstMin = boost::none;
     boost::optional<BSONObj> lastMax = boost::none;
 
     while (current != _chunkMap.cend()) {
-        const auto& firstChunkInRange = current->second;
+        const auto& firstChunkInRange = *current;
         const auto& currentRangeShardId = firstChunkInRange->getShardIdAt(boost::none);
 
         // Tracks the max shard version for the shard on which the current range will reside
         auto shardVersionIt = shardVersions.find(currentRangeShardId);
         if (shardVersionIt == shardVersions.end()) {
-            shardVersionIt = shardVersions.emplace(currentRangeShardId, epoch).first;
+            shardVersionIt =
+                shardVersions.emplace(currentRangeShardId, _collectionVersion.epoch()).first;
         }
 
         auto& maxShardVersion = shardVersionIt->second.shardVersion;
@@ -91,10 +139,7 @@ ShardVersionMap ChunkMap::constructShardVersionMap(const OID& epoch) const {
         current =
             std::find_if(current,
                          _chunkMap.cend(),
-                         [&currentRangeShardId,
-                          &maxShardVersion](const ChunkInfoMap::value_type& chunkMapEntry) {
-                             const auto& currentChunk = chunkMapEntry.second;
-
+                         [&currentRangeShardId, &maxShardVersion](const auto& currentChunk) {
                              if (currentChunk->getShardIdAt(boost::none) != currentRangeShardId)
                                  return true;
 
@@ -104,27 +149,23 @@ ShardVersionMap ChunkMap::constructShardVersionMap(const OID& epoch) const {
                              return false;
                          });
 
-        const auto rangeLast = std::prev(current);
+        const auto rangeLast = *std::prev(current);
 
         const auto& rangeMin = firstChunkInRange->getMin();
-        const auto& rangeMax = rangeLast->second->getMax();
+        const auto& rangeMax = rangeLast->getMax();
 
         // Check the continuity of the chunks map
         if (lastMax && !SimpleBSONObjComparator::kInstance.evaluate(*lastMax == rangeMin)) {
             if (SimpleBSONObjComparator::kInstance.evaluate(*lastMax < rangeMin))
                 uasserted(ErrorCodes::ConflictingOperationInProgress,
                           str::stream() << "Gap exists in the routing table between chunks "
-                                        << _chunkMap.at(ShardKeyPattern::toKeyString(*lastMax))
-                                               ->getRange()
-                                               .toString()
-                                        << " and " << rangeLast->second->getRange().toString());
+                                        << findIntersectingChunk(*lastMax)->getRange().toString()
+                                        << " and " << rangeLast->getRange().toString());
             else
                 uasserted(ErrorCodes::ConflictingOperationInProgress,
                           str::stream() << "Overlap exists in the routing table between chunks "
-                                        << _chunkMap.at(ShardKeyPattern::toKeyString(*lastMax))
-                                               ->getRange()
-                                               .toString()
-                                        << " and " << rangeLast->second->getRange().toString());
+                                        << findIntersectingChunk(*lastMax)->getRange().toString()
+                                        << " and " << rangeLast->getRange().toString());
         }
 
         if (!firstMin)
@@ -149,67 +190,110 @@ ShardVersionMap ChunkMap::constructShardVersionMap(const OID& epoch) const {
     return shardVersions;
 }
 
-void ChunkMap::addChunk(const ChunkType& chunk) {
-    const auto chunkMinKeyString = ShardKeyPattern::toKeyString(chunk.getMin());
-    const auto chunkMaxKeyString = ShardKeyPattern::toKeyString(chunk.getMax());
-
-    // Returns the first chunk with a max key that is > min - implies that the chunk overlaps
-    // min
-    const auto low = _chunkMap.upper_bound(chunkMinKeyString);
-
-    // Returns the first chunk with a max key that is > max - implies that the next chunk cannot
-    // not overlap max
-    const auto high = _chunkMap.upper_bound(chunkMaxKeyString);
-
-    // If we are in the middle of splitting a chunk, for the first few
-    // chunks inserted, low == high, because both lookups will point to the
-    // same chunk (the one being split). If we're inserting the last chunk
-    // for the current chunk being split, low will point to the chunk that
-    // we're splitting, and high will point to the next chunk past the one
-    // we're splitting (which could be chunkMap.end()). In this case,
-    // std::distance(low, high) == 1. Lastly, this does not apply during
-    // the creation of the original routing table, in which case the map is
-    // empty and the first chunk that is inserted will find that low ==
-    // high, but low == chunkMap.end(), and we aren't doing a split in that
-    // case.
-    auto foundSingleChunk =
-        ((low == high || std::distance(low, high) == 1) && low != _chunkMap.end());
-
-    auto newChunk = std::make_shared<ChunkInfo>(chunk);
-    if (foundSingleChunk) {
-        auto chunkBeingReplacedBySplit = low->second;
-        auto bytesInReplacedChunk =
-            chunkBeingReplacedBySplit->getWritesTracker()->getBytesWritten();
-        newChunk->getWritesTracker()->addBytesWritten(bytesInReplacedChunk);
-    }
-
-    // Erase all chunks from the map, which overlap the chunk we got from the persistent store
-    _chunkMap.erase(low, high);
+void ChunkMap::appendChunk(const std::shared_ptr<ChunkInfo>& chunk) {
+    appendChunkTo(_chunkMap, chunk);
 
-    // Insert only the chunk itself
-    _chunkMap.insert(std::make_pair(chunkMaxKeyString, newChunk));
+    _collectionVersion = std::max(_collectionVersion, chunk->getLastmod());
 }
 
 std::shared_ptr<ChunkInfo> ChunkMap::findIntersectingChunk(const BSONObj& shardKey) const {
     const auto it = _findIntersectingChunk(shardKey);
 
     if (it != _chunkMap.end())
-        return it->second;
+        return *it;
 
     return std::shared_ptr<ChunkInfo>();
 }
 
-ChunkMap::ChunkInfoMap::const_iterator ChunkMap::_findIntersectingChunk(
-    const BSONObj& shardKey) const {
-    return _chunkMap.upper_bound(ShardKeyPattern::toKeyString(shardKey));
+void validateChunk(const std::shared_ptr<ChunkInfo>& chunk, const ChunkVersion& version) {
+    uassert(ErrorCodes::ConflictingOperationInProgress,
+            str::stream() << "Changed chunk " << chunk->toString()
+                          << " has epoch different from that of the collection " << version.epoch(),
+            version.epoch() == chunk->getLastmod().epoch());
+
+    invariant(chunk->getLastmod() >= version);
+}
+
+ChunkMap ChunkMap::createMerged(const std::vector<std::shared_ptr<ChunkInfo>>& changedChunks) {
+    size_t chunkMapIndex = 0;
+    size_t changedChunkIndex = 0;
+
+    ChunkMap updatedChunkMap(getVersion().epoch(), _chunkMap.size() + changedChunks.size());
+
+    while (chunkMapIndex < _chunkMap.size() || changedChunkIndex < changedChunks.size()) {
+        if (chunkMapIndex >= _chunkMap.size()) {
+            validateChunk(changedChunks[changedChunkIndex], getVersion());
+            updatedChunkMap.appendChunk(changedChunks[changedChunkIndex++]);
+            continue;
+        }
+
+        if (changedChunkIndex >= changedChunks.size()) {
+            updatedChunkMap.appendChunk(_chunkMap[chunkMapIndex++]);
+            continue;
+        }
+
+        auto overlap = _chunkMap[chunkMapIndex]->getRange().overlaps(
+            changedChunks[changedChunkIndex]->getRange());
+
+        if (overlap) {
+            auto& changedChunk = changedChunks[changedChunkIndex++];
+            auto& chunkInfo = _chunkMap[chunkMapIndex];
+
+            auto bytesInReplacedChunk = chunkInfo->getWritesTracker()->getBytesWritten();
+            changedChunk->getWritesTracker()->addBytesWritten(bytesInReplacedChunk);
+
+            validateChunk(changedChunk, getVersion());
+            updatedChunkMap.appendChunk(changedChunk);
+        } else {
+            updatedChunkMap.appendChunk(_chunkMap[chunkMapIndex++]);
+        }
+    }
+
+    return updatedChunkMap;
+}
+
+BSONObj ChunkMap::toBSON() const {
+    BSONObjBuilder builder;
+
+    builder.append("startingVersion"_sd, getVersion().toBSON());
+    builder.append("chunkCount", static_cast<int64_t>(_chunkMap.size()));
+
+    {
+        BSONArrayBuilder arrayBuilder(builder.subarrayStart("chunks"_sd));
+        for (const auto& chunk : _chunkMap) {
+            arrayBuilder.append(chunk->toString());
+        }
+    }
+
+    return builder.obj();
+}
+
+ChunkMap::ChunkVector::const_iterator ChunkMap::_findIntersectingChunk(const BSONObj& shardKey,
+                                                                       bool isMaxInclusive) const {
+    auto shardKeyString = ShardKeyPattern::toKeyString(shardKey);
+
+    if (!isMaxInclusive) {
+        return std::lower_bound(_chunkMap.begin(),
+                                _chunkMap.end(),
+                                shardKey,
+                                [&shardKeyString](const auto& chunkInfo, const BSONObj& shardKey) {
+                                    return chunkInfo->getMaxKeyString() < shardKeyString;
+                                });
+    } else {
+        return std::upper_bound(_chunkMap.begin(),
+                                _chunkMap.end(),
+                                shardKey,
+                                [&shardKeyString](const BSONObj& shardKey, const auto& chunkInfo) {
+                                    return shardKeyString < chunkInfo->getMaxKeyString();
+                                });
+    }
 }
 
-std::pair<ChunkMap::ChunkInfoMap::const_iterator, ChunkMap::ChunkInfoMap::const_iterator>
+std::pair<ChunkMap::ChunkVector::const_iterator, ChunkMap::ChunkVector::const_iterator>
 ChunkMap::_overlappingBounds(const BSONObj& min, const BSONObj& max, bool isMaxInclusive) const {
-    const auto itMin = _chunkMap.upper_bound(ShardKeyPattern::toKeyString(min));
+    const auto itMin = _findIntersectingChunk(min);
     const auto itMax = [&]() {
-        auto it = isMaxInclusive ? _chunkMap.upper_bound(ShardKeyPattern::toKeyString(max))
-                                 : _chunkMap.lower_bound(ShardKeyPattern::toKeyString(max));
+        auto it = _findIntersectingChunk(max, isMaxInclusive);
         return it == _chunkMap.end() ? it : ++it;
     }();
 
@@ -224,8 +308,7 @@ RoutingTableHistory::RoutingTableHistory(NamespaceString nss,
                                          KeyPattern shardKeyPattern,
                                          std::unique_ptr<CollatorInterface> defaultCollator,
                                          bool unique,
-                                         ChunkMap chunkMap,
-                                         ChunkVersion collectionVersion)
+                                         ChunkMap chunkMap)
     : _sequenceNumber(nextCMSequenceNumber.addAndFetch(1)),
       _nss(std::move(nss)),
       _uuid(uuid),
@@ -233,8 +316,7 @@ RoutingTableHistory::RoutingTableHistory(NamespaceString nss,
       _defaultCollator(std::move(defaultCollator)),
       _unique(unique),
       _chunkMap(std::move(chunkMap)),
-      _collectionVersion(collectionVersion),
-      _shardVersions(_chunkMap.constructShardVersionMap(collectionVersion.epoch())) {}
+      _shardVersions(_chunkMap.constructShardVersionMap()) {}
 
 void RoutingTableHistory::setShardStale(const ShardId& shardId) {
     if (gEnableFinerGrainedCatalogCacheRefresh) {
@@ -580,7 +662,7 @@ ChunkVersion RoutingTableHistory::_getVersion(const ShardId& shardName,
     if (it == _shardVersions.end()) {
         // Shards without explicitly tracked shard versions (meaning they have no chunks) always
         // have a version of (0, 0, epoch)
-        return ChunkVersion(0, 0, _collectionVersion.epoch());
+        return ChunkVersion(0, 0, _chunkMap.getVersion().epoch());
     }
 
     if (throwOnStaleShard && gEnableFinerGrainedCatalogCacheRefresh) {
@@ -631,34 +713,15 @@ std::shared_ptr<RoutingTableHistory> RoutingTableHistory::makeNew(
                                std::move(shardKeyPattern),
                                std::move(defaultCollator),
                                std::move(unique),
-                               ChunkMap{},
-                               {0, 0, epoch})
+                               ChunkMap{epoch})
         .makeUpdated(chunks);
 }
 
 std::shared_ptr<RoutingTableHistory> RoutingTableHistory::makeUpdated(
     const std::vector<ChunkType>& changedChunks) {
 
-    const auto startingCollectionVersion = getVersion();
-    auto chunkMap = _chunkMap;
-
-    ChunkVersion collectionVersion = startingCollectionVersion;
-    for (const auto& chunk : changedChunks) {
-        const auto& chunkVersion = chunk.getVersion();
-
-        uassert(ErrorCodes::ConflictingOperationInProgress,
-                str::stream() << "Chunk with namespace " << chunk.getNS().ns() << " and min key "
-                              << chunk.getMin()
-                              << " has epoch different from that of the collection "
-                              << chunkVersion.epoch(),
-                collectionVersion.epoch() == chunkVersion.epoch());
-
-        // Chunks must always come in increasing sorted order
-        invariant(chunkVersion >= collectionVersion);
-        collectionVersion = chunkVersion;
-
-        chunkMap.addChunk(chunk);
-    }
+    auto changedChunkInfos = flatten(changedChunks);
+    auto chunkMap = _chunkMap.createMerged(changedChunkInfos);
 
     // If at least one diff was applied, the metadata is correct, but it might not have changed so
     // in this case there is no need to recreate the chunk manager.
@@ -667,7 +730,7 @@ std::shared_ptr<RoutingTableHistory> RoutingTableHistory::makeUpdated(
     // manager object, because the write commands' code relies on changes of the chunk manager's
     // sequence number to detect batch writes not making progress because of chunks moving across
     // shards too frequently.
-    if (collectionVersion == startingCollectionVersion) {
+    if (chunkMap.getVersion() == getVersion()) {
         return shared_from_this();
     }
 
@@ -677,8 +740,7 @@ std::shared_ptr<RoutingTableHistory> RoutingTableHistory::makeUpdated(
                                 KeyPattern(getShardKeyPattern().getKeyPattern()),
                                 CollatorInterface::cloneCollator(getDefaultCollator()),
                                 isUnique(),
-                                std::move(chunkMap),
-                                collectionVersion));
+                                std::move(chunkMap)));
 }
 
 }  // namespace mongo
diff --git a/src/mongo/s/chunk_manager.h b/src/mongo/s/chunk_manager.h
index dfd7d66cc4e..3a0858f6a3d 100644
--- a/src/mongo/s/chunk_manager.h
+++ b/src/mongo/s/chunk_manager.h
@@ -68,22 +68,28 @@ using ShardVersionMap = std::map<ShardId, ShardVersionTargetingInfo>;
 // provides a simpler, high-level interface for domain specific operations without exposing the
 // underlying implementation.
 class ChunkMap {
-    // Ordered map from the max for each chunk to an entry describing the chunk
-    using ChunkInfoMap = std::map<std::string, std::shared_ptr<ChunkInfo>>;
+    // Vector of chunks ordered by max key.
+    using ChunkVector = std::vector<std::shared_ptr<ChunkInfo>>;
 
 public:
-    ChunkMap() {}
+    explicit ChunkMap(OID epoch, size_t initialCapacity = 0) : _collectionVersion(0, 0, epoch) {
+        _chunkMap.reserve(initialCapacity);
+    }
 
     size_t size() const {
         return _chunkMap.size();
     }
 
+    ChunkVersion getVersion() const {
+        return _collectionVersion;
+    }
+
     template <typename Callable>
     void forEach(Callable&& handler, const BSONObj& shardKey = BSONObj()) const {
         auto it = shardKey.isEmpty() ? _chunkMap.begin() : _findIntersectingChunk(shardKey);
 
         for (; it != _chunkMap.end(); ++it) {
-            if (!handler(it->second))
+            if (!handler(*it))
                 break;
         }
     }
@@ -96,21 +102,29 @@ public:
         const auto bounds = _overlappingBounds(min, max, isMaxInclusive);
 
         for (auto it = bounds.first; it != bounds.second; ++it) {
-            if (!handler(it->second))
+            if (!handler(*it))
                 break;
         }
     }
 
-    ShardVersionMap constructShardVersionMap(const OID& epoch) const;
-    void addChunk(const ChunkType& chunk);
+    ShardVersionMap constructShardVersionMap() const;
     std::shared_ptr<ChunkInfo> findIntersectingChunk(const BSONObj& shardKey) const;
 
+    void appendChunk(const std::shared_ptr<ChunkInfo>& chunk);
+    ChunkMap createMerged(const std::vector<std::shared_ptr<ChunkInfo>>& changedChunks);
+
+    BSONObj toBSON() const;
+
 private:
-    ChunkInfoMap::const_iterator _findIntersectingChunk(const BSONObj& shardKey) const;
-    std::pair<ChunkInfoMap::const_iterator, ChunkInfoMap::const_iterator> _overlappingBounds(
+    ChunkVector::const_iterator _findIntersectingChunk(const BSONObj& shardKey,
+                                                       bool isMaxInclusive = true) const;
+    std::pair<ChunkVector::const_iterator, ChunkVector::const_iterator> _overlappingBounds(
         const BSONObj& min, const BSONObj& max, bool isMaxInclusive) const;
 
-    ChunkInfoMap _chunkMap;
+    ChunkVector _chunkMap;
+
+    // Max version across all chunks
+    ChunkVersion _collectionVersion;
 };
 
 /**
@@ -187,7 +201,7 @@ public:
     void setAllShardsRefreshed();
 
     ChunkVersion getVersion() const {
-        return _collectionVersion;
+        return _chunkMap.getVersion();
     }
 
     /**
@@ -256,8 +270,7 @@ private:
                         KeyPattern shardKeyPattern,
                         std::unique_ptr<CollatorInterface> defaultCollator,
                         bool unique,
-                        ChunkMap chunkMap,
-                        ChunkVersion collectionVersion);
+                        ChunkMap chunkMap);
 
     ChunkVersion _getVersion(const ShardId& shardName, bool throwOnStaleShard) const;
 
@@ -284,9 +297,6 @@ private:
     // ranges must cover the complete space from [MinKey, MaxKey).
     ChunkMap _chunkMap;
 
-    // Max version across all chunks
-    const ChunkVersion _collectionVersion;
-
     // The representation of shard versions and staleness indicators for this namespace. If a
     // shard does not exist, it will not have an entry in the map.
     // Note: this declaration must not be moved before _chunkMap since it is initialized by using
diff --git a/src/mongo/s/chunk_map_test.cpp b/src/mongo/s/chunk_map_test.cpp
index 96371906310..9bc1c70594a 100644
--- a/src/mongo/s/chunk_map_test.cpp
+++ b/src/mongo/s/chunk_map_test.cpp
@@ -55,37 +55,43 @@ TEST_F(ChunkMapTest, TestAddChunk) {
     const OID epoch = OID::gen();
     ChunkVersion version{1, 0, epoch};
 
-    auto chunk =
+    auto chunk = std::make_shared<ChunkInfo>(
         ChunkType{kNss,
                   ChunkRange{getShardKeyPattern().globalMin(), getShardKeyPattern().globalMax()},
                   version,
-                  kThisShard};
+                  kThisShard});
 
-    ChunkMap chunkMap{};
-    chunkMap.addChunk(chunk);
+    ChunkMap chunkMap{epoch};
+    auto newChunkMap = chunkMap.createMerged({chunk});
 
-    ASSERT_EQ(chunkMap.size(), 1);
+    ASSERT_EQ(newChunkMap.size(), 1);
 }
 
 TEST_F(ChunkMapTest, TestEnumerateAllChunks) {
-    ChunkMap chunkMap{};
-
     const OID epoch = OID::gen();
+    ChunkMap chunkMap{epoch};
     ChunkVersion version{1, 0, epoch};
 
-    chunkMap.addChunk(ChunkType{
-        kNss, ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)}, version, kThisShard});
+    auto newChunkMap = chunkMap.createMerged(
+        {std::make_shared<ChunkInfo>(
+             ChunkType{kNss,
+                       ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                       version,
+                       kThisShard}),
 
-    chunkMap.addChunk(
-        ChunkType{kNss, ChunkRange{BSON("a" << 0), BSON("a" << 100)}, version, kThisShard});
+         std::make_shared<ChunkInfo>(
+             ChunkType{kNss, ChunkRange{BSON("a" << 0), BSON("a" << 100)}, version, kThisShard}),
 
-    chunkMap.addChunk(ChunkType{
-        kNss, ChunkRange{BSON("a" << 100), getShardKeyPattern().globalMax()}, version, kThisShard});
+         std::make_shared<ChunkInfo>(ChunkType{
+             kNss,
+             ChunkRange{BSON("a" << 100), getShardKeyPattern().globalMax()},
+             version,
+             kThisShard})});
 
     int count = 0;
     auto lastMax = getShardKeyPattern().globalMin();
 
-    chunkMap.forEach([&](const auto& chunkInfo) {
+    newChunkMap.forEach([&](const auto& chunkInfo) {
         ASSERT(SimpleBSONObjComparator::kInstance.evaluate(chunkInfo->getMax() > lastMax));
         lastMax = chunkInfo->getMax();
         count++;
@@ -93,25 +99,31 @@ TEST_F(ChunkMapTest, TestEnumerateAllChunks) {
         return true;
     });
 
-    ASSERT_EQ(count, chunkMap.size());
+    ASSERT_EQ(count, newChunkMap.size());
 }
 
 TEST_F(ChunkMapTest, TestIntersectingChunk) {
-    ChunkMap chunkMap{};
-
     const OID epoch = OID::gen();
+    ChunkMap chunkMap{epoch};
     ChunkVersion version{1, 0, epoch};
 
-    chunkMap.addChunk(ChunkType{
-        kNss, ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)}, version, kThisShard});
+    auto newChunkMap = chunkMap.createMerged(
+        {std::make_shared<ChunkInfo>(
+             ChunkType{kNss,
+                       ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                       version,
+                       kThisShard}),
 
-    chunkMap.addChunk(
-        ChunkType{kNss, ChunkRange{BSON("a" << 0), BSON("a" << 100)}, version, kThisShard});
+         std::make_shared<ChunkInfo>(
+             ChunkType{kNss, ChunkRange{BSON("a" << 0), BSON("a" << 100)}, version, kThisShard}),
 
-    chunkMap.addChunk(ChunkType{
-        kNss, ChunkRange{BSON("a" << 100), getShardKeyPattern().globalMax()}, version, kThisShard});
+         std::make_shared<ChunkInfo>(ChunkType{
+             kNss,
+             ChunkRange{BSON("a" << 100), getShardKeyPattern().globalMax()},
+             version,
+             kThisShard})});
 
-    auto intersectingChunk = chunkMap.findIntersectingChunk(BSON("a" << 50));
+    auto intersectingChunk = newChunkMap.findIntersectingChunk(BSON("a" << 50));
 
     ASSERT(intersectingChunk);
     ASSERT(
@@ -121,25 +133,31 @@ TEST_F(ChunkMapTest, TestIntersectingChunk) {
 }
 
 TEST_F(ChunkMapTest, TestEnumerateOverlappingChunks) {
-    ChunkMap chunkMap{};
-
     const OID epoch = OID::gen();
+    ChunkMap chunkMap{epoch};
     ChunkVersion version{1, 0, epoch};
 
-    chunkMap.addChunk(ChunkType{
-        kNss, ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)}, version, kThisShard});
+    auto newChunkMap = chunkMap.createMerged(
+        {std::make_shared<ChunkInfo>(
+             ChunkType{kNss,
+                       ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                       version,
+                       kThisShard}),
 
-    chunkMap.addChunk(
-        ChunkType{kNss, ChunkRange{BSON("a" << 0), BSON("a" << 100)}, version, kThisShard});
+         std::make_shared<ChunkInfo>(
+             ChunkType{kNss, ChunkRange{BSON("a" << 0), BSON("a" << 100)}, version, kThisShard}),
 
-    chunkMap.addChunk(ChunkType{
-        kNss, ChunkRange{BSON("a" << 100), getShardKeyPattern().globalMax()}, version, kThisShard});
+         std::make_shared<ChunkInfo>(ChunkType{
+             kNss,
+             ChunkRange{BSON("a" << 100), getShardKeyPattern().globalMax()},
+             version,
+             kThisShard})});
 
     auto min = BSON("a" << -50);
     auto max = BSON("a" << 150);
 
     int count = 0;
-    chunkMap.forEachOverlappingChunk(min, max, true, [&](const auto& chunk) {
+    newChunkMap.forEachOverlappingChunk(min, max, true, [&](const auto& chunk) {
         count++;
         return true;
     });
diff --git a/src/mongo/s/routing_table_history_test.cpp b/src/mongo/s/routing_table_history_test.cpp
index 974437d512c..0d43cc9ba41 100644
--- a/src/mongo/s/routing_table_history_test.cpp
+++ b/src/mongo/s/routing_table_history_test.cpp
@@ -315,5 +315,339 @@ TEST_F(RoutingTableHistoryTestThreeInitialChunks,
                               expectedBytesInChunksNotSplit);
 }
 
+TEST_F(RoutingTableHistoryTest, TestSplits) {
+    const OID epoch = OID::gen();
+    ChunkVersion version{1, 0, epoch};
+
+    auto chunkAll =
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), getShardKeyPattern().globalMax()},
+                  version,
+                  kThisShard};
+
+    auto rt = RoutingTableHistory::makeNew(
+        kNss, UUID::gen(), getShardKeyPattern(), nullptr, false, epoch, {chunkAll});
+
+    std::vector<ChunkType> chunks1 = {
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                  ChunkVersion{2, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard}};
+
+    auto rt1 = rt->makeUpdated(chunks1);
+    auto v1 = ChunkVersion{2, 2, epoch};
+    ASSERT_EQ(v1, rt1->getVersion(kThisShard));
+
+    std::vector<ChunkType> chunks2 = {
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << -1)},
+                  ChunkVersion{3, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << -1), BSON("a" << 0)},
+                  ChunkVersion{3, 2, epoch},
+                  kThisShard}};
+
+    auto rt2 = rt1->makeUpdated(chunks2);
+    auto v2 = ChunkVersion{3, 2, epoch};
+    ASSERT_EQ(v2, rt2->getVersion(kThisShard));
+}
+
+TEST_F(RoutingTableHistoryTest, TestReplaceChunk) {
+    const OID epoch = OID::gen();
+    ChunkVersion version{2, 2, epoch};
+
+    std::vector<ChunkType> initialChunks = {
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                  ChunkVersion{2, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard}};
+
+    auto rt = RoutingTableHistory::makeNew(
+        kNss, UUID::gen(), getShardKeyPattern(), nullptr, false, epoch, {initialChunks});
+
+    std::vector<ChunkType> changedChunks = {
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard}};
+
+    auto rt1 = rt->makeUpdated(changedChunks);
+    auto v1 = ChunkVersion{2, 2, epoch};
+    ASSERT_EQ(v1, rt1->getVersion(kThisShard));
+    ASSERT_EQ(rt1->numChunks(), 2);
+    ASSERT_EQ(rt.get(), rt1.get());
+
+    std::shared_ptr<ChunkInfo> found;
+
+    rt1->forEachChunk(
+        [&](auto& chunkInfo) {
+            if (chunkInfo->getShardIdAt(boost::none) == kThisShard) {
+                found = chunkInfo;
+                return false;
+            }
+            return true;
+        },
+        BSON("a" << 0));
+    ASSERT(found);
+}
+
+TEST_F(RoutingTableHistoryTest, TestReplaceEmptyChunk) {
+    const OID epoch = OID::gen();
+
+    std::vector<ChunkType> initialChunks = {
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), getShardKeyPattern().globalMax()},
+                  ChunkVersion{1, 0, epoch},
+                  kThisShard}};
+
+    auto rt = RoutingTableHistory::makeNew(
+        kNss, UUID::gen(), getShardKeyPattern(), nullptr, false, epoch, initialChunks);
+
+    ASSERT_EQ(rt->numChunks(), 1);
+
+    std::vector<ChunkType> changedChunks = {
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                  ChunkVersion{2, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard}};
+
+    auto rt1 = rt->makeUpdated(changedChunks);
+    auto v1 = ChunkVersion{2, 2, epoch};
+    ASSERT_EQ(v1, rt1->getVersion(kThisShard));
+    ASSERT_EQ(rt1->numChunks(), 2);
+    ASSERT_NE(rt.get(), rt1.get());
+
+    std::shared_ptr<ChunkInfo> found;
+
+    rt1->forEachChunk(
+        [&](auto& chunkInfo) {
+            if (chunkInfo->getShardIdAt(boost::none) == kThisShard) {
+                found = chunkInfo;
+                return false;
+            }
+            return true;
+        },
+        BSON("a" << 0));
+    ASSERT(found);
+}
+
+TEST_F(RoutingTableHistoryTest, TestUseLatestVersions) {
+    const OID epoch = OID::gen();
+
+    std::vector<ChunkType> initialChunks = {
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), getShardKeyPattern().globalMax()},
+                  ChunkVersion{1, 0, epoch},
+                  kThisShard}};
+
+    auto rt = RoutingTableHistory::makeNew(
+        kNss, UUID::gen(), getShardKeyPattern(), nullptr, false, epoch, initialChunks);
+
+    ASSERT_EQ(rt->numChunks(), 1);
+
+    std::vector<ChunkType> changedChunks = {
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), getShardKeyPattern().globalMax()},
+                  ChunkVersion{1, 0, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                  ChunkVersion{2, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard}};
+
+    auto rt1 = rt->makeUpdated(changedChunks);
+    auto v1 = ChunkVersion{2, 2, epoch};
+    ASSERT_EQ(v1, rt1->getVersion(kThisShard));
+    ASSERT_EQ(rt1->numChunks(), 2);
+    ASSERT_NE(rt.get(), rt1.get());
+}
+
+TEST_F(RoutingTableHistoryTest, TestOutOfOrderVersion) {
+    const OID epoch = OID::gen();
+
+    std::vector<ChunkType> initialChunks = {
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                  ChunkVersion{2, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard}};
+
+    auto rt = RoutingTableHistory::makeNew(
+        kNss, UUID::gen(), getShardKeyPattern(), nullptr, false, epoch, initialChunks);
+
+    ASSERT_EQ(rt->numChunks(), 2);
+
+    std::vector<ChunkType> changedChunks = {
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), getShardKeyPattern().globalMax()},
+                  ChunkVersion{3, 0, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                  ChunkVersion{3, 1, epoch},
+                  kThisShard}};
+
+    auto rt1 = rt->makeUpdated(changedChunks);
+    auto v1 = ChunkVersion{3, 1, epoch};
+    ASSERT_EQ(v1, rt1->getVersion(kThisShard));
+    ASSERT_EQ(rt1->numChunks(), 2);
+    ASSERT_NE(rt.get(), rt1.get());
+
+    auto chunk1 = rt1->findIntersectingChunk(BSON("a" << 0));
+    ASSERT_EQ(chunk1->getLastmod(), ChunkVersion(3, 0, epoch));
+    ASSERT_EQ(chunk1->getMin().woCompare(BSON("a" << 0)), 0);
+    ASSERT_EQ(chunk1->getMax().woCompare(getShardKeyPattern().globalMax()), 0);
+}
+
+TEST_F(RoutingTableHistoryTest, TestMergeChunks) {
+    const OID epoch = OID::gen();
+
+    std::vector<ChunkType> initialChunks = {
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 0), BSON("a" << 10)},
+                  ChunkVersion{2, 0, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 0)},
+                  ChunkVersion{2, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 10), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard}};
+
+    auto rt = RoutingTableHistory::makeNew(
+        kNss, UUID::gen(), getShardKeyPattern(), nullptr, false, epoch, initialChunks);
+
+    ASSERT_EQ(rt->numChunks(), 3);
+    ASSERT_EQ(rt->getVersion(), ChunkVersion(2, 2, epoch));
+
+    std::vector<ChunkType> changedChunks = {
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 10), getShardKeyPattern().globalMax()},
+                  ChunkVersion{3, 0, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 10)},
+                  ChunkVersion{3, 1, epoch},
+                  kThisShard}};
+
+    auto rt1 = rt->makeUpdated(changedChunks);
+    auto v1 = ChunkVersion{3, 1, epoch};
+    ASSERT_EQ(v1, rt1->getVersion(kThisShard));
+    ASSERT_EQ(rt1->numChunks(), 2);
+    ASSERT_NE(rt.get(), rt1.get());
+}
+
+TEST_F(RoutingTableHistoryTest, TestMergeChunksOrdering) {
+    const OID epoch = OID::gen();
+
+    std::vector<ChunkType> initialChunks = {
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << -10), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 0, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << -500)},
+                  ChunkVersion{2, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << -500), BSON("a" << -10)},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard}};
+
+    auto rt = RoutingTableHistory::makeNew(
+        kNss, UUID::gen(), getShardKeyPattern(), nullptr, false, epoch, initialChunks);
+
+    ASSERT_EQ(rt->numChunks(), 3);
+    ASSERT_EQ(rt->getVersion(), ChunkVersion(2, 2, epoch));
+
+    std::vector<ChunkType> changedChunks = {
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << -500), BSON("a" << -10)},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << -10)},
+                  ChunkVersion{3, 1, epoch},
+                  kThisShard}};
+
+    auto rt1 = rt->makeUpdated(changedChunks);
+    auto v1 = ChunkVersion{3, 1, epoch};
+    ASSERT_EQ(v1, rt1->getVersion(kThisShard));
+    ASSERT_EQ(rt1->numChunks(), 2);
+    ASSERT_NE(rt.get(), rt1.get());
+
+    auto chunk1 = rt1->findIntersectingChunk(BSON("a" << -500));
+    ASSERT_EQ(chunk1->getLastmod(), ChunkVersion(3, 1, epoch));
+    ASSERT_EQ(chunk1->getMin().woCompare(getShardKeyPattern().globalMin()), 0);
+    ASSERT_EQ(chunk1->getMax().woCompare(BSON("a" << -10)), 0);
+}
+
+TEST_F(RoutingTableHistoryTest, TestFlatten) {
+    const OID epoch = OID::gen();
+
+    std::vector<ChunkType> initialChunks = {
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 10)},
+                  ChunkVersion{2, 0, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 10), BSON("a" << 20)},
+                  ChunkVersion{2, 1, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 20), getShardKeyPattern().globalMax()},
+                  ChunkVersion{2, 2, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), getShardKeyPattern().globalMax()},
+                  ChunkVersion{3, 0, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{getShardKeyPattern().globalMin(), BSON("a" << 10)},
+                  ChunkVersion{4, 0, epoch},
+                  kThisShard},
+        ChunkType{kNss,
+                  ChunkRange{BSON("a" << 10), getShardKeyPattern().globalMax()},
+                  ChunkVersion{4, 1, epoch},
+                  kThisShard},
+    };
+
+    auto rt = RoutingTableHistory::makeNew(
+        kNss, UUID::gen(), getShardKeyPattern(), nullptr, false, epoch, initialChunks);
+
+    ASSERT_EQ(rt->numChunks(), 2);
+    ASSERT_EQ(rt->getVersion(), ChunkVersion(4, 1, epoch));
+
+    auto chunk1 = rt->findIntersectingChunk(BSON("a" << 0));
+    ASSERT_EQ(chunk1->getLastmod(), ChunkVersion(4, 0, epoch));
+    ASSERT_EQ(chunk1->getMin().woCompare(getShardKeyPattern().globalMin()), 0);
+    ASSERT_EQ(chunk1->getMax().woCompare(BSON("a" << 10)), 0);
+}
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/s/shard_key_pattern.cpp b/src/mongo/s/shard_key_pattern.cpp
index 7fd4830afd1..94e6db3dbea 100644
--- a/src/mongo/s/shard_key_pattern.cpp
+++ b/src/mongo/s/shard_key_pattern.cpp
@@ -254,13 +254,12 @@ std::string ShardKeyPattern::toString() const {
 }
 
 std::string ShardKeyPattern::toKeyString(const BSONObj& shardKey) {
-    BSONObjBuilder strippedKeyValue;
-    for (const auto& elem : shardKey) {
-        strippedKeyValue.appendAs(elem, ""_sd);
-    }
+    KeyString::Builder ks(KeyString::Version::V1, Ordering::allAscending());
 
-    KeyString::Builder ks(
-        KeyString::Version::V1, strippedKeyValue.done(), Ordering::allAscending());
+    BSONObjIterator it(shardKey);
+    while (auto elem = it.next()) {
+        ks.appendBSONElement(elem);
+    }
 
     return {ks.getBuffer(), ks.getSize()};
 }