diff --git a/crates/api-core/src/handlers/machine_discovery.rs b/crates/api-core/src/handlers/machine_discovery.rs index 04c86e3e7f..ae973b58cd 100644 --- a/crates/api-core/src/handlers/machine_discovery.rs +++ b/crates/api-core/src/handlers/machine_discovery.rs @@ -116,6 +116,12 @@ pub(crate) async fn discover_machine( }; let mut txn = api.txn_begin().await?; + if hardware_info.is_dpu() { + // Discovery updates machine-interface rows before reconciliation. Take the segment locks + // first so the outer transaction preserves the allocator's advisory-lock-before-row-lock + // order instead of relying on reconciliation's inner savepoint to establish it later. + db::machine_interface::load_and_lock_all_admin_segments(&mut txn).await?; + } tracing::debug!( ?remote_ip, ?interface_id, diff --git a/crates/api-core/src/handlers/managed_host.rs b/crates/api-core/src/handlers/managed_host.rs index 9762f87deb..7bb03d7b87 100644 --- a/crates/api-core/src/handlers/managed_host.rs +++ b/crates/api-core/src/handlers/managed_host.rs @@ -283,6 +283,10 @@ async fn set_primary_interface_core( if current_primary_is_admin { db::machine_interface::reconcile_admin_addresses_for_host(&mut txn, &host_machine_id) .await?; + } else { + // The repair path has no current admin primary to reconcile, but it still must acquire + // segment locks before the primary-interface row updates below. + db::machine_interface::load_and_lock_all_admin_segments(&mut txn).await?; } // update the primary interface: clear the old primary (if any), then set the new. diff --git a/crates/api-core/src/tests/set_primary_interface.rs b/crates/api-core/src/tests/set_primary_interface.rs index 77c2a69c92..46bae4f71c 100644 --- a/crates/api-core/src/tests/set_primary_interface.rs +++ b/crates/api-core/src/tests/set_primary_interface.rs @@ -402,14 +402,64 @@ async fn test_set_primary_interface_repairs_dpu_host_with_no_admin_primary( .execute(&env.pool) .await?; - // Promoting the Admin interface must succeed (repair), not error after the BMC call. - env.api - .set_primary_interface(tonic::Request::new(forge::SetPrimaryInterfaceRequest { + // Hold the admin segment lock while the repair starts. The handler must wait for this lock + // before updating the interface row, or it can deadlock with DPU discovery's segment-first + // transaction. + let mut blocker_txn = db::Transaction::begin(&env.pool).await?; + let blocker_pid = sqlx::query_scalar::<_, i32>("SELECT pg_backend_pid()") + .fetch_one(blocker_txn.as_pgconn()) + .await?; + db::machine_interface::load_and_lock_all_admin_segments(&mut blocker_txn).await?; + + let api = env.api.clone(); + let set_primary_task = tokio::spawn(async move { + api.set_primary_interface(tonic::Request::new(forge::SetPrimaryInterfaceRequest { host_machine_id: Some(host_id), interface_id: Some(promote_id), reboot: false, })) - .await?; + .await + }); + + let wait_result = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS (\ + SELECT 1 FROM pg_locks \ + WHERE locktype = 'advisory' AND NOT granted \ + AND $1 = ANY(pg_blocking_pids(pid))\ + )", + ) + .bind(blocker_pid) + .fetch_one(&env.pool) + .await?; + if waiting { + return Ok::<(), sqlx::Error>(()); + } + tokio::task::yield_now().await; + } + }) + .await + .map_err(|_| { + std::io::Error::other("set-primary-interface did not wait for the admin segment lock") + }); + + let row_lock_result = + sqlx::query("SELECT id FROM machine_interfaces WHERE id = $1::uuid FOR UPDATE NOWAIT") + .bind(promote_id) + .fetch_one(&env.pool) + .await; + + blocker_txn.commit().await?; + let set_primary_result = set_primary_task.await?; + + wait_result??; + row_lock_result.map_err(|err| { + std::io::Error::other(format!( + "set-primary-interface locked the target row before the admin segment: {err}", + )) + })?; + set_primary_result?; // The promoted interface is now the only primary. let primaries_now: Vec<_> = { diff --git a/crates/api-db/src/machine_interface.rs b/crates/api-db/src/machine_interface.rs index 09e1ad1e53..e6d237704a 100644 --- a/crates/api-db/src/machine_interface.rs +++ b/crates/api-db/src/machine_interface.rs @@ -2627,7 +2627,7 @@ FOR UPDATE"#; /// This intentionally locks more broadly than the specific host touches so reconciliation follows /// the same high-level lock order as address allocation: segment advisory lock first, then /// machine interface/address row locks. -async fn load_and_lock_all_admin_segments( +pub async fn load_and_lock_all_admin_segments( txn: &mut Transaction<'_>, ) -> DatabaseResult> { let mut segment_ids =