RUST-1047 Ensure TopologyClosedEvent is the last SDAM event emitted (#485)

Author: patrickfreed

.evergreen/check-rustfmt.sh

@@ -4,3 +4,4 @@ set -o errexit
 
 . ~/.cargo/env
 rustfmt +nightly --unstable-features --check src/**/*.rs
+rustfmt +nightly --unstable-features --check src/*.rs

rustfmt.toml

@@ -8,3 +8,4 @@ use_try_shorthand = true
 wrap_comments = true
 imports_layout = "HorizontalVertical"
 imports_granularity = "Crate"
+ignore = ["src/lib.rs"]

src/cmap/establish/handshake/mod.rs

@@ -15,7 +15,7 @@ use crate::{
     event::sdam::SdamEventHandler,
     is_master::{is_master_command, run_is_master, IsMasterReply},
     options::{AuthMechanism, ClientOptions, Credential, DriverInfo, ServerApi},
-    sdam::WeakTopology,
+    sdam::Topology,
 };
 
 #[cfg(feature = "tokio-runtime")]
@@ -227,7 +227,7 @@ impl Handshaker {
     pub(crate) async fn handshake(
         &self,
         conn: &mut Connection,
-        topology: Option<&WeakTopology>,
+        topology: Option<&Topology>,
         handler: &Option<Arc<dyn SdamEventHandler>>,
     ) -> Result<HandshakeResult> {
         let mut command = self.command.clone();

src/is_master.rs

@@ -19,7 +19,7 @@ use crate::{
         ServerHeartbeatStartedEvent,
         ServerHeartbeatSucceededEvent,
     },
-    sdam::{ServerType, WeakTopology},
+    sdam::{ServerType, Topology},
     selection_criteria::TagSet,
 };
 
@@ -37,10 +37,15 @@ pub(crate) fn is_master_command(api: Option<&ServerApi>) -> Command {
     command
 }
 
+/// Execute an isMaster command, emiting events if a reference to the topology and a handler are
+/// provided.
+///
+/// A strong reference to the topology is used here to ensure it is still in scope and has not yet
+/// emitted a `TopologyClosedEvent`.
 pub(crate) async fn run_is_master(
     conn: &mut Connection,
     command: Command,
-    topology: Option<&WeakTopology>,
+    topology: Option<&Topology>,
     handler: &Option<Arc<dyn SdamEventHandler>>,
 ) -> Result<IsMasterReply> {
     emit_event(topology, handler, |handler| {
@@ -91,18 +96,13 @@ pub(crate) async fn run_is_master(
     }
 }
 
-fn emit_event<F>(
-    topology: Option<&WeakTopology>,
-    handler: &Option<Arc<dyn SdamEventHandler>>,
-    emit: F,
-) where
+fn emit_event<F>(topology: Option<&Topology>, handler: &Option<Arc<dyn SdamEventHandler>>, emit: F)
+where
     F: FnOnce(&Arc<dyn SdamEventHandler>),
 {
     if let Some(handler) = handler {
-        if let Some(topology) = topology {
-            if topology.is_alive() {
-                emit(handler);
-            }
+        if topology.is_some() {
+            emit(handler);
         }
     }
 }

src/sdam/mod.rs

@@ -26,6 +26,5 @@ pub(crate) use self::{
         server::{Server, ServerUpdate, ServerUpdateReceiver, ServerUpdateSender},
         HandshakePhase,
         Topology,
-        WeakTopology,
     },
 };

src/sdam/monitor.rs

@@ -160,7 +160,7 @@ impl HeartbeatMonitor {
     /// Returns true if the topology has changed and false otherwise.
     async fn check_server(&mut self, topology: &Topology, server: &Server) -> bool {
         let mut retried = false;
-        let check_result = match self.perform_is_master().await {
+        let check_result = match self.perform_is_master(topology).await {
             Ok(reply) => Ok(reply),
             Err(e) => {
                 let previous_description = topology.get_server_description(&server.address).await;
@@ -171,7 +171,7 @@ impl HeartbeatMonitor {
                 {
                     self.handle_error(e, topology, server).await;
                     retried = true;
-                    self.perform_is_master().await
+                    self.perform_is_master(topology).await
                 } else {
                     Err(e)
                 }
@@ -188,14 +188,14 @@ impl HeartbeatMonitor {
         }
     }
 
-    async fn perform_is_master(&mut self) -> Result<IsMasterReply> {
+    async fn perform_is_master(&mut self, topology: &Topology) -> Result<IsMasterReply> {
         let result = match self.connection {
             Some(ref mut conn) => {
                 let command = is_master_command(self.client_options.server_api.as_ref());
                 run_is_master(
                     conn,
                     command,
-                    Some(&self.topology),
+                    Some(topology),
                     &self.client_options.sdam_event_handler,
                 )
                 .await
@@ -212,7 +212,7 @@ impl HeartbeatMonitor {
                     .handshaker
                     .handshake(
                         &mut connection,
-                        Some(&self.topology),
+                        Some(topology),
                         &self.client_options.sdam_event_handler,
                     )
                     .await

src/sdam/state/mod.rs

@@ -81,6 +81,8 @@ struct TopologyState {
     http_client: HttpClient,
     description: TopologyDescription,
     servers: HashMap<ServerAddress, Arc<Server>>,
+    options: ClientOptions,
+    id: ObjectId,
 }
 
 impl Topology {
@@ -108,6 +110,8 @@ impl Topology {
             description,
             servers: Default::default(),
             http_client,
+            options: options.clone(),
+            id,
         };
 
         let state = Arc::new(RwLock::new(topology_state));
@@ -169,21 +173,6 @@ impl Topology {
 
     pub(crate) fn close(&self) {
         self.common.is_alive.store(false, Ordering::SeqCst);
-        if let Some(ref handler) = self.common.options.sdam_event_handler {
-            if self.common.options.load_balanced.unwrap_or(false) {
-                for host in &self.common.options.hosts {
-                    let event = ServerClosedEvent {
-                        address: host.clone(),
-                        topology_id: self.common.id,
-                    };
-                    handler.handle_server_closed_event(event);
-                }
-            }
-            let event = TopologyClosedEvent {
-                topology_id: self.common.id,
-            };
-            handler.handle_topology_closed_event(event);
-        }
     }
 
     /// Gets the addresses of the servers in the cluster.
@@ -485,6 +474,26 @@ impl Topology {
     }
 }
 
+impl Drop for TopologyState {
+    fn drop(&mut self) {
+        if let Some(ref handler) = self.options.sdam_event_handler {
+            if matches!(self.description.topology_type, TopologyType::LoadBalanced) {
+                for host in self.servers.keys() {
+                    let event = ServerClosedEvent {
+                        address: host.clone(),
+                        topology_id: self.id,
+                    };
+                    handler.handle_server_closed_event(event);
+                }
+            }
+            let event = TopologyClosedEvent {
+                topology_id: self.id,
+            };
+            handler.handle_topology_closed_event(event);
+        }
+    }
+}
+
 impl WeakTopology {
     /// Attempts to convert the WeakTopology to a strong reference.
     pub(crate) fn upgrade(&self) -> Option<Topology> {