RUST-242 Properly handle network errors during operations (#119)

Author: saghm

src/client/executor.rs

@@ -11,7 +11,7 @@ use crate::{
     error::Result,
     event::command::{CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent},
     operation::Operation,
-    sdam::{update_topology, ServerDescription},
+    sdam::{handle_post_handshake_error, handle_pre_handshake_error},
 };
 
 lazy_static! {
@@ -58,43 +58,26 @@ impl Client {
             Some(conn) => self.execute_operation_on_connection(op, conn),
             None => {
                 let server = self.select_server(op.selection_criteria())?;
-                let mut conn = server.checkout_connection()?;
 
-                let result = self.execute_operation_on_connection(op, &mut conn);
-
-                // If we encounter certain errors, we must update the topology as per the
-                // SDAM spec.
-                if let Err(ref e) = result {
-                    let update = || {
-                        let description =
-                            ServerDescription::new(conn.address().clone(), Some(Err(e.clone())));
-                        update_topology(self.topology(), description);
-                    };
+                let mut conn = match server.checkout_connection() {
+                    Ok(conn) => conn,
+                    Err(err) => {
+                        handle_pre_handshake_error(
+                            err.clone(),
+                            server.address.clone(),
+                            self.topology(),
+                        );
+                        return Err(err);
+                    }
+                };
 
-                    if e.is_non_timeout_network_error() {
-                        update();
-                    } else if e.is_recovering() || e.is_not_master() {
-                        update();
-
-                        // For "node is recovering" or "not master" errors, we must request a
-                        // topology check.
-                        server.request_topology_check();
-
-                        let wire_version = conn
-                            .stream_description()
-                            .map(|sd| sd.max_wire_version)
-                            .ok()
-                            .and_then(std::convert::identity)
-                            .unwrap_or(0);
-
-                        // in 4.2+, we only clear connection pool if we've received a
-                        // "node is shutting down" error. Otherwise, we always clear the pool.
-                        if wire_version < 8 || e.is_shutting_down() {
-                            server.clear_connection_pool();
-                        }
+                match self.execute_operation_on_connection(op, &mut conn) {
+                    Ok(result) => Ok(result),
+                    Err(err) => {
+                        handle_post_handshake_error(err.clone(), conn, server, self.topology());
+                        Err(err)
                     }
                 }
-                result
             }
         }
     }

src/error.rs

@@ -189,6 +189,13 @@ impl ErrorKind {
         }
     }
 
+    pub(crate) fn is_network_error(&self) -> bool {
+        match self {
+            ErrorKind::Io(..) => true,
+            _ => false,
+        }
+    }
+
     /// 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.
     fn code_and_message(&self) -> Option<(i32, &str)> {

src/sdam/mod.rs

@@ -5,8 +5,17 @@ mod state;
 
 pub use self::public::{ServerInfo, ServerType};
 
+#[cfg(test)]
+pub(crate) use self::description::server::ServerDescription;
 pub(crate) use self::{
-    description::{server::ServerDescription, topology::TopologyDescription},
+    description::topology::TopologyDescription,
     monitor::MIN_HEARTBEAT_FREQUENCY,
-    state::{server::Server, update_topology, Topology, TopologyUpdateCondvar},
+    state::{
+        handle_post_handshake_error,
+        handle_pre_handshake_error,
+        server::Server,
+        update_topology,
+        Topology,
+        TopologyUpdateCondvar,
+    },
 };

src/sdam/state/mod.rs

@@ -12,7 +12,7 @@ use self::server::Server;
 use super::TopologyDescription;
 use crate::{
     cmap::{Command, Connection},
-    error::Result,
+    error::{Error, Result},
     options::{ClientOptions, StreamAddress},
     sdam::{
         description::server::{ServerDescription, ServerType},
@@ -175,6 +175,54 @@ impl Topology {
     }
 }
 
+pub(crate) fn handle_pre_handshake_error(
+    error: Error,
+    address: StreamAddress,
+    topology: Arc<RwLock<Topology>>,
+) {
+    if error.is_network_error() {
+        mark_server_as_unknown(error, address, topology);
+    }
+}
+
+pub(crate) fn handle_post_handshake_error(
+    error: Error,
+    conn: Connection,
+    server: Arc<Server>,
+    topology: Arc<RwLock<Topology>>,
+) {
+    // If we encounter certain errors, we must update the topology as per the
+    // SDAM spec.
+    if error.is_non_timeout_network_error() {
+        mark_server_as_unknown(error, server.address.clone(), topology);
+        server.clear_connection_pool();
+    } else if error.is_recovering() || error.is_not_master() {
+        mark_server_as_unknown(error.clone(), server.address.clone(), topology);
+
+        // For "node is recovering" or "not master" errors, we must request a
+        // topology check.
+        server.request_topology_check();
+
+        let wire_version = conn
+            .stream_description()
+            .map(|sd| sd.max_wire_version)
+            .ok()
+            .and_then(std::convert::identity)
+            .unwrap_or(0);
+
+        // in 4.2+, we only clear connection pool if we've received a
+        // "node is shutting down" error. Otherwise, we always clear the pool.
+        if wire_version < 8 || error.is_shutting_down() {
+            server.clear_connection_pool();
+        }
+    }
+}
+
+fn mark_server_as_unknown(error: Error, address: StreamAddress, topology: Arc<RwLock<Topology>>) {
+    let description = ServerDescription::new(address, Some(Err(error)));
+    update_topology(topology, description);
+}
+
 /// Updates the provided topology in a minimally contentious way by cloning first.
 pub(crate) fn update_topology(
     topology: Arc<RwLock<Topology>>,