RUST-1100 Fix find_and_getmore_share_session test (#525)

patrickfreed · patrickfreed · commit bf0275e358d7 · 2021-12-02T14:27:45.000-05:00
This also fixes RUST-1113 and RUST-1114
diff --git a/src/client/mod.rs b/src/client/mod.rs
@@ -367,6 +367,11 @@ impl Client {
         }
     }
 
+    #[cfg(test)]
+    pub(crate) async fn topology_description(&self) -> crate::sdam::TopologyDescription {
+        self.inner.topology.description().await
+    }
+
     #[cfg(all(test, not(feature = "sync")))]
     pub(crate) async fn get_hosts(&self) -> Vec<String> {
         let servers = self.inner.topology.servers().await;
diff --git a/src/client/session/test/mod.rs b/src/client/session/test/mod.rs
@@ -1,15 +1,17 @@
 mod causal_consistency;
 
-use std::{future::Future, time::Duration};
+use std::{future::Future, sync::Arc, time::Duration};
 
 use bson::Document;
 use futures::stream::StreamExt;
 use tokio::sync::RwLockReadGuard;
 
 use crate::{
     bson::{doc, Bson},
+    coll::options::{CountOptions, InsertManyOptions},
     error::Result,
-    options::{Acknowledgment, FindOptions, InsertOneOptions, ReadPreference, WriteConcern},
+    options::{Acknowledgment, FindOptions, ReadConcern, ReadPreference, WriteConcern},
+    sdam::ServerInfo,
     selection_criteria::SelectionCriteria,
     test::{EventClient, TestClient, CLIENT_OPTIONS, LOCK},
     Collection,
@@ -449,22 +451,19 @@ async fn find_and_getmore_share_session() {
     let _guard: RwLockReadGuard<()> = LOCK.run_concurrently().await;
 
     let client = EventClient::new().await;
-    if !client.is_replica_set() {
+    if client.is_standalone() {
+        println!("skipping find_and_getmore_share_session due to unsupported topology: Standalone");
         return;
     }
 
     let coll = client
         .init_db_and_coll(function_name!(), function_name!())
         .await;
 
-    let options = InsertOneOptions::builder()
+    let options = InsertManyOptions::builder()
         .write_concern(WriteConcern::builder().w(Acknowledgment::Majority).build())
         .build();
-    for _ in 0..3 {
-        coll.insert_one(doc! {}, options.clone())
-            .await
-            .expect("insert should succeed");
-    }
+    coll.insert_many(vec![doc! {}; 3], options).await.unwrap();
 
     let read_preferences: Vec<ReadPreference> = vec![
         ReadPreference::Primary,
@@ -489,7 +488,8 @@ async fn find_and_getmore_share_session() {
     ) {
         let options = FindOptions::builder()
             .batch_size(2)
-            .selection_criteria(SelectionCriteria::ReadPreference(read_preference))
+            .selection_criteria(SelectionCriteria::ReadPreference(read_preference.clone()))
+            .read_concern(ReadConcern::local())
             .build();
 
         let mut cursor = coll
@@ -498,7 +498,21 @@ async fn find_and_getmore_share_session() {
             .expect("find should succeed");
 
         for _ in 0..3 {
-            assert!(matches!(cursor.next().await, Some(Ok(_))));
+            cursor
+                .next()
+                .await
+                .unwrap_or_else(|| {
+                    panic!(
+                        "should get result with read preference {:?}",
+                        read_preference
+                    )
+                })
+                .unwrap_or_else(|e| {
+                    panic!(
+                        "result should not be error with read preference {:?}, but got {:?}",
+                        read_preference, e
+                    )
+                });
         }
 
         let (find_started, _) = client.get_successful_command_execution("find");
@@ -516,6 +530,22 @@ async fn find_and_getmore_share_session() {
         assert_eq!(getmore_session_id, session_id);
     }
 
+    let topology_description = client.topology_description().await;
+    for (addr, server) in topology_description.servers {
+        if !server.server_type.is_data_bearing() {
+            continue;
+        }
+
+        let a = addr.clone();
+        let rp = Arc::new(move |si: &ServerInfo| si.address() == &a);
+        let options = CountOptions::builder()
+            .selection_criteria(SelectionCriteria::Predicate(rp))
+            .read_concern(ReadConcern::local())
+            .build();
+
+        while coll.count_documents(None, options.clone()).await.unwrap() != 3 {}
+    }
+
     for read_pref in read_preferences {
         run_test(&client, &coll, read_pref).await;
     }
diff --git a/src/cmap/worker.rs b/src/cmap/worker.rs
@@ -300,8 +300,11 @@ impl ConnectionPoolWorker {
                     } => {
                         self.clear(cause, service_id);
                     }
-                    PoolManagementRequest::MarkAsReady { .. } => {
+                    PoolManagementRequest::MarkAsReady {
+                        _completion_handler,
+                    } => {
                         self.mark_as_ready();
+                        _completion_handler.acknowledge(());
                     }
                     PoolManagementRequest::HandleConnectionSucceeded(conn) => {
                         self.handle_connection_succeeded(conn);
diff --git a/src/lib.rs b/src/lib.rs
@@ -354,7 +354,7 @@ pub(crate) use crate::{
     db::Database,
 };
 
-pub use {coll::Namespace, index::IndexModel, client::session::ClusterTime};
+pub use {coll::Namespace, index::IndexModel, client::session::ClusterTime, sdam::public::*};
 
 #[cfg(all(
     feature = "tokio-runtime",
diff --git a/src/sdam/description/server.rs b/src/sdam/description/server.rs
@@ -14,22 +14,41 @@ const DRIVER_MIN_DB_VERSION: &str = "3.6";
 const DRIVER_MIN_WIRE_VERSION: i32 = 6;
 const DRIVER_MAX_WIRE_VERSION: i32 = 14;
 
+/// Enum representing the possible types of servers that the driver can connect to.
 #[derive(Debug, Deserialize, Clone, Copy, Eq, PartialEq)]
 #[non_exhaustive]
 pub enum ServerType {
+    /// A single, non-replica set mongod.
     Standalone,
+
+    /// A router used in sharded deployments.
     Mongos,
+
+    /// The primary node in a replica set.
     #[serde(rename = "RSPrimary")]
     RsPrimary,
+
+    /// A secondary node in a replica set.
     #[serde(rename = "RSSecondary")]
     RsSecondary,
+
+    /// A non-data bearing node in a replica set which can participate in elections.
     #[serde(rename = "RSArbiter")]
     RsArbiter,
+
+    /// Hidden, starting up, or recovering nodes in a replica set.
     #[serde(rename = "RSOther")]
     RsOther,
+
+    /// A member of an uninitialized replica set or a member that has been removed from the replica
+    /// set config.
     #[serde(rename = "RSGhost")]
     RsGhost,
+
+    /// A load-balancing proxy between the driver and the MongoDB deployment.
     LoadBalancer,
+
+    /// A server that the driver hasn't yet communicated with or can't connect to.
     #[serde(alias = "PossiblePrimary")]
     Unknown,
 }
diff --git a/src/sdam/description/topology/server_selection/mod.rs b/src/sdam/description/topology/server_selection/mod.rs
@@ -119,7 +119,7 @@ impl TopologyDescription {
             SelectionCriteria::Predicate(ref filter) => self
                 .servers
                 .values()
-                .filter(|s| filter(&ServerInfo::new_borrowed(s)))
+                .filter(|s| s.server_type.is_data_bearing() && filter(&ServerInfo::new_borrowed(s)))
                 .collect(),
         };
 
diff --git a/src/sdam/description/topology/server_selection/test/in_window.rs b/src/sdam/description/topology/server_selection/test/in_window.rs
@@ -65,8 +65,7 @@ async fn run_test(test_file: TestFile) {
 
     let topology_description = test_file
         .topology_description
-        .into_topology_description(None)
-        .unwrap();
+        .into_topology_description(None);
 
     let read_pref = ReadPreference::Nearest {
         options: Default::default(),
diff --git a/src/sdam/description/topology/server_selection/test/logic.rs b/src/sdam/description/topology/server_selection/test/logic.rs
@@ -66,13 +66,9 @@ async fn run_test(test_file: TestFile) {
         None => return,
     };
 
-    let topology = match test_file
+    let topology = test_file
         .topology_description
-        .into_topology_description(test_file.heartbeat_frequency_ms.map(Duration::from_millis))
-    {
-        Some(t) => t,
-        None => return,
-    };
+        .into_topology_description(test_file.heartbeat_frequency_ms.map(Duration::from_millis));
 
     if let Some(ref expected_suitable_servers) = test_file.suitable_servers {
         let mut actual_servers: Vec<_> = topology.suitable_servers(&read_pref).unwrap();
diff --git a/src/sdam/description/topology/server_selection/test/mod.rs b/src/sdam/description/topology/server_selection/test/mod.rs
@@ -1,4 +1,4 @@
-use std::{collections::HashMap, time::Duration};
+use std::{collections::HashMap, sync::Arc, time::Duration};
 
 use serde::Deserialize;
 
@@ -12,7 +12,7 @@ use crate::{
         ServerType,
         TopologyDescription,
     },
-    selection_criteria::TagSet,
+    selection_criteria::{SelectionCriteria, TagSet},
 };
 
 mod in_window;
@@ -29,7 +29,7 @@ impl TestTopologyDescription {
     fn into_topology_description(
         self,
         heartbeat_frequency: Option<Duration>,
-    ) -> Option<TopologyDescription> {
+    ) -> TopologyDescription {
         let servers: HashMap<ServerAddress, ServerDescription> = self
             .servers
             .into_iter()
@@ -53,7 +53,6 @@ impl TestTopologyDescription {
             heartbeat_freq: heartbeat_frequency,
             servers,
         }
-        .into()
     }
 }
 
@@ -184,3 +183,66 @@ fn is_master_response_from_server_type(server_type: ServerType) -> Option<IsMast
 
     Some(response)
 }
+
+#[test]
+fn predicate_omits_unavailable() {
+    let criteria = SelectionCriteria::Predicate(Arc::new(|si| {
+        !matches!(si.server_type(), ServerType::RsPrimary)
+    }));
+
+    let desc = TestTopologyDescription {
+        topology_type: TopologyType::ReplicaSetWithPrimary,
+        servers: vec![
+            TestServerDescription {
+                address: "localhost:27017".to_string(),
+                avg_rtt_ms: Some(12.0),
+                server_type: TestServerType::RsPrimary,
+                tags: None,
+                last_update_time: None,
+                last_write: None,
+                _max_wire_version: None,
+            },
+            TestServerDescription {
+                address: "localhost:27018".to_string(),
+                avg_rtt_ms: Some(12.0),
+                server_type: TestServerType::Unknown,
+                tags: None,
+                last_update_time: None,
+                last_write: None,
+                _max_wire_version: None,
+            },
+            TestServerDescription {
+                address: "localhost:27019".to_string(),
+                avg_rtt_ms: Some(12.0),
+                server_type: TestServerType::RsArbiter,
+                tags: None,
+                last_update_time: None,
+                last_write: None,
+                _max_wire_version: None,
+            },
+            TestServerDescription {
+                address: "localhost:27020".to_string(),
+                avg_rtt_ms: Some(12.0),
+                server_type: TestServerType::RsGhost,
+                tags: None,
+                last_update_time: None,
+                last_write: None,
+                _max_wire_version: None,
+            },
+            TestServerDescription {
+                address: "localhost:27021".to_string(),
+                avg_rtt_ms: Some(12.0),
+                server_type: TestServerType::RsOther,
+                tags: None,
+                last_update_time: None,
+                last_write: None,
+                _max_wire_version: None,
+            },
+        ],
+    }
+    .into_topology_description(None);
+    pretty_assertions::assert_eq!(
+        desc.suitable_servers_in_latency_window(&criteria).unwrap(),
+        Vec::<&ServerDescription>::new()
+    );
+}

Original file line number	Diff line number	Diff line change
`@@ -367,6 +367,11 @@ impl Client {`
`367`	`367`	`}`
`368`	`368`	`}`
`369`	`369`
	`370`	`+ #[cfg(test)]`
	`371`	`+ pub(crate) async fn topology_description(&self) -> crate::sdam::TopologyDescription {`
	`372`	`+ self.inner.topology.description().await`
	`373`	`+ }`
	`374`	`+`
`370`	`375`	`#[cfg(all(test, not(feature = "sync")))]`
`371`	`376`	`pub(crate) async fn get_hosts(&self) -> Vec<String> {`
`372`	`377`	`let servers = self.inner.topology.servers().await;`
Original file line number	Diff line number	Diff line change
`@@ -300,8 +300,11 @@ impl ConnectionPoolWorker {`
`300`	`300`	`} => {`
`301`	`301`	`self.clear(cause, service_id);`
`302`	`302`	`}`
`303`		`- PoolManagementRequest::MarkAsReady { .. } => {`
	`303`	`+ PoolManagementRequest::MarkAsReady {`
	`304`	`+ _completion_handler,`
	`305`	`+ } => {`
`304`	`306`	`self.mark_as_ready();`
	`307`	`+ _completion_handler.acknowledge(());`
`305`	`308`	`}`
`306`	`309`	`PoolManagementRequest::HandleConnectionSucceeded(conn) => {`
`307`	`310`	`self.handle_connection_succeeded(conn);`