From 580caca66714b3e09fd3e0eb2e522d6f4aa5cc12 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 6 Nov 2026 09:28:55 +1100 Subject: [PATCH] fix(db): unlock the storage-accounting lease before its session closes (#9097) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary `StorageAccountingLeader` holds a session-level advisory lock on a detached connection. Dropping it only closes the socket, and PostgreSQL frees the lock when the backend exits, which can be after the next lock attempt. `release(self)` now unlocks on the leader's own session before closing it. The relay storage-sweep test releases its worker between iterations instead of dropping it, which fixes the `worker_publication_activates_reader_without_restart_or_mode_change` flake. The follow-up commit makes `release` report success only when the unlock actually ran on the session that held the lease. ### fix(db): unlock the storage-accounting lease before its session closes #### Problem The PostgreSQL lane's `storage_sweep::postgres_tests::worker_publication_activates_reader_without_restart_or_mode_change` failed with "called `Option::unwrap()` on a `None` value" at crates/buzz-relay/src/storage_sweep/postgres_tests.rs:66:74 after 5.53 s, on the unrelated branch elrond/cli-credential-chain: - https://github.com/block/buzz/actions/runs/41730818977/job/126253029211 #### Cause Line 66 is `db.try_lock_storage_accounting().await.unwrap().unwrap()` inside the loop over three metrics modes. `Ok(None)` means `pg_try_advisory_lock` returned false. A passing run takes about 15 s, roughly 5 s per iteration, so the failure is the start of iteration 2. Each test process gets its own database and advisory locks are per-database, so the holder is iteration 1's own worker. `StorageAccountingLeader` holds a session-level advisory lock on a detached connection. The test ended each iteration with `drop(worker)`, which only closes the socket. PostgreSQL frees session locks during backend exit, after it notices EOF, so the lock can outlive the drop on a loaded server. Iteration 2's lock attempt raced that exit. The db-crate test `lost_lock_session_cannot_overwrite_successor_snapshot` already polls for this release. The relay test didn't. #### Fix Add `StorageAccountingLeader::release(self)`. It runs `pg_advisory_unlock` on the leader's own session and then closes it, so the lease is free when it returns. This follows the unlock-then-close pattern in `migration.rs`. The relay test now releases the worker at the end of each iteration instead of dropping it. No assertions, timeouts or retries change. buzz-admin's `storage-snapshot` is left as is. It runs once per pod, and its CronJob uses `concurrencyPolicy: Forbid`, `backoffLimit: 0` and `restartPolicy: Never`, so the next lock attempt comes from a later scheduled run, not from the same process. `released_lease_admits_successor_before_old_session_exits` covers the new method. A closing backend drops its temp tables before it frees session locks. So the test locks a temp table owned by the leader's session from another transaction, which keeps the old backend alive after release. It then asserts that a fresh `try_lock` succeeds at once and that the old backend has not exited. Open PRs #9040 (edits this relay test and keeps `drop(worker)`) and #9043 (changes the `sqlx` import in storage_accounting.rs) will see small textual conflicts. #### Validation Local Linux workstation, throwaway `postgres:16-alpine` (CI uses `postgres:16`), desired-state schema applied as in the PostgreSQL lane: - Adversarial ordering, using temporary, uncommitted changes: the leader's connection was made public and the relay test held the worker backend's exit for 1 s with the same temp-table lock. Without the fix: 15/15 failed at postgres_tests.rs:66:74 after 5.1-5.3 s, the CI signature. With the fix: 10/10 passed. - Natural contention: Postgres, the test processes and busy loops were pinned to 4 CPUs. With 8 parallel runs and 4 busy loops, the unmodified test passed 24/24 without the fix. With 16 parallel runs and 16 busy loops, it passed 48/48 without the fix and 48/48 with it. Load alone did not open the window locally. - New db test: 50/50 passed. With the unlock removed from `release`, leaving only `close`, it fails with "a released lease must admit a successor at once". All four storage_accounting Postgres tests pass. - `cargo fmt --all -- --check`, `cargo clippy -p buzz-db -p buzz-relay -p buzz-admin --all-targets -- -D warnings`, `just file-size-check` and `git diff --check` are clean. A green run doesn't prove the flake is gone. Hosted CI covers broad validation. **Rebased onto origin/main `64d12e297` (this branch):** the relay test, run in a fresh database cloned from the desired-state template each time, passed 8/8 at the branch head, which includes this commit. The full rebased results are under the next commit. #### Rebase note #9040, which this commit's message lists as an open PR, has since merged to main. It removed the per-community storage metric, including the `buzz_community_storage_bytes` assertion just above `drop(worker)` in `crates/buzz-relay/src/storage_sweep/postgres_tests.rs`. The cherry-pick onto current main conflicted on that hunk. The resolution keeps main's version of the test and changes only `drop(worker)` to `worker.release().await.expect("release worker lease")` with its comment. Nothing else in the commit changed, and no semantic adjustment was needed. The commit message is unchanged, so it still describes #9040 as open. ### fix(db): report a storage-accounting release only once its unlock succeeded #### Problem Second review of the branch (nit) on 51c27212a ("fix(db): unlock the storage-accounting lease before its session closes"), `StorageAccountingLeader::release` in storage_accounting.rs. That commit cites the unlock-then-close pattern in runtime/migration.rs but handles errors differently: - A failed unlock returned early with `?`, skipping the graceful close. The dropped connection then closed only the socket. - A failed close after a successful unlock was returned as an error, although the lease was already free. - `pg_advisory_unlock` returning false, which means this session didn't hold the lock, was ignored and reported as success. So the doc comment's "Once this returns, the lease is already free" did not hold for every `Ok`. #### Cause `release` propagated the unlock's error with `?` before the close, propagated the close's error after a successful unlock, and ignored the boolean that `pg_advisory_unlock` returns. So a failed unlock skipped the graceful close, a failed close was reported although the lease was already free, and an unlock that returned false because this session no longer held the lock was reported as success. #### Fix Follow migration.rs's `with_schema_destruction_lock`. Run the unlock as `query_scalar::` and keep its result, always attempt the close and ignore its error, then return the unlock's error, or `InvalidData` if the unlock returned false. The doc comment now says that `Ok` means the unlock ran on the lease's own session, so the lease is free even if the close failed, and that an error means the unlock failed or the session no longer held the lease. `released_lease_admits_successor_before_old_session_exits` gains a final step: it unlocks the successor's lease on the successor's own session, then expects `release` to fail. Its successor check is now an `expect` with the same message. #### Validation Local Linux workstation, the `postgres:16` container used earlier on this branch, `BUZZ_TEST_DATABASE_URL` pointing at it: - `cargo test -p buzz-db --lib storage_accounting -- --ignored`: 4/4 storage_accounting Postgres tests pass. - With the unlock's boolean ignored, as before this change, the new step fails with "releasing a lease the session lost must fail". - `cargo fmt -p buzz-db -- --check`, `cargo clippy -p buzz-db --all-targets -- -D warnings` and `git diff --check` are clean. **Rebased onto origin/main `64d12e297` (this branch, both commits):** Linux workstation, private `postgres:16` container with the desired-state schema applied as in the PostgreSQL lane (pgschema plus the reconcile script). - `cargo test -p buzz-db`: 191 lib, 21 `observability_source` and 5 doc tests passed. The 674 Postgres-backed tests are `#[ignore]`d there. - `cargo test -p buzz-db --lib storage_accounting -- --ignored --test-threads=1`: 4/4 passed. `released_lease_admits_successor_before_old_session_exits` alone: 20/20 passed. - `storage_sweep::postgres_tests::worker_publication_activates_reader_without_restart_or_mode_change`, each run in a fresh database cloned from the desired-state template: 8/8 passed. - `cargo clippy -p buzz-db --all-targets -- -D warnings`, `cargo clippy -p buzz-relay --all-targets -- -D warnings`, `cargo fmt -p buzz-db -p buzz-relay -- --check` and `git diff --check` are clean. ## Coordination - #9043 (draft, "feat(relay): serve storage usage as an operator pull feed") edits `crates/buzz-db/src/store/storage_accounting.rs`. `git merge-tree` against its current head shows two small textual conflicts. First, both sides rewrite the `sqlx` import to the same `use sqlx::{postgres::PgConnection, Connection};`, and #9043 also drops `use uuid::Uuid;`. Second, #9043 replaces `usage_delivery_reads_snapshot_time_and_community_states` at the point where this branch adds `released_lease_admits_successor_before_old_session_exits`. Keeping both sides resolves each one. - #7899 also touches `storage_accounting.rs`. It already conflicts with main, and this branch adds no new conflicted file to it. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: mtoohey Co-authored-by: Claude Opus 5.5 --- .../buzz-db/src/store/storage_accounting.rs | 93 ++++++++++++++++++- .../src/storage_sweep/postgres_tests.rs | 4 +- 2 files changed, 94 insertions(+), 3 deletions(-) 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; }