RUST-107 Convenient transactions (#849)

Author: abr-egn

.evergreen/MSRV-Cargo.lock

@@ -110,6 +110,8 @@ pub struct ClientSession {
     pub(crate) transaction: Transaction,
     pub(crate) snapshot_time: Option<Timestamp>,
     pub(crate) operation_time: Option<Timestamp>,
+    #[cfg(test)]
+    pub(crate) convenient_transaction_timeout: Option<Duration>,
 }
 
 #[derive(Debug)]
@@ -216,6 +218,8 @@ impl ClientSession {
             transaction: Default::default(),
             snapshot_time: None,
             operation_time: None,
+            #[cfg(test)]
+            convenient_transaction_timeout: None,
         }
     }
 
@@ -561,13 +565,117 @@ impl ClientSession {
         }
     }
 
+    /// Starts a transaction, runs the given callback, and commits or aborts the transaction.
+    /// Transient transaction errors will cause the callback or the commit to be retried;
+    /// other errors will cause the transaction to be aborted and the error returned to the
+    /// caller.  If the callback needs to provide its own error information, the
+    /// [`Error::custom`](crate::error::Error::custom) method can accept an arbitrary payload that
+    /// can be retrieved via [`Error::get_custom`](crate::error::Error::get_custom).
+    ///
+    /// Because the callback can be repeatedly executed and because it returns a future, the rust
+    /// closure borrowing rules for captured values can be overly restrictive.  As a
+    /// convenience, `with_transaction` accepts a context argument that will be passed to the
+    /// callback along with the session:
+    ///
+    /// ```no_run
+    /// # use mongodb::{bson::{doc, Document}, error::Result, Client};
+    /// # use futures::FutureExt;
+    /// # async fn wrapper() -> Result<()> {
+    /// # let client = Client::with_uri_str("mongodb://example.com").await?;
+    /// # let mut session = client.start_session(None).await?;
+    /// let coll = client.database("mydb").collection::<Document>("mycoll");
+    /// let my_data = "my data".to_string();
+    /// // This works:
+    /// session.with_transaction(
+    ///     (&coll, &my_data),
+    ///     |session, (coll, my_data)| async move {
+    ///         coll.insert_one_with_session(doc! { "data": *my_data }, None, session).await
+    ///     }.boxed(),
+    ///     None,
+    /// ).await?;
+    /// /* This will not compile with a "variable moved due to use in generator" error:
+    /// session.with_transaction(
+    ///     (),
+    ///     |session, _| async move {
+    ///         coll.insert_one_with_session(doc! { "data": my_data }, None, session).await
+    ///     }.boxed(),
+    ///     None,
+    /// ).await?;
+    /// */
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn with_transaction<R, C, F>(
+        &mut self,
+        mut context: C,
+        mut callback: F,
+        options: impl Into<Option<TransactionOptions>>,
+    ) -> Result<R>
+    where
+        F: for<'a> FnMut(&'a mut ClientSession, &'a mut C) -> BoxFuture<'a, Result<R>>,
+    {
+        let options = options.into();
+        let timeout = Duration::from_secs(120);
+        #[cfg(test)]
+        let timeout = self.convenient_transaction_timeout.unwrap_or(timeout);
+        let start = Instant::now();
+
+        use crate::error::{TRANSIENT_TRANSACTION_ERROR, UNKNOWN_TRANSACTION_COMMIT_RESULT};
+
+        'transaction: loop {
+            self.start_transaction(options.clone()).await?;
+            let ret = match callback(self, &mut context).await {
+                Ok(v) => v,
+                Err(e) => {
+                    if matches!(
+                        self.transaction.state,
+                        TransactionState::Starting | TransactionState::InProgress
+                    ) {
+                        self.abort_transaction().await?;
+                    }
+                    if e.contains_label(TRANSIENT_TRANSACTION_ERROR) && start.elapsed() < timeout {
+                        continue 'transaction;
+                    }
+                    return Err(e);
+                }
+            };
+            if matches!(
+                self.transaction.state,
+                TransactionState::None
+                    | TransactionState::Aborted
+                    | TransactionState::Committed { .. }
+            ) {
+                return Ok(ret);
+            }
+            'commit: loop {
+                match self.commit_transaction().await {
+                    Ok(()) => return Ok(ret),
+                    Err(e) => {
+                        if e.is_max_time_ms_expired_error() || start.elapsed() >= timeout {
+                            return Err(e);
+                        }
+                        if e.contains_label(UNKNOWN_TRANSACTION_COMMIT_RESULT) {
+                            continue 'commit;
+                        }
+                        if e.contains_label(TRANSIENT_TRANSACTION_ERROR) {
+                            continue 'transaction;
+                        }
+                        return Err(e);
+                    }
+                }
+            }
+        }
+    }
+
     fn default_transaction_options(&self) -> Option<&TransactionOptions> {
         self.options
             .as_ref()
             .and_then(|options| options.default_transaction_options.as_ref())
     }
 }
 
+pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;
+
 struct DroppedClientSession {
     cluster_time: Option<ClusterTime>,
     server_session: ServerSession,
@@ -590,6 +698,8 @@ impl From<DroppedClientSession> for ClientSession {
             transaction: dropped_session.transaction,
             snapshot_time: dropped_session.snapshot_time,
             operation_time: dropped_session.operation_time,
+            #[cfg(test)]
+            convenient_transaction_timeout: None,
         }
     }
 }

