RUST-590 Pause connection pool when server is marked as Unknown (#278)

Author: patrickfreed

src/client/executor.rs

@@ -82,16 +82,7 @@ impl Client {
     ) -> Result<T::O> {
         let server = self.select_server(op.selection_criteria()).await?;
 
-        let mut conn = match server.checkout_connection().await {
-            Ok(conn) => conn,
-            Err(err) => {
-                self.inner
-                    .topology
-                    .handle_pre_handshake_error(err.clone(), server.address.clone())
-                    .await;
-                return Err(err);
-            }
-        };
+        let mut conn = server.pool.check_out().await?;
 
         let retryability = self.get_retryability(&conn, &op).await?;
 
@@ -153,15 +144,9 @@ impl Client {
             }
         };
 
-        let mut conn = match server.checkout_connection().await {
-            Ok(conn) => conn,
-            Err(err) => {
-                self.inner
-                    .topology
-                    .handle_pre_handshake_error(err.clone(), server.address.clone())
-                    .await;
-                return Err(first_error);
-            }
+        let mut conn = match server.pool.check_out().await {
+            Ok(c) => c,
+            Err(_) => return Err(first_error),
         };
 
         let retryability = self.get_retryability(&conn, &op).await?;

src/cmap/connection_requester.rs

@@ -1,12 +1,7 @@
 use tokio::sync::{mpsc, oneshot};
 
 use super::{worker::PoolWorkerHandle, Connection};
-use crate::{
-    error::{ErrorKind, Result},
-    options::StreamAddress,
-    runtime::AsyncJoinHandle,
-    RUNTIME,
-};
+use crate::{error::Result, options::StreamAddress, runtime::AsyncJoinHandle, RUNTIME};
 use std::time::Duration;
 
 /// Returns a new requester/receiver pair.
@@ -21,7 +16,7 @@ pub(super) fn channel(
             sender,
             handle,
         },
-        ConnectionRequestReceiver::new(receiver),
+        ConnectionRequestReceiver { receiver },
     )
 }
 
