RUST-978 Unified test format additions in support of load balanced tests (#465)

Author: abr-egn

src/client/mod.rs

@@ -360,4 +360,9 @@ impl Client {
             .map(|stream_address| format!("{}", stream_address))
             .collect()
     }
+
+    #[cfg(test)]
+    pub(crate) async fn sync_workers(&self) {
+        self.inner.topology.sync_workers().await;
+    }
 }

src/cmap/manager.rs

@@ -1,4 +1,6 @@
 use tokio::sync::mpsc;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 use super::Connection;
 use crate::{bson::oid::ObjectId, error::Error, runtime::AcknowledgedMessage};
@@ -72,6 +74,14 @@ impl PoolManager {
             .sender
             .send(PoolManagementRequest::HandleConnectionSucceeded(conn));
     }
+
+    /// Create a synchronization point for the pool's worker.
+    #[cfg(test)]
+    pub(super) fn sync_worker(&self) -> oneshot::Receiver<()> {
+        let (tx, rx) = oneshot::channel();
+        let _ = self.sender.send(PoolManagementRequest::Sync(tx));
+        rx
+    }
 }
 
 #[derive(Debug)]
@@ -108,6 +118,10 @@ pub(super) enum PoolManagementRequest {
     /// Update the pool after a successful connection, optionally populating the pool
     /// with the successful connection.
     HandleConnectionSucceeded(ConnectionSucceeded),
+
+    /// Synchronize the worker queue state with an external caller, i.e. a test.
+    #[cfg(test)]
+    Sync(oneshot::Sender<()>),
 }
 
 impl PoolManagementRequest {

src/cmap/mod.rs

@@ -12,6 +12,8 @@ mod worker;
 use std::sync::Arc;
 
 use derivative::Derivative;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 pub use self::conn::ConnectionInfo;
 pub(crate) use self::{
@@ -172,4 +174,9 @@ impl ConnectionPool {
     pub(crate) fn generation(&self) -> PoolGeneration {
         self.generation_subscriber.generation()
     }
+
+    #[cfg(test)]
+    pub(crate) fn sync_worker(&self) -> oneshot::Receiver<()> {
+        self.manager.sync_worker()
+    }
 }

src/cmap/worker.rs

@@ -313,6 +313,10 @@ impl ConnectionPoolWorker {
                 PoolTask::HandleManagementRequest(
                     PoolManagementRequest::HandleConnectionFailed,
                 ) => self.handle_connection_failed(),
+                #[cfg(test)]
+                PoolTask::HandleManagementRequest(PoolManagementRequest::Sync(tx)) => {
+                    let _ = tx.send(());
+                }
                 PoolTask::Maintenance => {
                     self.perform_maintenance();
                 }

src/cursor/common.rs

@@ -8,6 +8,8 @@ use std::{
 use derivative::Derivative;
 use futures_core::{Future, Stream};
 use serde::de::DeserializeOwned;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 use crate::{
     bson::Document,
@@ -303,13 +305,18 @@ pub(super) fn kill_cursor(
     ns: &Namespace,
     cursor_id: i64,
     pinned_conn: PinnedConnection,
+    #[cfg(test)] kill_watcher: Option<oneshot::Sender<()>>,
 ) {
     let coll = client
         .database(ns.db.as_str())
         .collection::<Document>(ns.coll.as_str());
     RUNTIME.execute(async move {
         if !pinned_conn.is_invalid() {
             let _ = coll.kill_cursor(cursor_id, pinned_conn.handle()).await;
+            #[cfg(test)]
+            if let Some(tx) = kill_watcher {
+                let _ = tx.send(());
+            }
         }
         pinned_conn.unpin().await;
     });

src/cursor/mod.rs

@@ -8,6 +8,8 @@ use std::{
 
 use futures_core::{future::BoxFuture, Stream};
 use serde::de::DeserializeOwned;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 use crate::{
     cmap::conn::PinnedConnectionHandle,
@@ -84,6 +86,8 @@ where
 {
     client: Client,
     wrapped_cursor: ImplicitSessionCursor<T>,
+    #[cfg(test)]
+    kill_watcher: Option<oneshot::Sender<()>>,
     _phantom: std::marker::PhantomData<T>,
 }
 
@@ -107,9 +111,24 @@ where
                 PinnedConnection::new(pin),
                 provider,
             ),
+            #[cfg(test)]
+            kill_watcher: None,
             _phantom: Default::default(),
         }
     }
+
+    /// Some tests need to be able to observe the events generated by `killCommand` execution;
+    /// however, because that happens asynchronously on `drop`, the test runner can conclude before
+    /// the event is published.  To fix that, tests can set a "kill watcher" on cursors - a
+    /// one-shot channel with a `()` value pushed after `killCommand` is run that the test can wait
+    /// on.
+    #[cfg(test)]
+    pub(crate) fn set_kill_watcher(&mut self, tx: oneshot::Sender<()>) {
+        if self.kill_watcher.is_some() {
+            panic!("cursor already has a kill_watcher");
+        }
+        self.kill_watcher = Some(tx);
+    }
 }
 
 impl<T> Stream for Cursor<T>
@@ -137,6 +156,8 @@ where
             self.wrapped_cursor.namespace(),
             self.wrapped_cursor.id(),
             self.wrapped_cursor.pinned_connection().replicate(),
+            #[cfg(test)]
+            self.kill_watcher.take(),
         );
     }
 }

src/cursor/session.rs

@@ -7,6 +7,8 @@ use std::{
 use futures_core::{future::BoxFuture, Stream};
 use futures_util::StreamExt;
 use serde::de::DeserializeOwned;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 use super::common::{
     kill_cursor,
@@ -64,6 +66,8 @@ where
     info: CursorInformation,
     buffer: VecDeque<T>,
     pinned_connection: PinnedConnection,
+    #[cfg(test)]
+    kill_watcher: Option<oneshot::Sender<()>>,
 }
 
 impl<T> SessionCursor<T>
@@ -83,6 +87,8 @@ where
             info: spec.info,
             buffer: spec.initial_buffer,
             pinned_connection: PinnedConnection::new(pinned),
+            #[cfg(test)]
+            kill_watcher: None,
         }
     }
 
