SERVER-29817 Move all ChunkManager construction logic into chunk_manager.cpp; hide implementation details.

3 files changed, 142 insertions, 93 deletions

diff --git a/src/mongo/s/catalog_cache.cpp b/src/mongo/s/catalog_cache.cpp
index 255ed96df69..831e7c4a264 100644
--- a/src/mongo/s/catalog_cache.cpp
+++ b/src/mongo/s/catalog_cache.cpp
@@ -36,6 +36,7 @@
 #include "mongo/base/status_with.h"
 #include "mongo/db/query/collation/collator_factory_interface.h"
 #include "mongo/db/repl/optime_with.h"
+#include "mongo/platform/unordered_set.h"
 #include "mongo/s/catalog/sharding_catalog_client.h"
 #include "mongo/s/catalog/type_collection.h"
 #include "mongo/s/catalog/type_database.h"
@@ -76,79 +77,36 @@ std::shared_ptr<ChunkManager> refreshCollectionRoutingInfo(
 
     const auto collectionAndChunks = uassertStatusOK(std::move(swCollectionAndChangedChunks));
 
-    // Check whether the collection epoch might have changed
-    ChunkVersion startingCollectionVersion;
-    ChunkMap chunkMap =
-        SimpleBSONObjComparator::kInstance.makeBSONObjIndexedMap<std::shared_ptr<Chunk>>();
-
-    if (!existingRoutingInfo) {
-        // If we don't have a basis chunk manager, do a full refresh
-        startingCollectionVersion = ChunkVersion(0, 0, collectionAndChunks.epoch);
-    } else if (existingRoutingInfo->getVersion().epoch() != collectionAndChunks.epoch) {
-        // If the collection's epoch has changed, do a full refresh
-        startingCollectionVersion = ChunkVersion(0, 0, collectionAndChunks.epoch);
-    } else {
-        startingCollectionVersion = existingRoutingInfo->getVersion();
-        chunkMap = existingRoutingInfo->chunkMap();
-    }
-
-    ChunkVersion collectionVersion = startingCollectionVersion;
-
-    for (const auto& chunk : collectionAndChunks.changedChunks) {
-        const auto& chunkVersion = chunk.getVersion();
-
-        uassert(ErrorCodes::ConflictingOperationInProgress,
-                str::stream() << "Chunk " << chunk.genID(nss.ns(), chunk.getMin())
-                              << " has epoch different from that of the collection "
-                              << chunkVersion.epoch(),
-                collectionVersion.epoch() == chunkVersion.epoch());
-
-        // Chunks must always come in incrementally sorted order
-        invariant(chunkVersion >= collectionVersion);
-        collectionVersion = chunkVersion;
-
-        // Ensure chunk references a valid shard and that the shard is available and loaded
-        uassertStatusOK(Grid::get(opCtx)->shardRegistry()->getShard(opCtx, chunk.getShard()));
-
-        // Returns the first chunk with a max key that is > min - implies that the chunk overlaps
-        // min
-        const auto low = chunkMap.upper_bound(chunk.getMin());
-
-        // 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(chunk.getMax());
+    auto chunkManager = [&] {
+        // If we have routing info already and it's for the same collection epoch, we're updating.
+        // Otherwise, we're making a whole new routing table.
+        if (existingRoutingInfo &&
+            existingRoutingInfo->getVersion().epoch() == collectionAndChunks.epoch) {
 
-        // Erase all chunks from the map, which overlap the chunk we got from the persistent store
-        chunkMap.erase(low, high);
-
-        // Insert only the chunk itself
-        chunkMap.insert(std::make_pair(chunk.getMax(), std::make_shared<Chunk>(chunk)));
-    }
-
-    // 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.
-    //
-    // NOTE: In addition to the above statement, it is also important that we return the same chunk
-    // 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) {
-        return existingRoutingInfo;
-    }
-
-    std::unique_ptr<CollatorInterface> defaultCollator;
-    if (!collectionAndChunks.defaultCollation.isEmpty()) {
-        // The collation should have been validated upon collection creation
-        defaultCollator = uassertStatusOK(CollatorFactoryInterface::get(opCtx->getServiceContext())
-                                              ->makeFromBSON(collectionAndChunks.defaultCollation));
+            return existingRoutingInfo->makeUpdated(collectionAndChunks.changedChunks);
+        }
+        auto defaultCollator = [&]() -> std::unique_ptr<CollatorInterface> {
+            if (!collectionAndChunks.defaultCollation.isEmpty()) {
+                // The collation should have been validated upon collection creation
+                return uassertStatusOK(CollatorFactoryInterface::get(opCtx->getServiceContext())
+                                           ->makeFromBSON(collectionAndChunks.defaultCollation));
+            }
+            return nullptr;
+        }();
+        return ChunkManager::makeNew(nss,
+                                     KeyPattern(collectionAndChunks.shardKeyPattern),
+                                     std::move(defaultCollator),
+                                     collectionAndChunks.shardKeyIsUnique,
+                                     collectionAndChunks.epoch,
+                                     collectionAndChunks.changedChunks);
+    }();
+
+    std::set<ShardId> shardIds;
+    chunkManager->getAllShardIds(&shardIds);
+    for (const auto& shardId : shardIds) {
+        uassertStatusOK(Grid::get(opCtx)->shardRegistry()->getShard(opCtx, shardId));
     }
-
-    return stdx::make_unique<ChunkManager>(nss,
-                                           KeyPattern(collectionAndChunks.shardKeyPattern),
-                                           std::move(defaultCollator),
-                                           collectionAndChunks.shardKeyIsUnique,
-                                           std::move(chunkMap),
-                                           collectionVersion);
+    return chunkManager;
 }
 
 }  // namespace