src/client/session/mod.rs

@@ -1,6 +1,7 @@
 //! Contains the `Error` and `Result` types that `mongodb` uses.
 
 use std::{
+    any::Any,
     collections::{HashMap, HashSet},
     fmt::{self, Debug},
     sync::Arc,
@@ -52,6 +53,22 @@ pub struct Error {
 }
 
 impl Error {
+    /// Create a new `Error` wrapping an arbitrary value.  Can be used to abort transactions in
+    /// callbacks for [`ClientSession::with_transaction`](crate::ClientSession::with_transaction).
+    pub fn custom(e: impl Any + Send + Sync) -> Self {
+        Self::new(ErrorKind::Custom(Arc::new(e)), None::<Option<String>>)
+    }
+
+    /// Retrieve a reference to a value provided to `Error::custom`.  Returns `None` if this is not
+    /// a custom error or if the payload types mismatch.
+    pub fn get_custom<E: Any>(&self) -> Option<&E> {
+        if let ErrorKind::Custom(c) = &*self.kind {
+            c.downcast_ref()
+        } else {
+            None
+        }
+    }
+
     pub(crate) fn new(kind: ErrorKind, labels: Option<impl IntoIterator<Item = String>>) -> Self {
         let mut labels: HashSet<String> = labels
             .map(|labels| labels.into_iter().collect())
@@ -140,6 +157,10 @@ impl Error {
         matches!(self.kind.as_ref(), ErrorKind::ServerSelection { .. })
     }
 
+    pub(crate) fn is_max_time_ms_expired_error(&self) -> bool {
+        self.code() == Some(50)
+    }
+
     /// Whether a read operation should be retried if this error occurs.
     pub(crate) fn is_read_retryable(&self) -> bool {
         if self.is_network_error() {
@@ -423,6 +444,7 @@ impl Error {
             | ErrorKind::IncompatibleServer { .. }
             | ErrorKind::MissingResumeToken
             | ErrorKind::Authentication { .. }
+            | ErrorKind::Custom(_)
             | ErrorKind::GridFs(_) => {}
             #[cfg(feature = "in-use-encryption-unstable")]
             ErrorKind::Encryption(_) => {}
@@ -578,6 +600,10 @@ pub enum ErrorKind {
     #[cfg(feature = "in-use-encryption-unstable")]
     #[error("An error occurred during client-side encryption: {0}")]
     Encryption(mongocrypt::error::Error),
+
+    /// A custom value produced by user code.
+    #[error("Custom user error")]
+    Custom(Arc<dyn Any + Send + Sync>),
 }
 
 impl ErrorKind {

src/error.rs

@@ -135,4 +135,72 @@ impl ClientSession {
     pub fn abort_transaction(&mut self) -> Result<()> {
         runtime::block_on(self.async_client_session.abort_transaction())
     }
+
+    /// Starts a transaction, runs the given callback, and commits or aborts the transaction.
+    /// Transient transaction errors will cause the callback or the commit to be retried;
+    /// other errors will cause the transaction to be aborted and the error returned to the
+    /// caller.  If the callback needs to provide its own error information, the
+    /// [`Error::custom`](crate::error::Error::custom) method can accept an arbitrary payload that
+    /// can be retrieved via [`Error::get_custom`](crate::error::Error::get_custom).
+    pub fn with_transaction<R, F>(
+        &mut self,
+        mut callback: F,
+        options: impl Into<Option<TransactionOptions>>,
+    ) -> Result<R>
+    where
+        F: for<'a> FnMut(&'a mut ClientSession) -> Result<R>,
+    {
+        let options = options.into();
+        let timeout = std::time::Duration::from_secs(120);
+        let start = std::time::Instant::now();
+
+        use crate::{
+            client::session::TransactionState,
+            error::{TRANSIENT_TRANSACTION_ERROR, UNKNOWN_TRANSACTION_COMMIT_RESULT},
+        };
+
+        'transaction: loop {
+            self.start_transaction(options.clone())?;
+            let ret = match callback(self) {
+                Ok(v) => v,
+                Err(e) => {
+                    if matches!(
+                        self.async_client_session.transaction.state,
+                        TransactionState::Starting | TransactionState::InProgress
+                    ) {
+                        self.abort_transaction()?;
+                    }
+                    if e.contains_label(TRANSIENT_TRANSACTION_ERROR) && start.elapsed() < timeout {
+                        continue 'transaction;
+                    }
+                    return Err(e);
+                }
+            };
+            if matches!(
+                self.async_client_session.transaction.state,
+                TransactionState::None
+                    | TransactionState::Aborted
+                    | TransactionState::Committed { .. }
+            ) {
+                return Ok(ret);
+            }
+            'commit: loop {
+                match self.commit_transaction() {
+                    Ok(()) => return Ok(ret),
+                    Err(e) => {
+                        if e.is_max_time_ms_expired_error() || start.elapsed() >= timeout {
+                            return Err(e);
+                        }
+                        if e.contains_label(UNKNOWN_TRANSACTION_COMMIT_RESULT) {
+                            continue 'commit;
+                        }
+                        if e.contains_label(TRANSIENT_TRANSACTION_ERROR) {
+                            continue 'transaction;
+                        }
+                        return Err(e);
+                    }
+                }
+            }
+        }
+    }
 }