diff --git a/crates/buzz-db/src/store/storage_accounting.rs b/crates/buzz-db/src/store/storage_accounting.rs index b34e83921c..72e9356e0d 100644 --- a/crates/buzz-db/src/store/storage_accounting.rs +++ b/crates/buzz-db/src/store/storage_accounting.rs @@ -2,10 +2,10 @@ use buzz_datastore_tracing::datastore_span; use chrono::{DateTime, Utc}; -use sqlx::postgres::PgConnection; +use sqlx::{postgres::PgConnection, Connection}; use uuid::Uuid; -use crate::{observability, Db, Result}; +use crate::{observability, Db, DbError, Result}; /// Deployment-global advisory lock for the run-once storage worker. pub const STORAGE_ACCOUNTING_LOCK_KEY: i64 = 0x4255_5a5a_5354_4f52; @@ -70,6 +70,31 @@ impl StorageAccountingLeader { .await?; Ok(()) } + + /// Unlock the lease on its own session, then close that session. + /// + /// Dropping the leader frees the lease only once PostgreSQL notices the + /// closed socket and ends the old backend, which can lag on a loaded + /// server. When this returns `Ok`, the unlock has run on the lease's own + /// session, so the lease is already free for a successor's + /// [`Db::try_lock_storage_accounting`], even if the close then failed. + /// An error means the unlock failed or this session no longer held the + /// lease. + #[datastore_span(name = "release_storage_accounting", system = "postgresql")] + pub async fn release(mut self) -> Result<()> { + let unlocked = sqlx::query_scalar::<_, bool>("SELECT pg_advisory_unlock($1)") + .bind(STORAGE_ACCOUNTING_LOCK_KEY) + .fetch_one(&mut self.connection) + .await; + let _ = self.connection.close().await; + if unlocked? { + Ok(()) + } else { + Err(DbError::InvalidData( + "storage-accounting session did not hold its lease".into(), + )) + } + } } impl Db { @@ -258,6 +283,70 @@ mod postgres_tests { assert_eq!(stored.code_sha, "b"); } + #[tokio::test] + #[ignore = "requires Postgres"] + async fn released_lease_admits_successor_before_old_session_exits() { + let (pool, _scratch) = + ScratchDb::migrated(&crate::test_support::database_url(), "storage_accounting").await; + let db = Db::from_pool(pool.clone()); + + let mut leader = db + .try_lock_storage_accounting() + .await + .expect("first lock") + .expect("first worker owns lock"); + // A closing backend drops its temp tables before it frees session + // locks. Locking the leader's temp table keeps that backend alive + // after release, so only an explicit unlock can admit a successor. + sqlx::query("CREATE TEMP TABLE exit_fence ()") + .execute(&mut leader.connection) + .await + .expect("create leader temp table"); + let (temp_schema, leader_pid) = sqlx::query_as::<_, (String, i32)>( + "SELECT pg_my_temp_schema()::regnamespace::text, pg_backend_pid()", + ) + .fetch_one(&mut leader.connection) + .await + .expect("load leader session"); + let mut exit_fence = pool.begin().await.expect("begin exit fence"); + sqlx::query(sqlx::AssertSqlSafe(format!( + "LOCK TABLE {temp_schema}.exit_fence IN ACCESS SHARE MODE" + ))) + .execute(&mut *exit_fence) + .await + .expect("hold the leader's temp-table cleanup"); + + leader.release().await.expect("release lease"); + let mut successor = db + .try_lock_storage_accounting() + .await + .expect("successor lock attempt") + .expect("a released lease must admit a successor at once"); + let leader_alive = sqlx::query_scalar::<_, bool>( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid = $1)", + ) + .bind(leader_pid) + .fetch_one(&pool) + .await + .expect("load leader backend"); + assert!(leader_alive, "the old backend must not have exited yet"); + exit_fence + .rollback() + .await + .expect("let the old backend exit"); + + // A session that no longer holds the lease must not report it freed. + sqlx::query("SELECT pg_advisory_unlock($1)") + .bind(STORAGE_ACCOUNTING_LOCK_KEY) + .execute(&mut successor.connection) + .await + .expect("unlock out from under the successor"); + successor + .release() + .await + .expect_err("releasing a lease the session lost must fail"); + } + #[tokio::test] #[ignore = "requires Postgres"] async fn usage_delivery_reads_snapshot_time_and_community_states() { diff --git a/crates/buzz-relay/src/storage_sweep/postgres_tests.rs b/crates/buzz-relay/src/storage_sweep/postgres_tests.rs index fccb237837..9434588849 100644 --- a/crates/buzz-relay/src/storage_sweep/postgres_tests.rs +++ b/crates/buzz-relay/src/storage_sweep/postgres_tests.rs @@ -145,7 +145,9 @@ async fn worker_publication_activates_reader_without_restart_or_mode_change() { values["buzz_total_storage_bytes"], 0.0, "a real empty snapshot may report zero" ); - drop(worker); + // A dropped worker keeps the lease until its backend exits, which can + // outlast the next iteration's lock attempt. + worker.release().await.expect("release worker lease"); } pool.close().await; }