RUST-51 Retryable Writes (#219)

Author: isabelatkinson

src/client/executor.rs

@@ -8,9 +8,9 @@ use time::PreciseTime;
 use crate::{
     bson::Document,
     cmap::Connection,
-    error::{ErrorKind, Result},
+    error::{Error, ErrorKind, Result},
     event::command::{CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent},
-    operation::Operation,
+    operation::{Operation, Retryability},
     options::SelectionCriteria,
     sdam::{Server, SessionSupportStatus},
 };
@@ -93,8 +93,17 @@ impl Client {
             }
         };
 
+        let retryability = self.get_retryability(&conn, &op).await?;
+
+        let txn_number = match session {
+            Some(ref mut session) if retryability == Retryability::Write => {
+                Some(session.get_and_increment_txn_number())
+            }
+            _ => None,
+        };
+
         let first_error = match self
-            .execute_operation_on_connection(&op, &mut conn, &mut session)
+            .execute_operation_on_connection(&op, &mut conn, &mut session, txn_number)
             .await
         {
             Ok(result) => {
@@ -103,16 +112,36 @@ impl Client {
             Err(err) => {
                 self.inner
                     .topology
-                    .handle_post_handshake_error(err.clone(), conn, server)
+                    .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()
+
+                // 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.as_ref() {
+                    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."
+                            .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,
+                };
+
+                // TODO RUST-90: Do not retry read if session is in a transaction
+                if retryability == Retryability::Read && err.is_read_retryable()
+                    || retryability == Retryability::Write && err.is_write_retryable()
                 {
-                    return Err(err);
-                } else {
                     err
+                } else {
+                    return Err(err);
                 }
             }
         };
@@ -135,17 +164,28 @@ impl Client {
             }
         };
 
+        let retryability = self.get_retryability(&conn, &op).await?;
+        if retryability == Retryability::None {
+            return Err(first_error);
+        }
+
         match self
