Merge pull request #4129 from Ten0/instrumentation_downcast_support

Add support for downcasting Instrumentation

diesel/Cargo.toml

@@ -44,6 +44,7 @@ bitflags = { version = "2.0.0", optional = true }
 r2d2 = { version = ">= 0.8.2, < 0.9.0", optional = true }
 itoa = { version = "1.0.0", optional = true }
 time = { version = "0.3.9", optional = true, features = ["macros"] }
+downcast-rs = "1.2.1"
 
 [dependencies.diesel_derives]
 version = "~2.2.0"

diesel/src/connection/instrumentation.rs

@@ -1,7 +1,7 @@
-use std::fmt::Debug;
-use std::fmt::Display;
+use downcast_rs::Downcast;
+use std::fmt::{Debug, Display};
 use std::num::NonZeroU32;
-use std::ops::DerefMut;
+use std::ops::{Deref, DerefMut};
 
 static GLOBAL_INSTRUMENTATION: std::sync::RwLock<fn() -> Option<Box<dyn Instrumentation>>> =
     std::sync::RwLock::new(|| None);
@@ -242,10 +242,11 @@ impl<'a> InstrumentationEvent<'a> {
 /// More complex usages and integrations with frameworks like
 /// `tracing` and `log` are supposed to be part of their own
 /// crates.
-pub trait Instrumentation: Send + 'static {
+pub trait Instrumentation: Downcast + Send + 'static {
     /// The function that is invoced for each event
     fn on_connection_event(&mut self, event: InstrumentationEvent<'_>);
 }
+downcast_rs::impl_downcast!(Instrumentation);
 
 /// Get an instance of the default [`Instrumentation`]
 ///
@@ -266,9 +267,11 @@ pub fn get_default_instrumentation() -> Option<Box<dyn Instrumentation>> {
 ///
 /// // a simple logger that prints all events to stdout
 /// fn simple_logger() -> Option<Box<dyn Instrumentation>> {
-///    // we need the explicit argument type there due
-///    // to bugs in rustc
-///    Some(Box::new(|event: InstrumentationEvent<'_>| println!("{event:?}")))
+///     // we need the explicit argument type there due
+///     // to bugs in rustc
+///     Some(Box::new(|event: InstrumentationEvent<'_>| {
+///         println!("{event:?}")
+///     }))
 /// }
 ///
 /// set_default_instrumentation(simple_logger);
@@ -313,3 +316,99 @@ where
         }
     }
 }
+
+#[diesel_derives::__diesel_public_if(
+    feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
+)]
+/// An optional dyn instrumentation.
+///
+/// For ease of use, this type implements [`Deref`] and [`DerefMut`] to `&dyn Instrumentation`,
+/// falling back to a no-op implementation if no instrumentation is set.
+///
+/// The DynInstrumentation type is useful because without it we actually did tend to return
+/// (accidentally) &mut Option<Box> as &mut dyn Instrumentation from connection.instrumentation(),
+/// so downcasting would have to be done in these two steps by the user, which is counter-intuitive.
+pub(crate) struct DynInstrumentation {
+    /// zst
+    no_instrumentation: NoInstrumentation,
+    inner: Option<Box<dyn Instrumentation>>,
+}
+
+impl Deref for DynInstrumentation {
+    type Target = dyn Instrumentation;
+
+    fn deref(&self) -> &Self::Target {
+        self.inner.as_deref().unwrap_or(&self.no_instrumentation)
+    }
+}
+
+impl DerefMut for DynInstrumentation {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.inner
+            .as_deref_mut()
+            .unwrap_or(&mut self.no_instrumentation)
+    }
+}
+
+impl DynInstrumentation {
+    #[diesel_derives::__diesel_public_if(
+        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
+    )]
+    pub(crate) fn default_instrumentation() -> Self {
+        Self {
+            inner: get_default_instrumentation(),
+            no_instrumentation: NoInstrumentation,
+        }
+    }
+
+    #[diesel_derives::__diesel_public_if(
+        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
+    )]
+    pub(crate) fn none() -> Self {
+        Self {
+            inner: None,
+            no_instrumentation: NoInstrumentation,
+        }
+    }
+
+    #[diesel_derives::__diesel_public_if(
+        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
+    )]
+    pub(crate) fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
+        // This implementation is not necessary to be able to call this method on this object
+        // because of the already existing Deref impl.
+        // However it allows avoiding the dynamic dispatch to the stub value
+        if let Some(inner) = self.inner.as_deref_mut() {
+            inner.on_connection_event(event)
+        }
+    }
+}
+
+impl<I: Instrumentation> From<I> for DynInstrumentation {
+    fn from(instrumentation: I) -> Self {
+        Self {
+            inner: Some(unpack_instrumentation(Box::new(instrumentation))),
+            no_instrumentation: NoInstrumentation,
+        }
+    }
+}
+
+struct NoInstrumentation;
+
+impl Instrumentation for NoInstrumentation {
+    fn on_connection_event(&mut self, _: InstrumentationEvent<'_>) {}
+}
+
+/// Unwrap unnecessary boxing levels
+fn unpack_instrumentation(
+    mut instrumentation: Box<dyn Instrumentation>,
+) -> Box<dyn Instrumentation> {
+    loop {
+        match instrumentation.downcast::<Box<dyn Instrumentation>>() {
+            Ok(extra_boxed_instrumentation) => instrumentation = *extra_boxed_instrumentation,
+            Err(not_extra_boxed_instrumentation) => {
+                break not_extra_boxed_instrumentation;
+            }
+        }
+    }
+}

diesel/src/mysql/connection/bind.rs

@@ -6,8 +6,7 @@ use std::ops::Index;
 use std::os::raw as libc;
 use std::ptr::NonNull;
 
-use super::stmt::MysqlFieldMetadata;
-use super::stmt::StatementUse;
+use super::stmt::{MysqlFieldMetadata, StatementUse};
 use crate::mysql::connection::stmt::StatementMetadata;
 use crate::mysql::types::date_and_time::MysqlTime;
 use crate::mysql::{MysqlType, MysqlValue};
@@ -870,7 +869,7 @@ mod tests {
             ),
             &mut conn.statement_cache,
             &mut conn.raw_connection,
-            &mut conn.instrumentation,
+            &mut *conn.instrumentation,
         ).unwrap();
 
         let metadata = stmt.metadata().unwrap();

