RUST-128 Retryable Reads (#179)

Author: isabelatkinson

src/client/executor.rs

@@ -46,7 +46,7 @@ impl Client {
             .into());
         }
         let mut implicit_session = self.start_implicit_session(&op).await?;
-        self.select_server_and_execute_operation(op, implicit_session.as_mut())
+        self.execute_operation_with_retry(op, implicit_session.as_mut())
             .await
     }
 
@@ -58,7 +58,7 @@ impl Client {
         op: T,
     ) -> Result<(T::O, Option<ClientSession>)> {
         let mut implicit_session = self.start_implicit_session(&op).await?;
-        self.select_server_and_execute_operation(op, implicit_session.as_mut())
+        self.execute_operation_with_retry(op, implicit_session.as_mut())
             .await
             .map(|result| (result, implicit_session))
     }
@@ -70,18 +70,15 @@ impl Client {
         op: T,
         session: &mut ClientSession,
     ) -> Result<T::O> {
-        self.select_server_and_execute_operation(op, Some(session))
-            .await
+        self.execute_operation_with_retry(op, Some(session)).await
     }
 
     /// Selects a server and executes the given operation on it, optionally using a provided
-    /// session.
-    ///
-    /// TODO: RUST-128: replace this with `execute_operation_with_retry` when implemented.
-    async fn select_server_and_execute_operation<T: Operation>(
+    /// session. Retries the operation upon failure if retryability is supported.
+    async fn execute_operation_with_retry<T: Operation>(
         &self,
         op: T,
-        session: Option<&mut ClientSession>,
+        mut session: Option<&mut ClientSession>,
     ) -> Result<T::O> {
         let server = self.select_server(op.selection_criteria()).await?;
 
@@ -96,8 +93,50 @@ impl Client {
             }
         };
 
+        let first_error = match self
+            .execute_operation_on_connection(&op, &mut conn, &mut session)
+            .await
+        {
+            Ok(result) => {
+                return Ok(result);
+            }
+            Err(err) => {
+                self.inner
+                    .topology
+                    .handle_post_handshake_error(err.clone(), conn, server)
+                    .await;
+                // TODO RUST-90: Do not retry if session is in a transaction
+                if self.inner.options.retry_reads == Some(false)
+                    || !op.is_read_retryable()
+                    || !err.is_read_retryable()
+                {
+                    return Err(err);
+                } else {
+                    err
+                }
+            }
+        };
+
+        let server = match self.select_server(op.selection_criteria()).await {
+            Ok(server) => server,
+            Err(err) => {
+                return Err(first_error);
+            }
+        };
+
+        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);
+            }
+        };
+
         match self
