[nexus] Make most transactions automatically retry

Cargo.toml

@@ -142,7 +142,9 @@ api_identity = { path = "api_identity" }
 approx = "0.5.1"
 assert_matches = "1.5.0"
 assert_cmd = "2.0.12"
-async-bb8-diesel = { git = "https://github.com/oxidecomputer/async-bb8-diesel", rev = "1446f7e0c1f05f33a0581abd51fa873c7652ab61" }
+# TODO: Patch before merging
+# async-bb8-diesel = { git = "https://github.com/oxidecomputer/async-bb8-diesel", rev = "1446f7e0c1f05f33a0581abd51fa873c7652ab61" }
+async-bb8-diesel = { git = "https://github.com/oxidecomputer/async-bb8-diesel", branch = "txn-retry" }
 async-trait = "0.1.74"
 atomicwrites = "0.4.2"
 authz-macros = { path = "nexus/authz-macros" }

nexus/db-model/src/sled.rs

@@ -154,7 +154,7 @@ impl DatastoreCollectionConfig<super::Service> for Sled {
 }
 
 /// A set of constraints that can be placed on operations that select a sled.
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct SledReservationConstraints {
     must_select_from: Vec<Uuid>,
 }

nexus/db-queries/Cargo.toml

@@ -32,6 +32,7 @@ oso.workspace = true
 paste.workspace = true
 # See omicron-rpaths for more about the "pq-sys" dependency.
 pq-sys = "*"
+rand.workspace = true
 ref-cast.workspace = true
 samael.workspace = true
 serde.workspace = true

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

@@ -563,12 +563,9 @@ where
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::db::{
-        self, error::TransactionError, identity::Resource as IdentityResource,
-    };
+    use crate::db::{self, identity::Resource as IdentityResource};
     use async_bb8_diesel::{
-        AsyncConnection, AsyncRunQueryDsl, AsyncSimpleConnection,
-        ConnectionManager,
+        AsyncRunQueryDsl, AsyncSimpleConnection, ConnectionManager,
     };
     use chrono::Utc;
     use db_macros::Resource;
@@ -999,22 +996,12 @@ mod test {
                 .set(resource::dsl::collection_id.eq(collection_id)),
         );
 
-        type TxnError =
-            TransactionError<AttachError<Resource, Collection, DieselError>>;
-        let result = conn
-            .transaction_async(|conn| async move {
-                attach_query.attach_and_get_result_async(&conn).await.map_err(
-                    |e| match e {
-                        AttachError::DatabaseError(e) => TxnError::from(e),
-                        e => TxnError::CustomError(e),
-                    },
-                )
-            })
-            .await;
-
         // "attach_and_get_result" should return the "attached" resource.
