RUST-575 Consider server load during server selection

Author: patrickfreed

Cargo.toml

@@ -41,7 +41,7 @@ lazy_static = "1.4.0"
 md-5 = "0.8.0"
 os_info = { version = "2.0.6", default-features = false }
 percent-encoding = "2.0.0"
-rand = "0.7.2"
+rand = { version = "0.7.2", features = ["small_rng"] }
 serde_with = "1.3.1"
 sha-1 = "0.8.1"
 sha2 = "0.8.0"

src/client/executor.rs

@@ -12,7 +12,7 @@ use crate::{
     event::command::{CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent},
     operation::{Operation, Retryability},
     options::SelectionCriteria,
-    sdam::{Server, SessionSupportStatus},
+    sdam::{SelectedServer, SessionSupportStatus},
 };
 
 lazy_static! {
@@ -356,7 +356,7 @@ impl Client {
                 let criteria = SelectionCriteria::Predicate(Arc::new(move |server_info| {
                     server_info.server_type().is_data_bearing()
                 }));
-                let _: Arc<Server> = self.select_server(Some(&criteria)).await?;
+                let _: SelectedServer = self.select_server(Some(&criteria)).await?;
                 Ok(self.inner.topology.session_support_status().await)
             }
             _ => Ok(initial_status),

src/client/mod.rs

@@ -24,7 +24,7 @@ use crate::{
         ReadPreference,
         SelectionCriteria,
     },
-    sdam::{Server, SessionSupportStatus, Topology},
+    sdam::{SelectedServer, SessionSupportStatus, Topology},
 };
 pub(crate) use session::{ClientSession, ClusterTime, SESSIONS_UNSUPPORTED_COMMANDS};
 use session::{ServerSession, ServerSessionPool};