diff --git a/src/mongo/s/chunk_manager.cpp b/src/mongo/s/chunk_manager.cpp
index 7b19211baaf..d07dc9a66eb 100644
--- a/src/mongo/s/chunk_manager.cpp
+++ b/src/mongo/s/chunk_manager.cpp
@@ -41,7 +41,6 @@
 #include "mongo/db/query/index_bounds_builder.h"
 #include "mongo/db/query/query_planner.h"
 #include "mongo/db/query/query_planner_common.h"
-#include "mongo/s/grid.h"
 #include "mongo/util/log.h"
 
 namespace mongo {
@@ -75,8 +74,6 @@ ChunkManager::ChunkManager(NamespaceString nss,
       _chunkMapViews(_constructChunkMapViews(collectionVersion.epoch(), _chunkMap)),
       _collectionVersion(collectionVersion) {}
 
-ChunkManager::~ChunkManager() = default;
-
 std::shared_ptr<Chunk> ChunkManager::findIntersectingChunk(const BSONObj& shardKey,
                                                            const BSONObj& collation) const {
     const bool hasSimpleCollation = (collation.isEmpty() && !_defaultCollator) ||
@@ -342,7 +339,6 @@ std::string ChunkManager::toString() const {
 
 ChunkManager::ChunkMapViews ChunkManager::_constructChunkMapViews(const OID& epoch,
                                                                   const ChunkMap& chunkMap) {
-    invariant(!chunkMap.empty());
 
     ChunkRangeMap chunkRangeMap =
         SimpleBSONObjComparator::kInstance.makeBSONObjIndexedMap<ShardAndChunkRange>();
@@ -409,13 +405,86 @@ ChunkManager::ChunkMapViews ChunkManager::_constructChunkMapViews(const OID& epo
         invariant(maxShardVersion.isSet());
     }
 
-    invariant(!chunkRangeMap.empty());
-    invariant(!shardVersions.empty());
+    if (!chunkMap.empty()) {
+        invariant(!chunkRangeMap.empty());
+        invariant(!shardVersions.empty());
 
-    checkAllElementsAreOfType(MinKey, chunkRangeMap.begin()->second.min());
-    checkAllElementsAreOfType(MaxKey, chunkRangeMap.rbegin()->first);
+        checkAllElementsAreOfType(MinKey, chunkRangeMap.begin()->second.min());
+        checkAllElementsAreOfType(MaxKey, chunkRangeMap.rbegin()->first);
+    }
 
     return {std::move(chunkRangeMap), std::move(shardVersions)};
 }
 
+std::shared_ptr<ChunkManager> ChunkManager::makeNew(
+    NamespaceString nss,
+    KeyPattern shardKeyPattern,
+    std::unique_ptr<CollatorInterface> defaultCollator,
+    bool unique,
+    OID epoch,
+    const std::vector<ChunkType>& chunks) {
+
+    return ChunkManager(
+               std::move(nss),
+               std::move(shardKeyPattern),
+               std::move(defaultCollator),
+               std::move(unique),
+               SimpleBSONObjComparator::kInstance.makeBSONObjIndexedMap<std::shared_ptr<Chunk>>(),
+               {0, 0, epoch})
+        .makeUpdated(chunks);
+}
+
+std::shared_ptr<ChunkManager> ChunkManager::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 " << chunk.genID(getns(), chunk.getMin())
+                              << " has epoch different from that of the collection "
+                              << chunkVersion.epoch(),
+                collectionVersion.epoch() == chunkVersion.epoch());
+
+        // Chunks must always come in incrementally sorted order
+        invariant(chunkVersion >= collectionVersion);
+        collectionVersion = chunkVersion;
+
+        // Returns the first chunk with a max key that is > min - implies that the chunk overlaps
+        // min
+        const auto low = chunkMap.upper_bound(chunk.getMin());
+
+        // 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(chunk.getMax());
+
+        // Erase all chunks from the map, which overlap the chunk we got from the persistent store
+        chunkMap.erase(low, high);
+
+        // Insert only the chunk itself
+        chunkMap.insert(std::make_pair(chunk.getMax(), std::make_shared<Chunk>(chunk)));
+    }
+
+    // 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.
+    //
+    // NOTE: In addition to the above statement, it is also important that we return the same chunk
+    // 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) {
+        return shared_from_this();
+    }
+
+    return std::shared_ptr<ChunkManager>(
+        new ChunkManager(_nss,
+                         KeyPattern(getShardKeyPattern().getKeyPattern()),
+                         CollatorInterface::cloneCollator(getDefaultCollator()),
+                         isUnique(),
+                         std::move(chunkMap),
+                         collectionVersion));
+}
 }  // namespace mongo
