RUST-232 SDAM Monitoring (#404)

Author: isabelatkinson

src/client/mod.rs

@@ -90,7 +90,7 @@ struct ClientInner {
 
 impl Drop for ClientInner {
     fn drop(&mut self) {
-        self.topology.mark_closed()
+        self.topology.close()
     }
 }
 

src/client/options/mod.rs

@@ -43,7 +43,7 @@ use crate::{
     client::auth::{AuthMechanism, Credential},
     concern::{Acknowledgment, ReadConcern, WriteConcern},
     error::{ErrorKind, Result},
-    event::{cmap::CmapEventHandler, command::CommandEventHandler},
+    event::{cmap::CmapEventHandler, command::CommandEventHandler, sdam::SdamEventHandler},
     options::ReadConcernLevel,
     sdam::MIN_HEARTBEAT_FREQUENCY,
     selection_criteria::{ReadPreference, SelectionCriteria, TagSet},
@@ -492,6 +492,13 @@ pub struct ClientOptions {
     #[builder(default)]
     pub retry_writes: Option<bool>,
 
+    /// The handler that should process all Server Discovery and Monitoring events. See the
+    /// [`SdamEventHandler`] type documentation for more details.
+    #[derivative(Debug = "ignore", PartialEq = "ignore")]
+    #[builder(default)]
+    #[serde(skip)]
+    pub sdam_event_handler: Option<Arc<dyn SdamEventHandler>>,
+
     /// The default selection criteria for operations performed on the Client. See the
     /// SelectionCriteria type documentation for more details.
     #[builder(default)]
@@ -934,6 +941,7 @@ impl From<ClientOptionsParser> for ClientOptions {
             heartbeat_freq_test: None,
             allow_load_balanced: false,
             load_balanced: parser.load_balanced,
+            sdam_event_handler: None,
         }
     }
 }

src/cmap/conn/command.rs

@@ -1,11 +1,14 @@
+use std::time::Duration;
+
 use serde::{de::DeserializeOwned, Serialize};
 
 use super::wire::Message;
 use crate::{
     bson::Document,
     client::{options::ServerApi, ClusterTime},
     error::{Error, ErrorKind, Result},
-    operation::{CommandErrorBody, CommandResponse},
+    is_master::{IsMasterCommandResponse, IsMasterReply},
+    operation::{CommandErrorBody, CommandResponse, Response},
     options::{ReadConcern, ServerAddress},
     selection_criteria::ReadPreference,
     ClientSession,
@@ -168,53 +171,33 @@ impl RawCommandResponse {
             .map_err(|_| Error::invalid_authentication_response(mechanism_name))
     }
 
-    /// Deserialize the raw bytes into a response backed by a `Document` for further processing.
-    pub(crate) fn into_document_response(self) -> Result<DocumentCommandResponse> {
-        let response: CommandResponse<Document> = self.body()?;
-        Ok(DocumentCommandResponse { response })
+    pub(crate) fn to_is_master_response(&self, round_trip_time: Duration) -> Result<IsMasterReply> {
+        match self.body::<CommandResponse<IsMasterCommandResponse>>() {
+            Ok(response) if response.is_success() => {
+                let server_address = self.source_address().clone();
+                let cluster_time = response.cluster_time().cloned();
+                Ok(IsMasterReply {
+                    server_address,
+                    command_response: response.body,
+                    round_trip_time,
+                    cluster_time,
+                })
+            }
+            _ => match self.body::<CommandResponse<CommandErrorBody>>() {
+                Ok(command_error_body) => Err(Error::new(
+                    ErrorKind::Command(command_error_body.body.command_error),
+                    command_error_body.body.error_labels,
+                )),
+                Err(_) => Err(ErrorKind::InvalidResponse {
+                    message: "invalid server response".into(),
+                }
+                .into()),
+            },
+        }
     }
 
     /// The address of the server that sent this response.
     pub(crate) fn source_address(&self) -> &ServerAddress {
         &self.source
     }
 }
-
-/// A command response backed by a `Document` rather than raw bytes.
-/// Use this for simple command responses where deserialization performance is not a high priority.
-pub(crate) struct DocumentCommandResponse {
-    response: CommandResponse<Document>,
-}
-
-impl DocumentCommandResponse {
-    /// Returns a result indicating whether this response corresponds to a command failure.
-    pub(crate) fn validate(&self) -> Result<()> {
-        if !self.response.is_success() {
-            let error_response: CommandErrorBody = bson::from_document(self.response.body.clone())
-                .map_err(|_| ErrorKind::InvalidResponse {
-                    message: "invalid server response".to_string(),
-                })?;
-            Err(Error::new(
-                ErrorKind::Command(error_response.command_error),
-                error_response.error_labels,
-            ))
-        } else {
-            Ok(())
-        }
-    }
-
-    /// Deserialize the body of the response.
-    pub(crate) fn body<T: DeserializeOwned>(self) -> Result<T> {
-        match bson::from_document(self.response.body) {
-            Ok(body) => Ok(body),
-            Err(e) => Err(ErrorKind::InvalidResponse {
-                message: format!("{}", e),
-            }
-            .into()),
-        }
-    }
-
-    pub(crate) fn cluster_time(&self) -> Option<&ClusterTime> {
-        self.response.cluster_time.as_ref()
-    }
-}

src/cmap/conn/mod.rs

@@ -49,7 +49,7 @@ pub struct ConnectionInfo {
 #[derivative(Debug)]
 pub(crate) struct Connection {
     pub(super) id: u32,
-    pub(super) address: ServerAddress,
+    pub(crate) address: ServerAddress,
     pub(crate) generation: ConnectionGeneration,
 
     /// The cached StreamDescription from the connection's handshake.

src/cmap/establish/handshake/mod.rs

@@ -1,6 +1,8 @@
 #[cfg(test)]
 mod test;
 
+use std::sync::Arc;
+
 use lazy_static::lazy_static;
 use os_info::{Type, Version};
 
@@ -9,8 +11,10 @@ use crate::{
     client::auth::{ClientFirst, FirstRound},
     cmap::{options::ConnectionPoolOptions, Command, Connection, StreamDescription},
     error::{ErrorKind, Result},
+    event::sdam::SdamEventHandler,
     is_master::{is_master_command, run_is_master, IsMasterReply},
     options::{AuthMechanism, ClientOptions, Credential, DriverInfo, ServerApi},
+    sdam::WeakTopology,
 };
 
 #[cfg(feature = "tokio-runtime")]
@@ -133,7 +137,7 @@ lazy_static! {
 }
 
 /// Contains the logic needed to handshake a connection.
-#[derive(Debug, Clone)]
+#[derive(Clone, Debug)]
 pub(crate) struct Handshaker {
     /// The `isMaster` command to send when handshaking. This will always be identical
     /// given the same pool options, so it can be created at the time the Handshaker is created.
@@ -192,12 +196,17 @@ impl Handshaker {
     }
 
     /// Handshakes a connection.
-    pub(crate) async fn handshake(&self, conn: &mut Connection) -> Result<HandshakeResult> {
+    pub(crate) async fn handshake(
+        &self,
+        conn: &mut Connection,
+        topology: Option<&WeakTopology>,
+        handler: &Option<Arc<dyn SdamEventHandler>>,
+    ) -> Result<HandshakeResult> {
         let mut command = self.command.clone();
 
         let client_first = set_speculative_auth_info(&mut command.body, self.credential.as_ref())?;
 
-        let mut is_master_reply = run_is_master(command, conn).await?;
+        let mut is_master_reply = run_is_master(conn, command, topology, handler).await?;
         if self.command.body.contains_key("loadBalanced")
             && is_master_reply.command_response.service_id.is_none()
         {

src/cmap/establish/mod.rs

@@ -52,7 +52,7 @@ impl ConnectionEstablisher {
 
         let handshake = self
             .handshaker
-            .handshake(&mut connection)
+            .handshake(&mut connection, None, &None)
             .await
             .map_err(|e| EstablishError::pre_hello(e, pool_gen.clone()))?;
         let service_id = handshake.is_master_reply.command_response.service_id;

src/cmap/establish/test.rs

@@ -1,8 +1,9 @@
 use tokio::sync::RwLockWriteGuard;
 
 use crate::{
-    bson::{doc, Bson},
+    bson::{doc, Bson, Document},
     cmap::{establish::Handshaker, Command, Connection, ConnectionPoolOptions},
+    operation::CommandResponse,
     options::{AuthMechanism, ClientOptions, Credential, ReadPreference},
     test::{TestClient, CLIENT_OPTIONS, LOCK},
 };
@@ -47,7 +48,11 @@ async fn speculative_auth_test(
         .await
         .unwrap();
 
-    let first_round = handshaker.handshake(&mut conn).await.unwrap().first_round;
+    let first_round = handshaker
+        .handshake(&mut conn, None, &None)
+        .await
+        .unwrap()
+        .first_round;
 
     // We expect that the server will return a response with the `speculativeAuthenticate` field if
     // and only if it's new enough to support it.
@@ -70,11 +75,9 @@ async fn speculative_auth_test(
     });
 
     let response = conn.send_command(command, None).await.unwrap();
-    let doc_response = response.into_document_response().unwrap();
+    let doc_response: CommandResponse<Document> = response.body().unwrap();
 
-    doc_response
-        .validate()
-        .expect("response should be successful");
+    assert!(doc_response.is_success());
 }
 
 #[cfg_attr(feature = "tokio-runtime", tokio::test)]

src/cmap/mod.rs

@@ -72,7 +72,9 @@ impl ConnectionPool {
             options.clone(),
         );
 
-        let event_handler = options.as_ref().and_then(|opts| opts.event_handler.clone());
+        let event_handler = options
+            .as_ref()
+            .and_then(|opts| opts.cmap_event_handler.clone());
 
         if let Some(ref handler) = event_handler {
             handler.handle_pool_created_event(PoolCreatedEvent {

src/cmap/options.rs

@@ -41,7 +41,7 @@ pub(crate) struct ConnectionPoolOptions {
     /// Processes all events generated by the pool.
     #[derivative(Debug = "ignore", PartialEq = "ignore")]
     #[serde(skip)]
-    pub(crate) event_handler: Option<Arc<dyn CmapEventHandler>>,
+    pub(crate) cmap_event_handler: Option<Arc<dyn CmapEventHandler>>,
 
     /// Interval between background thread maintenance runs (e.g. ensure minPoolSize).
     #[cfg(test)]
@@ -103,7 +103,7 @@ impl ConnectionPoolOptions {
             server_api: options.server_api.clone(),
             tls_options: options.tls_options(),
             credential: options.credential.clone(),
-            event_handler: options.cmap_event_handler.clone(),
+            cmap_event_handler: options.cmap_event_handler.clone(),
             #[cfg(test)]
             background_thread_interval: None,
             #[cfg(test)]
@@ -163,7 +163,7 @@ impl From<ConnectionPoolOptions> for ConnectionOptions {
         Self {
             connect_timeout: pool_options.connect_timeout,
             tls_options: pool_options.tls_options,
-            event_handler: pool_options.event_handler,
+            event_handler: pool_options.cmap_event_handler,
         }
     }
 }

src/cmap/test/integration.rs

@@ -6,9 +6,10 @@ use super::{
     EVENT_TIMEOUT,
 };
 use crate::{
-    bson::doc,
+    bson::{doc, Document},
     cmap::{options::ConnectionPoolOptions, Command, ConnectionPool},
     event::cmap::{CmapEventHandler, ConnectionClosedReason},
+    operation::CommandResponse,
     sdam::ServerUpdateSender,
     selection_criteria::ReadPreference,
     test::{FailCommandOptions, FailPoint, FailPointMode, TestClient, CLIENT_OPTIONS, LOCK},
@@ -55,13 +56,11 @@ async fn acquire_connection_and_send_command() {
     }
 
     let response = connection.send_command(cmd, None).await.unwrap();
-    let doc_response = response.into_document_response().unwrap();
+    let doc_response: CommandResponse<Document> = response.body().unwrap();
 
-    doc_response
-        .validate()
-        .expect("response should be successful");
+    assert!(doc_response.is_success());
 
-    let response: ListDatabasesResponse = doc_response.body().unwrap();
+    let response: ListDatabasesResponse = bson::from_document(doc_response.body).unwrap();
 
     let names: Vec<_> = response
         .databases
@@ -107,7 +106,7 @@ async fn concurrent_connections() {
     let handler = Arc::new(EventHandler::new());
     let client_options = CLIENT_OPTIONS.clone();
     let mut options = ConnectionPoolOptions::from_client_options(&client_options);
-    options.event_handler = Some(handler.clone() as Arc<dyn crate::cmap::CmapEventHandler>);
+    options.cmap_event_handler = Some(handler.clone() as Arc<dyn crate::cmap::CmapEventHandler>);
     options.ready = Some(true);
 
     let pool = ConnectionPool::new(
@@ -191,7 +190,7 @@ async fn connection_error_during_establishment() {
 
     let mut options = ConnectionPoolOptions::from_client_options(&client_options);
     options.ready = Some(true);
-    options.event_handler = Some(handler.clone() as Arc<dyn crate::cmap::CmapEventHandler>);
+    options.cmap_event_handler = Some(handler.clone() as Arc<dyn crate::cmap::CmapEventHandler>);
     let pool = ConnectionPool::new(
         client_options.hosts[0].clone(),
         Default::default(),