RUST-97 Support sharded transactions recovery token (#398)

Author: NBSquare

src/client/executor.rs

@@ -290,6 +290,7 @@ impl Client {
         }
 
         let stream_description = connection.stream_description()?;
+        let is_sharded = stream_description.initial_server_type == ServerType::Mongos;
         let mut cmd = op.build(stream_description)?;
         self.inner
             .topology
@@ -328,15 +329,22 @@ impl Client {
                         cmd.set_start_transaction();
                         cmd.set_autocommit();
                         cmd.set_txn_read_concern(*session)?;
-                        if stream_description.initial_server_type == ServerType::Mongos {
+                        if is_sharded {
                             session.pin_mongos(connection.address().clone());
                         }
                         session.transaction.state = TransactionState::InProgress;
                     }
-                    TransactionState::InProgress
-                    | TransactionState::Committed { .. }
-                    | TransactionState::Aborted => {
+                    TransactionState::InProgress => cmd.set_autocommit(),
+                    TransactionState::Committed { .. } | TransactionState::Aborted => {
                         cmd.set_autocommit();
+
+                        // Append the recovery token to the command if we are committing or aborting
+                        // on a sharded transaction.
+                        if is_sharded {
+                            if let Some(ref recovery_token) = session.transaction.recovery_token {
+                                cmd.set_recovery_token(recovery_token);
+                            }
+                        }
                     }
                     _ => {}
                 }
@@ -403,6 +411,9 @@ impl Client {
                     Ok(r) => {
                         self.update_cluster_time(&r, session).await;
                         if r.is_success() {
+                            // Retrieve recovery token from successful response.
+                            Client::update_recovery_token(is_sharded, &r, session).await;
+
                             Ok(CommandResult {
                                 raw: response,
                                 deserialized: r.into_body(),
@@ -447,7 +458,15 @@ impl Client {
                                         }))
                                     }
                                     // for ok: 1 just return the original deserialization error.
-                                    _ => Err(deserialize_error),
+                                    _ => {
+                                        Client::update_recovery_token(
+                                            is_sharded,
+                                            &error_response,
+                                            session,
+                                        )
+                                        .await;
+                                        Err(deserialize_error)
+                                    }
                                 }
                             }
                             // We failed to deserialize even that, so just return the original
@@ -626,6 +645,18 @@ impl Client {
             }
         }
     }
+
+    async fn update_recovery_token<T: Response>(
+        is_sharded: bool,
+        response: &T,
+        session: &mut Option<&mut ClientSession>,
+    ) {
+        if let Some(ref mut session) = session {
+            if is_sharded && session.in_transaction() {
+                session.transaction.recovery_token = response.recovery_token().cloned();
+            }
+        }
+    }
 }
 
 impl Error {

src/client/session/mod.rs

@@ -118,12 +118,14 @@ pub(crate) struct Transaction {
     pub(crate) state: TransactionState,
     pub(crate) options: Option<TransactionOptions>,
     pub(crate) pinned_mongos: Option<SelectionCriteria>,
+    pub(crate) recovery_token: Option<Document>,
 }
 
 impl Transaction {
     pub(crate) fn start(&mut self, options: Option<TransactionOptions>) {
         self.state = TransactionState::Starting;
         self.options = options;
+        self.recovery_token = None;
     }
 
     pub(crate) fn commit(&mut self, data_committed: bool) {
@@ -140,6 +142,7 @@ impl Transaction {
         self.state = TransactionState::None;
         self.options = None;
         self.pinned_mongos = None;
+        self.recovery_token = None;
     }
 }
 
@@ -149,6 +152,7 @@ impl Default for Transaction {
             state: TransactionState::None,
             options: None,
             pinned_mongos: None,
+            recovery_token: None,
         }
     }
 }

src/cmap/conn/command.rs

@@ -41,6 +41,10 @@ impl Command {
         }
     }
 
+    pub(crate) fn set_recovery_token(&mut self, recovery_token: &Document) {
+        self.body.insert("recoveryToken", recovery_token);
+    }
+
     pub(crate) fn set_txn_number(&mut self, txn_number: i64) {
         self.body.insert("txnNumber", txn_number);
     }

src/operation/mod.rs

@@ -153,6 +153,9 @@ pub(crate) trait Response: Sized {
     /// The `atClusterTime` field of the response.
     fn at_cluster_time(&self) -> Option<Timestamp>;
 
+    /// The `recoveryToken` field of the response.
+    fn recovery_token(&self) -> Option<&Document>;
+
     /// Convert into the body of the response.
     fn into_body(self) -> Self::Body;
 }
@@ -168,6 +171,8 @@ pub(crate) struct CommandResponse<T> {
 
     pub(crate) at_cluster_time: Option<Timestamp>,
 
+    pub(crate) recovery_token: Option<Document>,
+
     #[serde(flatten)]
     pub(crate) body: T,
 }
@@ -197,6 +202,10 @@ impl<T: DeserializeOwned> Response for CommandResponse<T> {
         self.at_cluster_time
     }
 
+    fn recovery_token(&self) -> Option<&Document> {
+        self.recovery_token.as_ref()
+    }
+
     fn into_body(self) -> Self::Body {
         self.body
     }
@@ -229,6 +238,10 @@ impl<T: DeserializeOwned> Response for CursorResponse<T> {
         self.response.body.cursor.at_cluster_time
     }
 
+    fn recovery_token(&self) -> Option<&Document> {
+        self.response.recovery_token()
+    }
+
     fn into_body(self) -> Self::Body {
         self.response.body
     }

src/operation/run_command/mod.rs

@@ -96,6 +96,7 @@ impl Operation for RunCommand {
 pub(crate) struct Response {
     doc: Document,
     cluster_time: Option<ClusterTime>,
+    recovery_token: Option<Document>,
 }
 
 impl super::Response for Response {
@@ -109,7 +110,13 @@ impl super::Response for Response {
             .ok()
             .and_then(|doc| bson::from_document(doc.clone()).ok());
 
-        Ok(Self { doc, cluster_time })
+        let recovery_token = doc.get_document("recoveryToken").ok().cloned();
+
+        Ok(Self {
+            doc,
+            cluster_time,
+            recovery_token,
+        })
     }
 
     fn ok(&self) -> Option<&Bson> {
@@ -131,6 +138,10 @@ impl super::Response for Response {
             .ok()
     }
 
+    fn recovery_token(&self) -> Option<&Document> {
+        self.recovery_token.as_ref()
+    }
+
     fn into_body(self) -> Self::Body {
         self.doc
     }