RUST-104 Poll SRV records for mongos discovery (#161)

Author: saghm

Cargo.toml

@@ -38,7 +38,7 @@ sha-1 = "0.8.1"
 sha2 = "0.8.0"
 stringprep = "0.1.2"
 time = "0.1.42"
-trust-dns-proto = "0.19.0"
+trust-dns-proto = "0.19.4"
 trust-dns-resolver = "0.19.0"
 typed-builder = "0.3.0"
 version_check = "0.9.1"

src/client/options/mod.rs

@@ -304,6 +304,9 @@ pub struct ClientOptions {
     #[builder(default)]
     pub(crate) zlib_compression: Option<i32>,
 
+    #[builder(default)]
+    original_srv_hostname: Option<String>,
+
     #[builder(default)]
     original_uri: Option<String>,
 }
@@ -492,12 +495,22 @@ impl From<ClientOptionsParser> for ClientOptions {
             credential: parser.credential,
             cmap_event_handler: None,
             command_event_handler: None,
+            original_srv_hostname: None,
             original_uri: Some(parser.original_uri),
         }
     }
 }
 
 impl ClientOptions {
+    /// Creates a new ClientOptions with the `original_srv_hostname` field set to the testing value
+    /// used in the SRV tests.
+    #[cfg(test)]
+    pub(crate) fn new_srv() -> Self {
+        let mut options = Self::default();
+        options.original_srv_hostname = Some("localhost.test.test.build.10gen.cc".into());
+        options
+    }
+
     /// Parses a MongoDB connection string into a ClientOptions struct. If the string is malformed
     /// or one of the options has an invalid value, an error will be returned.
     ///
@@ -555,11 +568,14 @@ impl ClientOptions {
         let mut options: Self = parser.into();
 
         if srv {
-            let resolver = SrvResolver::new().await?;
+            let mut resolver = SrvResolver::new().await?;
             let mut config = resolver
                 .resolve_client_options(&options.hosts[0].hostname)
                 .await?;
 
+            // Save the original SRV hostname to allow mongos polling.
+            options.original_srv_hostname = Some(options.hosts[0].hostname.clone());
+
             // Set the ClientOptions hosts to those found during the SRV lookup.
             options.hosts = config.hosts;
 
@@ -593,6 +609,11 @@ impl ClientOptions {
         Ok(options)
     }
 
+    /// Gets the original SRV hostname specified when this ClientOptions was parsed from a URI.
+    pub(crate) fn original_srv_hostname(&self) -> Option<&String> {
+        self.original_srv_hostname.as_ref()
+    }
+
     pub(crate) fn tls_options(&self) -> Option<TlsOptions> {
         match self.tls {
             Some(Tls::Enabled(ref opts)) => Some(opts.clone()),

src/sdam/description/topology/mod.rs

@@ -84,6 +84,27 @@ impl PartialEq for TopologyDescription {
 }
 
 impl TopologyDescription {
+    /// Creates a new TopologyDescription with the set of servers initialized to the addresses
+    /// specified in `hosts` and each other field set to its default value.
+    #[cfg(test)]
+    pub(crate) fn new_from_hosts(hosts: Vec<StreamAddress>) -> Self {
+        Self {
+            single_seed: false,
+            topology_type: TopologyType::Unknown,
+            set_name: None,
+            max_set_version: None,
+            max_election_id: None,
+            compatibility_error: None,
+            logical_session_timeout_minutes: None,
+            local_threshold: None,
+            heartbeat_freq: None,
+            servers: hosts
+                .into_iter()
+                .map(|address| (address.clone(), ServerDescription::new(address, None)))
+                .collect(),
+        }
+    }
+
     pub(crate) fn new(options: ClientOptions) -> Result<Self> {
         verify_max_staleness(
             options
@@ -124,6 +145,11 @@ impl TopologyDescription {
         })
     }
 
+    /// Gets the topology type of the cluster.
+    pub(crate) fn topology_type(&self) -> TopologyType {
+        self.topology_type
+    }
+
     pub(crate) fn server_addresses(&self) -> impl Iterator<Item = &StreamAddress> {
         self.servers.keys()
     }
@@ -272,6 +298,14 @@ impl TopologyDescription {
         })
     }
 
+    /// Syncs the set of servers in the description to those in `hosts`. Servers in the set not
+    /// already present in the cluster will be added, and servers in the cluster not present in the
+    /// set will be removed.
+    pub(crate) fn sync_hosts(&mut self, hosts: &HashSet<StreamAddress>) {
+        self.add_new_servers_from_addresses(hosts.iter());
+        self.servers.retain(|host, _| hosts.contains(host));
+    }
+
     /// Update the topology based on the new information about the topology contained by the
     /// ServerDescription.
     pub(crate) fn update(&mut self, mut server_description: ServerDescription) -> Result<()> {
@@ -532,17 +566,24 @@ impl TopologyDescription {
     }
 
     /// Create a new ServerDescription for each address and add it to the topology.
-    fn add_new_servers<'a>(&'a mut self, servers: impl Iterator<Item = &'a String>) -> Result<()> {
-        for server in servers {
-            let server = StreamAddress::parse(&server)?;
+    fn add_new_servers<'a>(&mut self, servers: impl Iterator<Item = &'a String>) -> Result<()> {
+        let servers: Result<Vec<_>> = servers.map(|server| StreamAddress::parse(server)).collect();
+
+        self.add_new_servers_from_addresses(servers?.iter());
+        Ok(())
+    }
 
+    /// Create a new ServerDescription for each address and add it to the topology.
+    fn add_new_servers_from_addresses<'a>(
+        &mut self,
+        servers: impl Iterator<Item = &'a StreamAddress>,
+    ) {
+        for server in servers {
             if !self.servers.contains_key(&server) {
                 self.servers
-                    .insert(server.clone(), ServerDescription::new(server, None));
+                    .insert(server.clone(), ServerDescription::new(server.clone(), None));
             }
         }
-
-        Ok(())
     }
 }
 

src/sdam/mod.rs

@@ -3,6 +3,7 @@ mod description;
 mod message_manager;
 mod monitor;
 pub mod public;
+mod srv_polling;
 mod state;
 
 pub use self::public::{ServerInfo, ServerType};

src/sdam/monitor.rs

@@ -18,7 +18,7 @@ use crate::{
     RUNTIME,
 };
 
-const DEFAULT_HEARTBEAT_FREQUENCY: Duration = Duration::from_secs(10);
+pub(super) const DEFAULT_HEARTBEAT_FREQUENCY: Duration = Duration::from_secs(10);
 
 pub(crate) const MIN_HEARTBEAT_FREQUENCY: Duration = Duration::from_millis(500);
 

src/sdam/srv_polling/mod.rs

@@ -0,0 +1,149 @@
+#[cfg(test)]
+mod test;
+
+use std::time::Duration;
+
+use super::{
+    monitor::DEFAULT_HEARTBEAT_FREQUENCY,
+    state::{Topology, TopologyState, WeakTopology},
+};
+use crate::{
+    error::{Error, Result},
+    options::{ClientOptions, StreamAddress},
+    srv::SrvResolver,
+    RUNTIME,
+};
+
+const DEFAULT_RESCAN_SRV_INTERVAL: Duration = Duration::from_secs(60);
+
+pub(crate) struct SrvPollingMonitor {
+    initial_hostname: String,
+    resolver: Option<SrvResolver>,
+    topology: WeakTopology,
+    rescan_interval: Option<Duration>,
+    client_options: ClientOptions,
+}
+
+struct LookupHosts {
+    hosts: Vec<StreamAddress>,
+    min_ttl: Option<Duration>,
+}
+
+impl SrvPollingMonitor {
+    pub(crate) fn new(topology: WeakTopology) -> Option<Self> {
+        let client_options = topology.client_options().clone();
+
+        let initial_hostname = match client_options.original_srv_hostname() {
+            Some(hostname) => hostname.clone(),
+            None => return None,
+        };
+
+        Some(Self {
+            initial_hostname,
+            resolver: None,
+            topology,
+            rescan_interval: None,
+            client_options,
+        })
+    }
+
+    /// Starts a monitoring task that periodically performs SRV record lookups to determine if the
+    /// set of mongos in the cluster have changed. A weak reference is used to ensure that the
+    /// monitoring task doesn't keep the topology alive after the client has been dropped.
+    pub(super) fn start(topology: WeakTopology) {
+        RUNTIME.execute(async move {
+            if let Some(mut monitor) = Self::new(topology) {
+                monitor.execute().await;
+            }
+        });
+    }
+
+    async fn execute(&mut self) {
+        while let Some(topology) = self.topology.upgrade() {
+            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;
+            }
+
+            RUNTIME
+                .delay_for(self.rescan_interval.unwrap_or(DEFAULT_RESCAN_SRV_INTERVAL))
+                .await;
+        }
+    }
+
+    async fn update_hosts(
+        &mut self,
+        lookup: Result<LookupHosts>,
+        topology: Topology,
+        mut topology_state: TopologyState,
+    ) {
+        let lookup = match lookup {
+            Ok(LookupHosts { hosts, .. }) if hosts.is_empty() => {
+                self.no_valid_hosts(None);
+
+                return;
+            }
+            Ok(lookup) => lookup,
+            Err(err) => {
+                self.no_valid_hosts(Some(err));
+
+                return;
+            }
+        };
+
+        self.rescan_interval = lookup.min_ttl;
+
+        let diff =
+            topology_state.update_hosts(&lookup.hosts.into_iter().collect(), &self.client_options);
+        topology.update_state(diff, topology_state).await;
+    }
+
+    async fn lookup_hosts(&mut self) -> Result<LookupHosts> {
+        let initial_hostname = self.initial_hostname.clone();
+        let resolver = self.get_or_create_srv_resolver().await?;
+        let mut new_hosts = Vec::new();
+
+        for host in resolver.get_srv_hosts(&initial_hostname).await? {
+            #[allow(clippy::single_match)]
+            match host {
+                Ok(host) => new_hosts.push(host),
+                Err(_) => {
+                    // TODO RUST-230: Log error with host that was returned.
+                }
+            }
+        }
+
+        Ok(LookupHosts {
+            hosts: new_hosts,
+            min_ttl: resolver
+                .min_ttl()
+                .map(|ttl| Duration::from_secs(ttl as u64)),
+        })
+    }
+
+    async fn get_or_create_srv_resolver(&mut self) -> Result<&mut SrvResolver> {
+        if let Some(ref mut resolver) = self.resolver {
+            return Ok(resolver);
+        }
+
+        let resolver = SrvResolver::new().await?;
+
+        // Since the connection was not `Some` above, this will always insert the new connection and
+        // return a reference to it.
+        Ok(self.resolver.get_or_insert(resolver))
+    }
+
+    fn no_valid_hosts(&mut self, _error: Option<Error>) {
+        // TODO RUST-230: Log error/lack of valid results.
+
+        self.rescan_interval = Some(self.heartbeat_freq());
+    }
+
+    fn heartbeat_freq(&self) -> Duration {
+        self.client_options
+            .heartbeat_freq
+            .unwrap_or(DEFAULT_HEARTBEAT_FREQUENCY)
+    }
+}