RUST-1505 3.0 event API: command events (#1004)

Author: abr-egn

src/client.rs

@@ -32,7 +32,7 @@ use crate::{
     concern::{ReadConcern, WriteConcern},
     db::Database,
     error::{Error, ErrorKind, Result},
-    event::command::{handle_command_event, CommandEvent},
+    event::command::CommandEvent,
     id_set::IdSet,
     options::{
         ClientOptions,
@@ -44,6 +44,7 @@ use crate::{
     },
     sdam::{server_selection, SelectedServer, Topology},
     tracking_arc::TrackingArc,
+    BoxFuture,
     ClientSession,
 };
 
@@ -257,27 +258,8 @@ impl Client {
         }
     }
 
-    #[cfg(not(feature = "tracing-unstable"))]
-    pub(crate) async fn emit_command_event(&self, generate_event: impl FnOnce() -> CommandEvent) {
-        let handler = self.inner.options.command_event_handler.as_ref();
-        let test_channel = self.test_command_event_channel();
-        if handler.is_none() && test_channel.is_none() {
-            return;
-        }
-
-        let event = generate_event();
-        if let Some(tx) = test_channel {
-            let (msg, ack) = crate::runtime::AcknowledgedMessage::package(event.clone());
-            let _ = tx.send(msg).await;
-            ack.wait_for_acknowledgment().await;
-        }
-        if let Some(handler) = handler {
-            handle_command_event(handler.as_ref(), event);
-        }
-    }
-
-    #[cfg(feature = "tracing-unstable")]
     pub(crate) async fn emit_command_event(&self, generate_event: impl FnOnce() -> CommandEvent) {
+        #[cfg(feature = "tracing-unstable")]
         let tracing_emitter = if trace_or_log_enabled!(
             target: COMMAND_TRACING_EVENT_TARGET,
             TracingOrLogLevel::Debug
@@ -289,9 +271,11 @@ impl Client {
         } else {
             None
         };
-        let apm_event_handler = self.inner.options.command_event_handler.as_ref();
         let test_channel = self.test_command_event_channel();
-        if !(tracing_emitter.is_some() || apm_event_handler.is_some() || test_channel.is_some()) {
+        let should_send = test_channel.is_some() || self.options().command_event_handler.is_some();
+        #[cfg(feature = "tracing-unstable")]
+        let should_send = should_send || tracing_emitter.is_some();
+        if !should_send {
             return;
         }
 
@@ -301,15 +285,12 @@ impl Client {
             let _ = tx.send(msg).await;
             ack.wait_for_acknowledgment().await;
         }
-        if let (Some(event_handler), Some(ref tracing_emitter)) =
-            (apm_event_handler, &tracing_emitter)
-        {
-            handle_command_event(event_handler.as_ref(), event.clone());
-            handle_command_event(tracing_emitter, event);
-        } else if let Some(event_handler) = apm_event_handler {
-            handle_command_event(event_handler.as_ref(), event);
-        } else if let Some(ref tracing_emitter) = tracing_emitter {
-            handle_command_event(tracing_emitter, event);
+        #[cfg(feature = "tracing-unstable")]
+        if let Some(ref tracing_emitter) = tracing_emitter {
+            tracing_emitter.handle(event.clone());
+        }
+        if let Some(handler) = &self.options().command_event_handler {
+            handler.handle(event);
         }
     }
 
@@ -693,7 +674,6 @@ impl Client {
         .ok()
     }
 
-    #[cfg(test)]
     pub(crate) fn options(&self) -> &ClientOptions {
         &self.inner.options
     }
@@ -731,6 +711,3 @@ impl AsyncDropToken {
         Self { tx: self.tx.take() }
     }
 }
-
-// TODO: merge this with other BoxFuture defs
-pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;

src/client/auth/oidc.rs

@@ -17,6 +17,7 @@ use crate::{
     },
     cmap::Connection,
     error::{Error, Result},
+    BoxFuture,
 };
 
 use super::{sasl::SaslContinue, Credential, MONGODB_OIDC_STR};
@@ -27,8 +28,6 @@ pub struct Callbacks {
     inner: Arc<CallbacksInner>,
 }
 
-pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;
-
 impl Callbacks {
     /// Create a new instance with a token request callback.
     pub fn new<F>(on_request: F) -> Self

src/client/options.rs

@@ -32,7 +32,7 @@ use crate::{
     compression::Compressor,
     concern::{Acknowledgment, ReadConcern, WriteConcern},
     error::{Error, ErrorKind, Result},
-    event::{cmap::CmapEventHandler, command::CommandEventHandler, sdam::SdamEventHandler},
+    event::{cmap::CmapEventHandler, sdam::SdamEventHandler},
     options::ReadConcernLevel,
     sdam::{verify_max_staleness, DEFAULT_HEARTBEAT_FREQUENCY, MIN_HEARTBEAT_FREQUENCY},
     selection_criteria::{ReadPreference, SelectionCriteria, TagSet},
@@ -416,9 +416,10 @@ pub struct ClientOptions {
     ///
     /// Note that monitoring command events may incur a performance penalty.
     #[derivative(Debug = "ignore", PartialEq = "ignore")]
-    #[builder(default)]
+    #[builder(default, setter(strip_option))]
     #[serde(skip)]
-    pub command_event_handler: Option<Arc<dyn CommandEventHandler>>,
+    pub command_event_handler:
+        Option<crate::event::EventHandler<crate::event::command::CommandEvent>>,
 
     /// The connect timeout passed to each underlying TcpStream when attemtping to connect to the
     /// server.

src/client/session.rs

@@ -20,6 +20,7 @@ use crate::{
     options::{SessionOptions, TransactionOptions},
     sdam::{ServerInfo, TransactionSupportStatus},
     selection_criteria::SelectionCriteria,
+    BoxFuture,
     Client,
 };
 pub use cluster_time::ClusterTime;
@@ -691,8 +692,6 @@ impl ClientSession {
     }
 }
 
-pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;
-
 struct DroppedClientSession {
     cluster_time: Option<ClusterTime>,
     server_session: ServerSession,

src/client/session/test.rs

@@ -305,7 +305,7 @@ async fn cluster_time_in_commands() {
     let handler = Arc::new(EventHandler::new());
     let mut options = get_client_options().await.clone();
     options.heartbeat_freq = Some(Duration::from_secs(1000));
-    options.command_event_handler = Some(handler.clone());
+    options.command_event_handler = Some(handler.clone().into());
     options.sdam_event_handler = Some(handler.clone());
 
     // Ensure we only connect to one server so the monitor checks from other servers

src/event.rs

@@ -3,3 +3,102 @@
 pub mod cmap;
 pub mod command;
 pub mod sdam;
+
+use std::sync::Arc;
+
+use futures_core::future::BoxFuture;
+
+use crate::event::command::CommandEvent;
+
+/// A destination for events.  Allows implicit conversion via [`From`] for concrete types for
+/// convenience with [`crate::options::ClientOptions`] construction:
+///
+/// ```rust
+/// # use mongodb::options::ClientOptions;
+/// # fn example() {
+/// let (tx, mut rx) = tokio::sync::mpsc::channel(100);
+/// tokio::spawn(async move {
+///     while let Some(ev) = rx.recv().await {
+///         println!("{:?}", ev);
+///     }
+/// });
+/// let options = ClientOptions::builder()
+///                 .command_event_handler(tx)
+///                 .build();
+/// # }
+/// ```
+///
+/// or explicit construction for `Fn` traits:
+///
+/// ```rust
+/// # use mongodb::options::ClientOptions;
+/// # use mongodb::event::EventHandler;
+/// # fn example() {
+/// let options = ClientOptions::builder()
+///                 .command_event_handler(EventHandler::callback(|ev| println!("{:?}", ev)))
+///                 .build();
+/// # }
+/// ```
+#[derive(Clone)]
+#[non_exhaustive]
+pub enum EventHandler<T> {
+    /// A callback.
+    Callback(Arc<dyn Fn(T) + Sync + Send>),
+    /// An async callback.
+    AsyncCallback(Arc<dyn Fn(T) -> BoxFuture<'static, ()> + Sync + Send>),
+    /// A `tokio` channel sender.
+    TokioMpsc(tokio::sync::mpsc::Sender<T>),
+}
+
+impl<T> std::fmt::Debug for EventHandler<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("EventHandler").finish()
+    }
+}
+
+impl<T> From<tokio::sync::mpsc::Sender<T>> for EventHandler<T> {
+    fn from(value: tokio::sync::mpsc::Sender<T>) -> Self {
+        Self::TokioMpsc(value)
+    }
+}
+
+#[allow(deprecated)]
+impl<T: crate::event::command::CommandEventHandler + 'static> From<Arc<T>>
+    for EventHandler<CommandEvent>
+{
+    fn from(value: Arc<T>) -> Self {
+        Self::callback(move |ev| match ev {
+            CommandEvent::Started(e) => value.handle_command_started_event(e),
+            CommandEvent::Succeeded(e) => value.handle_command_succeeded_event(e),
+            CommandEvent::Failed(e) => value.handle_command_failed_event(e),
+        })
+    }
+}
+
+impl<T: Send + Sync + 'static> EventHandler<T> {
+    /// Construct a new event handler with a callback.
+    pub fn callback(f: impl Fn(T) + Send + Sync + 'static) -> Self {
+        Self::Callback(Arc::new(f))
+    }
+
+    /// Construct a new event handler with an async callback.
+    pub fn async_callback(f: impl Fn(T) -> BoxFuture<'static, ()> + Send + Sync + 'static) -> Self {
+        Self::AsyncCallback(Arc::new(f))
+    }
+
+    pub(crate) fn handle(&self, event: T) {
+        match self {
+            // TODO RUST-1731 Use tokio's spawn_blocking
+            Self::Callback(cb) => (cb)(event),
+            Self::AsyncCallback(cb) => {
+                crate::runtime::spawn((cb)(event));
+            }
+            Self::TokioMpsc(sender) => {
+                let sender = sender.clone();
+                crate::runtime::spawn(async move {
+                    let _ = sender.send(event).await;
+                });
+            }
+        }
+    }
+}

src/event/command.rs

@@ -93,6 +93,9 @@ pub struct CommandFailedEvent {
     pub service_id: Option<ObjectId>,
 }
 
+/// Usage of this trait is deprecated.  Applications should use the simpler
+/// [`EventHandler`](crate::event::EventHandler) API.
+///
 /// Applications can implement this trait to specify custom logic to run on each command event sent
 /// by the driver.
 ///
@@ -121,7 +124,7 @@ pub struct CommandFailedEvent {
 /// }
 ///
 /// # fn do_stuff() -> Result<()> {
-/// let handler: Arc<dyn CommandEventHandler> = Arc::new(FailedCommandLogger);
+/// let handler = Arc::new(FailedCommandLogger);
 /// let options = ClientOptions::builder()
 ///                   .command_event_handler(handler)
 ///                   .build();
@@ -131,6 +134,7 @@ pub struct CommandFailedEvent {
 /// # Ok(())
 /// # }
 /// ```
+#[deprecated]
 pub trait CommandEventHandler: Send + Sync {
     /// A [`Client`](../../struct.Client.html) will call this method on each registered handler
     /// whenever a database command is initiated.
@@ -146,18 +150,11 @@ pub trait CommandEventHandler: Send + Sync {
 }
 
 #[derive(Clone, Debug, Serialize)]
+#[allow(missing_docs)]
 #[serde(untagged)]
-pub(crate) enum CommandEvent {
+#[non_exhaustive]
+pub enum CommandEvent {
     Started(CommandStartedEvent),
     Succeeded(CommandSucceededEvent),
     Failed(CommandFailedEvent),
 }
-
-/// Passes the specified event to the corresponding method on the provided handler.
-pub(crate) fn handle_command_event(handler: &dyn CommandEventHandler, event: CommandEvent) {
-    match event {
-        CommandEvent::Started(event) => handler.handle_command_started_event(event),
-        CommandEvent::Succeeded(event) => handler.handle_command_succeeded_event(event),
-        CommandEvent::Failed(event) => handler.handle_command_failed_event(event),
-    }
-}

src/lib.rs

@@ -363,9 +363,11 @@ pub use crate::{
     gridfs::{GridFsBucket, GridFsDownloadStream, GridFsUploadStream},
 };
 
-pub use {client::session::ClusterTime, coll::Namespace, index::IndexModel, sdam::public::*, search_index::SearchIndexModel};
+pub use {client::action, client::session::ClusterTime, coll::Namespace, index::IndexModel, sdam::public::*, search_index::SearchIndexModel};
+
+/// A boxed future.
+pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;
 
-pub use client::action;
 
 #[cfg(all(feature = "tokio-runtime", feature = "sync",))]
 compile_error!(

src/test/csfle.rs

@@ -32,12 +32,7 @@ use tokio::net::TcpListener;
 use crate::{
     client_encryption::{ClientEncryption, EncryptKey, MasterKey, RangeOptions},
     error::{ErrorKind, WriteError, WriteFailure},
-    event::command::{
-        CommandEventHandler,
-        CommandFailedEvent,
-        CommandStartedEvent,
-        CommandSucceededEvent,
-    },
+    event::command::{CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent},
     options::{
         CollectionOptions,
         CreateCollectionOptions,
@@ -568,7 +563,7 @@ async fn bson_size_limits() -> Result<()> {
     let mut opts = get_client_options().await.clone();
     let handler = Arc::new(EventHandler::new());
     let mut events = handler.subscribe();
-    opts.command_event_handler = Some(handler.clone());
+    opts.command_event_handler = Some(handler.clone().into());
     let client_encrypted =
         Client::encrypted_builder(opts, KV_NAMESPACE.clone(), LOCAL_KMS.clone())?
             .extra_options(EXTRA_OPTIONS.clone())
@@ -1627,7 +1622,7 @@ impl DeadlockTestCase {
         let mut encrypted_events = event_handler.subscribe();
         let mut opts = get_client_options().await.clone();
         opts.max_pool_size = Some(self.max_pool_size);
-        opts.command_event_handler = Some(event_handler.clone());
+        opts.command_event_handler = Some(event_handler.clone().into());
         opts.sdam_event_handler = Some(event_handler.clone());
         let client_encrypted =
             Client::encrypted_builder(opts, KV_NAMESPACE.clone(), LOCAL_KMS.clone())?
@@ -2701,7 +2696,7 @@ impl DecryptionEventsTestdata {
         let ev_handler = DecryptionEventsHandler::new();
         let mut opts = get_client_options().await.clone();
         opts.retry_reads = Some(false);
-        opts.command_event_handler = Some(ev_handler.clone());
+        opts.command_event_handler = Some(ev_handler.clone().into());
         let encrypted_client =
             Client::encrypted_builder(opts, KV_NAMESPACE.clone(), LOCAL_KMS.clone())?
                 .extra_options(EXTRA_OPTIONS.clone())
@@ -2737,7 +2732,8 @@ impl DecryptionEventsHandler {
     }
 }
 
-impl CommandEventHandler for DecryptionEventsHandler {
+#[allow(deprecated)]
+impl crate::event::command::CommandEventHandler for DecryptionEventsHandler {
     fn handle_command_succeeded_event(&self, event: CommandSucceededEvent) {
         if event.command_name == "aggregate" {
             *self.succeeded.lock().unwrap() = Some(event);