-            .execute_operation_on_connection(op, &mut conn, session)
+            .execute_operation_on_connection(&op, &mut conn, &mut session)
             .await
         {
             Ok(result) => Ok(result),
@@ -106,17 +145,21 @@ impl Client {
                     .topology
                     .handle_post_handshake_error(err.clone(), conn, server)
                     .await;
-                Err(err)
+                if err.is_server_error() || err.is_read_retryable() {
+                    Err(err)
+                } else {
+                    Err(first_error)
+                }
             }
         }
     }
 
     /// Executes an operation on a given connection, optionally using a provided session.
     async fn execute_operation_on_connection<T: Operation>(
         &self,
-        op: T,
+        op: &T,
         connection: &mut Connection,
-        mut session: Option<&mut ClientSession>,
+        session: &mut Option<&mut ClientSession>,
     ) -> Result<T::O> {
         if let Some(wc) = op.write_concern() {
             wc.validate()?;

src/client/options/mod.rs

@@ -304,8 +304,11 @@ pub struct ClientOptions {
     #[builder(default)]
     pub repl_set_name: Option<String>,
 
+    /// Whether or not the client should retry a read operation if the operation fails.
+    ///
+    /// The default value is true.
     #[builder(default)]
-    pub(crate) retry_reads: Option<bool>,
+    pub retry_reads: Option<bool>,
 
     #[builder(default)]
     pub(crate) retry_writes: Option<bool>,
@@ -607,7 +610,7 @@ impl ClientOptions {
     ///     tag set
     ///   * `replicaSet`: maps to the `repl_set_name` field
     ///   * `retryWrites`: not yet implemented
-    ///   * `retryReads`: not yet implemented
+    ///   * `retryReads`: maps to the `retry_reads` field
     ///   * `serverSelectionTimeoutMS`: maps to the `server_selection_timeout` field
     ///   * `socketTimeoutMS`: maps to the `socket_timeout` field
     ///   * `ssl`: an alias of the `tls` option

src/error.rs

@@ -13,6 +13,8 @@ lazy_static! {
     static ref RECOVERING_CODES: Vec<i32> = vec![11600, 11602, 13436, 189, 91];
     static ref NOTMASTER_CODES: Vec<i32> = vec![10107, 13435];
     static ref SHUTTING_DOWN_CODES: Vec<i32> = vec![11600, 91];
+    static ref RETRYABLE_READ_CODES: Vec<i32> =
+        vec![11600, 11602, 10107, 13435, 13436, 189, 91, 7, 6, 89, 9001];
 }
 
 /// The result type for all methods that can return an error in the `mongodb` crate.
@@ -69,6 +71,36 @@ impl Error {
             _ => false,
         }
     }
+
+    /// Whether a read operation should be retried if this error occurs
+    pub(crate) fn is_read_retryable(&self) -> bool {
+        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
+            }
+            None => false,
+        }
+    }
+
+    /// Whether an error originated from the server
+    pub(crate) fn is_server_error(&self) -> bool {
+        match self.kind.as_ref() {
+            ErrorKind::AuthenticationError { .. }
+            | ErrorKind::BulkWriteError(_)
+            | ErrorKind::CommandError(_)
+            | ErrorKind::WriteError(_) => true,
+            _ => false,
+        }
+    }
 }
 
 impl<E> From<E> for Error

src/operation/aggregate/mod.rs

@@ -92,6 +92,10 @@ impl Operation for Aggregate {
             .as_ref()
             .and_then(|opts| opts.write_concern.as_ref())
     }
+
+    fn is_read_retryable(&self) -> bool {
+        !self.is_out_or_merge()
+    }
 }
 
 impl Aggregate {

src/operation/count/mod.rs

@@ -62,6 +62,10 @@ impl Operation for Count {
         }
         None
     }
+
+    fn is_read_retryable(&self) -> bool {
+        true
+    }
 }
 
 #[derive(Debug, Deserialize)]

src/operation/distinct/mod.rs

@@ -80,6 +80,10 @@ impl Operation for Distinct {
         }
         None
     }
+
+    fn is_read_retryable(&self) -> bool {
+        true
+    }
 }
 
 #[derive(Debug, Deserialize)]

src/operation/find/mod.rs

@@ -113,4 +113,8 @@ impl Operation for Find {
             .as_ref()
             .and_then(|opts| opts.selection_criteria.as_ref())
     }
+
+    fn is_read_retryable(&self) -> bool {
+        true
+    }
 }

src/operation/list_collections/mod.rs

@@ -79,4 +79,8 @@ impl Operation for ListCollections {
     fn selection_criteria(&self) -> Option<&SelectionCriteria> {
         Some(SelectionCriteria::ReadPreference(ReadPreference::Primary)).as_ref()
     }
+
+    fn is_read_retryable(&self) -> bool {
+        true
+    }
 }

src/operation/list_databases/mod.rs

@@ -72,6 +72,10 @@ impl Operation for ListDatabases {
     fn selection_criteria(&self) -> Option<&SelectionCriteria> {
         Some(SelectionCriteria::ReadPreference(ReadPreference::Primary)).as_ref()
     }
+
+    fn is_read_retryable(&self) -> bool {
+        true
+    }
 }
 
 #[derive(Debug, Deserialize)]

src/operation/mod.rs

@@ -92,6 +92,11 @@ pub(crate) trait Operation {
     fn supports_sessions(&self) -> bool {
         true
     }
+
+    /// Whether or not the operation supports retryable reads
+    fn is_read_retryable(&self) -> bool {
+        false
+    }
 }
 
 /// Appends a serializable struct to the input document.