Adds database connection object with DTrace probes

nexus/Cargo.toml

@@ -11,10 +11,10 @@ path = "../rpaths"
 anyhow = "1.0"
 async-trait = "0.1.51"
 bb8 = "0.7.1"
-async-bb8-diesel = { git = "https://github.com/oxidecomputer/async-bb8-diesel", rev = "1a5c55d" }
+async-bb8-diesel = { git = "https://github.com/oxidecomputer/async-bb8-diesel", rev = "22c26ef" }
 cookie = "0.15"
 # Tracking pending 2.0 version.
-diesel = { git = "https://github.com/diesel-rs/diesel", rev = "a39dd2e", features = ["postgres", "r2d2", "chrono", "serde_json", "network-address", "uuid"] }
+diesel = { git = "https://github.com/diesel-rs/diesel", rev = "ce77c382", features = ["postgres", "r2d2", "chrono", "serde_json", "network-address", "uuid"] }
 futures = "0.3.18"
 hex = "0.4.3"
 http = "0.2.5"
@@ -40,6 +40,7 @@ sled-agent-client = { path = "../sled-agent-client" }
 structopt = "0.3"
 thiserror = "1.0"
 toml = "0.5.6"
+usdt = "0.2"
 
 [dependencies.api_identity]
 path = "../api_identity"
@@ -48,9 +49,14 @@ path = "../api_identity"
 version = "0.4"
 features = [ "serde" ]
 
+[dependencies.diesel-dtrace]
+git = "https://github.com/oxidecomputer/diesel-dtrace"
+branch = "main"
+
 [dependencies.dropshot]
 git = "https://github.com/oxidecomputer/dropshot"
 branch = "main"
+features = [ "usdt-probes" ]
 
 [dependencies.omicron-common]
 path = "../common"

nexus/src/config.rs

@@ -359,6 +359,7 @@ mod test {
             if_exists = "fail"
             [insecure]
             allow_any_request_to_spoof_authn_header = true
+            if_exists = "fail"
             "##,
         )
         .unwrap();

nexus/src/db/collection_insert.rs

@@ -9,6 +9,7 @@
 //! 2) updates the collection's child resource generation number
 //! 3) inserts the child resource row
 
+use super::pool::DbConnection;
 use async_bb8_diesel::{
     AsyncRunQueryDsl, ConnectionError, ConnectionManager, PoolError,
 };
@@ -204,11 +205,11 @@ where
     /// - Error(other diesel error)
     pub async fn insert_and_get_result_async(
         self,
-        pool: &bb8::Pool<ConnectionManager<PgConnection>>,
+        pool: &bb8::Pool<ConnectionManager<DbConnection>>,
     ) -> AsyncInsertIntoCollectionResult<ResourceType>
     where
         // We require this bound to ensure that "Self" is runnable as query.
