SERVER-32871 Invalidate CatalogCache entries when ShardRegistry reload discovers a shard has been removed

2 files changed, 29 insertions, 5 deletions

diff --git a/src/mongo/s/client/shard_registry.cpp b/src/mongo/s/client/shard_registry.cpp
index a5deb5da0f3..2c92c46e33e 100644
--- a/src/mongo/s/client/shard_registry.cpp
+++ b/src/mongo/s/client/shard_registry.cpp
@@ -87,8 +87,11 @@ const Seconds kRefreshPeriod(30);
 const ShardId ShardRegistry::kConfigServerShardId = ShardId("config");
 
 ShardRegistry::ShardRegistry(std::unique_ptr<ShardFactory> shardFactory,
-                             const ConnectionString& configServerCS)
-    : _shardFactory(std::move(shardFactory)), _initConfigServerCS(configServerCS) {}
+                             const ConnectionString& configServerCS,
+                             std::vector<ShardRemovalHook> shardRemovalHooks)
+    : _shardFactory(std::move(shardFactory)),
+      _initConfigServerCS(configServerCS),
+      _shardRemovalHooks(std::move(shardRemovalHooks)) {}
 
 ShardRegistry::~ShardRegistry() {
     shutdown();
@@ -334,6 +337,11 @@ bool ShardRegistry::reload(OperationContext* opCtx) {
 
         auto name = shard->getConnString().getSetName();
         ReplicaSetMonitor::remove(name);
+        for (auto& callback : _shardRemovalHooks) {
+            // Run callbacks asynchronously.
+            ExecutorFuture<void>(Grid::get(opCtx)->getExecutorPool()->getFixedExecutor())
+                .getAsync([=](const Status&) { callback(shardId); });
+        }
     }
 
     nextReloadState = ReloadState::Idle;
diff --git a/src/mongo/s/client/shard_registry.h b/src/mongo/s/client/shard_registry.h
index 5456f7bd06d..a5917c64413 100644
--- a/src/mongo/s/client/shard_registry.h
+++ b/src/mongo/s/client/shard_registry.h
@@ -147,13 +147,22 @@ public:
     static const ShardId kConfigServerShardId;
 
     /**
+     * A callback type for functions that can be called on shard removal.
+     */
+    using ShardRemovalHook = std::function<void(const ShardId&)>;
+
+    /**
      * Instantiates a new shard registry.
      *
-     * @param shardFactory Makes shards
-     * @param configServerCS ConnectionString used for communicating with the config servers
+     * @param shardFactory      Makes shards
+     * @param configServerCS    ConnectionString used for communicating with the config servers
+     * @param shardRemovalHooks A list of hooks that will be called when a shard is removed. The
+     *                          hook is expected not to throw. If it does throw, the process will be
+     *                          terminated.
      */
     ShardRegistry(std::unique_ptr<ShardFactory> shardFactory,
-                  const ConnectionString& configServerCS);
+                  const ConnectionString& configServerCS,
+                  std::vector<ShardRemovalHook> shardRemovalHooks = {});
 
     ~ShardRegistry();
     /**
@@ -282,6 +291,13 @@ private:
      * shard
      */
     ConnectionString _initConfigServerCS;
+
+    /**
+     * A list of callbacks to be called asynchronously when it has been discovered that a shard was
+     * removed.
+     */
+    std::vector<ShardRemovalHook> _shardRemovalHooks;
+
     void _internalReload(const executor::TaskExecutor::CallbackArgs& cbArgs);
     ShardRegistryData _data;