RUST-655 / RUST-691 Simplification of SDAM error handling (#312)

Author: patrickfreed

src/client/executor.rs

@@ -12,7 +12,7 @@ use crate::{
     event::command::{CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent},
     operation::{Operation, Retryability},
     options::SelectionCriteria,
-    sdam::{SelectedServer, SessionSupportStatus},
+    sdam::{HandshakePhase, SelectedServer, SessionSupportStatus},
 };
 
 lazy_static! {
@@ -100,30 +100,36 @@ impl Client {
                 return Ok(result);
             }
             Err(err) => {
-                self.inner
-                    .topology
-                    .handle_post_handshake_error(&err, &conn, server)
-                    .await;
+                // For a pre-4.4 connection, an error label should be added to any write-retryable
+                // error as long as the retry_writes client option is not set to false. For a 4.4+
+                // connection, a label should be added only to network errors.
+                let mut err = match retryability {
+                    Retryability::Write => get_error_with_retryable_write_label(&conn, err).await?,
+                    _ => err,
+                };
 
                 // Retryable writes are only supported by storage engines with document-level
                 // locking, so users need to disable retryable writes if using mmapv1.
-                if let ErrorKind::CommandError(ref err) = err.kind {
-                    if err.code == 20 && err.message.starts_with("Transaction numbers") {
-                        let mut err = err.clone();
-                        err.message = "This MongoDB deployment does not support retryable writes. \
-                                       Please add retryWrites=false to your connection string."
+                if let ErrorKind::CommandError(ref mut command_error) = err.kind {
+                    if command_error.code == 20
+                        && command_error.message.starts_with("Transaction numbers")
+                    {
+                        command_error.message = "This MongoDB deployment does not support \
+                                                 retryable writes. Please add retryWrites=false \
+                                                 to your connection string."
                             .to_string();
-                        return Err(ErrorKind::CommandError(err).into());
                     }
                 }
 
-                // For a pre-4.4 connection, an error label should be added to any write-retryable
-                // error as long as the retry_writes client option is not set to false. For a 4.4+
-                // connection, a label should be added only to network errors.
-                let err = match retryability {
-                    Retryability::Write => get_error_with_retryable_write_label(&conn, err).await?,
-                    _ => err,
-                };
+                self.inner
+                    .topology
+                    .handle_application_error(
+                        err.clone(),
+                        HandshakePhase::after_completion(conn),
+                        &server,
+                    )
+                    .await;
+                drop(server);
 
                 // TODO RUST-90: Do not retry read if session is in a transaction
                 if retryability == Retryability::Read && err.is_read_retryable()
@@ -159,15 +165,19 @@ impl Client {
         {
             Ok(result) => Ok(result),
             Err(err) => {
-                self.inner
-                    .topology
-                    .handle_post_handshake_error(&err, &conn, server)
-                    .await;
-
                 let err = match retryability {
                     Retryability::Write => get_error_with_retryable_write_label(&conn, err).await?,
                     _ => err,
                 };
+                self.inner
+                    .topology
+                    .handle_application_error(
+                        err.clone(),
+                        HandshakePhase::after_completion(conn),
+                        &server,
+                    )
+                    .await;
+                drop(server);
 
                 if err.is_server_error() || err.is_read_retryable() || err.is_write_retryable() {
                     Err(err)

src/cmap/conn/mod.rs

@@ -46,7 +46,7 @@ pub struct ConnectionInfo {
 pub(crate) struct Connection {
     pub(super) id: u32,
     pub(super) address: StreamAddress,
-    pub(super) generation: u32,
+    pub(crate) generation: u32,
 
     /// The cached StreamDescription from the connection's handshake.
     pub(super) stream_description: Option<StreamDescription>,

src/cmap/manager.rs

@@ -20,8 +20,15 @@ pub(super) struct PoolManager {
 
 impl PoolManager {
     /// Lazily clear the pool.
-    pub(super) fn clear(&self) {
-        let _ = self.sender.send(PoolManagementRequest::Clear);
+    pub(super) async fn clear(&self) {
+        let (message, acknowledgment_receiver) = AcknowledgedMessage::package(());
+        if self
+            .sender
+            .send(PoolManagementRequest::Clear(message))
+            .is_ok()
+        {
+            acknowledgment_receiver.wait_for_acknowledgment().await;
+        }
     }
 
     /// Mark the pool as "ready" as per the CMAP specification.
@@ -77,7 +84,7 @@ impl ManagementRequestReceiver {
 #[derive(Debug)]
 pub(super) enum PoolManagementRequest {
     /// Clear the pool, transitioning it to Paused.
-    Clear,
+    Clear(AcknowledgedMessage<()>),
 
     /// Mark the pool as Ready, allowing connections to be created and checked out.
     MarkAsReady {

src/cmap/mod.rs

@@ -176,20 +176,16 @@ 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) fn clear(&self) {
-        self.manager.clear();
+    pub(crate) async fn clear(&self) {
+        self.manager.clear().await
     }
 
     /// 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()
+    pub(crate) fn generation(&self) -> u32 {
+        self.generation_subscriber.generation()
     }
 }

src/cmap/status.rs

@@ -49,18 +49,4 @@ impl PoolGenerationSubscriber {
     pub(crate) fn generation(&self) -> u32 {
         self.receiver.borrow().generation
     }
-
-    #[cfg(test)]
-    pub(crate) async fn wait_for_generation_change(
-        &mut self,
-        timeout: std::time::Duration,
-    ) -> Option<u32> {
-        crate::RUNTIME
-            .timeout(timeout, self.receiver.changed())
-            .await
-            .ok()?
-            .ok()?;
-
-        Some(self.receiver.borrow().generation)
-    }
 }

src/cmap/test/mod.rs

@@ -164,7 +164,7 @@ impl Executor {
         RUNTIME.execute(async move {
             while let Some(update) = update_receiver.recv().await {
                 match update.into_message() {
-                    ServerUpdate::Error { .. } => manager.clear(),
+                    ServerUpdate::Error { .. } => manager.clear().await,
                 }
             }
         });
@@ -284,13 +284,7 @@ impl Operation {
             }
             Operation::Clear => {
                 if let Some(pool) = state.pool.read().await.as_ref() {
-                    let mut subscriber = pool.subscribe_to_generation_updates();
-                    pool.clear();
-
-                    subscriber
-                        .wait_for_generation_change(EVENT_TIMEOUT)
-                        .await
-                        .expect("generation did not change after clearing pool");
+                    pool.clear().await;
                 }
             }
             Operation::Ready => {

src/cmap/worker.rs

@@ -252,7 +252,9 @@ impl ConnectionPoolWorker {
                 PoolTask::HandleManagementRequest(PoolManagementRequest::CheckIn(connection)) => {
                     self.check_in(connection)
                 }
-                PoolTask::HandleManagementRequest(PoolManagementRequest::Clear) => self.clear(),
+                PoolTask::HandleManagementRequest(PoolManagementRequest::Clear(_message)) => {
+                    self.clear();
+                }
                 PoolTask::HandleManagementRequest(PoolManagementRequest::MarkAsReady {
                     completion_handler: _handler,
                 }) => {
@@ -562,7 +564,7 @@ async fn establish_connection(
                 };
                 handler.handle_connection_closed_event(event);
             }
-            server_updater.handle_error(&e, generation).await;
+            server_updater.handle_error(e.clone(), generation).await;
             manager.handle_connection_failed();
         }
         Ok(ref mut connection) => {

src/error.rs

@@ -8,7 +8,7 @@ use thiserror::Error;
 use crate::{bson::Document, options::StreamAddress};
 
 const RECOVERING_CODES: [i32; 5] = [11600, 11602, 13436, 189, 91];
-const NOTMASTER_CODES: [i32; 2] = [10107, 13435];
+const NOTMASTER_CODES: [i32; 3] = [10107, 13435, 10058];
 const SHUTTING_DOWN_CODES: [i32; 2] = [11600, 91];
 const RETRYABLE_READ_CODES: [i32; 11] =
     [11600, 11602, 10107, 13435, 13436, 189, 91, 7, 6, 89, 9001];
@@ -61,6 +61,22 @@ impl Error {
         Error::authentication_error(mechanism_name, "invalid server response")
     }
 
+    pub(crate) fn is_state_change_error(&self) -> bool {
+        self.is_recovering() || self.is_not_master()
+    }
+
+    pub(crate) fn is_auth_error(&self) -> bool {
+        matches!(self.kind, ErrorKind::AuthenticationError { .. })
+    }
+
+    pub(crate) fn is_command_error(&self) -> bool {
+        matches!(self.kind, ErrorKind::CommandError(_))
+    }
+
+    pub(crate) fn is_network_timeout(&self) -> bool {
+        matches!(self.kind, ErrorKind::Io(ref io_err) if io_err.kind() == std::io::ErrorKind::TimedOut)
+    }
+
     /// Whether this error is an "ns not found" error or not.
     pub(crate) fn is_ns_not_found(&self) -> bool {
         matches!(self.kind, ErrorKind::CommandError(ref err) if err.code == 26)
@@ -71,15 +87,9 @@ impl Error {
         if self.is_network_error() {
             return true;
         }
-        match &self.kind.code_and_message() {
-            Some((code, message)) => {
-                if RETRYABLE_READ_CODES.contains(&code) {
-                    return true;
-                }
-                if is_not_master(*code, message) || is_recovering(*code, message) {
-                    return true;
-                }
-                false
+        match &self.kind.code() {
+            Some(code) => {
+                RETRYABLE_READ_CODES.contains(&code)
             }
             None => false,
         }
@@ -100,8 +110,8 @@ impl Error {
         if self.is_network_error() {
             return true;
         }
-        match &self.kind.code_and_message() {
-            Some((code, _)) => RETRYABLE_WRITE_CODES.contains(&code),
+        match &self.kind.code() {
+            Some(code) => RETRYABLE_WRITE_CODES.contains(&code),
             None => false,
         }
     }
@@ -363,19 +373,48 @@ impl ErrorKind {
         )
     }
 
-    /// Gets the code/message tuple from this error, if applicable. In the case of write errors, the
-    /// code and message are taken from the write concern error, if there is one.
-    pub(crate) fn code_and_message(&self) -> Option<(i32, &str)> {
+    /// Gets the code from this error for performing SDAM updates, if applicable.
+    /// Any codes contained in WriteErrors are ignored.
+    pub(crate) fn code(&self) -> Option<i32> {
         match self {
-            ErrorKind::CommandError(ref cmd_err) => Some((cmd_err.code, cmd_err.message.as_str())),
-            ErrorKind::WriteError(WriteFailure::WriteConcernError(ref wc_err)) => {
-                Some((wc_err.code, wc_err.message.as_str()))
+            ErrorKind::CommandError(command_error) => {
+                Some(command_error.code)
+            },
+            // According to SDAM spec, write concern error codes MUST also be checked, and writeError codes
+            // MUST NOT be checked.
+            ErrorKind::BulkWriteError(BulkWriteFailure { write_concern_error: Some(wc_error), .. }) => {
+                Some(wc_error.code)
             }
-            ErrorKind::BulkWriteError(ref bwe) => bwe
-                .write_concern_error
-                .as_ref()
-                .map(|wc_err| (wc_err.code, wc_err.message.as_str())),
-            _ => None,
+            ErrorKind::WriteError(WriteFailure::WriteConcernError(wc_error)) => Some(wc_error.code),
+            _ => None
+        }
+    }
+
+    /// Gets the server's message for this error, if applicable, for use in testing.
+    /// If this error is a BulkWriteError, the messages are concatenated.
+    #[cfg(test)]
+    pub(crate) fn server_message(&self) -> Option<String> {
+        match self {
+            ErrorKind::CommandError(command_error) => {
+                Some(command_error.message.clone())
+            },
+            // since this is used primarily for errorMessageContains assertions in the unified runner, we just
+            // concatenate all the relevant server messages into one for bulk errors.
+            ErrorKind::BulkWriteError(BulkWriteFailure { write_concern_error, write_errors }) => {
+                let mut msg = "".to_string();
+                if let Some(wc_error) = write_concern_error {
+                    msg.push_str(wc_error.message.as_str());
+                }
+                if let Some(write_errors) = write_errors {
+                    for we in write_errors {
+                        msg.push_str(we.message.as_str());
+                    }
+                }
+                Some(msg)
+            }
+            ErrorKind::WriteError(WriteFailure::WriteConcernError(wc_error)) => Some(wc_error.message.clone()),
+            ErrorKind::WriteError(WriteFailure::WriteError(write_error)) => Some(write_error.message.clone()),
+            _ => None
         }
     }
 
@@ -398,42 +437,24 @@ impl ErrorKind {
 
     /// If this error corresponds to a "not master" error as per the SDAM spec.
     pub(crate) fn is_not_master(&self) -> bool {
-        self.code_and_message()
-            .map(|(code, msg)| is_not_master(code, msg))
-            .unwrap_or(false)
+        self.code().map(|code| NOTMASTER_CODES.contains(&code)).unwrap_or(false)
     }
 
     /// If this error corresponds to a "node is recovering" error as per the SDAM spec.
     pub(crate) fn is_recovering(&self) -> bool {
-        self.code_and_message()
-            .map(|(code, msg)| is_recovering(code, msg))
+        self.code()
+            .map(|code| RECOVERING_CODES.contains(&code))
             .unwrap_or(false)
     }
 
     /// If this error corresponds to a "node is shutting down" error as per the SDAM spec.
     pub(crate) fn is_shutting_down(&self) -> bool {
-        self.code_and_message()
-            .map(|(code, _)| SHUTTING_DOWN_CODES.contains(&code))
+        self.code()
+            .map(|code| SHUTTING_DOWN_CODES.contains(&code))
             .unwrap_or(false)
     }
 }
 
-fn is_not_master(code: i32, message: &str) -> bool {
-    if NOTMASTER_CODES.contains(&code) {
-        return true;
-    } else if is_recovering(code, message) {
-        return false;
-    }
-    message.contains("not master")
-}
-
-fn is_recovering(code: i32, message: &str) -> bool {
-    if RECOVERING_CODES.contains(&code) {
-        return true;
-    }
-    message.contains("not master or secondary") || message.contains("node is recovering")
-}
-
 /// An error that occurred due to a database command failing.
 #[derive(Clone, Debug, Deserialize)]
 #[non_exhaustive]
@@ -508,6 +529,7 @@ pub struct WriteError {
 #[non_exhaustive]
 pub struct BulkWriteError {
     /// Index into the list of operations that this error corresponds to.
+    #[serde(default)]
     pub index: usize,
 
     /// Identifies the type of write concern error.
@@ -526,7 +548,8 @@ pub struct BulkWriteError {
 }
 
 /// The set of errors that occurred during a write operation.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
 #[non_exhaustive]
 pub struct BulkWriteFailure {
     /// The error(s) that occurred on account of a non write concern failure.