RUST-122 Support mongos pinning for sharded transactions (#383)

Author: NBSquare

src/client/executor.rs

@@ -26,7 +26,13 @@ use crate::{
         Retryability,
     },
     options::SelectionCriteria,
-    sdam::{HandshakePhase, SelectedServer, SessionSupportStatus, TransactionSupportStatus},
+    sdam::{
+        HandshakePhase,
+        SelectedServer,
+        ServerType,
+        SessionSupportStatus,
+        TransactionSupportStatus,
+    },
     selection_criteria::ReadPreference,
 };
 
@@ -91,6 +97,7 @@ impl Client {
                         .into());
                     }
                 }
+
                 self.execute_operation_with_retry(op, Some(session)).await
             }
             None => {
@@ -135,18 +142,23 @@ impl Client {
             }
         }
 
-        let server = match self.select_server(op.selection_criteria()).await {
+        let selection_criteria = session
+            .as_ref()
+            .and_then(|s| s.transaction.pinned_mongos.as_ref())
+            .or_else(|| op.selection_criteria());
+
+        let server = match self.select_server(selection_criteria).await {
             Ok(server) => server,
             Err(mut err) => {
-                err.add_labels(None, &session, None)?;
+                err.add_labels_and_update_pin(None, &mut session, None)?;
                 return Err(err);
             }
         };
 
         let mut conn = match server.pool.check_out().await {
             Ok(conn) => conn,
             Err(mut err) => {
-                err.add_labels(None, &session, None)?;
+                err.add_labels_and_update_pin(None, &mut session, None)?;
 
                 if err.is_pool_cleared() {
                     return self.execute_retry(&mut op, &mut session, None, err).await;
@@ -229,6 +241,8 @@ impl Client {
         txn_number: Option<i64>,
         first_error: Error,
     ) -> Result<T::O> {
+        op.update_for_retry();
+
         let server = match self.select_server(op.selection_criteria()).await {
             Ok(server) => server,
             Err(_) => {
@@ -246,8 +260,6 @@ impl Client {
             return Err(first_error);
         }
 
-        op.update_for_retry();
-
         match self
             .execute_operation_on_connection(op, &mut conn, session, txn_number, &retryability)
             .await
@@ -286,7 +298,8 @@ impl Client {
             wc.validate()?;
         }
 
-        let mut cmd = op.build(connection.stream_description()?)?;
+        let stream_description = connection.stream_description()?;
+        let mut cmd = op.build(stream_description)?;
         self.inner
             .topology
             .update_command_with_read_pref(connection.address(), &mut cmd, op.selection_criteria())
@@ -324,6 +337,9 @@ impl Client {
                         cmd.set_start_transaction();
                         cmd.set_autocommit();
                         cmd.set_txn_read_concern(*session)?;
+                        if stream_description.initial_server_type == ServerType::Mongos {
+                            session.pin_mongos(connection.address().clone());
+                        }
                         session.transaction.state = TransactionState::InProgress;
                     }
                     TransactionState::InProgress
@@ -471,13 +487,13 @@ impl Client {
                     handler.handle_command_failed_event(command_failed_event);
                 });
 
-                if let Some(session) = session {
+                if let Some(ref mut session) = session {
                     if err.is_network_error() {
                         session.mark_dirty();
                     }
                 }
 
-                err.add_labels(Some(connection), session, Some(retryability))?;
+                err.add_labels_and_update_pin(Some(connection), session, Some(retryability))?;
                 op.handle_error(err)
             }
             Ok(response) => {
@@ -504,7 +520,11 @@ impl Client {
                 match op.handle_response(response.deserialized, connection.stream_description()?) {
                     Ok(response) => Ok(response),
                     Err(mut err) => {
-                        err.add_labels(Some(connection), session, Some(retryability))?;
+                        err.add_labels_and_update_pin(
+                            Some(connection),
+                            session,
+                            Some(retryability),
+                        )?;
                         Err(err)
                     }
                 }
@@ -618,7 +638,7 @@ impl Client {
 }
 
 impl Error {
-    /// Adds the necessary labels to this Error.
+    /// Adds the necessary labels to this Error, and unpins the session if needed.
     ///
     /// A TransientTransactionError label should be added if a transaction is in progress and the
     /// error is a network or server selection error.
@@ -628,10 +648,13 @@ impl Error {
     /// server version, a label should only be added if the `retry_writes` client option is not set
     /// to `false`, the operation during which the error occured is write-retryable, and a
     /// TransientTransactionError label has not already been added.
-    fn add_labels(
+    ///
+    /// If the TransientTransactionError or UnknownTransactionCommitResult labels are added, the
+    /// ClientSession should be unpinned.
+    fn add_labels_and_update_pin(
         &mut self,
         conn: Option<&Connection>,
-        session: &Option<&mut ClientSession>,
+        session: &mut Option<&mut ClientSession>,
         retryability: Option<&Retryability>,
     ) -> Result<()> {
         let transaction_state = session.as_ref().map_or(&TransactionState::None, |session| {
@@ -675,6 +698,15 @@ impl Error {
                 }
             }
         }
+
+        if let Some(ref mut session) = session {
+            if self.contains_label(TRANSIENT_TRANSACTION_ERROR)
+                || self.contains_label(UNKNOWN_TRANSACTION_COMMIT_RESULT)
+            {
+                session.unpin_mongos();
+            }
+        }
+
         Ok(())
     }
 }

src/client/session/mod.rs

@@ -5,6 +5,7 @@ mod test;
 
 use std::{
     collections::HashSet,
+    sync::Arc,
     time::{Duration, Instant},
 };
 
@@ -16,13 +17,16 @@ use crate::{
     error::{ErrorKind, Result},
     operation::{AbortTransaction, CommitTransaction, Operation},
     options::{SessionOptions, TransactionOptions},
-    sdam::TransactionSupportStatus,
+    sdam::{ServerInfo, TransactionSupportStatus},
+    selection_criteria::SelectionCriteria,
     Client,
     RUNTIME,
 };
 pub(crate) use cluster_time::ClusterTime;
 pub(super) use pool::ServerSessionPool;
 
+use super::options::ServerAddress;
+
 lazy_static! {
     pub(crate) static ref SESSIONS_UNSUPPORTED_COMMANDS: HashSet<&'static str> = {
         let mut hash_set = HashSet::new();
@@ -95,7 +99,7 @@ lazy_static! {
 ///     }
 /// }
 /// ```
-// TODO RUST-122 Remove this note and adjust the above description to indicate that sharded
+// TODO RUST-734 Remove this note and adjust the above description to indicate that sharded
 // transactions are supported on 4.2+
 /// Note: the driver does not currently support transactions on sharded clusters.
 #[derive(Clone, Debug)]
@@ -113,6 +117,7 @@ pub struct ClientSession {
 pub(crate) struct Transaction {
     pub(crate) state: TransactionState,
     pub(crate) options: Option<TransactionOptions>,
+    pub(crate) pinned_mongos: Option<SelectionCriteria>,
 }
 
 impl Transaction {
@@ -128,11 +133,13 @@ impl Transaction {
     pub(crate) fn abort(&mut self) {
         self.state = TransactionState::Aborted;
         self.options = None;
+        self.pinned_mongos = None;
     }
 
     pub(crate) fn reset(&mut self) {
         self.state = TransactionState::None;
         self.options = None;
+        self.pinned_mongos = None;
     }
 }
 
@@ -141,6 +148,7 @@ impl Default for Transaction {
         Self {
             state: TransactionState::None,
             options: None,
+            pinned_mongos: None,
         }
     }
 }
@@ -245,6 +253,17 @@ impl ClientSession {
         self.server_session.txn_number
     }
 
+    /// Pin mongos to session.
+    pub(crate) fn pin_mongos(&mut self, address: ServerAddress) {
+        self.transaction.pinned_mongos = Some(SelectionCriteria::Predicate(Arc::new(
+            move |server_info: &ServerInfo| *server_info.address() == address,
+        )));
+    }
+
+    pub(crate) fn unpin_mongos(&mut self) {
+        self.transaction.pinned_mongos = None;
+    }
+
     /// Whether this session is dirty.
     #[cfg(test)]
     pub(crate) fn is_dirty(&self) -> bool {
@@ -298,6 +317,9 @@ impl ClientSession {
                 }
                 .into());
             }
+            TransactionState::Committed { .. } => {
+                self.unpin_mongos(); // Unpin session if previous transaction is committed.
+            }
             _ => {}
         }
         match self.client.transaction_support_status().await? {
@@ -472,7 +494,8 @@ impl ClientSession {
                     .as_ref()
                     .and_then(|options| options.write_concern.as_ref())
                     .cloned();
-                let abort_transaction = AbortTransaction::new(write_concern);
+                let selection_criteria = self.transaction.pinned_mongos.clone();
+                let abort_transaction = AbortTransaction::new(write_concern, selection_criteria);
                 self.transaction.abort();
                 // Errors returned from running an abortTransaction command should be ignored.
                 let _result = self

src/concern/test.rs

@@ -179,7 +179,7 @@ async fn snapshot_read_concern() {
         .database(function_name!())
         .collection::<Document>(function_name!());
 
-    // TODO RUST-122 run this test on sharded clusters
+    // TODO RUST-734: Run this test on sharded clusters when transactions are complete.
     if client.is_replica_set() && client.server_version_gte(4, 0) {
         let mut session = client.start_session(None).await.unwrap();
         let options = TransactionOptions::builder()

src/operation/abort_transaction/mod.rs

@@ -4,17 +4,25 @@ use crate::{
     error::Result,
     operation::{Operation, Retryability},
     options::WriteConcern,
+    selection_criteria::SelectionCriteria,
 };
 
 use super::{CommandResponse, Response, WriteConcernOnlyBody};
 
 pub(crate) struct AbortTransaction {
     write_concern: Option<WriteConcern>,
+    selection_criteria: Option<SelectionCriteria>,
 }
 
 impl AbortTransaction {
-    pub(crate) fn new(write_concern: Option<WriteConcern>) -> Self {
-        Self { write_concern }
+    pub(crate) fn new(
+        write_concern: Option<WriteConcern>,
+        selection_criteria: Option<SelectionCriteria>,
+    ) -> Self {
+        Self {
+            write_concern,
+            selection_criteria,
+        }
     }
 }
 
@@ -47,11 +55,20 @@ impl Operation for AbortTransaction {
         response.validate()
     }
 
+    fn selection_criteria(&self) -> Option<&SelectionCriteria> {
+        self.selection_criteria.as_ref()
+    }
+
     fn write_concern(&self) -> Option<&WriteConcern> {
         self.write_concern.as_ref()
     }
 
     fn retryability(&self) -> Retryability {
         Retryability::Write
     }
+
+    fn update_for_retry(&mut self) {
+        // The session must be "unpinned" before server selection for a retry.
+        self.selection_criteria = None;
+    }
 }

src/sdam/description/topology/mod.rs

@@ -326,21 +326,15 @@ impl TopologyDescription {
             self.transaction_support_status = TransactionSupportStatus::Unsupported;
         }
         if let Ok(Some(max_wire_version)) = server_description.max_wire_version() {
-            match self.topology_type {
-                TopologyType::Sharded => {
-                    // TODO RUST-122: support transactions on sharded clusters
-                    self.transaction_support_status = TransactionSupportStatus::Unsupported;
-                }
-                _ => {
-                    if max_wire_version < 7 {
-                        self.transaction_support_status = TransactionSupportStatus::Unsupported;
-                    } else {
-                        self.transaction_support_status = TransactionSupportStatus::Supported;
-                    }
-                }
+            // TODO RUST-734: Evaluate whether we should permanently support sharded transactions.
+            //                If we leave the feature as unsupported, we should revert this code.
+            self.transaction_support_status = if max_wire_version < 7
+                || (max_wire_version < 8 && self.topology_type == TopologyType::Sharded)
+            {
+                TransactionSupportStatus::Unsupported
+            } else {
+                TransactionSupportStatus::Supported
             }
-        } else {
-            self.transaction_support_status = TransactionSupportStatus::Unsupported;
         }
     }
 
@@ -749,7 +743,7 @@ pub(crate) enum TransactionSupportStatus {
 
     /// Transactions are supported by this topology. A topology supports transactions if it
     /// supports sessions and its maxWireVersion >= 7. Transactions are not currently supported
-    /// on sharded clusters (TODO RUST-122).
+    /// on sharded clusters (TODO RUST-734).
     ///
     /// Note that meeting these conditions does not guarantee that a deployment
     /// supports transactions; any other missing qualification will be reported by the server.

src/sync/test.rs

@@ -202,7 +202,7 @@ fn transactions() {
 
     let should_skip = RUNTIME.block_on(async {
         let test_client = AsyncTestClient::new().await;
-        // TODO RUST-122: Unskip this test on sharded clusters
+        // TODO RUST-734: Unskip this test on sharded clusters when transactions are complete.
         !test_client.is_replica_set() || test_client.server_version_lt(4, 0)
     });
     if should_skip {