fix(db): unlock the storage-accounting lease before its session closes (#9097)

## 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::<bool>` 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 <mtoohey@squareup.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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() {

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;
 }