SERVER-33797: Reload ShardRegistry on shard collection command

16 files changed, 44 insertions, 30 deletions

diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index be50d5c8f60..52699c273a2 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -467,7 +467,7 @@ bool Balancer::_checkOIDs(OperationContext* opCtx) {
     auto shardingContext = Grid::get(opCtx);
 
     vector<ShardId> all;
-    shardingContext->shardRegistry()->getAllShardIds(&all);
+    shardingContext->shardRegistry()->getAllShardIdsNoReload(&all);
 
     // map of OID machine ID => shardId
     map<int, ShardId> oids;
diff --git a/src/mongo/db/s/config/configsvr_drop_database_command.cpp b/src/mongo/db/s/config/configsvr_drop_database_command.cpp
index fe4b658dca2..7d4cfcbd5a4 100644
--- a/src/mongo/db/s/config/configsvr_drop_database_command.cpp
+++ b/src/mongo/db/s/config/configsvr_drop_database_command.cpp
@@ -154,7 +154,7 @@ public:
         // Drop the database from each of the remaining shards.
         {
             std::vector<ShardId> allShardIds;
-            Grid::get(opCtx)->shardRegistry()->getAllShardIds(&allShardIds);
+            Grid::get(opCtx)->shardRegistry()->getAllShardIdsNoReload(&allShardIds);
 
             for (const ShardId& shardId : allShardIds) {
                 _dropDatabaseFromShard(opCtx, shardId, dbname);
diff --git a/src/mongo/db/s/config/configsvr_shard_collection_command.cpp b/src/mongo/db/s/config/configsvr_shard_collection_command.cpp
index b88978aecd7..6732c7e9678 100644
--- a/src/mongo/db/s/config/configsvr_shard_collection_command.cpp
+++ b/src/mongo/db/s/config/configsvr_shard_collection_command.cpp
@@ -745,6 +745,7 @@ public:
         auto const catalogManager = ShardingCatalogManager::get(opCtx);
         auto const catalogCache = Grid::get(opCtx)->catalogCache();
         auto const catalogClient = Grid::get(opCtx)->catalogClient();
+        auto shardRegistry = Grid::get(opCtx)->shardRegistry();
 
         // Make the distlocks boost::optional so that they can be released by being reset below.
         boost::optional<DistLockManager::ScopedDistLock> dbDistLock(
@@ -773,7 +774,7 @@ public:
         ShardKeyPattern shardKeyPattern(proposedKey);
 
         std::vector<ShardId> shardIds;
-        Grid::get(opCtx)->shardRegistry()->getAllShardIds(&shardIds);
+        shardRegistry->getAllShardIds(opCtx, &shardIds);
         const int numShards = shardIds.size();
 
         uassert(ErrorCodes::IllegalOperation,
@@ -788,8 +789,7 @@ public:
                     "only special collections in the config db may be sharded",
                     nss.ns() == SessionsCollection::kSessionsFullNS || getTestCommandsEnabled());
 
-            auto configShard = uassertStatusOK(
-                Grid::get(opCtx)->shardRegistry()->getShard(opCtx, dbType.getPrimary()));
+            auto configShard = uassertStatusOK(shardRegistry->getShard(opCtx, dbType.getPrimary()));
             ScopedDbConnection configConn(configShard->getConnString());
             ON_BLOCK_EXIT([&configConn] { configConn.done(); });
 
@@ -810,8 +810,7 @@ public:
             }
         }();
 
-        auto primaryShard =
-            uassertStatusOK(Grid::get(opCtx)->shardRegistry()->getShard(opCtx, primaryShardId));
+        auto primaryShard = uassertStatusOK(shardRegistry->getShard(opCtx, primaryShardId));
         ScopedDbConnection conn(primaryShard->getConnString());
         ON_BLOCK_EXIT([&conn] { conn.done(); });
 
diff --git a/src/mongo/db/s/config/sharding_catalog_manager_collection_operations.cpp b/src/mongo/db/s/config/sharding_catalog_manager_collection_operations.cpp
index 511b7bdfae9..574a3e5a9c1 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager_collection_operations.cpp
+++ b/src/mongo/db/s/config/sharding_catalog_manager_collection_operations.cpp
@@ -191,7 +191,7 @@ ChunkVersion ShardingCatalogManager::_createFirstChunks(OperationContext* opCtx,
         [&opCtx, primaryShardId, &shardIds, primaryDraining]() {
             if (primaryDraining) {
                 vector<ShardId> allShardIds;
-                Grid::get(opCtx)->shardRegistry()->getAllShardIds(&allShardIds);
+                Grid::get(opCtx)->shardRegistry()->getAllShardIdsNoReload(&allShardIds);
 
                 auto dbShardId = allShardIds[0];
                 if (allShardIds[0] == primaryShardId && allShardIds.size() > 1) {
@@ -238,7 +238,7 @@ ChunkVersion ShardingCatalogManager::_createFirstChunks(OperationContext* opCtx,
         // If docs already exist for the collection, must use primary shard,
         // otherwise defer to passed-in distribution option.
         if (numObjects == 0 && distributeInitialChunks) {
-            Grid::get(opCtx)->shardRegistry()->getAllShardIds(&shardIds);
+            Grid::get(opCtx)->shardRegistry()->getAllShardIdsNoReload(&shardIds);
             if (primaryDraining && shardIds.size() > 1) {
                 shardIds.erase(std::remove(shardIds.begin(), shardIds.end(), primaryShardId),
                                shardIds.end());
@@ -259,7 +259,7 @@ ChunkVersion ShardingCatalogManager::_createFirstChunks(OperationContext* opCtx,
         }
 
         if (distributeInitialChunks) {
-            Grid::get(opCtx)->shardRegistry()->getAllShardIds(&shardIds);
+            Grid::get(opCtx)->shardRegistry()->getAllShardIdsNoReload(&shardIds);
             if (primaryDraining) {
                 shardIds.erase(std::remove(shardIds.begin(), shardIds.end(), primaryShardId),
                                shardIds.end());
@@ -655,7 +655,6 @@ void ShardingCatalogManager::createCollection(OperationContext* opCtx,
     createCmdBuilder.append("create", ns.coll());
     collOptions.appendBSON(&createCmdBuilder);
     createCmdBuilder.append(kWriteConcernField, opCtx->getWriteConcern().toBSON());
-
     auto swResponse = primaryShard->runCommandWithFixedRetryAttempts(
         opCtx,
         ReadPreferenceSetting{ReadPreference::PrimaryOnly},
diff --git a/src/mongo/db/s/config/sharding_catalog_manager_shard_operations.cpp b/src/mongo/db/s/config/sharding_catalog_manager_shard_operations.cpp
index 34263b7cb31..2ee60c2011d 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager_shard_operations.cpp
+++ b/src/mongo/db/s/config/sharding_catalog_manager_shard_operations.cpp
@@ -887,14 +887,9 @@ StatusWith<ShardId> ShardingCatalogManager::_selectShardForNewDatabase(
     OperationContext* opCtx, ShardRegistry* shardRegistry) {
     std::vector<ShardId> allShardIds;
 
-    shardRegistry->getAllShardIds(&allShardIds);
+    shardRegistry->getAllShardIds(opCtx, &allShardIds);
     if (allShardIds.empty()) {
-        shardRegistry->reload(opCtx);
-        shardRegistry->getAllShardIds(&allShardIds);
-
-        if (allShardIds.empty()) {
-            return Status(ErrorCodes::ShardNotFound, "No shards found");
-        }
+        return Status(ErrorCodes::ShardNotFound, "No shards found");
     }
 
     ShardId candidateShardId = allShardIds[0];
diff --git a/src/mongo/s/client/shard_connection.cpp b/src/mongo/s/client/shard_connection.cpp
index 2aa50893ecf..9d117a706af 100644
--- a/src/mongo/s/client/shard_connection.cpp
+++ b/src/mongo/s/client/shard_connection.cpp
@@ -217,7 +217,7 @@ public:
         auto const shardRegistry = Grid::get(opCtx)->shardRegistry();
 
         std::vector<ShardId> all;
-        shardRegistry->getAllShardIds(&all);
+        shardRegistry->getAllShardIdsNoReload(&all);
 
         // Don't report exceptions here as errors in GetLastError
         LastError::Disabled ignoreForGLE(&LastError::get(cc()));
diff --git a/src/mongo/s/client/shard_registry.cpp b/src/mongo/s/client/shard_registry.cpp
index c97279a8403..ff43db37b72 100644
--- a/src/mongo/s/client/shard_registry.cpp
+++ b/src/mongo/s/client/shard_registry.cpp
@@ -163,12 +163,26 @@ shared_ptr<Shard> ShardRegistry::lookupRSName(const string& name) const {
     return _data.findByRSName(name);
 }
 
-void ShardRegistry::getAllShardIds(vector<ShardId>* all) const {
+void ShardRegistry::getAllShardIdsNoReload(vector<ShardId>* all) const {
     std::set<ShardId> seen;
     _data.getAllShardIds(seen);
     all->assign(seen.begin(), seen.end());
 }
 
+void ShardRegistry::getAllShardIds(OperationContext* opCtx, vector<ShardId>* all) {
+    getAllShardIdsNoReload(all);
+    if (all->empty()) {
+        bool didReload = reload(opCtx);
+        getAllShardIdsNoReload(all);
+        // If we didn't do the reload ourselves, we should retry to ensure
+        // that the reload is actually initiated while we're executing this
+        if (!didReload && all->empty()) {
+            reload(opCtx);
+            getAllShardIdsNoReload(all);
+        }
+    }
+}
+
 int ShardRegistry::getNumShards() const {
     std::set<ShardId> seen;
     _data.getAllShardIds(seen);
diff --git a/src/mongo/s/client/shard_registry.h b/src/mongo/s/client/shard_registry.h
index af1f11e8f06..ab27d97ac3c 100644
--- a/src/mongo/s/client/shard_registry.h
+++ b/src/mongo/s/client/shard_registry.h
@@ -232,7 +232,14 @@ public:
      */
     std::shared_ptr<Shard> lookupRSName(const std::string& name) const;
 
-    void getAllShardIds(std::vector<ShardId>* all) const;
+    void getAllShardIdsNoReload(std::vector<ShardId>* all) const;
+
+    /**
+     * Like getAllShardIdsNoReload(), but does a reload internally in the case that
+     * getAllShardIdsNoReload() comes back empty
+     */
+    void getAllShardIds(OperationContext* opCtx, std::vector<ShardId>* all);
+
     int getNumShards() const;
 
     void toBSON(BSONObjBuilder* result) const;
diff --git a/src/mongo/s/commands/cluster_aggregate.cpp b/src/mongo/s/commands/cluster_aggregate.cpp
index 3249c0a740c..cd977d6b7fa 100644
--- a/src/mongo/s/commands/cluster_aggregate.cpp
+++ b/src/mongo/s/commands/cluster_aggregate.cpp
@@ -167,7 +167,7 @@ StatusWith<CachedCollectionRoutingInfo> getExecutionNsRoutingInfo(OperationConte
     // further check that at least one shard exists if the aggregation is collectionless.
     if (swRoutingInfo.isOK() && execNss.isCollectionlessAggregateNS()) {
         std::vector<ShardId> shardIds;
-        Grid::get(opCtx)->shardRegistry()->getAllShardIds(&shardIds);
+        Grid::get(opCtx)->shardRegistry()->getAllShardIdsNoReload(&shardIds);
 
         if (shardIds.size() == 0) {
             return {ErrorCodes::NamespaceNotFound, "No shards are present in the cluster"};
@@ -186,7 +186,7 @@ std::set<ShardId> getTargetedShards(OperationContext* opCtx,
     if (mustRunOnAllShards(nss, routingInfo, litePipe)) {
         // The pipeline begins with a stage which must be run on all shards.
         std::vector<ShardId> shardIds;
-        Grid::get(opCtx)->shardRegistry()->getAllShardIds(&shardIds);
+        Grid::get(opCtx)->shardRegistry()->getAllShardIdsNoReload(&shardIds);
         return {shardIds.begin(), shardIds.end()};
     }
 
diff --git a/src/mongo/s/commands/cluster_commands_helpers.cpp b/src/mongo/s/commands/cluster_commands_helpers.cpp
index 7e880f95f11..4037e67bb7e 100644
--- a/src/mongo/s/commands/cluster_commands_helpers.cpp
+++ b/src/mongo/s/commands/cluster_commands_helpers.cpp
@@ -112,7 +112,7 @@ std::vector<AsyncRequestsSender::Request> buildUnversionedRequestsForShards(
 std::vector<AsyncRequestsSender::Request> buildUnversionedRequestsForAllShards(
     OperationContext* opCtx, const BSONObj& cmdObj) {
     std::vector<ShardId> shardIds;
-    Grid::get(opCtx)->shardRegistry()->getAllShardIds(&shardIds);
+    Grid::get(opCtx)->shardRegistry()->getAllShardIdsNoReload(&shardIds);
     return buildUnversionedRequestsForShards(std::move(shardIds), cmdObj);
 }
 
diff --git a/src/mongo/s/commands/cluster_fsync_cmd.cpp b/src/mongo/s/commands/cluster_fsync_cmd.cpp
index ecb4297c1d2..ee7bdd14afd 100644
--- a/src/mongo/s/commands/cluster_fsync_cmd.cpp
+++ b/src/mongo/s/commands/cluster_fsync_cmd.cpp
@@ -83,7 +83,7 @@ public:
         int numFiles = 0;
 
         std::vector<ShardId> shardIds;
-        grid.shardRegistry()->getAllShardIds(&shardIds);
+        grid.shardRegistry()->getAllShardIdsNoReload(&shardIds);
 
         for (const ShardId& shardId : shardIds) {
             auto shardStatus = grid.shardRegistry()->getShard(opCtx, shardId);
diff --git a/src/mongo/s/commands/cluster_list_databases_cmd.cpp b/src/mongo/s/commands/cluster_list_databases_cmd.cpp
index 8fa0b0e8927..6b7dd77d8d6 100644
--- a/src/mongo/s/commands/cluster_list_databases_cmd.cpp
+++ b/src/mongo/s/commands/cluster_list_databases_cmd.cpp
@@ -87,7 +87,7 @@ public:
         std::map<std::string, std::unique_ptr<BSONObjBuilder>> dbShardInfo;
 
         std::vector<ShardId> shardIds;
-        grid.shardRegistry()->getAllShardIds(&shardIds);
+        grid.shardRegistry()->getAllShardIdsNoReload(&shardIds);
         shardIds.emplace_back(ShardRegistry::kConfigServerShardId);
 
         auto filteredCmd = CommandHelpers::filterCommandRequestForPassthrough(cmdObj);
diff --git a/src/mongo/s/commands/cluster_multicast.cpp b/src/mongo/s/commands/cluster_multicast.cpp
index 105893fb7d4..d80979bf03b 100644
--- a/src/mongo/s/commands/cluster_multicast.cpp
+++ b/src/mongo/s/commands/cluster_multicast.cpp
@@ -48,7 +48,7 @@ std::vector<HostAndPort> getAllClusterHosts(OperationContext* opCtx) {
     auto registry = Grid::get(opCtx)->shardRegistry();
 
     std::vector<ShardId> shardIds;
-    registry->getAllShardIds(&shardIds);
+    registry->getAllShardIds(opCtx, &shardIds);
 
     std::vector<HostAndPort> servers;
     for (const auto& shardId : shardIds) {
diff --git a/src/mongo/s/commands/kill_sessions_remote.cpp b/src/mongo/s/commands/kill_sessions_remote.cpp
index 78db3419f10..735a70cd954 100644
--- a/src/mongo/s/commands/kill_sessions_remote.cpp
+++ b/src/mongo/s/commands/kill_sessions_remote.cpp
@@ -62,7 +62,7 @@ std::vector<HostAndPort> getAllClusterHosts(OperationContext* opCtx) {
     auto registry = Grid::get(opCtx)->shardRegistry();
 
     std::vector<ShardId> shardIds;
-    registry->getAllShardIds(&shardIds);
+    registry->getAllShardIds(opCtx, &shardIds);
 
     std::vector<HostAndPort> servers;
     for (const auto& shardId : shardIds) {
diff --git a/src/mongo/s/query/router_stage_update_on_add_shard.cpp b/src/mongo/s/query/router_stage_update_on_add_shard.cpp
index 3d9ddb3428c..451a1ee9699 100644
--- a/src/mongo/s/query/router_stage_update_on_add_shard.cpp
+++ b/src/mongo/s/query/router_stage_update_on_add_shard.cpp
@@ -97,7 +97,7 @@ RouterStageUpdateOnAddShard::establishShardCursorsOnNewShards(std::vector<ShardI
     }
 
     std::vector<ShardId> shardIds, newShardIds;
-    shardRegistry->getAllShardIds(&shardIds);
+    shardRegistry->getAllShardIdsNoReload(&shardIds);
     std::sort(existingShardIds.begin(), existingShardIds.end());
     std::sort(shardIds.begin(), shardIds.end());
     std::set_difference(shardIds.begin(),
diff --git a/src/mongo/s/write_ops/chunk_manager_targeter.cpp b/src/mongo/s/write_ops/chunk_manager_targeter.cpp
index 04b0d3a1e21..925c50d95df 100644
--- a/src/mongo/s/write_ops/chunk_manager_targeter.cpp
+++ b/src/mongo/s/write_ops/chunk_manager_targeter.cpp
@@ -601,7 +601,7 @@ StatusWith<std::vector<ShardEndpoint>> ChunkManagerTargeter::targetAllShards(
     }
 
     std::vector<ShardId> shardIds;
-    Grid::get(opCtx)->shardRegistry()->getAllShardIds(&shardIds);
+    Grid::get(opCtx)->shardRegistry()->getAllShardIdsNoReload(&shardIds);
 
     std::vector<ShardEndpoint> endpoints;
     for (auto&& shardId : shardIds) {