-        let (returned_collection, returned_resource) =
-            result.expect("Attach should have worked");
+        let (returned_collection, returned_resource) = attach_query
+            .attach_and_get_result_async(&conn)
+            .await
+            .expect("Attach should have worked");
+
         assert_eq!(
             returned_resource.collection_id.expect("Expected a collection ID"),
             collection_id

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

@@ -479,12 +479,9 @@ where
 mod test {
     use super::*;
     use crate::db::collection_attach::DatastoreAttachTarget;
-    use crate::db::{
-        self, error::TransactionError, identity::Resource as IdentityResource,
-    };
+    use crate::db::{self, identity::Resource as IdentityResource};
     use async_bb8_diesel::{
-        AsyncConnection, AsyncRunQueryDsl, AsyncSimpleConnection,
-        ConnectionManager,
+        AsyncRunQueryDsl, AsyncSimpleConnection, ConnectionManager,
     };
     use chrono::Utc;
     use db_macros::Resource;
@@ -919,21 +916,12 @@ mod test {
                 .set(resource::dsl::collection_id.eq(Option::<Uuid>::None)),
         );
 
-        type TxnError =
-            TransactionError<DetachManyError<Collection, DieselError>>;
-        let result = conn
-            .transaction_async(|conn| async move {
-                detach_query.detach_and_get_result_async(&conn).await.map_err(
-                    |e| match e {
-                        DetachManyError::DatabaseError(e) => TxnError::from(e),
-                        e => TxnError::CustomError(e),
-                    },
-                )
-            })
-            .await;
-
         // "detach_and_get_result" should return the "detached" resource.
-        let returned_collection = result.expect("Detach should have worked");
+        let returned_collection = detach_query
+            .detach_and_get_result_async(&conn)
+            .await
+            .expect("Detach should have worked");
+
         // The returned values should be the latest value in the DB.
         assert_eq!(
             returned_collection,

nexus/db-queries/src/db/datastore/address_lot.rs

@@ -13,6 +13,7 @@ use crate::db::error::TransactionError;
 use crate::db::model::Name;
 use crate::db::model::{AddressLot, AddressLotBlock, AddressLotReservedBlock};
 use crate::db::pagination::paginated;
+use crate::transaction_retry::RetryHelper;
 use async_bb8_diesel::{AsyncConnection, AsyncRunQueryDsl, Connection};
 use chrono::Utc;
 use diesel::result::Error as DieselError;
@@ -28,6 +29,7 @@ use omicron_common::api::external::{
 };
 use ref_cast::RefCast;
 use serde::{Deserialize, Serialize};
+use std::sync::{Arc, OnceLock};
 use uuid::Uuid;
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
@@ -45,51 +47,61 @@ impl DataStore {
         use db::schema::address_lot::dsl as lot_dsl;
         use db::schema::address_lot_block::dsl as block_dsl;
 
+        let retry_helper = RetryHelper::new(
+            &self.transaction_retry_producer,
+            "address_lot_create",
+        );
         self.pool_connection_authorized(opctx)
             .await?
             // TODO https://github.com/oxidecomputer/omicron/issues/2811
             // Audit external networking database transaction usage
-            .transaction_async(|conn| async move {
-                let lot = AddressLot::new(&params.identity, params.kind.into());
-
-                let db_lot: AddressLot =
-                    diesel::insert_into(lot_dsl::address_lot)
-                        .values(lot)
-                        .returning(AddressLot::as_returning())
-                        .get_result_async(&conn)
-                        .await?;
-
-                let blocks: Vec<AddressLotBlock> = params
-                    .blocks
-                    .iter()
-                    .map(|b| {
-                        AddressLotBlock::new(
-                            db_lot.id(),
-                            b.first_address.into(),
-                            b.last_address.into(),
-                        )
+            .transaction_async_with_retry(
+                |conn| async move {
+                    let lot =
+                        AddressLot::new(&params.identity, params.kind.into());
+
+                    let db_lot: AddressLot =
+                        diesel::insert_into(lot_dsl::address_lot)
+                            .values(lot)
+                            .returning(AddressLot::as_returning())
+                            .get_result_async(&conn)
+                            .await?;
+
+                    let blocks: Vec<AddressLotBlock> = params
+                        .blocks
+                        .iter()
+                        .map(|b| {
+                            AddressLotBlock::new(
+                                db_lot.id(),
+                                b.first_address.into(),
+                                b.last_address.into(),
+                            )
+                        })
+                        .collect();
+
+                    let db_blocks =
+                        diesel::insert_into(block_dsl::address_lot_block)
+                            .values(blocks)
+                            .returning(AddressLotBlock::as_returning())
+                            .get_results_async(&conn)
+                            .await?;
+
+                    Ok(AddressLotCreateResult {
+                        lot: db_lot,
+                        blocks: db_blocks,
                     })
-                    .collect();
-
-                let db_blocks =
-                    diesel::insert_into(block_dsl::address_lot_block)
-                        .values(blocks)
-                        .returning(AddressLotBlock::as_returning())
-                        .get_results_async(&conn)
-                        .await?;
-
-                Ok(AddressLotCreateResult { lot: db_lot, blocks: db_blocks })
-            })
+                },
+                retry_helper.as_callback(),
+            )
             .await
-            .map_err(|e| match e {
-                DieselError::DatabaseError(_, _) => public_error_from_diesel(
+            .map_err(|e| {
+                public_error_from_diesel(
                     e,
                     ErrorHandler::Conflict(
                         ResourceType::AddressLot,
                         &params.identity.name.as_str(),
                     ),
-                ),
-                _ => public_error_from_diesel(e, ErrorHandler::Server),
+                )
             })
     }
 
@@ -113,46 +125,60 @@ impl DataStore {
             LotInUse,
         }
 
-        type TxnError = TransactionError<AddressLotDeleteError>;
+        let err = Arc::new(OnceLock::new());
+        let retry_helper = RetryHelper::new(
+            &self.transaction_retry_producer,
+            "address_lot_delete",
+        );
 
         // TODO https://github.com/oxidecomputer/omicron/issues/2811
         // Audit external networking database transaction usage
-        conn.transaction_async(|conn| async move {
-            let rsvd: Vec<AddressLotReservedBlock> =
-                rsvd_block_dsl::address_lot_rsvd_block
-                    .filter(rsvd_block_dsl::address_lot_id.eq(id))
-                    .select(AddressLotReservedBlock::as_select())
-                    .limit(1)
-                    .load_async(&conn)
-                    .await?;
-
-            if !rsvd.is_empty() {
-                Err(TxnError::CustomError(AddressLotDeleteError::LotInUse))?;
-            }
-
-            let now = Utc::now();
-            diesel::update(lot_dsl::address_lot)
-                .filter(lot_dsl::time_deleted.is_null())
-                .filter(lot_dsl::id.eq(id))
-                .set(lot_dsl::time_deleted.eq(now))
-                .execute_async(&conn)
-                .await?;
+        conn.transaction_async_with_retry(
+            |conn| {
+                let err = err.clone();
+                async move {
+                    let rsvd: Vec<AddressLotReservedBlock> =
+                        rsvd_block_dsl::address_lot_rsvd_block
+                            .filter(rsvd_block_dsl::address_lot_id.eq(id))
+                            .select(AddressLotReservedBlock::as_select())
+                            .limit(1)
+                            .load_async(&conn)
+                            .await?;
+
+                    if !rsvd.is_empty() {
+                        err.set(AddressLotDeleteError::LotInUse).unwrap();
+                        return Err(DieselError::RollbackTransaction);
+                    }
+
+                    let now = Utc::now();
+                    diesel::update(lot_dsl::address_lot)
+                        .filter(lot_dsl::time_deleted.is_null())
+                        .filter(lot_dsl::id.eq(id))
+                        .set(lot_dsl::time_deleted.eq(now))
+                        .execute_async(&conn)
+                        .await?;
 
-            diesel::delete(block_dsl::address_lot_block)
-                .filter(block_dsl::address_lot_id.eq(id))
-                .execute_async(&conn)
-                .await?;
+                    diesel::delete(block_dsl::address_lot_block)
+                        .filter(block_dsl::address_lot_id.eq(id))
+                        .execute_async(&conn)
+                        .await?;
 
-            Ok(())
-        })
+                    Ok(())
+                }
+            },
+            retry_helper.as_callback(),
+        )
         .await
-        .map_err(|e| match e {
-            TxnError::Database(e) => {
+        .map_err(|e| {
+            if let Some(err) = err.get() {
+                match err {
+                    AddressLotDeleteError::LotInUse => {
+                        Error::invalid_request("lot is in use")
+                    }
+                }
+            } else {
                 public_error_from_diesel(e, ErrorHandler::Server)
             }
-            TxnError::CustomError(AddressLotDeleteError::LotInUse) => {
-                Error::invalid_request("lot is in use")
-            }
         })
     }