SERVER-32901 Pull the CollectionShardingState map out of ShardingState

5 files changed, 69 insertions, 55 deletions

diff --git a/src/mongo/db/s/collection_sharding_state.cpp b/src/mongo/db/s/collection_sharding_state.cpp
index 9b4d84756aa..e94547ccfe9 100644
--- a/src/mongo/db/s/collection_sharding_state.cpp
+++ b/src/mongo/db/s/collection_sharding_state.cpp
@@ -32,6 +32,7 @@
 
 #include "mongo/db/s/collection_sharding_state.h"
 
+#include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/bson/util/bson_extract.h"
 #include "mongo/db/catalog/catalog_raii.h"
 #include "mongo/db/client.h"
@@ -54,7 +55,6 @@
 #include "mongo/s/catalog/type_shard_collection.h"
 #include "mongo/s/catalog_cache.h"
 #include "mongo/s/catalog_cache_loader.h"
-#include "mongo/s/chunk_version.h"
 #include "mongo/s/cluster_identity_loader.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/stale_exception.h"
@@ -66,6 +66,61 @@ namespace {
 // How long to wait before starting cleanup of an emigrated chunk range
 MONGO_EXPORT_SERVER_PARAMETER(orphanCleanupDelaySecs, int, 900);  // 900s = 15m
 
+// This map matches 1:1 with the set of collections in the storage catalog. It is not safe to
+// look-up values from this map without holding some form of collection lock. It is only safe to
+// add/remove values when holding X lock on the respective namespace.
+class CollectionShardingStateMap {
+    MONGO_DISALLOW_COPYING(CollectionShardingStateMap);
+
+public:
+    CollectionShardingStateMap() = default;
+
+    CollectionShardingState& getOrCreate(OperationContext* opCtx, const std::string& ns) {
+        stdx::lock_guard<stdx::mutex> lg(_mutex);
+
+        auto it = _collections.find(ns);
+        if (it == _collections.end()) {
+            auto inserted =
+                _collections.emplace(ns,
+                                     std::make_unique<CollectionShardingState>(
+                                         opCtx->getServiceContext(), NamespaceString(ns)));
+            invariant(inserted.second);
+            it = std::move(inserted.first);
+        }
+
+        return *it->second;
+    }
+
+    void report(BSONObjBuilder* builder) {
+        BSONObjBuilder versionB(builder->subobjStart("versions"));
+
+        {
+            stdx::lock_guard<stdx::mutex> lg(_mutex);
+
+            for (auto& coll : _collections) {
+                ScopedCollectionMetadata metadata = coll.second->getMetadata();
+                if (metadata) {
+                    versionB.appendTimestamp(coll.first, metadata->getShardVersion().toLong());
+                } else {
+                    versionB.appendTimestamp(coll.first, ChunkVersion::UNSHARDED().toLong());
+                }
+            }
+        }
+
+        versionB.done();
+    }
+
+private:
+    mutable stdx::mutex _mutex;
+
+    using CollectionsMap =
+        stdx::unordered_map<std::string, std::unique_ptr<CollectionShardingState>>;
+    CollectionsMap _collections;
+};
+
+const auto collectionShardingStateMap =
+    ServiceContext::declareDecoration<CollectionShardingStateMap>();
+
 /**
  * Used to perform shard identity initialization once it is certain that the document is committed.
  */
@@ -145,8 +200,13 @@ CollectionShardingState* CollectionShardingState::get(OperationContext* opCtx,
     // Collection lock must be held to have a reference to the collection's sharding state
     dassert(opCtx->lockState()->isCollectionLockedForMode(ns, MODE_IS));
 
-    ShardingState* const shardingState = ShardingState::get(opCtx);
-    return shardingState->getNS(ns, opCtx);
+    auto& collectionsMap = collectionShardingStateMap(opCtx->getServiceContext());
+    return &collectionsMap.getOrCreate(opCtx, ns);
+}
+
+void CollectionShardingState::report(OperationContext* opCtx, BSONObjBuilder* builder) {
+    auto& collectionsMap = collectionShardingStateMap(opCtx->getServiceContext());
+    collectionsMap.report(builder);
 }
 
 ScopedCollectionMetadata CollectionShardingState::getMetadata() {
diff --git a/src/mongo/db/s/collection_sharding_state.h b/src/mongo/db/s/collection_sharding_state.h
index 7af282f3463..30ba8a295d0 100644
--- a/src/mongo/db/s/collection_sharding_state.h
+++ b/src/mongo/db/s/collection_sharding_state.h
@@ -34,7 +34,6 @@
 #include "mongo/base/disallow_copying.h"
 #include "mongo/base/string_data.h"
 #include "mongo/db/namespace_string.h"
-#include "mongo/db/s/collection_range_deleter.h"
 #include "mongo/db/s/metadata_manager.h"
 #include "mongo/util/concurrency/notification.h"
 
@@ -42,6 +41,7 @@ namespace mongo {
 
 class BalancerConfiguration;
 class BSONObj;
+class BSONObjBuilder;
 struct ChunkVersion;
 class CollectionMetadata;
 class MigrationSourceManager;
@@ -93,6 +93,8 @@ public:
     static CollectionShardingState* get(OperationContext* opCtx, const NamespaceString& nss);
     static CollectionShardingState* get(OperationContext* opCtx, const std::string& ns);
 
+    static void report(OperationContext* opCtx, BSONObjBuilder* builder);
+
     /**
      * Returns the chunk metadata for the collection. The metadata it represents lives as long as
      * the object itself, and the collection, exist. After dropping the collection lock, the
diff --git a/src/mongo/db/s/sharding_state.cpp b/src/mongo/db/s/sharding_state.cpp
index 77298091f12..5ea59edb813 100644
--- a/src/mongo/db/s/sharding_state.cpp
+++ b/src/mongo/db/s/sharding_state.cpp
@@ -45,7 +45,6 @@
 #include "mongo/db/ops/update_lifecycle_impl.h"
 #include "mongo/db/repl/optime.h"
 #include "mongo/db/repl/replication_coordinator.h"
-#include "mongo/db/s/collection_sharding_state.h"
 #include "mongo/db/s/operation_sharding_state.h"
 #include "mongo/db/s/sharded_connection_info.h"
 #include "mongo/db/s/sharding_initialization_mongod.h"
@@ -185,21 +184,6 @@ Status ShardingState::updateConfigServerOpTimeFromMetadata(OperationContext* opC
     return Status::OK();
 }
 
-CollectionShardingState* ShardingState::getNS(const std::string& ns, OperationContext* opCtx) {
-    stdx::lock_guard<stdx::mutex> lk(_mutex);
-    CollectionShardingStateMap::iterator it = _collections.find(ns);
-    if (it == _collections.end()) {
-        auto inserted =
-            _collections.insert(make_pair(ns,
-                                          stdx::make_unique<CollectionShardingState>(
-                                              opCtx->getServiceContext(), NamespaceString(ns))));
-        invariant(inserted.second);
-        it = std::move(inserted.first);
-    }
-
-    return it->second.get();
-}
-
 ChunkSplitter* ShardingState::getChunkSplitter() {
     return _chunkSplitter.get();
 }
@@ -565,20 +549,6 @@ void ShardingState::appendInfo(OperationContext* opCtx, BSONObjBuilder& builder)
                    Grid::get(opCtx)->shardRegistry()->getConfigServerConnectionString().toString());
     builder.append("shardName", _shardName);
     builder.append("clusterId", _clusterId);
-
-    BSONObjBuilder versionB(builder.subobjStart("versions"));
-    for (CollectionShardingStateMap::const_iterator it = _collections.begin();
-         it != _collections.end();
-         ++it) {
-        ScopedCollectionMetadata metadata = it->second->getMetadata();
-        if (metadata) {
-            versionB.appendTimestamp(it->first, metadata->getShardVersion().toLong());
-        } else {
-            versionB.appendTimestamp(it->first, ChunkVersion::UNSHARDED().toLong());
-        }
-    }
-
-    versionB.done();
 }
 
 bool ShardingState::needCollectionMetadata(OperationContext* opCtx, const string& ns) {
diff --git a/src/mongo/db/s/sharding_state.h b/src/mongo/db/s/sharding_state.h
index 55605d957f0..f14815502b3 100644
--- a/src/mongo/db/s/sharding_state.h
+++ b/src/mongo/db/s/sharding_state.h
@@ -36,25 +36,20 @@
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/s/active_migrations_registry.h"
 #include "mongo/db/s/chunk_splitter.h"
-#include "mongo/db/s/collection_range_deleter.h"
+#include "mongo/db/s/collection_sharding_state.h"
 #include "mongo/db/s/migration_destination_manager.h"
 #include "mongo/executor/task_executor.h"
 #include "mongo/executor/thread_pool_task_executor.h"
 #include "mongo/stdx/functional.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/stdx/mutex.h"
-#include "mongo/stdx/unordered_map.h"
 
 namespace mongo {
 
 class BSONObj;
 class BSONObjBuilder;
-struct ChunkVersion;
-class CollectionMetadata;
-class CollectionShardingState;
 class ConnectionString;
 class OperationContext;
-class ScopedCollectionMetadata;
 class ServiceContext;
 class ShardIdentityType;
 class Status;
@@ -141,8 +136,6 @@ public:
      */
     Status updateConfigServerOpTimeFromMetadata(OperationContext* opCtx);
 
-    CollectionShardingState* getNS(const std::string& ns, OperationContext* opCtx);
-
     ChunkSplitter* getChunkSplitter();
 
     /**
@@ -277,10 +270,6 @@ public:
     executor::TaskExecutor* getRangeDeleterTaskExecutor();
 
 private:
-    // Map from a namespace into the sharding state for each collection we have
-    typedef stdx::unordered_map<std::string, std::unique_ptr<CollectionShardingState>>
-        CollectionShardingStateMap;
-
     // Progress of the sharding state initialization
     enum class InitializationState : uint32_t {
         // Initial state. The server must be under exclusive lock when this state is entered. No
@@ -338,11 +327,6 @@ private:
     // Sets the shard name for this host (comes through setShardVersion)
     std::string _shardName;
 
-    // Cache of collection metadata on this shard. It is not safe to look-up values from this map
-    // without holding some form of collection lock. It is only safe to add/remove values when
-    // holding X lock on the respective namespace.
-    CollectionShardingStateMap _collections;
-
     // The id for the cluster this shard belongs to.
     OID _clusterId;
 
diff --git a/src/mongo/db/s/sharding_state_command.cpp b/src/mongo/db/s/sharding_state_command.cpp
index 4a2e546ac76..0e3aea8b652 100644
--- a/src/mongo/db/s/sharding_state_command.cpp
+++ b/src/mongo/db/s/sharding_state_command.cpp
@@ -32,13 +32,10 @@
 
 #include "mongo/db/auth/action_set.h"
 #include "mongo/db/auth/action_type.h"
-#include "mongo/db/auth/authorization_manager.h"
-#include "mongo/db/auth/authorization_session.h"
 #include "mongo/db/auth/privilege.h"
 #include "mongo/db/commands.h"
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/util/log.h"
-#include "mongo/util/stringutils.h"
 
 namespace mongo {
 namespace {
@@ -47,7 +44,7 @@ class ShardingStateCmd : public BasicCommand {
 public:
     ShardingStateCmd() : BasicCommand("shardingState") {}
 
-    virtual bool supportsWriteConcern(const BSONObj& cmd) const override {
+    bool supportsWriteConcern(const BSONObj& cmd) const override {
         return false;
     }
 
@@ -72,6 +69,7 @@ public:
              const BSONObj& cmdObj,
              BSONObjBuilder& result) override {
         ShardingState::get(opCtx)->appendInfo(opCtx, result);
+        CollectionShardingState::report(opCtx, &result);
         return true;
     }