SERVER-42455 Make ReplicaSetChangeNotifier::onConfirmedSet safe during shutdown

5 files changed, 20 insertions, 29 deletions

diff --git a/src/mongo/client/replica_set_change_notifier.cpp b/src/mongo/client/replica_set_change_notifier.cpp
index 00523770718..5d16973cdc6 100644
--- a/src/mongo/client/replica_set_change_notifier.cpp
+++ b/src/mongo/client/replica_set_change_notifier.cpp
@@ -39,20 +39,13 @@
 
 namespace mongo {
 
-void ReplicaSetChangeNotifier::_addListener(Listener* listener) {
+void ReplicaSetChangeNotifier::_addListener(std::shared_ptr<Listener> listener) {
     stdx::lock_guard lk(_mutex);
 
     listener->init(this);
     _listeners.push_back(listener);
 }
 
-void ReplicaSetChangeNotifier::_removeListener(Listener* listener) {
-    stdx::lock_guard lk(_mutex);
-
-    auto& listeners = _listeners;
-    listeners.erase(std::remove(listeners.begin(), listeners.end(), listener), listeners.end());
-}
-
 void ReplicaSetChangeNotifier::onFoundSet(const std::string& name) noexcept {
     LOGV2_DEBUG(20158, 2, "Signaling found set {name}", "name"_attr = name);
 
@@ -64,7 +57,9 @@ void ReplicaSetChangeNotifier::onFoundSet(const std::string& name) noexcept {
     lk.unlock();
 
     for (auto listener : listeners) {
-        listener->onFoundSet(name);
+        if (auto l = listener.lock()) {
+            l->onFoundSet(name);
+        }
     };
 }
 
@@ -92,7 +87,9 @@ void ReplicaSetChangeNotifier::onPossibleSet(ConnectionString connectionString)
     lk.unlock();
 
     for (auto listener : listeners) {
-        listener->onPossibleSet(state);
+        if (auto l = listener.lock()) {
+            l->onPossibleSet(state);
+        }
     };
 }
 
@@ -123,7 +120,9 @@ void ReplicaSetChangeNotifier::onConfirmedSet(ConnectionString connectionString,
     lk.unlock();
 
     for (auto listener : listeners) {
-        listener->onConfirmedSet(state);
+        if (auto l = listener.lock()) {
+            l->onConfirmedSet(state);
+        }
     };
 }
 
@@ -144,7 +143,9 @@ void ReplicaSetChangeNotifier::onDroppedSet(const std::string& name) noexcept {
     lk.unlock();
 
     for (auto listener : listeners) {
-        listener->onDroppedSet(name);
+        if (auto l = listener.lock()) {
+            l->onDroppedSet(name);
+        }
     };
 }
 
diff --git a/src/mongo/client/replica_set_change_notifier.h b/src/mongo/client/replica_set_change_notifier.h
index 3ab802e7147..4408cd46d53 100644
--- a/src/mongo/client/replica_set_change_notifier.h
+++ b/src/mongo/client/replica_set_change_notifier.h
@@ -49,7 +49,6 @@ class ReplicaSetChangeNotifier {
 public:
     using Key = std::string;
     class Listener;
-    using ListenerHandle = std::unique_ptr<Listener, unique_function<void(Listener*)>>;
     struct State;
 
 public:
@@ -88,24 +87,17 @@ public:
               typename... Args,
               typename = std::enable_if_t<std::is_constructible_v<DerivedT, Args...>>>
     auto makeListener(Args&&... args) {
-        auto deleter = [this](auto listener) {
-            _removeListener(listener);
-            delete listener;
-        };
-        auto ptr = new DerivedT(std::forward<Args>(args)...);
-
+        auto ptr = std::make_shared<DerivedT>(std::forward<Args>(args)...);
         _addListener(ptr);
-
-        return ListenerHandle(ptr, std::move(deleter));
+        return ptr;
     }
 
 private:
-    void _addListener(Listener* listener);
-    void _removeListener(Listener* listener);
+    void _addListener(std::shared_ptr<Listener> listener);
 
     Mutex _mutex =
         MONGO_MAKE_LATCH(HierarchicalAcquisitionLevel(0), "ReplicaSetChangeNotifier::_mutex");
-    std::vector<Listener*> _listeners;
+    std::vector<std::weak_ptr<Listener>> _listeners;
     stdx::unordered_map<Key, State> _replicaSetStates;
 };
 
@@ -168,8 +160,6 @@ private:
     Notifier* _notifier = nullptr;
 };
 
-using ReplicaSetChangeListenerHandle = ReplicaSetChangeNotifier::ListenerHandle;
-
 struct ReplicaSetChangeNotifier::State {
     ConnectionString connStr;
     HostAndPort primary;
diff --git a/src/mongo/client/scanning_replica_set_monitor_scan_test.cpp b/src/mongo/client/scanning_replica_set_monitor_scan_test.cpp
index 451b20af3e0..d294cd859d4 100644
--- a/src/mongo/client/scanning_replica_set_monitor_scan_test.cpp
+++ b/src/mongo/client/scanning_replica_set_monitor_scan_test.cpp
@@ -1550,7 +1550,7 @@ public:
     }
 
 protected:
-    decltype(_notifier)::ListenerHandle _listener = _notifier.makeListener<Listener>();
+    std::shared_ptr<decltype(_notifier)::Listener> _listener = _notifier.makeListener<Listener>();
 
     std::shared_ptr<SetState> _state = makeState(basicUri);
 };
diff --git a/src/mongo/db/s/sharding_initialization_mongod.h b/src/mongo/db/s/sharding_initialization_mongod.h
index 2eaefd22fbe..a8e8eec3c72 100644
--- a/src/mongo/db/s/sharding_initialization_mongod.h
+++ b/src/mongo/db/s/sharding_initialization_mongod.h
@@ -120,7 +120,7 @@ private:
     // Function for initializing the sharding environment components (i.e. everything on the Grid)
     ShardingEnvironmentInitFunc _initFunc;
 
-    ReplicaSetChangeListenerHandle _replicaSetChangeListener;
+    std::shared_ptr<ReplicaSetChangeNotifier::Listener> _replicaSetChangeListener;
 };
 
 /**
diff --git a/src/mongo/s/sharding_task_executor_pool_controller.h b/src/mongo/s/sharding_task_executor_pool_controller.h
index 187fbba55de..718c2b0cd0a 100644
--- a/src/mongo/s/sharding_task_executor_pool_controller.h
+++ b/src/mongo/s/sharding_task_executor_pool_controller.h
@@ -193,7 +193,7 @@ private:
         boost::optional<PoolId> maybeId;
     };
 
-    ReplicaSetChangeListenerHandle _listener;
+    std::shared_ptr<ReplicaSetChangeNotifier::Listener> _listener;
 
     Mutex _mutex = MONGO_MAKE_LATCH("ShardingTaskExecutorPoolController::_mutex");