Skip to content

Commit

Permalink
Make Instrumentation downcasting work correctly by avoiding additiona…
Browse files Browse the repository at this point in the history
…l dyn dispatch level

we actually returned &mut Option<Box<dyn Instrumentation>> as &mut dyn Instrumentation, so downcasting would have to be done in these two steps, which is counter-intuitive and doesn't seem ideal inside diesel itself either.
  • Loading branch information
Ten0 committed Jul 25, 2024
1 parent 034826a commit 5500029
Show file tree
Hide file tree
Showing 5 changed files with 120 additions and 50 deletions.
85 changes: 84 additions & 1 deletion diesel/src/connection/instrumentation.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
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);
Expand Down Expand Up @@ -316,3 +316,86 @@ 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.
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<'_>) {
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;
}
}
}
}
19 changes: 9 additions & 10 deletions diesel/src/mysql/connection/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
});
Expand All @@ -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
Expand All @@ -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 {
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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)?;
Expand Down
1 change: 0 additions & 1 deletion diesel/src/pg/connection/cursor.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down
49 changes: 19 additions & 30 deletions diesel/src/pg/connection/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,14 @@ mod result;
mod row;
mod stmt;

use self::copy::CopyFromSink;
use self::copy::CopyToBuffer;
use self::copy::{CopyFromSink, CopyToBuffer};
use self::cursor::*;
use self::private::ConnectionAndTransactionManager;
use self::raw::{PgTransactionStatus, RawConnection};
use self::stmt::Statement;
use crate::connection::instrumentation::DebugQuery;
use crate::connection::instrumentation::Instrumentation;
use crate::connection::instrumentation::StrQueryHelper;
use crate::connection::instrumentation::{
DebugQuery, DynInstrumentation, Instrumentation, StrQueryHelper,
};
use crate::connection::statement_cache::{MaybeCached, StatementCache};
use crate::connection::*;
use crate::expression::QueryMetadata;
Expand All @@ -29,9 +28,7 @@ use std::ffi::CString;
use std::fmt::Debug;
use std::os::raw as libc;

use super::query_builder::copy::CopyFromExpression;
use super::query_builder::copy::CopyTarget;
use super::query_builder::copy::CopyToCommand;
use super::query_builder::copy::{CopyFromExpression, CopyTarget, CopyToCommand};

pub(super) use self::result::PgResult;

Expand Down Expand Up @@ -178,7 +175,7 @@ impl Connection for PgConnection {
type TransactionManager = AnsiTransactionManager;

fn establish(database_url: &str) -> ConnectionResult<PgConnection> {
let mut instrumentation = crate::connection::instrumentation::get_default_instrumentation();
let mut instrumentation = DynInstrumentation::default_instrumentation();
instrumentation.on_connection_event(InstrumentationEvent::StartEstablishConnection {
url: database_url,
});
Expand All @@ -187,7 +184,7 @@ impl Connection for PgConnection {
connection_and_transaction_manager: ConnectionAndTransactionManager {
raw_connection: raw_conn,
transaction_state: AnsiTransactionManager::default(),
instrumentation: None,
instrumentation: DynInstrumentation::none(),
},
statement_cache: StatementCache::new(),
metadata_cache: PgMetadataCache::new(),
Expand Down Expand Up @@ -233,11 +230,11 @@ impl Connection for PgConnection {
}

fn instrumentation(&mut self) -> &mut dyn Instrumentation {
&mut self.connection_and_transaction_manager.instrumentation
&mut *self.connection_and_transaction_manager.instrumentation
}

fn set_instrumentation(&mut self, instrumentation: impl Instrumentation) {
self.connection_and_transaction_manager.instrumentation = Some(Box::new(instrumentation));
self.connection_and_transaction_manager.instrumentation = instrumentation.into();
}
}

Expand Down Expand Up @@ -517,7 +514,7 @@ impl PgConnection {
};
Statement::prepare(conn, sql, query_name.as_deref(), &metadata)
},
&mut self.connection_and_transaction_manager.instrumentation,
&mut *self.connection_and_transaction_manager.instrumentation,
);
if !execute_returning_count {
if let Err(ref e) = query {
Expand Down Expand Up @@ -557,7 +554,7 @@ mod private {
pub struct ConnectionAndTransactionManager {
pub(super) raw_connection: RawConnection,
pub(super) transaction_state: AnsiTransactionManager,
pub(super) instrumentation: Option<Box<dyn Instrumentation>>,
pub(super) instrumentation: DynInstrumentation,
}

pub trait PgLoadingMode<B> {
Expand Down Expand Up @@ -800,8 +797,7 @@ mod tests {

#[test]
fn transaction_manager_returns_an_error_when_attempting_to_commit_outside_of_a_transaction() {
use crate::connection::AnsiTransactionManager;
use crate::connection::TransactionManager;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::result::Error;
use crate::PgConnection;

Expand All @@ -818,8 +814,7 @@ mod tests {

#[test]
fn transaction_manager_returns_an_error_when_attempting_to_rollback_outside_of_a_transaction() {
use crate::connection::AnsiTransactionManager;
use crate::connection::TransactionManager;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::result::Error;
use crate::PgConnection;

Expand All @@ -838,8 +833,7 @@ mod tests {
fn postgres_transaction_is_rolled_back_upon_syntax_error() {
use std::num::NonZeroU32;

use crate::connection::AnsiTransactionManager;
use crate::connection::TransactionManager;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::pg::connection::raw::PgTransactionStatus;
use crate::*;
let conn = &mut crate::test_helpers::pg_connection_no_transaction();
Expand Down Expand Up @@ -883,8 +877,7 @@ mod tests {
fn nested_postgres_transaction_is_rolled_back_upon_syntax_error() {
use std::num::NonZeroU32;

use crate::connection::AnsiTransactionManager;
use crate::connection::TransactionManager;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::pg::connection::raw::PgTransactionStatus;
use crate::*;
let conn = &mut crate::test_helpers::pg_connection_no_transaction();
Expand Down Expand Up @@ -1165,8 +1158,7 @@ mod tests {

#[test]
fn postgres_transaction_is_rolled_back_upon_deferred_constraint_failure() {
use crate::connection::AnsiTransactionManager;
use crate::connection::TransactionManager;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::pg::connection::raw::PgTransactionStatus;
use crate::result::Error;
use crate::*;
Expand Down Expand Up @@ -1215,8 +1207,7 @@ mod tests {

#[test]
fn postgres_transaction_is_rolled_back_upon_deferred_trigger_failure() {
use crate::connection::AnsiTransactionManager;
use crate::connection::TransactionManager;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::pg::connection::raw::PgTransactionStatus;
use crate::result::Error;
use crate::*;
Expand Down Expand Up @@ -1291,8 +1282,7 @@ mod tests {

#[test]
fn nested_postgres_transaction_is_rolled_back_upon_deferred_trigger_failure() {
use crate::connection::AnsiTransactionManager;
use crate::connection::TransactionManager;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::pg::connection::raw::PgTransactionStatus;
use crate::result::Error;
use crate::*;
Expand Down Expand Up @@ -1374,8 +1364,7 @@ mod tests {

#[test]
fn nested_postgres_transaction_is_rolled_back_upon_deferred_constraint_failure() {
use crate::connection::AnsiTransactionManager;
use crate::connection::TransactionManager;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::pg::connection::raw::PgTransactionStatus;
use crate::result::Error;
use crate::*;
Expand Down
Loading

0 comments on commit 5500029

Please sign in to comment.