-        Self: query_methods::LoadQuery<PgConnection, ResourceType>,
+        Self: query_methods::LoadQuery<DbConnection, ResourceType>,
     {
         self.get_result_async::<ResourceType>(pool)
             .await
@@ -223,11 +224,11 @@ where
     /// - Error(other diesel error)
     pub async fn insert_and_get_results_async(
         self,
-        pool: &bb8::Pool<ConnectionManager<PgConnection>>,
+        pool: &bb8::Pool<ConnectionManager<DbConnection>>,
     ) -> AsyncInsertIntoCollectionResult<Vec<ResourceType>>
     where
         // We require this bound to ensure that "Self" is runnable as query.
-        Self: query_methods::LoadQuery<PgConnection, ResourceType>,
+        Self: query_methods::LoadQuery<DbConnection, ResourceType>,
     {
         self.get_results_async::<ResourceType>(pool)
             .await
@@ -242,11 +243,11 @@ where
     /// - Error(other diesel error)
     pub fn insert_and_get_result(
         self,
-        conn: &mut PgConnection,
+        conn: &mut DbConnection,
     ) -> SyncInsertIntoCollectionResult<ResourceType>
     where
         // We require this bound to ensure that "Self" is runnable as query.
-        Self: query_methods::LoadQuery<PgConnection, ResourceType>,
+        Self: query_methods::LoadQuery<DbConnection, ResourceType>,
     {
         self.get_result::<ResourceType>(conn)
             .map_err(Self::translate_sync_error)
@@ -260,11 +261,11 @@ where
     /// - Error(other diesel error)
     pub fn insert_and_get_results(
         self,
-        conn: &mut PgConnection,
+        conn: &mut DbConnection,
     ) -> SyncInsertIntoCollectionResult<Vec<ResourceType>>
     where
         // We require this bound to ensure that "Self" is runnable as query.
-        Self: query_methods::LoadQuery<PgConnection, ResourceType>,
+        Self: query_methods::LoadQuery<DbConnection, ResourceType>,
     {
         self.get_results::<ResourceType>(conn)
             .map_err(Self::translate_sync_error)
@@ -323,7 +324,7 @@ where
     type SqlType = SelectableSqlType<ResourceType>;
 }
 
-impl<ResourceType, ISR, C> RunQueryDsl<PgConnection>
+impl<ResourceType, ISR, C> RunQueryDsl<DbConnection>
     for InsertIntoCollectionStatement<ResourceType, ISR, C>
 where
     ResourceTable<ResourceType, C>: Table + Copy + Debug,

nexus/src/db/datastore.rs

@@ -27,6 +27,7 @@ use super::collection_insert::{
 };
 use super::error::diesel_pool_result_optional;
 use super::identity::{Asset, Resource};
+use super::pool::DbConnection;
 use super::Pool;
 use crate::authz;
 use crate::context::OpContext;
@@ -81,15 +82,14 @@ impl DataStore {
     // the database.  Eventually, this function should only be used for doing
     // authentication in the first place (since we can't do an authz check in
     // that case).
-    fn pool(&self) -> &bb8::Pool<ConnectionManager<diesel::PgConnection>> {
+    fn pool(&self) -> &bb8::Pool<ConnectionManager<DbConnection>> {
         self.pool.pool()
     }
 
     fn pool_authorized(
         &self,
         opctx: &OpContext,
-    ) -> Result<&bb8::Pool<ConnectionManager<diesel::PgConnection>>, Error>
-    {
+    ) -> Result<&bb8::Pool<ConnectionManager<DbConnection>>, Error> {
         opctx.authorize(authz::Action::Query, authz::DATABASE)?;
         Ok(self.pool.pool())
     }

nexus/src/db/pool.rs

@@ -31,24 +31,27 @@
 use super::Config as DbConfig;
 use async_bb8_diesel::ConnectionManager;
 use diesel::PgConnection;
+use diesel_dtrace::DTraceConnection;
+
+pub type DbConnection = DTraceConnection<PgConnection>;
 
 /// Wrapper around a database connection pool.
 ///
 /// Expected to be used as the primary interface to the database.
 pub struct Pool {
-    pool: bb8::Pool<ConnectionManager<diesel::PgConnection>>,
+    pool: bb8::Pool<ConnectionManager<DbConnection>>,
 }
 
 impl Pool {
     pub fn new(db_config: &DbConfig) -> Self {
         let manager =
-            ConnectionManager::<PgConnection>::new(&db_config.url.url());
+            ConnectionManager::<DbConnection>::new(&db_config.url.url());
         let pool = bb8::Builder::new().build_unchecked(manager);
         Pool { pool }
     }
 
     /// Returns a reference to the underlying pool.
-    pub fn pool(&self) -> &bb8::Pool<ConnectionManager<diesel::PgConnection>> {
+    pub fn pool(&self) -> &bb8::Pool<ConnectionManager<DbConnection>> {
         &self.pool
     }
 }

nexus/src/db/update_and_check.rs

@@ -4,6 +4,7 @@
 
 //! CTE implementation for "UPDATE with extended return status".
 
+use super::pool::DbConnection;
 use async_bb8_diesel::{AsyncRunQueryDsl, ConnectionManager, PoolError};
 use diesel::associations::HasTable;
 use diesel::helper_types::*;
@@ -142,11 +143,11 @@ where
     /// - Error (row doesn't exist, or other diesel error)
     pub async fn execute_and_check(
         self,
-        pool: &bb8::Pool<ConnectionManager<PgConnection>>,
+        pool: &bb8::Pool<ConnectionManager<DbConnection>>,
     ) -> Result<UpdateAndQueryResult<Q>, PoolError>
     where
         // We require this bound to ensure that "Self" is runnable as query.
-        Self: LoadQuery<PgConnection, (Option<K>, Option<K>, Q)>,
+        Self: LoadQuery<DbConnection, (Option<K>, Option<K>, Q)>,
     {
         let (id0, id1, found) =
             self.get_result_async::<(Option<K>, Option<K>, Q)>(pool).await?;
@@ -176,7 +177,7 @@ where
     );
 }
 
-impl<US, K, Q> RunQueryDsl<PgConnection> for UpdateAndQueryStatement<US, K, Q>
+impl<US, K, Q> RunQueryDsl<DbConnection> for UpdateAndQueryStatement<US, K, Q>
 where
     US: UpdateStatementExt,
     US::Table: Table,

nexus/src/lib.rs

@@ -167,6 +167,13 @@ pub async fn run_server(config: &Config) -> Result<(), String> {
         .log
         .to_logger("nexus")
         .map_err(|message| format!("initializing logger: {}", message))?;
+    if let Err(e) = usdt::register_probes() {
+        let msg = format!("failed to register DTrace probes: {}", e);
+        error!(log, "{}", msg);
+        return Err(msg);
+    } else {
+        debug!(log, "registered DTrace probes");
+    }
     let rack_id = Uuid::new_v4();
     let server = Server::start(config, &rack_id, &log).await?;
     server.register_as_producer().await;