-            .execute_operation_on_connection(&op, &mut conn, &mut session)
+            .execute_operation_on_connection(&op, &mut conn, &mut session, txn_number)
             .await
         {
             Ok(result) => Ok(result),
             Err(err) => {
                 self.inner
                     .topology
-                    .handle_post_handshake_error(err.clone(), conn, server)
+                    .handle_post_handshake_error(err.clone(), &conn, server)
                     .await;
-                if err.is_server_error() || err.is_read_retryable() {
+
+                let err = match retryability {
+                    Retryability::Write => get_error_with_retryable_write_label(&conn, err).await?,
+                    _ => err,
+                };
+
+                if err.is_server_error() || err.is_read_retryable() || err.is_write_retryable() {
                     Err(err)
                 } else {
                     Err(first_error)
@@ -160,6 +200,7 @@ impl Client {
         op: &T,
         connection: &mut Connection,
         session: &mut Option<&mut ClientSession>,
+        txn_number: Option<u64>,
     ) -> Result<T::O> {
         if let Some(wc) = op.write_concern() {
             wc.validate()?;
@@ -174,6 +215,9 @@ impl Client {
         match session {
             Some(ref mut session) if op.supports_sessions() && op.is_acknowledged() => {
                 cmd.set_session(session);
+                if let Some(txn_number) = txn_number {
+                    cmd.set_txn_number(txn_number);
+                }
                 session.update_last_use();
             }
             Some(ref session) if !op.supports_sessions() && !session.is_implicit() => {
@@ -318,4 +362,37 @@ impl Client {
             _ => Ok(initial_status),
         }
     }
+
+    /// Returns the retryability level for the execution of this operation.
+    async fn get_retryability<T: Operation>(
+        &self,
+        conn: &Connection,
+        op: &T,
+    ) -> Result<Retryability> {
+        match op.retryability() {
+            Retryability::Read if self.inner.options.retry_reads != Some(false) => {
+                Ok(Retryability::Read)
+            }
+            Retryability::Write
+                if self.inner.options.retry_writes != Some(false)
+                    && conn.stream_description()?.supports_retryable_writes() =>
+            {
+                Ok(Retryability::Write)
+            }
+            _ => Ok(Retryability::None),
+        }
+    }
+}
+
+/// Returns an Error with a "RetryableWriteError" label added if necessary. On a pre-4.4
+/// connection, a label should be added to any write-retryable error. On a 4.4+ connection, a
+/// label should only be added to network errors. Regardless of server version, a label should
+/// only be added if the `retry_writes` client option is not set to `false`.
+async fn get_error_with_retryable_write_label(conn: &Connection, err: Error) -> Result<Error> {
+    if let Some(max_wire_version) = conn.stream_description()?.max_wire_version {
+        if err.should_add_retryable_write_label(max_wire_version) {
+            return Ok(err.with_label("RetryableWriteError"));
+        }
+    }
+    Ok(err)
 }

src/client/options/mod.rs

@@ -316,8 +316,11 @@ pub struct ClientOptions {
     #[builder(default)]
     pub retry_reads: Option<bool>,
 
+    /// Whether or not the client should retry a write operation if the operation fails.
+    ///
+    /// The default value is true.
     #[builder(default)]
-    pub(crate) retry_writes: Option<bool>,
+    pub retry_writes: Option<bool>,
 
     /// The default selection criteria for operations performed on the Client. See the
     /// SelectionCriteria type documentation for more details.
@@ -370,6 +373,10 @@ pub struct ClientOptions {
     #[builder(default)]
     #[serde(skip)]
     pub(crate) resolver_config: Option<ResolverConfig>,
+
+    /// Used by tests to override MIN_HEARTBEAT_FREQUENCY.
+    #[builder(default)]
+    pub(crate) heartbeat_freq_test: Option<Duration>,
 }
 
 fn default_hosts() -> Vec<StreamAddress> {
@@ -585,6 +592,7 @@ impl From<ClientOptionsParser> for ClientOptions {
             original_srv_hostname: None,
             original_uri: Some(parser.original_uri),
             resolver_config: None,
+            heartbeat_freq_test: None,
         }
     }
 }
@@ -759,6 +767,43 @@ impl ClientOptions {
         }
         Ok(())
     }
+
+    /// Applies the options in other to these options if a value is not already present
+    #[cfg(test)]
+    pub(crate) fn merge(&mut self, other: ClientOptions) {
+        merge_options!(
+            other,
+            self,
+            [
+                app_name,
+                compressors,
+                cmap_event_handler,
+                command_event_handler,
+                connect_timeout,
+                credential,
+                direct_connection,
+                driver_info,
+                heartbeat_freq,
+                local_threshold,
+                max_idle_time,
+                max_pool_size,
+                min_pool_size,
+                read_concern,
+                repl_set_name,
+                retry_reads,
+                retry_writes,
+                selection_criteria,
+                server_selection_timeout,
+                socket_timeout,
+                tls,
+                wait_queue_timeout,
+                write_concern,
+                zlib_compression,
+                original_srv_hostname,
+                original_uri
+            ]
+        );
+    }
 }
 
 /// Splits a string into a section before a given index and a section exclusively after the index.

src/client/session/mod.rs

@@ -83,6 +83,12 @@ impl ClientSession {
     pub(crate) fn update_last_use(&mut self) {
         self.server_session.last_use = Instant::now();
     }
+
+    /// Increments the txn_number and returns the new value.
+    pub(crate) fn get_and_increment_txn_number(&mut self) -> u64 {
+        self.server_session.txn_number += 1;
+        self.server_session.txn_number
+    }
 }
 
 impl Drop for ClientSession {
@@ -92,6 +98,7 @@ impl Drop for ClientSession {
             id: self.server_session.id.clone(),
             last_use: self.server_session.last_use,
             dirty: self.server_session.dirty,
+            txn_number: self.server_session.txn_number,
         };
 
         RUNTIME.execute(async move {
@@ -112,6 +119,9 @@ pub(crate) struct ServerSession {
 
     /// Whether a network error was encountered while using this session.
     dirty: bool,
+
+    /// A monotonically increasing transaction number for this session.
+    txn_number: u64,
 }
 
 impl ServerSession {
@@ -126,6 +136,7 @@ impl ServerSession {
             id: doc! { "id": binary },
             last_use: Instant::now(),
             dirty: false,
+            txn_number: 0,
         }
     }
 

src/cmap/conn/command.rs

@@ -56,6 +56,10 @@ impl Command {
             self.body.insert("$clusterTime", doc);
         }
     }
+
+    pub(crate) fn set_txn_number(&mut self, txn_number: u64) {
+        self.body.insert("txnNumber", txn_number);
+    }
 }
 
 #[derive(Debug, Clone)]

src/cmap/conn/stream_description.rs

@@ -38,6 +38,13 @@ impl StreamDescription {
         }
     }
 
+    /// Whether this StreamDescription supports retryable writes.
+    pub(crate) fn supports_retryable_writes(&self) -> bool {
+        self.initial_server_type != ServerType::Standalone
+            && self.logical_session_timeout.is_some()
+            && self.max_wire_version.map_or(false, |version| version >= 6)
+    }
+
     /// Gets a description of a stream for a 4.2 connection.
     #[cfg(test)]
     pub(crate) fn new_testing() -> Self {

src/coll/mod.rs

@@ -287,6 +287,11 @@ impl Collection {
     }
 
     /// Deletes up to one document found matching `query`.
+    ///
+    /// This operation will retry once upon failure if the connection and encountered error support
+    /// retryability. See the documentation
+    /// [here](https://docs.mongodb.com/manual/core/retryable-writes/) for more information on
+    /// retryable writes.
     pub async fn delete_one(
         &self,
         query: Document,
@@ -349,6 +354,11 @@ impl Collection {
     }
 
     /// Atomically finds up to one document in the collection matching `filter` and deletes it.
+    ///
+    /// This operation will retry once upon failure if the connection and encountered error support
+    /// retryability. See the documentation
+    /// [here](https://docs.mongodb.com/manual/core/retryable-writes/) for more information on
+    /// retryable writes.
     pub async fn find_one_and_delete(
         &self,
         filter: Document,
@@ -363,6 +373,11 @@ impl Collection {
 
     /// Atomically finds up to one document in the collection matching `filter` and replaces it with
     /// `replacement`.
+    ///
+    /// This operation will retry once upon failure if the connection and encountered error support
+    /// retryability. See the documentation
+    /// [here](https://docs.mongodb.com/manual/core/retryable-writes/) for more information on
+    /// retryable writes.
     pub async fn find_one_and_replace(
         &self,
         filter: Document,
@@ -380,6 +395,11 @@ impl Collection {
     /// Both `Document` and `Vec<Document>` implement `Into<UpdateModifications>`, so either can be
     /// passed in place of constructing the enum case. Note: pipeline updates are only supported
     /// in MongoDB 4.2+.
+    ///
+    /// This operation will retry once upon failure if the connection and encountered error support
+    /// retryability. See the documentation
+    /// [here](https://docs.mongodb.com/manual/core/retryable-writes/) for more information on
+    /// retryable writes.
     pub async fn find_one_and_update(
         &self,
         filter: Document,
@@ -395,6 +415,11 @@ impl Collection {
     }
 
     /// Inserts the documents in `docs` into the collection.
+    ///
+    /// This operation will retry once upon failure if the connection and encountered error support
+    /// retryability. See the documentation
+    /// [here](https://docs.mongodb.com/manual/core/retryable-writes/) for more information on
+    /// retryable writes.
     pub async fn insert_many(
         &self,
         docs: impl IntoIterator<Item = Document>,
@@ -477,6 +502,11 @@ impl Collection {
     }
 
     /// Inserts `doc` into the collection.
+    ///
+    /// This operation will retry once upon failure if the connection and encountered error support
+    /// retryability. See the documentation
+    /// [here](https://docs.mongodb.com/manual/core/retryable-writes/) for more information on
+    /// retryable writes.
     pub async fn insert_one(
         &self,
         doc: Document,
@@ -498,6 +528,11 @@ impl Collection {
     }
 
     /// Replaces up to one document matching `query` in the collection with `replacement`.
+    ///
+    /// This operation will retry once upon failure if the connection and encountered error support
+    /// retryability. See the documentation
+    /// [here](https://docs.mongodb.com/manual/core/retryable-writes/) for more information on
+    /// retryable writes.
     pub async fn replace_one(
         &self,
         query: Document,
@@ -550,6 +585,11 @@ impl Collection {
     /// passed in place of constructing the enum case. Note: pipeline updates are only supported
     /// in MongoDB 4.2+. See the official MongoDB
     /// [documentation](https://docs.mongodb.com/manual/reference/command/update/#behavior) for more information on specifying updates.
+    ///
+    /// This operation will retry once upon failure if the connection and encountered error support
+    /// retryability. See the documentation
+    /// [here](https://docs.mongodb.com/manual/core/retryable-writes/) for more information on
+    /// retryable writes.
     pub async fn update_one(
         &self,
         query: Document,