@@ -246,7 +246,7 @@ impl Client {
 
     /// Select a server using the provided criteria. If none is provided, a primary read preference
     /// will be used instead.
-    async fn select_server(&self, criteria: Option<&SelectionCriteria>) -> Result<Arc<Server>> {
+    async fn select_server(&self, criteria: Option<&SelectionCriteria>) -> Result<SelectedServer> {
         let criteria =
             criteria.unwrap_or(&SelectionCriteria::ReadPreference(ReadPreference::Primary));
 

src/client/options/mod.rs

@@ -22,7 +22,7 @@ use rustls::{
     ServerCertVerifier,
     TLSError,
 };
-use serde::Deserialize;
+use serde::{de::Error, Deserialize, Deserializer};
 use strsim::jaro_winkler;
 pub use trust_dns_resolver::config::ResolverConfig;
 use typed_builder::TypedBuilder;
@@ -91,7 +91,7 @@ lazy_static! {
 }
 
 /// A hostname:port address pair.
-#[derive(Clone, Debug, Deserialize, Eq)]
+#[derive(Clone, Debug, Eq)]
 pub struct StreamAddress {
     /// The hostname of the address.
     pub hostname: String,
@@ -102,6 +102,16 @@ pub struct StreamAddress {
     pub port: Option<u16>,
 }
 
+impl<'de> Deserialize<'de> for StreamAddress {
+    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let s: String = Deserialize::deserialize(deserializer)?;
+        Self::parse(s.as_str()).map_err(|e| D::Error::custom(format!("{}", e)))
+    }
+}
+
 impl Default for StreamAddress {
     fn default() -> Self {
         Self {

src/cmap/connection_requester.rs

@@ -9,6 +9,22 @@ use crate::{
 };
 use std::time::Duration;
 
+/// Returns a new requester/receiver pair.
+pub(super) fn channel(
+    address: StreamAddress,
+    handle: PoolWorkerHandle,
+) -> (ConnectionRequester, ConnectionRequestReceiver) {
+    let (sender, receiver) = mpsc::unbounded_channel();
+    (
+        ConnectionRequester {
+            address,
+            sender,
+            handle,
+        },
+        ConnectionRequestReceiver::new(receiver),
+    )
+}
+
 /// Handle for requesting Connections from the pool.
 /// This requester will keep the pool alive. Once all requesters have been dropped,
 /// the pool will stop servicing requests, drop its available connections, and close.
@@ -20,22 +36,6 @@ pub(super) struct ConnectionRequester {
 }
 
 impl ConnectionRequester {
-    /// Returns a new requester/receiver pair.
-    pub(super) fn new(
-        address: StreamAddress,
-        handle: PoolWorkerHandle,
-    ) -> (Self, ConnectionRequestReceiver) {
-        let (sender, receiver) = mpsc::unbounded_channel();
-        (
-            Self {
-                address,
-                sender,
-                handle,
-            },
-            ConnectionRequestReceiver::new(receiver),
-        )
-    }
-
     /// Request a connection from the pool that owns the receiver end of this requester.
     /// Returns an error if it takes longer than wait_queue_timeout before either a connection is
     /// received or an establishment begins.

src/cmap/manager.rs

@@ -3,19 +3,22 @@ use tokio::sync::mpsc;
 use super::Connection;
 use crate::error::Error;
 
+pub(super) fn channel() -> (PoolManager, ManagementRequestReceiver) {
+    let (sender, receiver) = mpsc::unbounded_channel();
+    (
+        PoolManager { sender },
+        ManagementRequestReceiver { receiver },
+    )
+}
+
 /// Struct used to make management requests to the pool (e.g. checking in a connection).
 /// A PoolManager will NOT keep a pool from going out of scope and closing.
 #[derive(Clone, Debug)]
 pub(super) struct PoolManager {
-    pub(super) sender: mpsc::UnboundedSender<PoolManagementRequest>,
+    sender: mpsc::UnboundedSender<PoolManagementRequest>,
 }
 
 impl PoolManager {
-    pub(super) fn new() -> (PoolManager, ManagementRequestReceiver) {
-        let (sender, receiver) = mpsc::unbounded_channel();
-        (Self { sender }, ManagementRequestReceiver { receiver })
-    }
-
     /// Lazily clear the pool.
     pub(super) fn clear(&self) {
         let _ = self.sender.send(PoolManagementRequest::Clear);

src/cmap/mod.rs

@@ -31,6 +31,9 @@ use connection_requester::ConnectionRequester;
 use manager::PoolManager;
 use worker::ConnectionPoolWorker;
 
+#[cfg(test)]
+use self::worker::PoolWorkerHandle;
+
 const DEFAULT_MAX_POOL_SIZE: u32 = 100;
 
 /// A pool of connections implementing the CMAP spec. All state is kept internally in an `Arc`, and
@@ -41,6 +44,7 @@ pub(crate) struct ConnectionPool {
     address: StreamAddress,
     manager: PoolManager,
     connection_requester: ConnectionRequester,
+
     wait_queue_timeout: Option<Duration>,
 
     #[derivative(Debug = "ignore")]
@@ -75,6 +79,21 @@ impl ConnectionPool {
         }
     }
 
+    #[cfg(test)]
+    pub(crate) fn new_mocked(address: StreamAddress) -> Self {
+        let (manager, _) = manager::channel();
+        let handle = PoolWorkerHandle::new_mocked();
+        let (connection_requester, _) = connection_requester::channel(Default::default(), handle);
+
+        Self {
+            address,
+            manager,
+            connection_requester,
+            wait_queue_timeout: None,
+            event_handler: None,
+        }
+    }
+
     fn emit_event<F>(&self, emit: F)
     where
         F: FnOnce(&Arc<dyn CmapEventHandler>),

src/cmap/test/integration.rs

@@ -72,7 +72,11 @@ async fn acquire_connection_and_send_command() {
 async fn concurrent_connections() {
     let _guard = LOCK.run_exclusively().await;
 
-    let client = TestClient::new().await;
+    let mut options = CLIENT_OPTIONS.clone();
+    options.direct_connection = Some(true);
+    options.hosts.drain(1..);
+
+    let client = TestClient::with_options(Some(options)).await;
     let version = VersionReq::parse(">= 4.2.9").unwrap();
     // blockConnection failpoint option only supported in 4.2.9+.
     if !version.matches(&client.server_version) {
@@ -165,7 +169,7 @@ async fn connection_error_during_establishment() {
 
     let options = FailCommandOptions::builder().error_code(1234).build();
     let failpoint = FailPoint::fail_command(&["isMaster"], FailPointMode::Times(10), Some(options));
-    let _fp_guard = client.enable_failpoint(failpoint).await.unwrap();
+    let _fp_guard = client.enable_failpoint(failpoint, None).await.unwrap();
 
     let handler = Arc::new(EventHandler::new());
     let mut subscriber = handler.subscribe();
@@ -213,7 +217,7 @@ async fn connection_error_during_operation() {
 
     let options = FailCommandOptions::builder().close_connection(true).build();
     let failpoint = FailPoint::fail_command(&["ping"], FailPointMode::Times(10), Some(options));
-    let _fp_guard = client.enable_failpoint(failpoint).await.unwrap();
+    let _fp_guard = client.enable_failpoint(failpoint, None).await.unwrap();
 
     let mut subscriber = handler.subscribe();
 

src/cmap/worker.rs

@@ -2,13 +2,15 @@ use derivative::Derivative;
 
 use super::{
     conn::PendingConnection,
+    connection_requester,
     connection_requester::{
         ConnectionRequest,
         ConnectionRequestReceiver,
         ConnectionRequester,
         RequestedConnection,
     },
     establish::ConnectionEstablisher,
+    manager,
     manager::{ManagementRequestReceiver, PoolManagementRequest, PoolManager},
     options::{ConnectionOptions, ConnectionPoolOptions},
     Connection,
@@ -132,10 +134,10 @@ impl ConnectionPoolWorker {
             .as_ref()
             .map(|pool_options| ConnectionOptions::from(pool_options.clone()));
 
-        let (handle_listener, handle) = HandleListener::new();
+        let (handle, handle_listener) = handle_channel();
         let (connection_requester, request_receiver) =
-            ConnectionRequester::new(address.clone(), handle);
-        let (manager, management_receiver) = PoolManager::new();
+            connection_requester::channel(address.clone(), handle);
+        let (manager, management_receiver) = manager::channel();
 
         let worker = ConnectionPoolWorker {
             address,
@@ -498,25 +500,35 @@ impl From<PoolManagementRequest> for PoolTask {
     }
 }
 
+/// Constructs a new channel for for monitoring whether this pool still has references
+/// to it.
+fn handle_channel() -> (PoolWorkerHandle, HandleListener) {
+    let (sender, receiver) = mpsc::channel(1);
+    (PoolWorkerHandle { sender }, HandleListener { receiver })
+}
+
 /// Handle to the worker. Once all handles have been dropped, the worker
 /// will stop waiting for new requests and drop the pool itself.
 #[derive(Debug, Clone)]
 pub(super) struct PoolWorkerHandle {
     sender: mpsc::Sender<()>,
 }
 
+impl PoolWorkerHandle {
+    #[cfg(test)]
+    pub(super) fn new_mocked() -> Self {
+        let (s, _) = handle_channel();
+        s
+    }
+}
+
 /// Listener used to determine when all handles have been dropped.
 #[derive(Debug)]
 struct HandleListener {
     receiver: mpsc::Receiver<()>,
 }
 
 impl HandleListener {
-    fn new() -> (Self, PoolWorkerHandle) {
-        let (sender, receiver) = mpsc::channel(1);
-        (Self { receiver }, PoolWorkerHandle { sender })
-    }
-
     /// Listen until all handles are dropped.
     /// This will not return until all handles are dropped, so make sure to only poll this via
     /// select or with a timeout.

src/sdam/description/topology/mod.rs

@@ -1,4 +1,4 @@
-mod server_selection;
+pub(crate) mod server_selection;
 #[cfg(test)]
 mod test;