Qorb integration as connection pool for database

Cargo.toml

@@ -219,13 +219,12 @@ api_identity = { path = "api_identity" }
 approx = "0.5.1"
 assert_matches = "1.5.0"
 assert_cmd = "2.0.14"
-async-bb8-diesel = { git = "https://github.com/oxidecomputer/async-bb8-diesel", rev = "ed7ab5ef0513ba303d33efd41d3e9e381169d59b" }
+async-bb8-diesel = "0.2"
 async-trait = "0.1.80"
 atomicwrites = "0.4.3"
 authz-macros = { path = "nexus/authz-macros" }
 backoff = { version = "0.4.0", features = [ "tokio" ] }
 base64 = "0.22.1"
-bb8 = "0.8.3"
 bcs = "0.1.6"
 bincode = "1.3.3"
 bootstore = { path = "bootstore" }
@@ -398,6 +397,7 @@ bhyve_api = { git = "https://github.com/oxidecomputer/propolis", rev = "6d7ed9a0
 propolis-client = { git = "https://github.com/oxidecomputer/propolis", rev = "6d7ed9a033babc054db9eff5b59dee978d2b0d76" }
 propolis-mock-server = { git = "https://github.com/oxidecomputer/propolis", rev = "6d7ed9a033babc054db9eff5b59dee978d2b0d76" }
 proptest = "1.4.0"
+qorb = { git = "https://github.com/oxidecomputer/qorb", branch = "master" }
 quote = "1.0"
 rand = "0.8.5"
 rand_core = "0.6.4"
@@ -475,7 +475,7 @@ textwrap = "0.16.1"
 test-strategy = "0.3.1"
 thiserror = "1.0"
 tofino = { git = "http://github.com/oxidecomputer/tofino", branch = "main" }
-tokio = "1.37.0"
+tokio = "1.38.0"
 tokio-postgres = { version = "0.7", features = [ "with-chrono-0_4", "with-uuid-1" ] }
 tokio-stream = "0.1.15"
 tokio-tungstenite = "0.20"

dev-tools/omdb/src/bin/omdb/db.rs

@@ -231,7 +231,8 @@ impl DbUrlOptions {
         eprintln!("note: using database URL {}", &db_url);
 
         let db_config = db::Config { url: db_url.clone() };
-        let pool = Arc::new(db::Pool::new(&log.clone(), &db_config));
+        let pool =
+            Arc::new(db::Pool::new_single_host(&log.clone(), &db_config));
 
         // Being a dev tool, we want to try this operation even if the schema
         // doesn't match what we expect.  So we use `DataStore::new_unchecked()`
@@ -3069,7 +3070,7 @@ async fn cmd_db_inventory(
 }
 
 async fn cmd_db_inventory_baseboard_ids(
-    conn: &DataStoreConnection<'_>,
+    conn: &DataStoreConnection,
     limit: NonZeroU32,
 ) -> Result<(), anyhow::Error> {
     #[derive(Tabled)]
@@ -3106,7 +3107,7 @@ async fn cmd_db_inventory_baseboard_ids(
 }
 
 async fn cmd_db_inventory_cabooses(
-    conn: &DataStoreConnection<'_>,
+    conn: &DataStoreConnection,
     limit: NonZeroU32,
 ) -> Result<(), anyhow::Error> {
     #[derive(Tabled)]
@@ -3147,7 +3148,7 @@ async fn cmd_db_inventory_cabooses(
 }
 
 async fn cmd_db_inventory_physical_disks(
-    conn: &DataStoreConnection<'_>,
+    conn: &DataStoreConnection,
     limit: NonZeroU32,
     args: PhysicalDisksArgs,
 ) -> Result<(), anyhow::Error> {
@@ -3204,7 +3205,7 @@ async fn cmd_db_inventory_physical_disks(
 }
 
 async fn cmd_db_inventory_rot_pages(
-    conn: &DataStoreConnection<'_>,
+    conn: &DataStoreConnection,
     limit: NonZeroU32,
 ) -> Result<(), anyhow::Error> {
     #[derive(Tabled)]
@@ -3239,7 +3240,7 @@ async fn cmd_db_inventory_rot_pages(
 }
 
 async fn cmd_db_inventory_collections_list(
-    conn: &DataStoreConnection<'_>,
+    conn: &DataStoreConnection,
     limit: NonZeroU32,
 ) -> Result<(), anyhow::Error> {
     #[derive(Tabled)]

nexus-config/src/postgres_config.rs

@@ -5,6 +5,7 @@
 //! Common objects used for configuration
 
 use std::fmt;
+use std::net::SocketAddr;
 use std::ops::Deref;
 use std::str::FromStr;
 
@@ -32,6 +33,29 @@ impl PostgresConfigWithUrl {
     pub fn url(&self) -> String {
         self.url_raw.clone()
     }
+
+    /// Accesses the first ip / port pair within the URL.
+    ///
+    /// # Panics
+    ///
+    /// This method makes the assumption that the hostname has at least one
+    /// "host IP / port" pair which can be extracted. If the supplied URL
+    /// does not have such a pair, this function will panic.
+    // Yes, panicking in the above scenario sucks. But this type is already
+    // pretty ubiquitous within Omicron, and integration with the qorb
+    // connection pooling library requires access to database by SocketAddr.
+    pub fn address(&self) -> SocketAddr {
+        let tokio_postgres::config::Host::Tcp(host) =
+            &self.config.get_hosts()[0]
+        else {
+            panic!("Non-TCP hostname");
+        };
+        let ip: std::net::IpAddr =
+            host.parse().expect("Failed to parse host as IP address");
+
+        let port = self.config.get_ports()[0];
+        SocketAddr::new(ip, port)
+    }
 }
 
 impl FromStr for PostgresConfigWithUrl {

nexus/db-queries/Cargo.toml

@@ -14,21 +14,22 @@ omicron-rpaths.workspace = true
 anyhow.workspace = true
 async-bb8-diesel.workspace = true
 async-trait.workspace = true
-bb8.workspace = true
 camino.workspace = true
 chrono.workspace = true
 const_format.workspace = true
 diesel.workspace = true
 diesel-dtrace.workspace = true
 dropshot.workspace = true
 futures.workspace = true
+internal-dns.workspace = true
 ipnetwork.workspace = true
 macaddr.workspace = true
 once_cell.workspace = true
 oxnet.workspace = true
 paste.workspace = true
 # See omicron-rpaths for more about the "pq-sys" dependency.
 pq-sys = "*"
+qorb = { workspace = true, features = [ "qtop" ] }
 rand.workspace = true
 ref-cast.workspace = true
 schemars.workspace = true

nexus/db-queries/src/db/collection_attach.rs

@@ -564,9 +564,7 @@ where
 mod test {
     use super::*;
     use crate::db::{self, identity::Resource as IdentityResource};
-    use async_bb8_diesel::{
-        AsyncRunQueryDsl, AsyncSimpleConnection, ConnectionManager,
-    };
+    use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection};
     use chrono::Utc;
     use db_macros::Resource;
     use diesel::expression_methods::ExpressionMethods;
@@ -603,8 +601,8 @@ mod test {
 
     async fn setup_db(
         pool: &crate::db::Pool,
-    ) -> bb8::PooledConnection<ConnectionManager<DbConnection>> {
-        let connection = pool.pool().get().await.unwrap();
+    ) -> crate::db::datastore::DataStoreConnection {
+        let connection = pool.claim().await.unwrap();
         (*connection)
             .batch_execute_async(
                 "CREATE SCHEMA IF NOT EXISTS test_schema; \
@@ -859,7 +857,7 @@ mod test {
             dev::test_setup_log("test_attach_missing_collection_fails");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -888,7 +886,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_missing_resource_fails");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -925,7 +923,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_once");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -973,7 +971,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_once_synchronous");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -1022,7 +1020,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_multiple_times");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -1078,7 +1076,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_beyond_capacity_fails");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -1142,7 +1140,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_while_already_attached");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -1249,7 +1247,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_once");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -1304,7 +1302,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_deleted_resource_fails");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -1349,7 +1347,7 @@ mod test {
         let logctx = dev::test_setup_log("test_attach_without_update_filter");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 

nexus/db-queries/src/db/collection_detach.rs

@@ -482,9 +482,7 @@ mod test {
     use super::*;
     use crate::db::collection_attach::DatastoreAttachTarget;
     use crate::db::{self, identity::Resource as IdentityResource};
-    use async_bb8_diesel::{
-        AsyncRunQueryDsl, AsyncSimpleConnection, ConnectionManager,
-    };
+    use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection};
     use chrono::Utc;
     use db_macros::Resource;
     use diesel::expression_methods::ExpressionMethods;
@@ -521,8 +519,8 @@ mod test {
 
     async fn setup_db(
         pool: &crate::db::Pool,
-    ) -> bb8::PooledConnection<ConnectionManager<DbConnection>> {
-        let connection = pool.pool().get().await.unwrap();
+    ) -> crate::db::datastore::DataStoreConnection {
+        let connection = pool.claim().await.unwrap();
         (*connection)
             .batch_execute_async(
                 "CREATE SCHEMA IF NOT EXISTS test_schema; \
@@ -786,7 +784,7 @@ mod test {
             dev::test_setup_log("test_detach_missing_collection_fails");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -814,7 +812,7 @@ mod test {
         let logctx = dev::test_setup_log("test_detach_missing_resource_fails");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -850,7 +848,7 @@ mod test {
         let logctx = dev::test_setup_log("test_detach_once");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -890,7 +888,7 @@ mod test {
         let logctx = dev::test_setup_log("test_detach_while_already_detached");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -954,7 +952,7 @@ mod test {
         let logctx = dev::test_setup_log("test_detach_deleted_resource_fails");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;
 
@@ -998,7 +996,7 @@ mod test {
         let logctx = dev::test_setup_log("test_detach_without_update_filter");
         let mut db = test_setup_database(&logctx.log).await;
         let cfg = db::Config { url: db.pg_config().clone() };
-        let pool = db::Pool::new(&logctx.log, &cfg);
+        let pool = db::Pool::new_single_host(&logctx.log, &cfg);
 
         let conn = setup_db(&pool).await;