RUST-2074 Fix retryability bug, add disabled test (#1427)

Author: isabelatkinson

src/client/executor.rs

@@ -349,12 +349,19 @@ impl Client {
                         err.add_label(RETRYABLE_WRITE_ERROR);
                     }
 
-                    let op_retry = match self.get_op_retryability(op, &session) {
-                        Retryability::Read => err.is_read_retryable(),
-                        Retryability::Write => err.is_write_retryable(),
-                        _ => false,
+                    let retryability = op.retryability().with_options(self.options());
+                    let can_retry = match retryability {
+                        // Read-retryable operations should be retried on pool cleared errors during
+                        // connection checkout regardless of transaction status.
+                        Retryability::Read if err.is_pool_cleared() => true,
+                        _ => {
+                            retryability.can_retry_error(&err)
+                                && !session
+                                    .as_ref()
+                                    .is_some_and(|session| session.in_transaction())
+                        }
                     };
-                    if err.is_pool_cleared() || op_retry {
+                    if can_retry {
                         retry = Some(ExecutionRetry {
                             prior_txn_number: None,
                             first_error: err,
@@ -380,7 +387,7 @@ impl Client {
                 session = implicit_session.as_mut();
             }
 
-            let retryability = self.get_retryability(&conn, op, &session)?;
+            let retryability = self.get_retryability(op, &session, conn.stream_description()?);
             if retryability == Retryability::None {
                 retry.first_error()?;
             }
@@ -449,9 +456,7 @@ impl Client {
                         } else {
                             return Err(r.first_error);
                         }
-                    } else if retryability == Retryability::Read && err.is_read_retryable()
-                        || retryability == Retryability::Write && err.is_write_retryable()
-                    {
+                    } else if retryability.can_retry_error(&err) {
                         retry = Some(ExecutionRetry {
                             prior_txn_number: txn_number,
                             first_error: err,
@@ -910,52 +915,37 @@ impl Client {
         }
     }
 
-    /// Returns the retryability level for the execution of this operation.
-    fn get_op_retryability<T: Operation>(
+    /// Returns the retryability level for the execution of this operation with the given session
+    /// and connection stream description.
+    fn get_retryability<T: Operation>(
         &self,
         op: &T,
         session: &Option<&mut ClientSession>,
+        stream_description: &StreamDescription,
     ) -> Retryability {
+        // commitTransaction and abortTransaction are always retried, regardless of the value of
+        // retry_writes.
+        if op.name() == CommitTransaction::NAME || op.name() == AbortTransaction::NAME {
+            return Retryability::Write;
+        }
+
         if session
             .as_ref()
-            .map(|session| session.in_transaction())
-            .unwrap_or(false)
+            .is_some_and(|session| session.in_transaction())
         {
             return Retryability::None;
         }
-        match op.retryability() {
-            Retryability::Read if self.inner.options.retry_reads != Some(false) => {
-                Retryability::Read
-            }
-            // commitTransaction and abortTransaction should be retried regardless of the
-            // value for retry_writes set on the Client
-            Retryability::Write
-                if op.name() == CommitTransaction::NAME
-                    || op.name() == AbortTransaction::NAME
-                    || self.inner.options.retry_writes != Some(false) =>
-            {
+
+        match op.retryability().with_options(self.options()) {
+            Retryability::Write if stream_description.supports_retryable_writes() => {
                 Retryability::Write
             }
+            // All servers compatible with the driver support retryable reads.
+            Retryability::Read => Retryability::Read,
             _ => Retryability::None,
         }
     }
 
-    /// Returns the retryability level for the execution of this operation on this connection.
-    fn get_retryability<T: Operation>(
-        &self,
-        conn: &PooledConnection,
-        op: &T,
-        session: &Option<&mut ClientSession>,
-    ) -> Result<Retryability> {
-        match self.get_op_retryability(op, session) {
-            Retryability::Read => Ok(Retryability::Read),
-            Retryability::Write if conn.stream_description()?.supports_retryable_writes() => {
-                Ok(Retryability::Write)
-            }
-            _ => Ok(Retryability::None),
-        }
-    }
-
     async fn update_cluster_time(
         &self,
         cluster_time: Option<ClusterTime>,

src/cmap/test.rs

@@ -486,3 +486,74 @@ async fn cmap_spec_tests() {
     )
     .await;
 }
+
+// TODO RUST-2074: investigate why this test is flaky
+// #[tokio::test(flavor = "multi_thread")]
+// async fn pool_cleared_error_has_transient_transaction_error_label() {
+//     if !block_connection_supported().await {
+//         log_uncaptured(
+//             "skipping pool_cleared_error_has_transient_transaction_error_label: block connection
+// \              unsupported",
+//         );
+//         return;
+//     }
+//     if !transactions_supported().await {
+//         log_uncaptured(
+//             "skipping pool_cleared_error_has_transient_transaction_error_label: transactions \
+//              unsupported",
+//         );
+//         return;
+//     }
+//     if topology_is_load_balanced().await {
+//         log_uncaptured(
+//             "skipping pool_cleared_error_has_transient_transaction_error_label: load balanced \
+//              topology",
+//         );
+//     }
+
+//     let app_name = "pool_cleared_error_has_transient_transaction_error_label";
+
+//     let mut client_options = get_client_options().await.clone();
+//     if topology_is_sharded().await {
+//         client_options.hosts.drain(1..);
+//     }
+//     client_options.connect_timeout = Some(Duration::from_millis(500));
+//     client_options.heartbeat_freq = Some(Duration::from_millis(500));
+//     client_options.app_name = Some(app_name.to_string());
+//     let client = Client::for_test()
+//         .options(client_options)
+//         .monitor_events()
+//         .await;
+
+//     let mut session = client.start_session().await.unwrap();
+//     session.start_transaction().await.unwrap();
+
+//     let fail_point = FailPoint::fail_command(&["insert"], FailPointMode::Times(1))
+//         .block_connection(Duration::from_secs(15))
+//         .app_name(app_name);
+//     let _guard = client.enable_fail_point(fail_point).await.unwrap();
+
+//     let insert_client = client.clone();
+//     let insert_handle = tokio::spawn(async move {
+//         insert_client
+//             .database("db")
+//             .collection("coll")
+//             .insert_one(doc! { "x": 1 })
+//             .session(&mut session)
+//             .await
+//     });
+
+//     let fail_point = FailPoint::fail_command(
+//         &["hello", LEGACY_HELLO_COMMAND_NAME],
+//         // The RTT hellos may encounter this failpoint, so use FailPointMode::AlwaysOn to ensure
+//         // that the server monitors hit it as well.
+//         FailPointMode::AlwaysOn,
+//     )
+//     .block_connection(Duration::from_millis(1500))
+//     .app_name(app_name);
+//     let _guard = client.enable_fail_point(fail_point).await.unwrap();
+
+//     let insert_error = insert_handle.await.unwrap().unwrap_err();
+//     assert!(insert_error.is_pool_cleared(), "{:?}", insert_error);
+//     assert!(insert_error.contains_label(TRANSIENT_TRANSACTION_ERROR));
+// }

src/operation.rs

@@ -52,7 +52,7 @@ use crate::{
         WriteConcernError,
         WriteFailure,
     },
-    options::WriteConcern,
+    options::{ClientOptions, WriteConcern},
     selection_criteria::SelectionCriteria,
     BoxFuture,
     ClientSession,
@@ -100,6 +100,26 @@ pub(crate) enum Retryability {
     None,
 }
 
+impl Retryability {
+    /// Returns this level of retryability in tandem with the client options.
+    pub(crate) fn with_options(&self, options: &ClientOptions) -> Self {
+        match self {
+            Self::Write if options.retry_writes != Some(false) => Self::Write,
+            Self::Read if options.retry_reads != Some(false) => Self::Read,
+            _ => Self::None,
+        }
+    }
+
+    /// Whether this level of retryability can retry the given error.
+    pub(crate) fn can_retry_error(&self, error: &Error) -> bool {
+        match self {
+            Self::Write => error.is_write_retryable(),
+            Self::Read => error.is_read_retryable(),
+            Self::None => false,
+        }
+    }
+}
+
 /// A trait modeling the behavior of a server side operation.
 ///
 /// No methods in this trait should have default behaviors to ensure that wrapper operations

src/test/util/fail_point.rs

@@ -147,6 +147,16 @@ impl Drop for FailPointGuard {
         // multi-threaded runtime.
         let result = tokio::task::block_in_place(|| {
             futures::executor::block_on(async move {
+                let client = if client.options().app_name.is_some() {
+                    // Create a fresh client with no app name to avoid issues when disabling a
+                    // failpoint configured on the "hello" command.
+                    let mut options = client.options().clone();
+                    options.app_name = None;
+                    Client::for_test().options(options).await.into_client()
+                } else {
+                    client
+                };
+
                 client
                     .database("admin")
                     .run_command(