diff --git a/src/mongo/s/chunk_manager.h b/src/mongo/s/chunk_manager.h
index 374fe929234..bb2580f8b4a 100644
--- a/src/mongo/s/chunk_manager.h
+++ b/src/mongo/s/chunk_manager.h
@@ -53,7 +53,10 @@ using ChunkMap = BSONObjIndexedMap<std::shared_ptr<Chunk>>;
 // Map from a shard is to the max chunk version on that shard
 using ShardVersionMap = std::map<ShardId, ChunkVersion>;
 
-class ChunkManager {
+/**
+ * In-memory representation of the routing table for a single sharded collection.
+ */
+class ChunkManager : public std::enable_shared_from_this<ChunkManager> {
     MONGO_DISALLOW_COPYING(ChunkManager);
 
 public:
@@ -100,14 +103,32 @@ public:
         ConstChunkIterator _end;
     };
 
-    ChunkManager(NamespaceString nss,
-                 KeyPattern shardKeyPattern,
-                 std::unique_ptr<CollatorInterface> defaultCollator,
-                 bool unique,
-                 ChunkMap chunkMap,
-                 ChunkVersion collectionVersion);
+    /**
+     * Makes an instance with a routing table for collection "nss", sharded on
+     * "shardKeyPattern".
+     *
+     * "defaultCollator" is the default collation for the collection, "unique" indicates whether
+     * or not the shard key for each document will be globally unique, and "epoch" is the globally
+     * unique identifier for this version of the collection.
+     *
+     * The "chunks" vector must contain the chunk routing information sorted in ascending order by
+     * chunk version, and adhere to the requirements of the routing table update algorithm.
+     */
+    static std::shared_ptr<ChunkManager> makeNew(NamespaceString nss,
+                                                 KeyPattern shardKeyPattern,
+                                                 std::unique_ptr<CollatorInterface> defaultCollator,
+                                                 bool unique,
+                                                 OID epoch,
+                                                 const std::vector<ChunkType>& chunks);
 
-    ~ChunkManager();
+    /**
+     * Constructs a new instance with a routing table updated according to the changes described
+     * in "changedChunks".
+     *
+     * The changes in "changedChunks" must be sorted in ascending order by chunk version, and adhere
+     * to the requirements of the routing table update algorithm.
+     */
+    std::shared_ptr<ChunkManager> makeUpdated(const std::vector<ChunkType>& changedChunks);
 
     /**
      * Returns an increasing number of the reload sequence number of this chunk manager.
@@ -142,10 +163,6 @@ public:
         return {ConstChunkIterator{_chunkMap.cbegin()}, ConstChunkIterator{_chunkMap.cend()}};
     }
 
-    const ChunkMap& chunkMap() const {
-        return _chunkMap;
-    }
-
     int numChunks() const {
         return _chunkMap.size();
     }
@@ -219,8 +236,6 @@ public:
     std::string toString() const;
 
 private:
-    friend class CollectionRoutingDataLoader;
-
     /**
      * Represents a range of chunk keys [getMin(), getMax()) and the id of the shard on which they
      * reside according to the metadata.
@@ -259,6 +274,13 @@ private:
      */
     static ChunkMapViews _constructChunkMapViews(const OID& epoch, const ChunkMap& chunkMap);
 
+    ChunkManager(NamespaceString nss,
+                 KeyPattern shardKeyPattern,
+                 std::unique_ptr<CollatorInterface> defaultCollator,
+                 bool unique,
+                 ChunkMap chunkMap,
+                 ChunkVersion collectionVersion);
+
     // The shard versioning mechanism hinges on keeping track of the number of times we reload
     // ChunkManagers.
     const unsigned long long _sequenceNumber;