RUST-2203 Allow client metadata to be appended post-construction (#1471)

Co-authored-by: Isabel Atkinson <[email protected]>

Author: abr-egn

src/client.rs

@@ -36,7 +36,14 @@ use crate::{
     event::command::CommandEvent,
     id_set::IdSet,
     operation::OverrideCriteriaFn,
-    options::{ClientOptions, DatabaseOptions, ReadPreference, SelectionCriteria, ServerAddress},
+    options::{
+        ClientOptions,
+        DatabaseOptions,
+        DriverInfo,
+        ReadPreference,
+        SelectionCriteria,
+        ServerAddress,
+    },
     sdam::{
         server_selection::{self, attempt_to_select_server},
         SelectedServer,
@@ -380,6 +387,16 @@ impl Client {
             .map(|db_name| self.database(db_name))
     }
 
+    /// Append new information to the metadata of the handshake with the server.
+    pub fn append_metadata(&self, driver_info: DriverInfo) -> Result<()> {
+        self.inner
+            .topology
+            .metadata
+            .write()
+            .unwrap()
+            .append(driver_info)
+    }
+
     pub(crate) fn register_async_drop(&self) -> AsyncDropToken {
         let (cleanup_tx, cleanup_rx) = tokio::sync::oneshot::channel::<BoxFuture<'static, ()>>();
         let (id_tx, id_rx) = tokio::sync::oneshot::channel::<crate::id_set::Id>();
@@ -418,7 +435,7 @@ impl Client {
     /// Check in a server session to the server session pool. The session will be discarded if it is
     /// expired or dirty.
     pub(crate) async fn check_in_server_session(&self, session: ServerSession) {
-        let timeout = self.inner.topology.logical_session_timeout();
+        let timeout = self.inner.topology.watcher().logical_session_timeout();
         self.inner.session_pool.check_in(session, timeout).await;
     }
 
@@ -481,12 +498,12 @@ impl Client {
             timeout,
         );
         #[cfg(feature = "tracing-unstable")]
-        event_emitter.emit_started_event(self.inner.topology.watch().observe_latest().description);
+        event_emitter.emit_started_event(self.inner.topology.latest().description.clone());
         // We only want to emit this message once per operation at most.
         #[cfg(feature = "tracing-unstable")]
         let mut emitted_waiting_message = false;
 
-        let mut watcher = self.inner.topology.watch();
+        let mut watcher = self.inner.topology.watcher().clone();
         loop {
             let state = watcher.observe_latest();
             let override_slot;
@@ -550,8 +567,7 @@ impl Client {
 
     #[cfg(all(test, feature = "dns-resolver"))]
     pub(crate) fn get_hosts(&self) -> Vec<String> {
-        let watcher = self.inner.topology.watch();
-        let state = watcher.peek_latest();
+        let state = self.inner.topology.latest();
 
         state
             .servers()
@@ -562,17 +578,12 @@ impl Client {
 
     #[cfg(test)]
     pub(crate) async fn sync_workers(&self) {
-        self.inner.topology.sync_workers().await;
+        self.inner.topology.updater().sync_workers().await;
     }
 
     #[cfg(test)]
     pub(crate) fn topology_description(&self) -> crate::sdam::TopologyDescription {
-        self.inner
-            .topology
-            .watch()
-            .peek_latest()
-            .description
-            .clone()
+        self.inner.topology.latest().description.clone()
     }
 
     #[cfg(test)]
@@ -588,7 +599,7 @@ impl Client {
             .options
             .server_selection_timeout
             .unwrap_or(DEFAULT_SERVER_SELECTION_TIMEOUT);
-        let mut watcher = self.inner.topology.watch();
+        let mut watcher = self.inner.topology.watcher().clone();
         loop {
             let topology = watcher.observe_latest();
             if let Some(desc) = topology.description.primary() {
@@ -628,7 +639,7 @@ impl Client {
         // The maximum number of session IDs that should be sent in a single endSessions command.
         const MAX_END_SESSIONS_BATCH_SIZE: usize = 10_000;
 
-        let mut watcher = self.inner.topology.watch();
+        let mut watcher = self.inner.topology.watcher().clone();
         let selection_criteria =
             SelectionCriteria::from(ReadPreference::PrimaryPreferred { options: None });
 

src/client/action/perf.rs

@@ -15,6 +15,6 @@ impl<'a> Action for crate::action::WarmConnectionPool<'a> {
             // No-op when min_pool_size is zero.
             return;
         }
-        self.client.inner.topology.warm_pool().await;
+        self.client.inner.topology.updater().fill_pool().await;
     }
 }

src/client/action/shutdown.rs

@@ -28,7 +28,7 @@ impl Action for crate::action::Shutdown {
         if !self.client.inner.shutdown.executed.load(Ordering::SeqCst) {
             self.client.end_all_sessions().await;
         }
-        self.client.inner.topology.shutdown().await;
+        self.client.inner.topology.updater().shutdown().await;
         // This has to happen last to allow pending cleanup to execute commands.
         self.client
             .inner

src/client/executor.rs

@@ -433,6 +433,7 @@ impl Client {
 
                     self.inner
                         .topology
+                        .updater()
                         .handle_application_error(
                             server_addr.clone(),
                             err.clone(),
@@ -651,7 +652,7 @@ impl Client {
         let stream_description = connection.stream_description()?;
         let is_sharded = stream_description.initial_server_type == ServerType::Mongos;
         let mut cmd = op.build(stream_description)?;
-        self.inner.topology.update_command_with_read_pref(
+        self.inner.topology.watcher().update_command_with_read_pref(
             connection.address(),
             &mut cmd,
             &effective_criteria,
@@ -749,7 +750,7 @@ impl Client {
         }
 
         let session_cluster_time = session.as_ref().and_then(|session| session.cluster_time());
-        let client_cluster_time = self.inner.topology.cluster_time();
+        let client_cluster_time = self.inner.topology.watcher().cluster_time();
         let max_cluster_time = std::cmp::max(session_cluster_time, client_cluster_time.as_ref());
         if let Some(cluster_time) = max_cluster_time {
             cmd.set_cluster_time(cluster_time);
@@ -883,7 +884,7 @@ impl Client {
     }
 
     async fn select_data_bearing_server(&self, operation_name: &str) -> Result<()> {
-        let topology_type = self.inner.topology.topology_type();
+        let topology_type = self.inner.topology.watcher().topology_type();
         let criteria = SelectionCriteria::Predicate(Arc::new(move |server_info| {
             let server_type = server_info.server_type();
             (matches!(topology_type, TopologyType::Single) && server_type.is_available())
@@ -899,15 +900,15 @@ impl Client {
     /// topology supports transactions, this method will perform a server selection that will force
     /// that determination to be made.
     pub(crate) async fn transaction_support_status(&self) -> Result<TransactionSupportStatus> {
-        let initial_status = self.inner.topology.transaction_support_status();
+        let initial_status = self.inner.topology.watcher().transaction_support_status();
 
         // Need to guarantee that we're connected to at least one server that can determine if
         // sessions are supported or not.
         match initial_status {
             TransactionSupportStatus::Undetermined => {
                 self.select_data_bearing_server("Check transactions support status")
                     .await?;
-                Ok(self.inner.topology.transaction_support_status())
+                Ok(self.inner.topology.watcher().transaction_support_status())
             }
             _ => Ok(initial_status),
         }
@@ -953,6 +954,7 @@ impl Client {
         if let Some(ref cluster_time) = cluster_time {
             self.inner
                 .topology
+                .updater()
                 .advance_cluster_time(cluster_time.clone())
                 .await;
             if let Some(ref mut session) = session {

src/client/options.rs

@@ -656,6 +656,9 @@ pub(crate) struct TestOptions {
 
     /// Async-capable command event listener.
     pub(crate) async_event_listener: Option<TestEventSender>,
+
+    /// Callback to receive hello commands.
+    pub(crate) hello_cb: Option<EventHandler<crate::cmap::Command>>,
 }
 
 pub(crate) type TestEventSender = tokio::sync::mpsc::Sender<
@@ -1112,7 +1115,7 @@ impl TlsOptions {
 
 /// Extra information to append to the driver version in the metadata of the handshake with the
 /// server. This should be used by libraries wrapping the driver, e.g. ODMs.
-#[derive(Clone, Debug, Deserialize, TypedBuilder, PartialEq)]
+#[derive(Clone, Debug, Deserialize, TypedBuilder, Eq)]
 #[builder(field_defaults(default, setter(into)))]
 #[non_exhaustive]
 pub struct DriverInfo {
@@ -1127,6 +1130,32 @@ pub struct DriverInfo {
     pub platform: Option<String>,
 }
 
+impl DriverInfo {
+    pub(crate) fn spec_version(&self) -> &str {
+        self.version.as_deref().unwrap_or("")
+    }
+
+    pub(crate) fn spec_platform(&self) -> &str {
+        self.platform.as_deref().unwrap_or("")
+    }
+}
+
+impl PartialEq for DriverInfo {
+    fn eq(&self, other: &Self) -> bool {
+        self.name == other.name
+            && self.spec_version() == other.spec_version()
+            && self.spec_platform() == other.spec_platform()
+    }
+}
+
+impl Hash for DriverInfo {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.spec_version().hash(state);
+        self.spec_platform().hash(state);
+    }
+}
+
 impl ClientOptions {
     /// Creates a new ClientOptions with the `original_srv_hostname` field set to the testing value
     /// used in the SRV tests.

src/client/session.rs

@@ -212,7 +212,7 @@ impl ClientSession {
         options: Option<SessionOptions>,
         is_implicit: bool,
     ) -> Self {
-        let timeout = client.inner.topology.logical_session_timeout();
+        let timeout = client.inner.topology.watcher().logical_session_timeout();
         let server_session = client.inner.session_pool.check_out(timeout).await;
         Self {
             drop_token: client.register_async_drop(),

src/cmap/establish.rs

@@ -13,15 +13,17 @@ use super::{
     Connection,
     PoolGeneration,
 };
+#[cfg(test)]
+use crate::options::ClientOptions;
 use crate::{
     client::{
         auth::Credential,
-        options::{ClientOptions, ServerAddress, TlsOptions},
+        options::{ServerAddress, TlsOptions},
     },
     error::{Error as MongoError, ErrorKind, Result},
     hello::HelloReply,
     runtime::{self, stream::DEFAULT_CONNECT_TIMEOUT, AsyncStream, TlsConfig},
-    sdam::HandshakePhase,
+    sdam::{topology::TopologySpec, HandshakePhase},
 };
 
 /// Contains the logic to establish a connection, including handshaking, authenticating, and
@@ -48,18 +50,25 @@ pub(crate) struct EstablisherOptions {
     pub(crate) test_patch_reply: Option<fn(&mut Result<HelloReply>)>,
 }
 
-impl From<&ClientOptions> for EstablisherOptions {
-    fn from(opts: &ClientOptions) -> Self {
+impl From<&TopologySpec> for EstablisherOptions {
+    fn from(spec: &TopologySpec) -> Self {
         Self {
-            handshake_options: HandshakerOptions::from(opts),
-            tls_options: opts.tls_options(),
-            connect_timeout: opts.connect_timeout,
+            handshake_options: HandshakerOptions::from(spec),
+            tls_options: spec.options.tls_options(),
+            connect_timeout: spec.options.connect_timeout,
             #[cfg(test)]
             test_patch_reply: None,
         }
     }
 }
 
+#[cfg(test)]
+impl From<&ClientOptions> for EstablisherOptions {
+    fn from(options: &ClientOptions) -> Self {
+        Self::from(&TopologySpec::try_from(options.clone()).unwrap())
+    }
+}
+
 impl ConnectionEstablisher {
     /// Creates a new ConnectionEstablisher from the given options.
     pub(crate) fn new(options: EstablisherOptions) -> Result<Self> {