@@ -31,112 +26,87 @@ pub(super) fn channel(
 #[derive(Clone, Debug)]
 pub(super) struct ConnectionRequester {
     address: StreamAddress,
-    sender: mpsc::UnboundedSender<oneshot::Sender<RequestedConnection>>,
+    sender: mpsc::UnboundedSender<oneshot::Sender<ConnectionRequestResult>>,
     handle: PoolWorkerHandle,
 }
 
 impl ConnectionRequester {
     /// 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.
-    pub(super) async fn request(&self, wait_queue_timeout: Option<Duration>) -> Result<Connection> {
+    /// Returns None if it takes longer than wait_queue_timeout before the pool returns a result.
+    pub(super) async fn request(
+        &self,
+        wait_queue_timeout: Option<Duration>,
+    ) -> Option<ConnectionRequestResult> {
         let (sender, receiver) = oneshot::channel();
 
         // this only errors if the receiver end is dropped, which can't happen because
         // we own a handle to the worker, keeping it alive.
         self.sender.send(sender).unwrap();
 
-        let response = match wait_queue_timeout {
+        match wait_queue_timeout {
             Some(timeout) => RUNTIME
                 .timeout(timeout, receiver)
                 .await
                 .map(|r| r.unwrap()) // see comment below as to why this is safe
-                .map_err(|_| {
-                    ErrorKind::WaitQueueTimeoutError {
-                        address: self.address.clone(),
-                    }
-                    .into()
-                }),
+                .ok(),
 
             // similarly, the receiver only returns an error if the sender is dropped, which
             // can't happen due to the handle.
-            None => Ok(receiver.await.unwrap()),
-        };
-
-        match response {
-            Ok(RequestedConnection::Pooled(c)) => Ok(c),
-            Ok(RequestedConnection::Establishing(task)) => task.await,
-            Err(e) => Err(e),
+            None => Some(receiver.await.unwrap()),
         }
     }
 }
 
 /// Receiving end of a given ConnectionRequester.
 #[derive(Debug)]
 pub(super) struct ConnectionRequestReceiver {
-    receiver: mpsc::UnboundedReceiver<oneshot::Sender<RequestedConnection>>,
-    cache: Option<ConnectionRequest>,
+    receiver: mpsc::UnboundedReceiver<oneshot::Sender<ConnectionRequestResult>>,
 }
 
 impl ConnectionRequestReceiver {
-    pub(super) fn new(
-        receiver: mpsc::UnboundedReceiver<oneshot::Sender<RequestedConnection>>,
-    ) -> Self {
-        Self {
-            receiver,
-            cache: None,
-        }
-    }
-
     pub(super) async fn recv(&mut self) -> Option<ConnectionRequest> {
-        match self.cache.take() {
-            Some(request) => Some(request),
-            None => self
-                .receiver
-                .recv()
-                .await
-                .map(|sender| ConnectionRequest { sender }),
-        }
-    }
-
-    /// Put a request back into the receiver. Next call to `recv` will immediately
-    /// return this value.
-    pub(super) fn cache_request(&mut self, request: ConnectionRequest) {
-        self.cache = Some(request);
+        self.receiver
+            .recv()
+            .await
+            .map(|sender| ConnectionRequest { sender })
     }
 }
 
 /// Struct encapsulating a request for a connection.
 #[derive(Debug)]
 pub(super) struct ConnectionRequest {
-    sender: oneshot::Sender<RequestedConnection>,
+    sender: oneshot::Sender<ConnectionRequestResult>,
 }
 
 impl ConnectionRequest {
     /// Respond to the connection request, either with a pooled connection or one that is
     /// establishing asynchronously.
     pub(super) fn fulfill(
         self,
-        conn: RequestedConnection,
-    ) -> std::result::Result<(), RequestedConnection> {
-        self.sender.send(conn)
+        result: ConnectionRequestResult,
+    ) -> std::result::Result<(), ConnectionRequestResult> {
+        self.sender.send(result)
     }
 }
 
 #[derive(Debug)]
-pub(super) enum RequestedConnection {
+pub(super) enum ConnectionRequestResult {
     /// A connection that was already established and was simply checked out of the pool.
     Pooled(Connection),
 
     /// A new connection in the process of being established.
     /// The handle can be awaited upon to receive the established connection.
     Establishing(AsyncJoinHandle<Result<Connection>>),
+
+    /// The request was rejected because the pool was cleared before it could
+    /// be fulfilled.
+    PoolCleared,
 }
 
-impl RequestedConnection {
+impl ConnectionRequestResult {
     pub(super) fn unwrap_pooled_connection(self) -> Connection {
         match self {
-            RequestedConnection::Pooled(c) => c,
+            ConnectionRequestResult::Pooled(c) => c,
             _ => panic!("attempted to unwrap pooled connection when was establishing"),
         }
     }

src/cmap/manager.rs

@@ -1,7 +1,7 @@
 use tokio::sync::mpsc;
 
 use super::Connection;
-use crate::error::Error;
+use crate::runtime::AcknowledgedMessage;
 
 pub(super) fn channel() -> (PoolManager, ManagementRequestReceiver) {
     let (sender, receiver) = mpsc::unbounded_channel();
@@ -24,6 +24,20 @@ impl PoolManager {
         let _ = self.sender.send(PoolManagementRequest::Clear);
     }
 
+    /// Mark the pool as "ready" as per the CMAP specification.
+    pub(super) async fn mark_as_ready(&self) {
+        let (message, acknowledgment_receiver) = AcknowledgedMessage::package(());
+        if self
+            .sender
+            .send(PoolManagementRequest::MarkAsReady {
+                completion_handler: message,
+            })
+            .is_ok()
+        {
+            acknowledgment_receiver.wait_for_acknowledgment().await;
+        }
+    }
+
     /// Check in the given connection to the pool.
     /// This returns an error containing the connection if the pool has been dropped already.
     pub(crate) fn check_in(&self, connection: Connection) -> std::result::Result<(), Connection> {
@@ -35,10 +49,10 @@ impl PoolManager {
     }
 
     /// Notify the pool that establishing a connection failed.
-    pub(super) fn handle_connection_failed(&self, error: Error) {
+    pub(super) fn handle_connection_failed(&self) {
         let _ = self
             .sender
-            .send(PoolManagementRequest::HandleConnectionFailed(error));
+            .send(PoolManagementRequest::HandleConnectionFailed);
     }
 
     /// Notify the pool that establishing a connection succeeded.
@@ -62,9 +76,22 @@ impl ManagementRequestReceiver {
 
 #[derive(Debug)]
 pub(super) enum PoolManagementRequest {
+    /// Clear the pool, transitioning it to Paused.
     Clear,
+
+    /// Mark the pool as Ready, allowing connections to be created and checked out.
+    MarkAsReady {
+        completion_handler: AcknowledgedMessage<()>,
+    },
+
+    /// Check in the given connection.
     CheckIn(Connection),
-    HandleConnectionFailed(Error),
+
+    /// Update the pool based on the given establishment error.
+    HandleConnectionFailed,
+
+    /// Update the pool after a successful connection, optionally populating the pool
+    /// with the successful connection.
     HandleConnectionSucceeded(Option<Connection>),
 }
 

src/cmap/mod.rs

@@ -1,25 +1,27 @@
 #[cfg(test)]
-mod test;
+pub(crate) mod test;
 
 pub(crate) mod conn;
 mod connection_requester;
 mod establish;
 mod manager;
 pub(crate) mod options;
+mod status;
 mod worker;
 
 use std::{sync::Arc, time::Duration};
 
 use derivative::Derivative;
 
 pub use self::conn::ConnectionInfo;
-use self::options::ConnectionPoolOptions;
 pub(crate) use self::{
     conn::{Command, CommandResponse, Connection, StreamDescription},
     establish::handshake::{is_master, Handshaker},
+    status::PoolGenerationSubscriber,
 };
+use self::{connection_requester::ConnectionRequestResult, options::ConnectionPoolOptions};
 use crate::{
-    error::{ErrorKind, Result},
+    error::{Error, ErrorKind, Result},
     event::cmap::{
         CmapEventHandler,
         ConnectionCheckoutFailedEvent,
@@ -29,6 +31,7 @@ use crate::{
     },
     options::StreamAddress,
     runtime::HttpClient,
+    sdam::ServerUpdateSender,
 };
 use connection_requester::ConnectionRequester;
 use manager::PoolManager;
@@ -47,6 +50,7 @@ pub(crate) struct ConnectionPool {
     address: StreamAddress,
     manager: PoolManager,
     connection_requester: ConnectionRequester,
+    generation_subscriber: PoolGenerationSubscriber,
 
     wait_queue_timeout: Option<Duration>,
 
@@ -58,10 +62,15 @@ impl ConnectionPool {
     pub(crate) fn new(
         address: StreamAddress,
         http_client: HttpClient,
+        server_updater: ServerUpdateSender,
         options: Option<ConnectionPoolOptions>,
     ) -> Self {
-        let (manager, connection_requester) =
-            ConnectionPoolWorker::start(address.clone(), http_client, options.clone());
+        let (manager, connection_requester, generation_subscriber) = ConnectionPoolWorker::start(
+            address.clone(),
+            http_client,
+            server_updater,
+            options.clone(),
+        );
 
         let event_handler = options.as_ref().and_then(|opts| opts.event_handler.clone());
         let wait_queue_timeout = options.as_ref().and_then(|opts| opts.wait_queue_timeout);
@@ -79,6 +88,7 @@ impl ConnectionPool {
             connection_requester,
             wait_queue_timeout,
             event_handler,
+            generation_subscriber,
         }
     }
 
@@ -87,11 +97,13 @@ impl ConnectionPool {
         let (manager, _) = manager::channel();
         let handle = PoolWorkerHandle::new_mocked();
         let (connection_requester, _) = connection_requester::channel(Default::default(), handle);
+        let (_, generation_subscriber) = status::channel();
 
         Self {
             address,
             manager,
             connection_requester,
+            generation_subscriber,
             wait_queue_timeout: None,
             event_handler: None,
         }
@@ -120,11 +132,23 @@ impl ConnectionPool {
             handler.handle_connection_checkout_started_event(event);
         });
 
-        let conn = self
+        let response = self
             .connection_requester
             .request(self.wait_queue_timeout)
             .await;
 
+        let conn = match response {
+            Some(ConnectionRequestResult::Pooled(c)) => Ok(c),
+            Some(ConnectionRequestResult::Establishing(task)) => task.await,
+            Some(ConnectionRequestResult::PoolCleared) => {
+                Err(Error::pool_cleared_error(&self.address))
+            }
+            None => Err(ErrorKind::WaitQueueTimeoutError {
+                address: self.address.clone(),
+            }
+            .into()),
+        };
+
         match conn {
             Ok(ref conn) => {
                 self.emit_event(|handler| {
@@ -156,4 +180,17 @@ impl ConnectionPool {
     pub(crate) fn clear(&self) {
         self.manager.clear();
     }
+
+    /// Mark the pool as "ready", allowing connections to be created and checked out.
+    pub(crate) async fn mark_as_ready(&self) {
+        self.manager.mark_as_ready().await;
+    }
+
+    /// Subscribe to updates to the pool's generation.
+    ///
+    /// This can be used to listen for errors that occur during connection
+    /// establishment or to get the current generation of the pool.
+    pub(crate) fn subscribe_to_generation_updates(&self) -> PoolGenerationSubscriber {
+        self.generation_subscriber.clone()
+    }
 }