RUST-544 Ensure correct termination of monitoring tasks (#242)

Co-authored-by: Saghm Rossi <[email protected]>

Author: saghm

src/client/mod.rs

@@ -82,6 +82,12 @@ struct ClientInner {
     session_pool: ServerSessionPool,
 }
 
+impl Drop for ClientInner {
+    fn drop(&mut self) {
+        self.topology.mark_closed()
+    }
+}
+
 impl Client {
     /// Creates a new `Client` connected to the cluster specified by `uri`. `uri` must be a valid
     /// MongoDB connection string.

src/lib.rs

@@ -79,6 +79,7 @@
     )
 )]
 #![cfg_attr(docsrs, feature(doc_cfg))]
+#![cfg_attr(test, type_length_limit = "6000000")]
 
 macro_rules! define_if_single_runtime_enabled {
     ( $( $def:item )+ ) => {

src/sdam/monitor.rs

@@ -52,7 +52,11 @@ impl Monitor {
             .heartbeat_freq
             .unwrap_or(DEFAULT_HEARTBEAT_FREQUENCY);
 
-        while let Some(_server) = self.server.upgrade() {
+        while self.topology.is_alive() {
+            if self.server.upgrade().is_none() {
+                break;
+            }
+
             let topology = match self.topology.upgrade() {
                 Some(topology) => topology,
                 None => break,

src/sdam/srv_polling/mod.rs

@@ -59,14 +59,21 @@ impl SrvPollingMonitor {
     }
 
     async fn execute(&mut self) {
-        while let Some(topology) = self.topology.upgrade() {
+        while self.topology.is_alive() {
+            let topology = match self.topology.upgrade() {
+                Some(topology) => topology,
+                None => break,
+            };
+
             let state = topology.clone_state().await;
 
             if state.is_sharded() || state.is_unknown() {
                 let hosts = self.lookup_hosts().await;
-                self.update_hosts(hosts, topology, state).await;
+                self.update_hosts(hosts, topology.clone(), state).await;
             }
 
+            std::mem::drop(topology);
+
             RUNTIME
                 .delay_for(self.rescan_interval.unwrap_or(DEFAULT_RESCAN_SRV_INTERVAL))
                 .await;

src/sdam/state/mod.rs

@@ -2,7 +2,11 @@ pub(super) mod server;
 
 use std::{
     collections::{HashMap, HashSet},
-    sync::{Arc, Weak},
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+        Weak,
+    },
     time::Duration,
 };
 
@@ -47,6 +51,7 @@ pub(crate) struct WeakTopology {
 /// manager and the client options.
 #[derive(Clone, Debug)]
 struct Common {
+    is_alive: Arc<AtomicBool>,
     message_manager: TopologyMessageManager,
     options: ClientOptions,
 }
@@ -71,6 +76,7 @@ impl Topology {
         let description = TopologyDescription::new_from_hosts(hosts.clone());
 
         let common = Common {
+            is_alive: Arc::new(AtomicBool::new(true)),
             message_manager: TopologyMessageManager::new(),
             options: ClientOptions::new_srv(),
         };
@@ -106,6 +112,7 @@ impl Topology {
         let hosts: Vec<_> = options.hosts.drain(..).collect();
 
         let common = Common {
+            is_alive: Arc::new(AtomicBool::new(true)),
             message_manager: TopologyMessageManager::new(),
             options: options.clone(),
         };
@@ -132,6 +139,10 @@ impl Topology {
         Ok(topology)
     }
 
+    pub(crate) fn mark_closed(&self) {
+        self.common.is_alive.store(false, Ordering::SeqCst);
+    }
+
     /// Gets the addresses of the servers in the cluster.
     #[cfg(test)]
     pub(super) async fn servers(&self) -> HashSet<StreamAddress> {
@@ -360,6 +371,10 @@ impl WeakTopology {
         })
     }
 
+    pub(crate) fn is_alive(&self) -> bool {
+        self.common.is_alive.load(Ordering::SeqCst)
+    }
+
     pub(crate) fn client_options(&self) -> &ClientOptions {
         &self.common.options
     }