RUST-935 direct retries to different mongos (#968)

Author: abr-egn

src/client.rs

@@ -19,8 +19,6 @@ use derivative::Derivative;
 use futures_core::{future::BoxFuture, Future};
 use futures_util::{future::join_all, FutureExt};
 
-#[cfg(test)]
-use crate::options::ServerAddress;
 #[cfg(feature = "tracing-unstable")]
 use crate::trace::{
     command::CommandTracingEventEmitter,
@@ -49,6 +47,7 @@ use crate::{
         ListDatabasesOptions,
         ReadPreference,
         SelectionCriteria,
+        ServerAddress,
         SessionOptions,
     },
     results::DatabaseSpecification,
@@ -654,17 +653,20 @@ impl Client {
         &self,
         criteria: Option<&SelectionCriteria>,
     ) -> Result<ServerAddress> {
-        let server = self.select_server(criteria, "Test select server").await?;
+        let server = self
+            .select_server(criteria, "Test select server", None)
+            .await?;
         Ok(server.address.clone())
     }
 
     /// Select a server using the provided criteria. If none is provided, a primary read preference
     /// will be used instead.
-    #[allow(unused_variables)] // we only use the operation_name for tracing.
     async fn select_server(
         &self,
         criteria: Option<&SelectionCriteria>,
+        #[allow(unused_variables)] // we only use the operation_name for tracing.
         operation_name: &str,
+        deprioritized: Option<&ServerAddress>,
     ) -> Result<SelectedServer> {
         let criteria =
             criteria.unwrap_or(&SelectionCriteria::ReadPreference(ReadPreference::Primary));
@@ -698,6 +700,7 @@ impl Client {
                 criteria,
                 &state.description,
                 &state.servers(),
+                deprioritized,
             );
             match result {
                 Err(error) => {

src/client/executor.rs

@@ -12,7 +12,7 @@ use std::{
     time::Instant,
 };
 
-use super::{session::TransactionState, Client, ClientSession};
+use super::{options::ServerAddress, session::TransactionState, Client, ClientSession};
 use crate::{
     bson::Document,
     change_stream::{
@@ -313,7 +313,14 @@ impl Client {
                 .and_then(|s| s.transaction.pinned_mongos())
                 .or_else(|| op.selection_criteria());
 
-            let server = match self.select_server(selection_criteria, op.name()).await {
+            let server = match self
+                .select_server(
+                    selection_criteria,
+                    op.name(),
+                    retry.as_ref().map(|r| &r.first_server),
+                )
+                .await
+            {
                 Ok(server) => server,
                 Err(mut err) => {
                     retry.first_error()?;
@@ -322,6 +329,7 @@ impl Client {
                     return Err(err);
                 }
             };
+            let server_addr = server.address.clone();
 
             let mut conn = match get_connection(&session, &op, &server.pool).await {
                 Ok(c) => c,
@@ -342,6 +350,7 @@ impl Client {
                         retry = Some(ExecutionRetry {
                             prior_txn_number: None,
                             first_error: err,
+                            first_server: server_addr.clone(),
                         });
                         continue;
                     } else {
@@ -407,7 +416,7 @@ impl Client {
                     self.inner
                         .topology
                         .handle_application_error(
-                            server.address.clone(),
+                            server_addr.clone(),
                             err.clone(),
                             HandshakePhase::after_completion(&conn),
                         )
@@ -433,6 +442,7 @@ impl Client {
                         retry = Some(ExecutionRetry {
                             prior_txn_number: txn_number,
                             first_error: err,
+                            first_server: server_addr.clone(),
                         });
                         continue;
                     } else {
@@ -806,7 +816,9 @@ impl Client {
             (matches!(topology_type, TopologyType::Single) && server_type.is_available())
                 || server_type.is_data_bearing()
         }));
-        let _: SelectedServer = self.select_server(Some(&criteria), operation_name).await?;
+        let _: SelectedServer = self
+            .select_server(Some(&criteria), operation_name, None)
+            .await?;
         Ok(())
     }
 
@@ -1019,6 +1031,7 @@ struct ExecutionDetails<T: Operation> {
 struct ExecutionRetry {
     prior_txn_number: Option<i64>,
     first_error: Error,
+    first_server: ServerAddress,
 }
 
 trait RetryHelper {

src/sdam/description/topology/server_selection.rs

@@ -61,13 +61,20 @@ pub(crate) fn attempt_to_select_server<'a>(
     criteria: &'a SelectionCriteria,
     topology_description: &'a TopologyDescription,
     servers: &'a HashMap<ServerAddress, Arc<Server>>,
+    deprioritized: Option<&ServerAddress>,
 ) -> Result<Option<SelectedServer>> {
-    let in_window = topology_description.suitable_servers_in_latency_window(criteria)?;
+    let mut in_window = topology_description.suitable_servers_in_latency_window(criteria)?;
+    if let Some(addr) = deprioritized {
+        if in_window.len() > 1 {
+            in_window.retain(|d| &d.address != addr);
+        }
+    }
     let in_window_servers = in_window
         .into_iter()
         .flat_map(|desc| servers.get(&desc.address))
         .collect();
-    Ok(select_server_in_latency_window(in_window_servers).map(SelectedServer::new))
+    let selected = select_server_in_latency_window(in_window_servers);
+    Ok(selected.map(SelectedServer::new))
 }
 
 /// Choose a server from several suitable choices within the latency window according to

src/sdam/description/topology/server_selection/test/in_window.rs

@@ -77,10 +77,14 @@ async fn run_test(test_file: TestFile) {
     .into();
 
     for _ in 0..test_file.iterations {
-        let selection =
-            server_selection::attempt_to_select_server(&read_pref, &topology_description, &servers)
-                .expect("selection should not fail")
-                .expect("a server should have been selected");
+        let selection = server_selection::attempt_to_select_server(
+            &read_pref,
+            &topology_description,
+            &servers,
+            None,
+        )
+        .expect("selection should not fail")
+        .expect("a server should have been selected");
         *tallies.entry(selection.address.clone()).or_insert(0) += 1;
     }
 

src/test/spec/retryable_reads.rs

@@ -6,7 +6,7 @@ use crate::{
     error::Result,
     event::{
         cmap::{CmapEvent, CmapEventHandler, ConnectionCheckoutFailedReason},
-        command::CommandEventHandler,
+        command::{CommandEvent, CommandEventHandler},
     },
     runtime,
     runtime::AsyncJoinHandle,
@@ -21,6 +21,7 @@ use crate::{
         TestClient,
         CLIENT_OPTIONS,
     },
+    Client,
 };
 
 #[cfg_attr(feature = "tokio-runtime", tokio::test(flavor = "multi_thread"))]
@@ -159,3 +160,123 @@ async fn retry_read_pool_cleared() {
 
     assert_eq!(handler.get_command_started_events(&["find"]).len(), 3);
 }
+
+// Retryable Reads Are Retried on a Different mongos if One is Available
+#[cfg_attr(feature = "tokio-runtime", tokio::test(flavor = "multi_thread"))]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn retry_read_different_mongos() {
+    let mut client_options = CLIENT_OPTIONS.get().await.clone();
+    if client_options.repl_set_name.is_some() || client_options.hosts.len() < 2 {
+        log_uncaptured(
+            "skipping retry_read_different_mongos: requires sharded cluster with at least two \
+             hosts",
+        );
+        return;
+    }
+    client_options.hosts.drain(2..);
+    client_options.retry_reads = Some(true);
+
+    let mut guards = vec![];
+    for ix in [0, 1] {
+        let mut opts = client_options.clone();
+        opts.hosts.remove(ix);
+        opts.direct_connection = Some(true);
+        let client = Client::test_builder().options(opts).build().await;
+        if !client.supports_fail_command() {
+            log_uncaptured("skipping retry_read_different_mongos: requires failCommand");
+            return;
+        }
+        let fail_opts = FailCommandOptions::builder()
+            .error_code(6)
+            .close_connection(true)
+            .build();
+        let fp = FailPoint::fail_command(&["find"], FailPointMode::Times(1), Some(fail_opts));
+        guards.push(client.enable_failpoint(fp, None).await.unwrap());
+    }
+
+    let client = Client::test_builder()
+        .options(client_options)
+        .event_client()
+        .build()
+        .await;
+    let result = client
+        .database("test")
+        .collection::<bson::Document>("retry_read_different_mongos")
+        .find(doc! {}, None)
+        .await;
+    assert!(result.is_err());
+    let events = client.get_command_events(&["find"]);
+    assert!(
+        matches!(
+            &events[..],
+            &[
+                CommandEvent::Started(_),
+                CommandEvent::Failed(_),
+                CommandEvent::Started(_),
+                CommandEvent::Failed(_),
+            ]
+        ),
+        "unexpected events: {:#?}",
+        events,
+    );
+
+    drop(guards); // enforce lifetime
+}
+
+// Retryable Reads Are Retried on the Same mongos if No Others are Available
+#[cfg_attr(feature = "tokio-runtime", tokio::test(flavor = "multi_thread"))]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn retry_read_same_mongos() {
+    let init_client = Client::test_builder().build().await;
+    if !init_client.supports_fail_command() {
+        log_uncaptured("skipping retry_read_same_mongos: requires failCommand");
+        return;
+    }
+    if !init_client.is_sharded() {
+        log_uncaptured("skipping retry_read_same_mongos: requires sharded cluster");
+        return;
+    }
+
+    let mut client_options = CLIENT_OPTIONS.get().await.clone();
+    client_options.hosts.drain(1..);
+    client_options.retry_reads = Some(true);
+    let fp_guard = {
+        let mut client_options = client_options.clone();
+        client_options.direct_connection = Some(true);
+        let client = Client::test_builder().options(client_options).build().await;
+        let fail_opts = FailCommandOptions::builder()
+            .error_code(6)
+            .close_connection(true)
+            .build();
+        let fp = FailPoint::fail_command(&["find"], FailPointMode::Times(1), Some(fail_opts));
+        client.enable_failpoint(fp, None).await.unwrap()
+    };
+
+    let client = Client::test_builder()
+        .options(client_options)
+        .event_client()
+        .build()
+        .await;
+    let result = client
+        .database("test")
+        .collection::<bson::Document>("retry_read_same_mongos")
+        .find(doc! {}, None)
+        .await;
+    assert!(result.is_ok(), "{:?}", result);
+    let events = client.get_command_events(&["find"]);
+    assert!(
+        matches!(
+            &events[..],
+            &[
+                CommandEvent::Started(_),
+                CommandEvent::Failed(_),
+                CommandEvent::Started(_),
+                CommandEvent::Succeeded(_),
+            ]
+        ),
+        "unexpected events: {:#?}",
+        events,
+    );
+
+    drop(fp_guard); // enforce lifetime
+}