RUST-828 Reduce flakiness of cluster_time_in_commands test (#740)

Author: patrickfreed

src/client/session/test/mod.rs

@@ -14,7 +14,17 @@ use crate::{
     runtime,
     sdam::ServerInfo,
     selection_criteria::SelectionCriteria,
-    test::{log_uncaptured, EventClient, TestClient, CLIENT_OPTIONS, LOCK},
+    test::{
+        log_uncaptured,
+        Event,
+        EventClient,
+        EventHandler,
+        SdamEvent,
+        TestClient,
+        CLIENT_OPTIONS,
+        LOCK,
+    },
+    Client,
     Collection,
 };
 
@@ -238,19 +248,22 @@ async fn pool_is_lifo() {
 async fn cluster_time_in_commands() {
     let _guard: RwLockReadGuard<()> = LOCK.run_concurrently().await;
 
-    let client = TestClient::new().await;
-    if client.is_standalone() {
+    let test_client = TestClient::new().await;
+    if test_client.is_standalone() {
+        log_uncaptured("skipping cluster_time_in_commands test due to standalone topology");
         return;
     }
 
-    async fn cluster_time_test<F, G, R>(command_name: &str, operation: F)
-    where
-        F: Fn(EventClient) -> G,
+    async fn cluster_time_test<F, G, R>(
+        command_name: &str,
+        client: &Client,
+        event_handler: &EventHandler,
+        operation: F,
+    ) where
+        F: Fn(Client) -> G,
         G: Future<Output = Result<R>>,
     {
-        let mut options = CLIENT_OPTIONS.get().await.clone();
-        options.heartbeat_freq = Some(Duration::from_secs(1000));
-        let client = EventClient::with_options(options).await;
+        let mut subscriber = event_handler.subscribe();
 
         operation(client.clone())
             .await
@@ -260,16 +273,31 @@ async fn cluster_time_in_commands() {
             .await
             .expect("operation should succeed");
 
-        let (first_command_started, first_command_succeeded) =
-            client.get_successful_command_execution(command_name);
+        let (first_command_started, first_command_succeeded) = subscriber
+            .wait_for_successful_command_execution(Duration::from_secs(5), command_name)
+            .await
+            .unwrap_or_else(|| {
+                panic!(
+                    "did not see command started and succeeded events for {}",
+                    command_name
+                )
+            });
 
         assert!(first_command_started.command.get("$clusterTime").is_some());
         let response_cluster_time = first_command_succeeded
             .reply
             .get("$clusterTime")
             .expect("should get cluster time from command response");
 
-        let (second_command_started, _) = client.get_successful_command_execution(command_name);
+        let (second_command_started, _) = subscriber
+            .wait_for_successful_command_execution(Duration::from_secs(5), command_name)
+            .await
+            .unwrap_or_else(|| {
+                panic!(
+                    "did not see command started and succeeded events for {}",
+                    command_name
+                )
+            });
 
         assert_eq!(
             response_cluster_time,
@@ -282,24 +310,73 @@ async fn cluster_time_in_commands() {
         );
     }
 
-    cluster_time_test("ping", |client| async move {
+    let handler = Arc::new(EventHandler::new());
+    let mut options = CLIENT_OPTIONS.get().await.clone();
+    options.heartbeat_freq = Some(Duration::from_secs(1000));
+    options.command_event_handler = Some(handler.clone());
+    options.sdam_event_handler = Some(handler.clone());
+
+    // Ensure we only connect to one server so the monitor checks from other servers
+    // don't affect the TopologyDescription's clusterTime value between commands.
+    if options.load_balanced != Some(true) {
+        options.direct_connection = Some(true);
+
+        // Since we need to run an insert below, ensure the single host is a primary
+        // if we're connected to a replica set.
+        if let Some(primary) = test_client.primary() {
+            options.hosts = vec![primary];
+        } else {
+            options.hosts.drain(1..);
+        }
+    }
+
+    let mut subscriber = handler.subscribe();
+
+    let client = Client::with_options(options).unwrap();
+
+    // Wait for initial monitor check to complete and discover the server.
+    subscriber
+        .wait_for_event(Duration::from_secs(5), |event| match event {
+            Event::Sdam(SdamEvent::ServerDescriptionChanged(e)) => {
+                !e.previous_description.server_type().is_available()
+                    && e.new_description.server_type().is_available()
+            }
+            _ => false,
+        })
+        .await
+        .expect("server should be discovered");
+
+    // LoadBalanced topologies don't have monitors, so the client needs to get a clusterTime from
+    // a command invocation.
+    client
+        .database("admin")
+        .run_command(doc! { "ping": 1 }, None)
+        .await
+        .unwrap();
+
+    cluster_time_test("ping", &client, handler.as_ref(), |client| async move {
         client
             .database(function_name!())
             .run_command(doc! { "ping": 1 }, None)
             .await
     })
     .await;
 
-    cluster_time_test("aggregate", |client| async move {
-        client
-            .database(function_name!())
-            .collection::<Document>(function_name!())
-            .aggregate(vec![doc! { "$match": { "x": 1 } }], None)
-            .await
-    })
+    cluster_time_test(
+        "aggregate",
+        &client,
+        handler.as_ref(),
+        |client| async move {
+            client
+                .database(function_name!())
+                .collection::<Document>(function_name!())
+                .aggregate(vec![doc! { "$match": { "x": 1 } }], None)
+                .await
+        },
+    )
     .await;
 
-    cluster_time_test("find", |client| async move {
+    cluster_time_test("find", &client, handler.as_ref(), |client| async move {
         client
             .database(function_name!())
             .collection::<Document>(function_name!())
@@ -308,7 +385,7 @@ async fn cluster_time_in_commands() {
     })
     .await;
 
-    cluster_time_test("insert", |client| async move {
+    cluster_time_test("insert", &client, handler.as_ref(), |client| async move {
         client
             .database(function_name!())
             .collection::<Document>(function_name!())

src/test/util/event.rs

@@ -481,17 +481,49 @@ impl EventSubscriber<'_> {
     ) -> Option<Event>
     where
         F: FnMut(&Event) -> bool,
+    {
+        self.filter_map_event(timeout, |e| if filter(&e) { Some(e) } else { None })
+            .await
+    }
+
+    pub(crate) async fn collect_events<F>(&mut self, timeout: Duration, mut filter: F) -> Vec<Event>
+    where
+        F: FnMut(&Event) -> bool,
+    {
+        let mut events = Vec::new();
+        let _ = runtime::timeout(timeout, async {
+            while let Some(event) = self.wait_for_event(timeout, &mut filter).await {
+                events.push(event);
+            }
+        })
+        .await;
+        events
+    }
+
+    /// Consume and pass events to the provided closure until it returns Some or the timeout is hit.
+    pub(crate) async fn filter_map_event<F, T>(
+        &mut self,
+        timeout: Duration,
+        mut filter_map: F,
+    ) -> Option<T>
+    where
+        F: FnMut(Event) -> Option<T>,
     {
         runtime::timeout(timeout, async {
             loop {
                 match self.receiver.recv().await {
-                    Ok(event) if filter(&event) => return event.into(),
+                    Ok(event) => {
+                        if let Some(e) = filter_map(event) {
+                            return Some(e);
+                        } else {
+                            continue;
+                        }
+                    }
                     // the channel hit capacity and missed some events.
                     Err(RecvError::Lagged(amount_skipped)) => {
                         panic!("receiver lagged and skipped {} events", amount_skipped)
                     }
                     Err(_) => return None,
-                    _ => continue,
                 }
             }
         })
@@ -500,18 +532,52 @@ impl EventSubscriber<'_> {
         .flatten()
     }
 
-    pub(crate) async fn collect_events<F>(&mut self, timeout: Duration, mut filter: F) -> Vec<Event>
-    where
-        F: FnMut(&Event) -> bool,
-    {
-        let mut events = Vec::new();
-        let _ = runtime::timeout(timeout, async {
-            while let Some(event) = self.wait_for_event(timeout, &mut filter).await {
-                events.push(event);
-            }
+    /// Waits for the next CommandStartedEvent/CommandFailedEvent pair.
+    /// If the next CommandStartedEvent is associated with a CommandFailedEvent, this method will
+    /// panic.
+    pub(crate) async fn wait_for_successful_command_execution(
+        &mut self,
+        timeout: Duration,
+        command_name: impl AsRef<str>,
+    ) -> Option<(CommandStartedEvent, CommandSucceededEvent)> {
+        runtime::timeout(timeout, async {
+            let started = self
+                .filter_map_event(Duration::MAX, |event| match event {
+                    Event::Command(CommandEvent::Started(s))
+                        if s.command_name == command_name.as_ref() =>
+                    {
+                        Some(s)
+                    }
+                    _ => None,
+                })
+                .await
+                .unwrap();
+
+            let succeeded = self
+                .filter_map_event(Duration::MAX, |event| match event {
+                    Event::Command(CommandEvent::Succeeded(s))
+                        if s.request_id == started.request_id =>
+                    {
+                        Some(s)
+                    }
+                    Event::Command(CommandEvent::Failed(f))
+                        if f.request_id == started.request_id =>
+                    {
+                        panic!(
+                            "expected {} to succeed but it failed: {:#?}",
+                            command_name.as_ref(),
+                            f
+                        )
+                    }
+                    _ => None,
+                })
+                .await
+                .unwrap();
+
+            (started, succeeded)
         })
-        .await;
-        events
+        .await
+        .ok()
     }
 }
 

src/test/util/mod.rs

@@ -14,6 +14,7 @@ use std::{fmt::Debug, sync::Arc, time::Duration};
 
 use crate::{
     bson::{doc, Bson},
+    client::options::ServerAddress,
     hello::{hello_command, HelloCommandResponse},
     selection_criteria::SelectionCriteria,
 };
@@ -386,6 +387,13 @@ impl TestClient {
         .to_string()
     }
 
+    pub(crate) fn primary(&self) -> Option<ServerAddress> {
+        self.server_info
+            .primary
+            .as_ref()
+            .map(|s| ServerAddress::parse(s).unwrap())
+    }
+
     pub(crate) async fn options_for_multiple_mongoses(
         options: Option<ClientOptions>,
         use_multiple_mongoses: bool,