@@ -177,6 +183,19 @@ where
     pub async fn next(&mut self, session: &mut ClientSession) -> Option<Result<T>> {
         self.stream(session).next().await
     }
+
+    /// Some tests need to be able to observe the events generated by `killCommand` execution;
+    /// however, because that happens asynchronously on `drop`, the test runner can conclude before
+    /// the event is published.  To fix that, tests can set a "kill watcher" on cursors - a
+    /// one-shot channel with a `()` value pushed after `killCommand` is run that the test can wait
+    /// on.
+    #[cfg(test)]
+    pub(crate) fn set_kill_watcher(&mut self, tx: oneshot::Sender<()>) {
+        if self.kill_watcher.is_some() {
+            panic!("cursor already has a kill_watcher");
+        }
+        self.kill_watcher = Some(tx);
+    }
 }
 
 impl<T> Drop for SessionCursor<T>
@@ -193,6 +212,8 @@ where
             &self.info.ns,
             self.info.id,
             self.pinned_connection.replicate(),
+            #[cfg(test)]
+            self.kill_watcher.take(),
         );
     }
 }

src/operation/count/mod.rs

@@ -47,8 +47,10 @@ impl Operation for Count {
     const NAME: &'static str = "count";
 
     fn build(&mut self, description: &StreamDescription) -> Result<Command> {
+        let mut name = Self::NAME.to_string();
         let mut body = match description.max_wire_version {
             Some(v) if v >= SERVER_4_9_0_WIRE_VERSION => {
+                name = "aggregate".to_string();
                 doc! {
                     "aggregate": self.ns.coll.clone(),
                     "pipeline": [
@@ -74,11 +76,7 @@ impl Operation for Count {
 
         append_options(&mut body, self.options.as_ref())?;
 
-        Ok(Command::new(
-            Self::NAME.to_string(),
-            self.ns.db.clone(),
-            body,
-        ))
+        Ok(Command::new(name, self.ns.db.clone(), body))
     }
 
     fn handle_response(

src/sdam/state/mod.rs

@@ -9,6 +9,8 @@ use std::{
     },
 };
 
+#[cfg(test)]
+use futures_util::stream::{FuturesUnordered, StreamExt};
 use tokio::sync::{RwLock, RwLockWriteGuard};
 
 use self::server::Server;
@@ -532,6 +534,11 @@ impl Topology {
             .map(|(addr, server)| (addr.clone(), Arc::downgrade(server)))
             .collect()
     }
+
+    #[cfg(test)]
+    pub(crate) async fn sync_workers(&self) {
+        self.state.read().await.sync_workers().await;
+    }
 }
 
 impl WeakTopology {
@@ -680,6 +687,16 @@ impl TopologyState {
 
         self.servers.retain(|host, _| hosts.contains(host));
     }
+
+    #[cfg(test)]
+    async fn sync_workers(&self) {
+        let rxen: FuturesUnordered<_> = self
+            .servers
+            .values()
+            .map(|v| v.pool.sync_worker())
+            .collect();
+        let _: Vec<_> = rxen.collect().await;
+    }
 }
 
 /// Enum describing a point in time during an operation's execution relative to when the MongoDB

src/test/mod.rs

@@ -66,6 +66,10 @@ lazy_static! {
     };
     pub(crate) static ref SERVERLESS: bool =
         matches!(std::env::var("SERVERLESS"), Ok(s) if s == "serverless");
+    pub(crate) static ref LOAD_BALANCED_SINGLE_URI: Option<String> =
+        std::env::var("MONGODB_LOAD_BALANCED_SINGLE_URI").ok();
+    pub(crate) static ref LOAD_BALANCED_MULTIPLE_URI: Option<String> =
+        std::env::var("MONGODB_LOAD_BALANCED_MULTIPLE_URI").ok();
 }
 
 fn get_default_uri() -> String {