diesel/src/mysql/connection/mod.rs

@@ -8,8 +8,7 @@ use self::stmt::iterator::StatementIterator;
 use self::stmt::Statement;
 use self::url::ConnectionOptions;
 use super::backend::Mysql;
-use crate::connection::instrumentation::DebugQuery;
-use crate::connection::instrumentation::StrQueryHelper;
+use crate::connection::instrumentation::{DebugQuery, DynInstrumentation, StrQueryHelper};
 use crate::connection::statement_cache::{MaybeCached, StatementCache};
 use crate::connection::*;
 use crate::expression::QueryMetadata;
@@ -111,7 +110,7 @@ pub struct MysqlConnection {
     raw_connection: RawConnection,
     transaction_state: AnsiTransactionManager,
     statement_cache: StatementCache<Mysql, Statement>,
-    instrumentation: Option<Box<dyn Instrumentation>>,
+    instrumentation: DynInstrumentation,
 }
 
 // mysql connection can be shared between threads according to libmysqlclients documentation
@@ -156,7 +155,7 @@ impl Connection for MysqlConnection {
     /// * `ssl_mode` expects a value defined for MySQL client command option `--ssl-mode`
     ///   See <https://dev.mysql.com/doc/refman/5.7/en/connection-options.html#option_general_ssl-mode>
     fn establish(database_url: &str) -> ConnectionResult<Self> {
-        let mut instrumentation = crate::connection::instrumentation::get_default_instrumentation();
+        let mut instrumentation = DynInstrumentation::default_instrumentation();
         instrumentation.on_connection_event(InstrumentationEvent::StartEstablishConnection {
             url: database_url,
         });
@@ -181,7 +180,7 @@ impl Connection for MysqlConnection {
                 &source,
                 &mut self.statement_cache,
                 &mut self.raw_connection,
-                &mut self.instrumentation,
+                &mut *self.instrumentation,
             )
             .and_then(|stmt| {
                 // we have not called result yet, so calling `execute` is
@@ -200,19 +199,19 @@ impl Connection for MysqlConnection {
     }
 
     fn instrumentation(&mut self) -> &mut dyn Instrumentation {
-        &mut self.instrumentation
+        &mut *self.instrumentation
     }
 
     fn set_instrumentation(&mut self, instrumentation: impl Instrumentation) {
-        self.instrumentation = Some(Box::new(instrumentation));
+        self.instrumentation = instrumentation.into();
     }
 }
 
 #[inline(always)]
 fn update_transaction_manager_status<T>(
     query_result: QueryResult<T>,
     transaction_manager: &mut AnsiTransactionManager,
-    instrumentation: &mut Option<Box<dyn Instrumentation>>,
+    instrumentation: &mut DynInstrumentation,
     query: &dyn DebugQuery,
 ) -> QueryResult<T> {
     if let Err(Error::DatabaseError(DatabaseErrorKind::SerializationFailure, _)) = query_result {
@@ -244,7 +243,7 @@ impl LoadConnection<DefaultLoadingMode> for MysqlConnection {
                 &source,
                 &mut self.statement_cache,
                 &mut self.raw_connection,
-                &mut self.instrumentation,
+                &mut *self.instrumentation,
             )
             .and_then(|stmt| {
                 let mut metadata = Vec::new();
@@ -329,7 +328,7 @@ impl MysqlConnection {
             raw_connection,
             transaction_state: AnsiTransactionManager::default(),
             statement_cache: StatementCache::new(),
-            instrumentation: None,
+            instrumentation: DynInstrumentation::none(),
         };
         conn.set_config_options()
             .map_err(CouldntSetupConfiguration)?;

diesel/src/pg/connection/cursor.rs

@@ -1,7 +1,6 @@
 use super::raw::RawConnection;
 use super::result::PgResult;
 use super::row::PgRow;
-use crate::connection::Instrumentation;
 use crate::pg::Pg;
 use crate::query_builder::QueryFragment;
 use std::rc::Rc;