RUST-360 Streaming monitoring protocol

src/client/executor.rs

@@ -601,6 +601,7 @@ impl Client {
         let raw_cmd = RawCommand {
             name: cmd_name.clone(),
             target_db,
+            exhaust_allowed: false,
             bytes: serialized,
         };
 

src/client/mod.rs

@@ -409,7 +409,7 @@ impl Client {
             if let Some(server) = server_selection::attempt_to_select_server(
                 criteria,
                 &state.description,
-                &state.servers,
+                &state.servers(),
             )? {
                 return Ok(server);
             }
@@ -432,13 +432,14 @@ impl Client {
         }
     }
 
-    #[cfg(all(test, not(feature = "sync"), not(feature = "tokio-sync")))]
-    pub(crate) async fn get_hosts(&self) -> Vec<String> {
+    #[cfg(test)]
+    pub(crate) fn get_hosts(&self) -> Vec<String> {
         let watcher = self.inner.topology.watch();
         let state = watcher.peek_latest();
 
-        let servers = state.servers.keys();
-        servers
+        state
+            .servers()
+            .keys()
             .map(|stream_address| format!("{}", stream_address))
             .collect()
     }
@@ -458,6 +459,11 @@ impl Client {
             .clone()
     }
 
+    #[cfg(test)]
+    pub(crate) fn topology(&self) -> &crate::sdam::Topology {
+        &self.inner.topology
+    }
+
     #[cfg(feature = "csfle")]
     pub(crate) fn weak(&self) -> WeakClient {
         WeakClient {
@@ -475,11 +481,6 @@ impl Client {
         })
         .ok()
     }
-
-    #[cfg(test)]
-    pub(crate) fn topology(&self) -> &Topology {
-        &self.inner.topology
-    }
 }
 
 #[cfg(feature = "csfle")]

src/cmap/conn/command.rs

@@ -20,6 +20,8 @@ use crate::{
 pub(crate) struct RawCommand {
     pub(crate) name: String,
     pub(crate) target_db: String,
+    /// Whether or not the server may respond to this command multiple times via the moreToComeBit.
+    pub(crate) exhaust_allowed: bool,
     pub(crate) bytes: Vec<u8>,
 }
 
@@ -38,6 +40,9 @@ pub(crate) struct Command<T = Document> {
     #[serde(skip)]
     pub(crate) name: String,
 
+    #[serde(skip)]
+    pub(crate) exhaust_allowed: bool,
+
     #[serde(flatten)]
     pub(crate) body: T,
 
@@ -71,6 +76,7 @@ impl<T> Command<T> {
         Self {
             name,
             target_db,
+            exhaust_allowed: false,
             body,
             lsid: None,
             cluster_time: None,
@@ -93,6 +99,7 @@ impl<T> Command<T> {
         Self {
             name,
             target_db,
+            exhaust_allowed: false,
             body,
             lsid: None,
             cluster_time: None,
@@ -244,7 +251,7 @@ impl RawCommandResponse {
             .map_err(|_| Error::invalid_authentication_response(mechanism_name))
     }
 
-    pub(crate) fn into_hello_reply(self, round_trip_time: Duration) -> Result<HelloReply> {
+    pub(crate) fn into_hello_reply(self, round_trip_time: Option<Duration>) -> Result<HelloReply> {
         match self.body::<CommandResponse<HelloCommandResponse>>() {
             Ok(response) if response.is_success() => {
                 let server_address = self.source_address().clone();

src/cmap/conn/mod.rs

@@ -11,7 +11,7 @@ use derivative::Derivative;
 use serde::Serialize;
 use tokio::sync::{mpsc, Mutex};
 
-use self::wire::Message;
+use self::wire::{Message, MessageFlags};
 use super::manager::PoolManager;
 use crate::{
     bson::oid::ObjectId,
@@ -86,6 +86,12 @@ pub(crate) struct Connection {
     /// into a pool.
     error: bool,
 
+    /// Whether the most recently received message included the moreToCome flag, indicating the
+    /// server may send more responses without any additional requests. Attempting to send new
+    /// messages on this connection while this value is true will return an error. This value
+    /// will remain true until a server response does not include the moreToComeFlag.
+    more_to_come: bool,
+
     stream: AsyncStream,
 
     /// Compressor that the client will use before sending messages.
@@ -131,6 +137,7 @@ impl Connection {
             error: false,
             pinned_sender: None,
             compressor: None,
+            more_to_come: false,
         };
 
         Ok(conn)
@@ -156,15 +163,14 @@ impl Connection {
     /// Construct and connect a new connection used for monitoring.
     pub(crate) async fn connect_monitoring(
         address: ServerAddress,
-        connect_timeout: Option<Duration>,
         tls_options: Option<TlsOptions>,
     ) -> Result<Self> {
         Self::new(
             0,
             address,
             0,
             Some(ConnectionOptions {
-                connect_timeout,
+                connect_timeout: None, // handled by the monitor
                 tls_options,
                 event_handler: None,
             }),
@@ -272,6 +278,13 @@ impl Connection {
         message: Message,
         to_compress: bool,
     ) -> Result<RawCommandResponse> {
+        if self.more_to_come {
+            return Err(Error::internal(format!(
+                "attempted to send a new message to {} but moreToCome bit was set",
+                self.address()
+            )));
+        }
+
         self.command_executing = true;
 
         // If the client has agreed on a compressor with the server, and the command
@@ -298,7 +311,10 @@ impl Connection {
         self.command_executing = false;
         self.error = response_message_result.is_err();
 
-        RawCommandResponse::new(self.address.clone(), response_message_result?)
+        let response_message = response_message_result?;
+        self.more_to_come = response_message.flags.contains(MessageFlags::MORE_TO_COME);
+
+        RawCommandResponse::new(self.address.clone(), response_message)
     }
 
     /// Executes a `Command` and returns a `CommandResponse` containing the result from the server.
@@ -332,6 +348,34 @@ impl Connection {
         self.send_message(message, to_compress).await
     }
 
+    /// Receive the next message from the connection.
+    /// This will return an error if the previous response on this connection did not include the
+    /// moreToCome flag.
+    pub(crate) async fn receive_message(&mut self) -> Result<RawCommandResponse> {
+        if !self.more_to_come {
+            return Err(Error::internal(format!(
+                "attempted to stream response from connection to {} but moreToCome bit was not set",
+                self.address()
+            )));
+        }
+
+        self.command_executing = true;
+        let response_message_result = Message::read_from(
+            &mut self.stream,
+            self.stream_description
+                .as_ref()
+                .map(|d| d.max_message_size_bytes),
+        )
+        .await;
+        self.command_executing = false;
+        self.error = response_message_result.is_err();
+
+        let response_message = response_message_result?;
+        self.more_to_come = response_message.flags.contains(MessageFlags::MORE_TO_COME);
+
+        RawCommandResponse::new(self.address.clone(), response_message)
+    }
+
     /// Gets the connection's StreamDescription.
     pub(crate) fn stream_description(&self) -> Result<&StreamDescription> {
         self.stream_description.as_ref().ok_or_else(|| {
@@ -394,8 +438,15 @@ impl Connection {
             ready_and_available_time: None,
             pinned_sender: self.pinned_sender.clone(),
             compressor: self.compressor.clone(),
+            more_to_come: false,
         }
     }
+
+    /// Whether or not the previous command response indicated that the server may send
+    /// more responses without another request.
+    pub(crate) fn is_streaming(&self) -> bool {
+        self.more_to_come
+    }
 }
 
 impl Drop for Connection {

src/cmap/conn/wire/message.rs

@@ -37,6 +37,7 @@ impl Message {
                 bytes,
                 target_db: command.target_db,
                 name: command.name,
+                exhaust_allowed: command.exhaust_allowed,
             },
             request_id,
         ))
@@ -46,9 +47,14 @@ impl Message {
     ///
     /// Note that `response_to` will need to be set manually.
     pub(crate) fn with_raw_command(command: RawCommand, request_id: Option<i32>) -> Self {
+        let mut flags = MessageFlags::empty();
+        if command.exhaust_allowed {
+            flags |= MessageFlags::EXHAUST_ALLOWED;
+        }
+
         Self {
             response_to: 0,
-            flags: MessageFlags::empty(),
+            flags,
             sections: vec![MessageSection::Document(command.bytes)],
             checksum: None,
             request_id,

src/cmap/conn/wire/mod.rs

@@ -2,4 +2,7 @@ mod header;
 mod message;
 mod util;
 
-pub(crate) use self::{message::Message, util::next_request_id};
+pub(crate) use self::{
+    message::{Message, MessageFlags},
+    util::next_request_id,
+};

src/cmap/establish/handshake/mod.rs

@@ -1,8 +1,6 @@
 #[cfg(test)]
 mod test;
 
-use std::sync::Arc;
-
 use lazy_static::lazy_static;
 use os_info::{Type, Version};
 
@@ -12,10 +10,8 @@ use crate::{
     cmap::{options::ConnectionPoolOptions, Command, Connection, StreamDescription},
     compression::Compressor,
     error::{ErrorKind, Result},
-    event::sdam::SdamEventHandler,
     hello::{hello_command, run_hello, HelloReply},
     options::{AuthMechanism, ClientOptions, Credential, DriverInfo, ServerApi},
-    sdam::Topology,
 };
 
 #[cfg(all(feature = "tokio-runtime", not(feature = "tokio-sync")))]
@@ -158,6 +154,7 @@ impl Handshaker {
             options.as_ref().and_then(|opts| opts.server_api.as_ref()),
             options.as_ref().and_then(|opts| opts.load_balanced.into()),
             None,
+            None,
         );
 
         if let Some(options) = options {
@@ -214,17 +211,12 @@ impl Handshaker {
     }
 
     /// Handshakes a connection.
-    pub(crate) async fn handshake(
-        &self,
-        conn: &mut Connection,
-        topology: Option<&Topology>,
-        handler: &Option<Arc<dyn SdamEventHandler>>,
-    ) -> Result<HandshakeResult> {
+    pub(crate) async fn handshake(&self, conn: &mut Connection) -> Result<HandshakeResult> {
         let mut command = self.command.clone();
 
         let client_first = set_speculative_auth_info(&mut command.body, self.credential.as_ref())?;
 
-        let mut hello_reply = run_hello(conn, command, topology, handler).await?;
+        let mut hello_reply = run_hello(conn, command).await?;
 
         if self.command.body.contains_key("loadBalanced")
             && hello_reply.command_response.service_id.is_none()

src/cmap/establish/mod.rs

@@ -52,7 +52,7 @@ impl ConnectionEstablisher {
 
         let handshake = self
             .handshaker
-            .handshake(&mut connection, None, &None)
+            .handshake(&mut connection)
             .await
             .map_err(|e| EstablishError::pre_hello(e, pool_gen.clone()))?;
         let service_id = handshake.hello_reply.command_response.service_id;

src/cmap/establish/test.rs

@@ -62,11 +62,7 @@ async fn speculative_auth_test(
         .await
         .unwrap();
 
-    let first_round = handshaker
-        .handshake(&mut conn, None, &None)
-        .await
-        .unwrap()
-        .first_round;
+    let first_round = handshaker.handshake(&mut conn).await.unwrap().first_round;
 
     // We expect that the server will return a response with the `speculativeAuthenticate` field if
     // and only if it's new enough to support it.

src/cmap/manager.rs

@@ -38,17 +38,11 @@ impl PoolManager {
     }
 
     /// Mark the pool as "ready" as per the CMAP specification.
-    pub(super) async fn mark_as_ready(&self) {
-        let (message, acknowledgment_receiver) = AcknowledgedMessage::package(());
-        if self
-            .sender
-            .send(PoolManagementRequest::MarkAsReady {
-                _completion_handler: message,
-            })
-            .is_ok()
-        {
-            acknowledgment_receiver.wait_for_acknowledgment().await;
-        }
+    ///
+    /// Since management requests are treated with the highest priority by the pool and will be
+    /// processed before all others, there's no need to wait for acknowledgment.
+    pub(super) fn mark_as_ready(&self) {
+        let _ = self.sender.send(PoolManagementRequest::MarkAsReady);
     }
 
     /// Check in the given connection to the pool.
@@ -108,9 +102,7 @@ pub(super) enum PoolManagementRequest {
     },
 
     /// Mark the pool as Ready, allowing connections to be created and checked out.
-    MarkAsReady {
-        _completion_handler: AcknowledgedMessage<()>,
-    },
+    MarkAsReady,
 
     /// Check in the given connection.
     CheckIn(Box<Connection>),

src/cmap/mod.rs

@@ -46,8 +46,9 @@ use crate::runtime::WorkerHandle;
 
 const DEFAULT_MAX_POOL_SIZE: u32 = 10;
 
-/// A pool of connections implementing the CMAP spec. All state is kept internally in an `Arc`, and
-/// internal state that is mutable is additionally wrapped by a lock.
+/// A pool of connections implementing the CMAP spec.
+/// This type is actually a handle to task that manages the connections and is cheap to clone and
+/// pass around.
 #[derive(Clone, Derivative)]
 #[derivative(Debug)]
 pub(crate) struct ConnectionPool {
@@ -166,9 +167,9 @@ impl ConnectionPool {
         self.manager.clear(cause, service_id).await
     }
 
-    /// Mark the pool as "ready", allowing connections to be created and checked out.
-    pub(crate) async fn mark_as_ready(&self) {
-        self.manager.mark_as_ready().await;
+    /// Mark the pool as "ready" as per the CMAP specification.
+    pub(crate) fn mark_as_ready(&self) {
+        self.manager.mark_as_ready()
     }
 
     pub(crate) fn generation(&self) -> PoolGeneration {

src/cmap/test/mod.rs

@@ -297,7 +297,7 @@ impl Operation {
             }
             Operation::Ready => {
                 if let Some(pool) = state.pool.read().await.deref() {
-                    pool.mark_as_ready().await;
+                    pool.mark_as_ready();
                 }
             }
             Operation::Close => {