RUST-695 Retry ConnectionPoolCleared errors (#365)

Author: patrickfreed

src/client/executor.rs

@@ -131,7 +131,12 @@ impl Client {
             Ok(conn) => conn,
             Err(mut err) => {
                 err.add_labels(None, &session, None)?;
-                return Err(err);
+
+                if err.is_pool_cleared() {
+                    return self.execute_retry(&mut op, &mut session, None, err).await;
+                } else {
+                    return Err(err);
+                }
             }
         };
 
@@ -151,7 +156,7 @@ impl Client {
             None => None,
         };
 
-        let first_error = match self
+        match self
             .execute_operation_on_connection(
                 &op,
                 &mut conn,
@@ -161,9 +166,7 @@ impl Client {
             )
             .await
         {
-            Ok(result) => {
-                return Ok(result);
-            }
+            Ok(result) => Ok(result),
             Err(mut err) => {
                 // Retryable writes are only supported by storage engines with document-level
                 // locking, so users need to disable retryable writes if using mmapv1.
@@ -194,13 +197,22 @@ impl Client {
                 if retryability == Retryability::Read && err.is_read_retryable()
                     || retryability == Retryability::Write && err.is_write_retryable()
                 {
-                    err
+                    self.execute_retry(&mut op, &mut session, txn_number, err)
+                        .await
                 } else {
-                    return Err(err);
+                    Err(err)
                 }
             }
-        };
+        }
+    }
 
+    async fn execute_retry<T: Operation>(
+        &self,
+        op: &mut T,
+        session: &mut Option<&mut ClientSession>,
+        txn_number: Option<u64>,
+        first_error: Error,
+    ) -> Result<T::O> {
         let server = match self.select_server(op.selection_criteria()).await {
             Ok(server) => server,
             Err(_) => {
@@ -213,21 +225,15 @@ impl Client {
             Err(_) => return Err(first_error),
         };
 
-        let retryability = self.get_retryability(&conn, &op, &session).await?;
+        let retryability = self.get_retryability(&conn, op, &session).await?;
         if retryability == Retryability::None {
             return Err(first_error);
         }
 
         op.update_for_retry();
 
         match self
-            .execute_operation_on_connection(
-                &op,
-                &mut conn,
-                &mut session,
-                txn_number,
-                &retryability,
-            )
+            .execute_operation_on_connection(op, &mut conn, session, txn_number, &retryability)
             .await
         {
             Ok(result) => Ok(result),

src/cmap/connection_requester.rs

@@ -1,7 +1,11 @@
 use tokio::sync::{mpsc, oneshot};
 
 use super::{worker::PoolWorkerHandle, Connection};
-use crate::{error::Result, options::ServerAddress, runtime::AsyncJoinHandle};
+use crate::{
+    error::{Error, Result},
+    options::ServerAddress,
+    runtime::AsyncJoinHandle,
+};
 
 /// Returns a new requester/receiver pair.
 pub(super) fn channel(
@@ -86,8 +90,8 @@ pub(super) enum ConnectionRequestResult {
     Establishing(AsyncJoinHandle<Result<Connection>>),
 
     /// The request was rejected because the pool was cleared before it could
-    /// be fulfilled.
-    PoolCleared,
+    /// be fulfilled. The error that caused the pool to be cleared is returned.
+    PoolCleared(Error),
 }
 
 impl ConnectionRequestResult {

src/cmap/manager.rs

@@ -1,7 +1,7 @@
 use tokio::sync::mpsc;
 
 use super::Connection;
-use crate::runtime::AcknowledgedMessage;
+use crate::{error::Error, runtime::AcknowledgedMessage};
 
 pub(super) fn channel() -> (PoolManager, ManagementRequestReceiver) {
     let (sender, receiver) = mpsc::unbounded_channel();
@@ -20,11 +20,14 @@ pub(super) struct PoolManager {
 
 impl PoolManager {
     /// Lazily clear the pool.
-    pub(super) async fn clear(&self) {
+    pub(super) async fn clear(&self, cause: Error) {
         let (message, acknowledgment_receiver) = AcknowledgedMessage::package(());
         if self
             .sender
-            .send(PoolManagementRequest::Clear(message))
+            .send(PoolManagementRequest::Clear {
+                completion_handler: message,
+                cause,
+            })
             .is_ok()
         {
             acknowledgment_receiver.wait_for_acknowledgment().await;
@@ -84,7 +87,10 @@ impl ManagementRequestReceiver {
 #[derive(Debug)]
 pub(super) enum PoolManagementRequest {
     /// Clear the pool, transitioning it to Paused.
-    Clear(AcknowledgedMessage<()>),
+    Clear {
+        completion_handler: AcknowledgedMessage<()>,
+        cause: Error,
+    },
 
     /// Mark the pool as Ready, allowing connections to be created and checked out.
     MarkAsReady {

src/cmap/mod.rs

@@ -130,7 +130,9 @@ impl ConnectionPool {
         let conn = match response {
             ConnectionRequestResult::Pooled(c) => Ok(c),
             ConnectionRequestResult::Establishing(task) => task.await,
-            ConnectionRequestResult::PoolCleared => Err(Error::pool_cleared_error(&self.address)),
+            ConnectionRequestResult::PoolCleared(e) => {
+                Err(Error::pool_cleared_error(&self.address, &e))
+            }
         };
 
         match conn {
@@ -154,8 +156,8 @@ impl ConnectionPool {
 
     /// Increments the generation of the pool. Rather than eagerly removing stale connections from
     /// the pool, they are left for the background thread to clean up.
-    pub(crate) async fn clear(&self) {
-        self.manager.clear().await
+    pub(crate) async fn clear(&self, cause: Error) {
+        self.manager.clear(cause).await
     }
 
     /// Mark the pool as "ready", allowing connections to be created and checked out.

src/cmap/test/mod.rs

@@ -13,7 +13,7 @@ use self::{
 
 use crate::{
     cmap::{Connection, ConnectionPool, ConnectionPoolOptions},
-    error::{Error, Result},
+    error::{Error, ErrorKind, Result},
     event::cmap::ConnectionPoolOptions as EventOptions,
     options::TlsOptions,
     runtime::AsyncJoinHandle,
@@ -168,7 +168,7 @@ impl Executor {
         RUNTIME.execute(async move {
             while let Some(update) = update_receiver.recv().await {
                 match update.into_message() {
-                    ServerUpdate::Error { .. } => manager.clear().await,
+                    ServerUpdate::Error { error, .. } => manager.clear(error).await,
                 }
             }
         });
@@ -285,7 +285,13 @@ impl Operation {
             }
             Operation::Clear => {
                 if let Some(pool) = state.pool.read().await.as_ref() {
-                    pool.clear().await;
+                    pool.clear(
+                        ErrorKind::Internal {
+                            message: "test error".to_string(),
+                        }
+                        .into(),
+                    )
+                    .await;
                 }
             }
             Operation::Ready => {

src/cmap/worker.rs

@@ -19,7 +19,7 @@ use super::{
     DEFAULT_MAX_POOL_SIZE,
 };
 use crate::{
-    error::Result,
+    error::{Error, ErrorKind, Result},
     event::cmap::{
         CmapEventHandler,
         ConnectionClosedEvent,
@@ -172,7 +172,7 @@ impl ConnectionPoolWorker {
         {
             PoolState::Ready
         } else {
-            PoolState::Paused
+            PoolState::New
         };
         #[cfg(test)]
         let maintenance_frequency = options
@@ -181,7 +181,7 @@ impl ConnectionPoolWorker {
             .unwrap_or(MAINTENACE_FREQUENCY);
 
         #[cfg(not(test))]
-        let (state, maintenance_frequency) = (PoolState::Paused, MAINTENACE_FREQUENCY);
+        let (state, maintenance_frequency) = (PoolState::New, MAINTENACE_FREQUENCY);
 
         let worker = ConnectionPoolWorker {
             address,
@@ -244,16 +244,27 @@ impl ConnectionPoolWorker {
                     PoolState::Ready => {
                         self.wait_queue.push_back(request);
                     }
-                    PoolState::Paused => {
+                    PoolState::Paused(ref e) => {
                         // if receiver doesn't listen to error that's ok.
-                        let _ = request.fulfill(ConnectionRequestResult::PoolCleared);
+                        let _ = request.fulfill(ConnectionRequestResult::PoolCleared(e.clone()));
+                    }
+                    PoolState::New => {
+                        let _ = request.fulfill(ConnectionRequestResult::PoolCleared(
+                            ErrorKind::Internal {
+                                message: "check out attempted from new pool".to_string(),
+                            }
+                            .into(),
+                        ));
                     }
                 },
                 PoolTask::HandleManagementRequest(PoolManagementRequest::CheckIn(connection)) => {
                     self.check_in(connection)
                 }
-                PoolTask::HandleManagementRequest(PoolManagementRequest::Clear(_message)) => {
-                    self.clear();
+                PoolTask::HandleManagementRequest(PoolManagementRequest::Clear {
+                    completion_handler: _,
+                    cause,
+                }) => {
+                    self.clear(cause);
                 }
                 PoolTask::HandleManagementRequest(PoolManagementRequest::MarkAsReady {
                     completion_handler: _handler,
@@ -428,12 +439,12 @@ impl ConnectionPoolWorker {
         }
     }
 
-    fn clear(&mut self) {
+    fn clear(&mut self, cause: Error) {
         self.generation += 1;
-        let previous_state = std::mem::replace(&mut self.state, PoolState::Paused);
+        let previous_state = std::mem::replace(&mut self.state, PoolState::Paused(cause.clone()));
         self.generation_publisher.publish(self.generation);
 
-        if !matches!(previous_state, PoolState::Paused) {
+        if matches!(previous_state, PoolState::Ready) {
             self.emit_event(|handler| {
                 let event = PoolClearedEvent {
                     address: self.address.clone(),
@@ -446,7 +457,7 @@ impl ConnectionPoolWorker {
                 // an error means the other end hung up already, which is okay because we were
                 // returning an error anyways
                 let _: std::result::Result<_, _> =
-                    request.fulfill(ConnectionRequestResult::PoolCleared);
+                    request.fulfill(ConnectionRequestResult::PoolCleared(cause.clone()));
             }
         }
     }
@@ -583,8 +594,11 @@ async fn establish_connection(
 /// once it goes out of scope and cannot be manually closed before then.
 #[derive(Debug)]
 enum PoolState {
+    /// Same as Paused, but only for a new pool, not one that has been cleared due to an error.
+    New,
+
     /// Connections may not be checked out nor created in the background to satisfy minPoolSize.
-    Paused,
+    Paused(Error),
 
     /// Pool is operational.
     Ready,

src/error.rs

@@ -57,11 +57,12 @@ impl Error {
         }
     }
 
-    pub(crate) fn pool_cleared_error(address: &ServerAddress) -> Self {
+    pub(crate) fn pool_cleared_error(address: &ServerAddress, cause: &Error) -> Self {
         ErrorKind::ConnectionPoolCleared {
             message: format!(
-                "Connection pool for {} cleared during operation execution",
-                address
+                "Connection pool for {} cleared because another operation failed with: {}",
+                address,
+                cause
             ),
         }
         .into()
@@ -292,6 +293,10 @@ impl Error {
             .map(|code| SHUTTING_DOWN_CODES.contains(&code))
             .unwrap_or(false)
     }
+
+    pub(crate) fn is_pool_cleared(&self) -> bool {
+        matches!(self.kind.as_ref(), ErrorKind::ConnectionPoolCleared { .. })
+    }
 }
 
 impl<E> From<E> for Error

src/sdam/state/mod.rs

@@ -257,7 +257,7 @@ impl Topology {
                 .await;
 
             if updated && (error.is_shutting_down() || handshake.wire_version().unwrap_or(0) < 8) {
-                server.pool.clear().await;
+                server.pool.clear(error).await;
             }
             self.request_topology_check();
 
@@ -272,7 +272,7 @@ impl Topology {
                 .mark_server_as_unknown(error.to_string(), server, state_lock)
                 .await;
             if updated {
-                server.pool.clear().await;
+                server.pool.clear(error).await;
             }
             updated
         } else {
@@ -286,7 +286,7 @@ impl Topology {
             .mark_server_as_unknown(error.to_string(), server, state_lock)
             .await;
         if updated {
-            server.pool.clear().await;
+            server.pool.clear(error).await;
         }
         updated
     }