RUST-981 Merge new load balancer spec tests (#480)

Author: abr-egn

src/client/executor.rs

@@ -7,7 +7,13 @@ use std::{collections::HashSet, sync::Arc, time::Instant};
 use super::{session::TransactionState, Client, ClientSession};
 use crate::{
     bson::Document,
-    cmap::{conn::PinnedConnectionHandle, Connection, RawCommand, RawCommandResponse},
+    cmap::{
+        conn::PinnedConnectionHandle,
+        Connection,
+        ConnectionPool,
+        RawCommand,
+        RawCommandResponse,
+    },
     cursor::{session::SessionCursor, Cursor, CursorSpecification},
     error::{
         Error,
@@ -158,11 +164,13 @@ impl Client {
         Op: Operation<O = CursorSpecification<T>>,
         T: DeserializeOwned + Unpin + Send + Sync,
     {
-        let mut details = self.execute_operation_with_details(op, session).await?;
-        let pinned = if details.output.connection.is_pinned() {
-            // Cursor operations on load-balanced transactions will be pinned via the transaction
-            // pin.
-            None
+        let mut details = self
+            .execute_operation_with_details(op, &mut *session)
+            .await?;
+
+        let pinned = if let Some(handle) = session.transaction.pinned_connection() {
+            // Cursor operations on a transaction share the same pinned connection.
+            Some(handle.replicate())
         } else {
             self.pin_connection_for_cursor(&mut details.output)?
         };
@@ -225,20 +233,17 @@ impl Client {
             }
         };
 
-        let mut conn = match op.pinned_connection() {
-            Some(l) => l.take_connection().await?,
-            None => match server.pool.check_out().await {
-                Ok(c) => c,
-                Err(mut err) => {
-                    err.add_labels_and_update_pin(None, &mut session, None)?;
+        let mut conn = match get_connection(&session, &op, &server.pool).await {
+            Ok(c) => c,
+            Err(mut err) => {
+                err.add_labels_and_update_pin(None, &mut session, None)?;
 
-                    if err.is_pool_cleared() {
-                        return self.execute_retry(&mut op, &mut session, None, err).await;
-                    } else {
-                        return Err(err);
-                    }
+                if err.is_pool_cleared() {
+                    return self.execute_retry(&mut op, &mut session, None, err).await;
+                } else {
+                    return Err(err);
                 }
-            },
+            }
         };
 
         let retryability = self.get_retryability(&conn, &op, &session).await?;
@@ -326,23 +331,9 @@ impl Client {
             }
         };
 
-        let session_pinned = session
-            .as_ref()
-            .and_then(|s| s.transaction.pinned_connection());
-        let mut conn = match (session_pinned, op.pinned_connection()) {
-            (Some(c), None) | (None, Some(c)) => c.take_connection().await?,
-            (Some(c), Some(_)) => {
-                // An operation executing in a transaction should never have a pinned connection,
-                // but in case it does, prefer the transaction's pin.
-                if cfg!(debug_assertions) {
-                    panic!("pinned operation executing in pinned transaction");
-                }
-                c.take_connection().await?
-            }
-            (None, None) => match server.pool.check_out().await {
-                Ok(c) => c,
-                Err(_) => return Err(first_error),
-            },
+        let mut conn = match get_connection(session, op, &server.pool).await {
+            Ok(c) => c,
+            Err(_) => return Err(first_error),
         };
 
         let retryability = self.get_retryability(&conn, op, session).await?;
@@ -776,6 +767,25 @@ impl Client {
     }
 }
 
+async fn get_connection<T: Operation>(
+    session: &Option<&mut ClientSession>,
+    op: &T,
+    pool: &ConnectionPool,
+) -> Result<Connection> {
+    let session_pinned = session
+        .as_ref()
+        .and_then(|s| s.transaction.pinned_connection());
+    match (session_pinned, op.pinned_connection()) {
+        (Some(c), None) | (None, Some(c)) => c.take_connection().await,
+        (Some(session_handle), Some(op_handle)) => {
+            // An operation executing in a transaction should be sharing the same pinned connection.
+            debug_assert_eq!(session_handle.id(), op_handle.id());
+            session_handle.take_connection().await
+        }
+        (None, None) => pool.check_out().await,
+    }
+}
+
 impl Error {
     /// Adds the necessary labels to this Error, and unpins the session if needed.
     ///

src/cmap/conn/mod.rs

@@ -345,11 +345,6 @@ impl Connection {
         })
     }
 
-    /// Whether this connection has a live `PinnedConnectionHandle`.
-    pub(crate) fn is_pinned(&self) -> bool {
-        self.pinned_sender.is_some()
-    }
-
     /// Close this connection, emitting a `ConnectionClosedEvent` with the supplied reason.
     pub(super) fn close_and_drop(mut self, reason: ConnectionClosedReason) {
         self.close(reason);
@@ -470,12 +465,8 @@ impl PinnedConnectionHandle {
         })
     }
 
-    /// Return the pinned connection to the normal connection pool.
-    pub(crate) async fn unpin_connection(&self) {
-        let mut receiver = self.receiver.lock().await;
-        receiver.close();
-        // Ensure any connections buffered in the channel are dropped, returning them to the pool.
-        while receiver.recv().await.is_some() {}
+    pub(crate) fn id(&self) -> u32 {
+        self.id
     }
 }
 

src/cmap/test/mod.rs

@@ -39,6 +39,8 @@ const TEST_DESCRIPTIONS_TO_SKIP: &[&str] = &[
     // WaitQueueTimeoutMS is not supported
     "must aggressively timeout threads enqueued longer than waitQueueTimeoutMS",
     "waiting on maxConnecting is limited by WaitQueueTimeoutMS",
+    // TODO DRIVERS-1785 remove this skip when test event order is fixed
+    "error during minPoolSize population clears pool",
 ];
 
 /// Many different types of CMAP events are emitted from tasks spawned in the drop
@@ -165,9 +167,12 @@ impl Executor {
         let manager = pool.manager.clone();
         RUNTIME.execute(async move {
             while let Some(update) = update_receiver.recv().await {
-                match update.into_message() {
-                    ServerUpdate::Error { error, .. } => manager.clear(error.cause, None).await,
+                match update.message() {
+                    ServerUpdate::Error { error, .. } => {
+                        manager.clear(error.cause.clone(), None).await
+                    }
                 }
+                drop(update);
             }
         });
 
@@ -426,6 +431,7 @@ async fn cmap_spec_tests() {
 
         let mut options = CLIENT_OPTIONS.clone();
         if options.load_balanced.unwrap_or(false) {
+            println!("skipping due to load balanced topology");
             return;
         }
         options.hosts.drain(1..);
@@ -434,6 +440,7 @@ async fn cmap_spec_tests() {
         if let Some(ref run_on) = test_file.run_on {
             let can_run_on = run_on.iter().any(|run_on| run_on.can_run_on(&client));
             if !can_run_on {
+                println!("skipping due to runOn requirements");
                 return;
             }
         }

src/cmap/worker.rs

@@ -646,6 +646,7 @@ async fn establish_connection(
 
     match establish_result {
         Err(ref e) => {
+            server_updater.handle_error(e.clone()).await;
             if let Some(handler) = event_handler {
                 let event = ConnectionClosedEvent {
                     address,
@@ -654,7 +655,6 @@ async fn establish_connection(
                 };
                 handler.handle_connection_closed_event(event);
             }
-            server_updater.handle_error(e.clone()).await;
             manager.handle_connection_failed();
         }
         Ok(ref mut connection) => {

src/coll/options.rs

@@ -920,7 +920,7 @@ pub struct ListIndexesOptions {
     pub max_time: Option<Duration>,
 
     /// The number of indexes the server should return per cursor batch.
-    #[serde(default, serialize_with = "bson_util::serialize_u32_option_as_i32")]
+    #[serde(default, skip_serializing)]
     pub batch_size: Option<u32>,
 }
 

src/cursor/common.rs

@@ -283,12 +283,6 @@ impl PinnedConnection {
         matches!(self, Self::Invalid(_))
     }
 
-    async fn unpin(&self) {
-        if let Some(h) = self.handle() {
-            h.unpin_connection().await;
-        }
-    }
-
     fn invalidate(&mut self) {
         take_mut::take(self, |self_| {
             if let Self::Valid(c) = self_ {
@@ -318,6 +312,5 @@ pub(super) fn kill_cursor(
                 let _ = tx.send(());
             }
         }
-        pinned_conn.unpin().await;
     });
 }

src/operation/list_indexes/mod.rs

@@ -48,6 +48,9 @@ impl Operation for ListIndexes {
         let mut body = doc! {
             "listIndexes": self.ns.coll.clone(),
         };
+        if let Some(size) = self.options.as_ref().and_then(|o| o.batch_size) {
+            body.insert("cursor", doc! { "batchSize": size });
+        }
         append_options(&mut body, self.options.as_ref())?;
 
         Ok(Command::new(

src/operation/list_indexes/test.rs

@@ -33,7 +33,9 @@ async fn build() {
         doc! {
             "listIndexes": "test_coll",
             "maxTimeMS": 42,
-            "batchSize": 4,
+            "cursor": doc! {
+                "batchSize": 4,
+            },
         }
     );
 }

src/runtime/acknowledged_message.rs

@@ -20,9 +20,9 @@ impl<M, R> AcknowledgedMessage<M, R> {
         )
     }
 
-    /// Get the message.
-    pub(crate) fn into_message(self) -> M {
-        self.message
+    /// Borrow the message.
+    pub(crate) fn message(&self) -> &M {
+        &self.message
     }
 
     /// Send acknowledgement to the receiver.

src/sdam/description/topology/mod.rs

@@ -261,12 +261,8 @@ impl TopologyDescription {
         match read_preference {
             ReadPreference::Secondary { .. }
             | ReadPreference::PrimaryPreferred { .. }
-            | ReadPreference::Nearest { .. } => {
-                command.set_read_preference(read_preference.clone())
-            }
-            ReadPreference::SecondaryPreferred { ref options }
-                if options.max_staleness.is_some() || options.tag_sets.is_some() =>
-            {
+            | ReadPreference::Nearest { .. }
+            | ReadPreference::SecondaryPreferred { .. } => {
                 command.set_read_preference(read_preference.clone())
             }
             _ => {}