diesel_dtrace/
lib.rs

1#![doc = include_str!("../README.md")]
2// Copyright 2024 Oxide Computer Company
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use diesel::backend::Backend;
17use diesel::connection::{
18    AnsiTransactionManager, LoadConnection, SimpleConnection, TransactionManager,
19    TransactionManagerStatus,
20};
21use diesel::debug_query;
22use diesel::expression::QueryMetadata;
23use diesel::prelude::*;
24use diesel::query_builder::{AsQuery, QueryFragment, QueryId};
25use diesel::r2d2::R2D2Connection;
26use std::ops::{Deref, DerefMut};
27use usdt::UniqueId;
28use uuid::Uuid;
29
30#[usdt::provider(provider = "diesel_db")]
31pub mod probes {
32    /// Fires right before we attempt to establish a connection.
33    pub fn connection__establish__start(_: &UniqueId, conn_id: Uuid, url: &str) {}
34    /// Fires when we finish establishing a connection, with a flag indicating
35    /// whether it succeeded or failed.
36    pub fn connection__establish__done(_: &UniqueId, conn_id: Uuid, success: u8) {}
37    /// Fires just before issuing a SQL query.
38    pub fn query__start(_: &UniqueId, conn_id: Uuid, query: &str) {}
39    /// Fires when a query completes.
40    pub fn query__done(_: &UniqueId, conn_id: Uuid) {}
41    /// Fires when we start a transaction.
42    ///
43    /// This includes the connection ID as well as the depth of the transaction.
44    /// As transactions can be nested, _both_ of these are required to unique ID
45    /// a transaction in full.
46    ///
47    /// The depth is `0` if there is no outstanding transaction, meaning this is
48    /// not nested inside another transaction. Querying the transaction status
49    /// may fail, in which case `depth == -1`.
50    pub fn transaction__start(conn_id: Uuid, depth: i64) {}
51    /// Fires when a transaction completes.
52    ///
53    /// This includes the connection ID as well as the depth of the transaction.
54    /// As transactions can be nested, _both_ of these are required to uniquely
55    /// ID a transaction in full.
56    ///
57    /// The depth is `0` if there is no outstanding transaction, meaning this is
58    /// not nested inside another transaction. Querying the transaction status
59    /// may fail, in which case `depth == -1`.
60    ///
61    /// This also includes a flag indicating whether the transaction was
62    /// committed (`committed == 1`) or rolled back (`committed == 0`).
63    pub fn transaction__done(conn_id: Uuid, depth: i64, committed: u8) {}
64}
65
66/// A [`Connection`] wrapper that inserts DTrace probe points.
67///
68/// See the module-level documentation for more details.
69#[derive(Debug)]
70pub struct DTraceConnection<C: Connection> {
71    inner: C,
72    id: Uuid,
73}
74
75impl<C: Connection> DTraceConnection<C> {
76    pub fn id(&self) -> Uuid {
77        self.id
78    }
79}
80
81impl<C: Connection> Deref for DTraceConnection<C> {
82    type Target = C;
83    fn deref(&self) -> &Self::Target {
84        &self.inner
85    }
86}
87
88impl<C: Connection> DerefMut for DTraceConnection<C> {
89    fn deref_mut(&mut self) -> &mut Self::Target {
90        &mut self.inner
91    }
92}
93
94impl<C: Connection> SimpleConnection for DTraceConnection<C> {
95    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
96        let id = UniqueId::new();
97        probes::query__start!(|| (&id, self.id, query));
98        let result = self.inner.batch_execute(query);
99        probes::query__done!(|| (&id, self.id));
100        result
101    }
102}
103
104impl<C> LoadConnection for DTraceConnection<C>
105where
106    C: Connection<TransactionManager = AnsiTransactionManager> + LoadConnection,
107    C::Backend: Default,
108    <C::Backend as Backend>::QueryBuilder: Default,
109{
110    type Cursor<'conn, 'query>
111        = C::Cursor<'conn, 'query>
112    where
113        Self: 'conn;
114    type Row<'conn, 'query>
115        = C::Row<'conn, 'query>
116    where
117        Self: 'conn;
118
119    fn load<'conn, 'query, T>(
120        &'conn mut self,
121        source: T,
122    ) -> QueryResult<Self::Cursor<'conn, 'query>>
123    where
124        T: diesel::query_builder::Query + QueryFragment<Self::Backend> + QueryId + 'query,
125        Self::Backend: QueryMetadata<T::SqlType>,
126    {
127        let query = source.as_query();
128        let id = UniqueId::new();
129        probes::query__start!(|| (
130            &id,
131            self.id,
132            debug_query::<Self::Backend, _>(&query).to_string()
133        ));
134        let result = self.inner.load(query);
135        probes::query__done!(|| (&id, self.id));
136        result
137    }
138}
139
140impl<C> Connection for DTraceConnection<C>
141where
142    C: Connection<TransactionManager = AnsiTransactionManager>,
143    C::Backend: Default,
144    <C::Backend as Backend>::QueryBuilder: Default,
145{
146    type Backend = C::Backend;
147    type TransactionManager = DTraceTransactionManager<C>;
148
149    fn establish(database_url: &str) -> ConnectionResult<Self> {
150        let id = UniqueId::new();
151        let conn_id = Uuid::new_v4();
152        probes::connection__establish__start!(|| (&id, conn_id, database_url));
153        let conn = C::establish(database_url);
154        probes::connection__establish__done!(|| (&id, conn_id, u8::from(conn.is_ok())));
155        let inner = conn?;
156        Ok(DTraceConnection { inner, id: conn_id })
157    }
158
159    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
160    where
161        T: QueryFragment<Self::Backend> + QueryId,
162    {
163        let id = UniqueId::new();
164        probes::query__start!(|| (
165            &id,
166            self.id,
167            debug_query::<Self::Backend, _>(&source).to_string()
168        ));
169        let result = self.inner.execute_returning_count(source);
170        probes::query__done!(|| (&id, self.id));
171        result
172    }
173
174    fn transaction_state(
175        &mut self,
176    ) -> &mut <DTraceTransactionManager<C> as TransactionManager<DTraceConnection<C>>>::TransactionStateData{
177        self.inner.transaction_state()
178    }
179
180    fn instrumentation(&mut self) -> &mut dyn diesel::connection::Instrumentation {
181        self.inner.instrumentation()
182    }
183
184    fn set_instrumentation(&mut self, instrumentation: impl diesel::connection::Instrumentation) {
185        self.inner.set_instrumentation(instrumentation)
186    }
187
188    fn set_prepared_statement_cache_size(&mut self, size: diesel::connection::CacheSize) {
189        self.inner.set_prepared_statement_cache_size(size);
190    }
191}
192
193impl<C> diesel::connection::ConnectionSealed for DTraceConnection<C>
194where
195    C: Connection<TransactionManager = AnsiTransactionManager>,
196    C::Backend: Default,
197    <C::Backend as Backend>::QueryBuilder: Default,
198{
199}
200
201impl<C> R2D2Connection for DTraceConnection<C>
202where
203    C: R2D2Connection + Connection<TransactionManager = AnsiTransactionManager>,
204    C::Backend: Default,
205    <C::Backend as Backend>::QueryBuilder: Default,
206{
207    fn ping(&mut self) -> QueryResult<()> {
208        self.inner.ping()
209    }
210}
211
212/// A [`TransactionManager`] for a [`DTraceConnection`].
213///
214/// This manager is responsible for the probes `transaction-start` and
215/// `transaction-done`. See the module-level documentation for more details on
216/// these probes.
217pub struct DTraceTransactionManager<C> {
218    _data: std::marker::PhantomData<C>,
219}
220
221impl<C> DTraceTransactionManager<C>
222where
223    C: Connection<TransactionManager = AnsiTransactionManager>,
224{
225    /// Compute the current transaction depth for the DTrace probes.
226    fn depth(conn: &mut DTraceConnection<C>) -> i64 {
227        let status = AnsiTransactionManager::transaction_manager_status_mut(&mut conn.inner);
228        match status.transaction_depth() {
229            Ok(Some(depth)) => i64::from(depth.get()),
230            Ok(None) => 0,
231            Err(_) => -1,
232        }
233    }
234}
235
236impl<C> TransactionManager<DTraceConnection<C>> for DTraceTransactionManager<C>
237where
238    C: Connection<TransactionManager = AnsiTransactionManager>,
239    C::Backend: Default,
240    <C::Backend as Backend>::QueryBuilder: Default,
241{
242    type TransactionStateData = AnsiTransactionManager;
243
244    fn begin_transaction(conn: &mut DTraceConnection<C>) -> QueryResult<()> {
245        // TODO-performance: We're unconditionally computing the transaction
246        // depth here, even if the probe is not enabled.
247        //
248        // This ultimately comes from the interaction of a few things. These
249        // trait methods don't make it easy to store state -- since they take
250        // the mutable connection, not `&mut self`, we have to store everything
251        // on the connection type, but that interacts with the
252        // `Connection::transaction_state()` method weirdly. Second, even if we
253        // could do that, there's no good way to keep that in sync if the probes
254        // are disabled _while_ a transaction is outstanding (we'd end up
255        // thinking we were still in a transaction, when we're not anymore).
256        //
257        // Last, this interacts with a `Clone` bound on the `usdt` crate's
258        // argument closure that we pass to the probe macro itself. That is
259        // required today so that we can accurately type check the return value
260        // of the closure. Still, there are probably ways around that which
261        // still give nice error messages. See
262        // https://github.com/oxidecomputer/usdt/issues/136 for some more
263        // background and context.
264        //
265        // In any case, it is probably "fine" to pay this cost all the time,
266        // even though it's antithetical to the "zero disabled-probe effect"
267        // ethos of DTrace. These methods really just take a pointer to a field
268        // of `AnsiTransactionManager`, and destructure a few enums. It should
269        // be in the noise for any realistic database application.
270        let depth = Self::depth(conn);
271        probes::transaction__start!(|| (&conn.id, depth));
272        AnsiTransactionManager::begin_transaction(&mut conn.inner)
273    }
274
275    fn rollback_transaction(conn: &mut DTraceConnection<C>) -> QueryResult<()> {
276        let result = AnsiTransactionManager::rollback_transaction(&mut conn.inner);
277        let depth = Self::depth(conn);
278        probes::transaction__done!(|| (&conn.id, depth, 0));
279        result
280    }
281
282    fn commit_transaction(conn: &mut DTraceConnection<C>) -> QueryResult<()> {
283        let result = AnsiTransactionManager::commit_transaction(&mut conn.inner);
284        let depth = Self::depth(conn);
285        probes::transaction__done!(|| (&conn.id, depth, 1));
286        result
287    }
288
289    fn transaction_manager_status_mut(
290        conn: &mut DTraceConnection<C>,
291    ) -> &mut TransactionManagerStatus {
292        AnsiTransactionManager::transaction_manager_status_mut(&mut conn.inner)
293    }
294}
diesel_dtrace/lib.rs

diesel_dtrace/
lib.rs
