fix(admin): allow cold storage worker DB startup

Justfile

@@ -390,7 +390,7 @@ test-unit:
         cargo nextest run -p buzz-media --lib \
             -E 'test(=bucket_index::tests::bucket_snapshot_json_round_trip_preserves_community_keys)'
         cargo nextest run -p buzz-admin \
-            -E 'test(=storage_snapshot_tests::failed_fold_never_invokes_snapshot_persistence)'
+            -E 'test(storage_snapshot)'
         # Multi-tenant conformance gate (buzz-conformance): the independent
         # replay checker + golden fixtures. No infra — pure in-process trace
         # replay — so it belongs in the unit job. Run all targets (lib + the

crates/buzz-admin/Cargo.toml

@@ -37,3 +37,6 @@ sqlx = { workspace = true }
 url = { workspace = true }
 uuid = { workspace = true }
 clap = { version = "4", features = ["derive"] }
+
+[dev-dependencies]
+tokio = { workspace = true, features = ["test-util"] }

crates/buzz-admin/src/main.rs

@@ -21,6 +21,7 @@
 //! the guard against parallel adds (e.g. `xargs -P`).
 
 mod deletions;
+mod storage_snapshot_startup;
 
 use std::future::Future;
 use std::sync::atomic::{AtomicU64, Ordering};
@@ -186,7 +187,7 @@ async fn cmd_storage_snapshot(max_objects: u64) -> Result<i32> {
         return Err(anyhow::anyhow!("--max-objects must be greater than zero"));
     }
 
-    let db = connect_db().await?;
+    let db = storage_snapshot_startup::connect_db().await?;
     let mut leader = db.try_lock_storage_accounting().await?.ok_or_else(|| {
         anyhow::anyhow!("another storage-snapshot worker already holds the lease")
     })?;
@@ -576,17 +577,17 @@ async fn connect_member_services() -> Result<(Db, Arc<PubSubManager>, Keys)> {
 }
 
 async fn connect_db() -> Result<Db> {
+    Ok(Db::new(&db_config_from_env()).await?)
+}
+
+fn db_config_from_env() -> DbConfig {
     let db_url = std::env::var("DATABASE_URL")
         .unwrap_or_else(|_| "postgres://buzz:buzz_dev@localhost:5432/buzz".to_string());
-    let db = Db::new(
-        &DbConfig {
-            database_url: db_url,
-            ..DbConfig::default()
-        }
-        .with_session_timeouts_from_env(),
-    )
-    .await?;
-    Ok(db)
+    DbConfig {
+        database_url: db_url,
+        ..DbConfig::default()
+    }
+    .with_session_timeouts_from_env()
 }
 
 /// Resolve the deployment's tenant from the configured `RELAY_URL` host.

crates/buzz-admin/src/storage_snapshot_startup.rs

@@ -0,0 +1,155 @@
+//! Bounded cold database startup for the run-once storage accounting worker.
+
+use std::future::Future;
+use std::io::ErrorKind;
+use std::time::Duration;
+
+use anyhow::{Context, Result};
+use buzz_db::{Db, DbConfig, DbError};
+use tokio::time::{sleep, timeout, Instant};
+
+const ACQUIRE_TIMEOUT: Duration = Duration::from_secs(30);
+const RETRY_DELAYS: [Duration; 2] = [Duration::from_secs(2), Duration::from_secs(5)];
+
+/// Establish the worker's database pool before it acquires its accounting lease.
+pub(super) async fn connect_db() -> Result<Db> {
+    let config = worker_config(crate::db_config_from_env());
+    connect_with_retry(|| Db::new(&config)).await
+}
+
+fn worker_config(config: DbConfig) -> DbConfig {
+    DbConfig {
+        max_connections: 1,
+        // The worker detaches its lock-owning session. Do not open idle
+        // replacements while that session scans S3 and publishes the result.
+        min_connections: 0,
+        acquire_timeout_secs: ACQUIRE_TIMEOUT.as_secs(),
+        ..config
+    }
+}
+
+// Only initial connection establishment is retried. Once the command owns
+// the advisory lock, reconnecting would lose its publication fence.
+async fn connect_with_retry<T, Connect, Attempt>(mut connect: Connect) -> Result<T>
+where
+    Connect: FnMut() -> Attempt,
+    Attempt: Future<Output = buzz_db::Result<T>>,
+{
+    let started = Instant::now();
+    // Three attempts of at most 30s, plus 2s and 5s backoffs: at most 97s.
+    for attempt in 0..=RETRY_DELAYS.len() {
+        let attempt_started = Instant::now();
+        eprintln!(
+            "{}",
+            serde_json::json!({
+                "event": "storage_snapshot_db_connect_started",
+                "stage": "db_connect",
+                "attempt": attempt + 1,
+                "timeout_ms": ACQUIRE_TIMEOUT.as_millis(),
+            })
+        );
+        // Bound the entire initialization future, including session setup.
+        let result = timeout(ACQUIRE_TIMEOUT, connect())
+            .await
+            .unwrap_or_else(|_| Err(sqlx::Error::PoolTimedOut.into()));
+        match result {
+            Ok(db) => {
+                eprintln!(
+                    "{}",
+                    serde_json::json!({
+                        "event": "storage_snapshot_db_connect_completed",
+                        "stage": "db_connect",
+                        "attempt": attempt + 1,
+                        "attempt_elapsed_ms": attempt_started.elapsed().as_millis(),
+                        "elapsed_ms": started.elapsed().as_millis(),
+                    })
+                );
+                return Ok(db);
+            }
+            Err(error) => {
+                let delay = RETRY_DELAYS
+                    .get(attempt)
+                    .copied()
+                    .filter(|_| retryable(&error));
+                let class = error_class(&error);
+                // Avoid raw connection errors/URLs: configuration and driver
+                // messages can contain credentials. Keep the source on the
+                // returned error, but print only this bounded classification.
+                eprintln!(
+                    "{}",
+                    serde_json::json!({
+                        "event": "storage_snapshot_db_connect_failed",
+                        "stage": "db_connect",
+                        "attempt": attempt + 1,
+                        "attempt_elapsed_ms": attempt_started.elapsed().as_millis(),
+                        "elapsed_ms": started.elapsed().as_millis(),
+                        "error_class": class,
+                        "io_kind": match &error {
+                            DbError::Sqlx(sqlx::Error::Io(error)) => Some(format!("{:?}", error.kind())),
+                            _ => None,
+                        },
+                        "sqlstate": match &error {
+                            DbError::Sqlx(sqlx::Error::Database(error)) => error.code(),
+                            _ => None,
+                        },
+                        "retry_in_ms": delay.map(|value| value.as_millis()),
+                    })
+                );
+                match delay {
+                    Some(delay) => sleep(delay).await,
+                    None => {
+                        return Err(error).with_context(|| {
+                            format!(
+                                "storage snapshot database startup failed after {} attempt(s) and {}ms ({class})",
+                                attempt + 1,
+                                started.elapsed().as_millis(),
+                            )
+                        });
+                    }
+                }
+            }
+        }
+    }
+    unreachable!("the last connection attempt always returns")
+}
+
+fn retryable(error: &DbError) -> bool {
+    match error {
+        DbError::Sqlx(sqlx::Error::PoolTimedOut) => true,
+        // SQLx already backs off on connection refusal and transient server
+        // errors. Other transport/resolver failures may also be transient;
+        // allow only the bounded startup retries, excluding local input and
+        // permission errors. Unknown resolver errors are not labeled as DNS.
+        DbError::Sqlx(sqlx::Error::Io(error)) => !matches!(
+            error.kind(),
+            ErrorKind::InvalidInput
+                | ErrorKind::InvalidData
+                | ErrorKind::PermissionDenied
+                | ErrorKind::NotFound
+                | ErrorKind::Unsupported
+        ),
+        // Includes authentication, TLS, URL configuration and protocol errors.
+        _ => false,
+    }
+}
+
+fn error_class(error: &DbError) -> &'static str {
+    match error {
+        DbError::Sqlx(sqlx::Error::PoolTimedOut) => "timeout",
+        DbError::Sqlx(sqlx::Error::Io(_)) => "io",
+        DbError::Sqlx(sqlx::Error::Tls(_)) => "tls",
+        DbError::Sqlx(sqlx::Error::Configuration(_)) => "configuration",
+        DbError::Sqlx(sqlx::Error::Database(error)) => {
+            if error.code().is_some_and(|code| code.starts_with("28")) {
+                "authentication"
+            } else {
+                "database"
+            }
+        }
+        DbError::Sqlx(sqlx::Error::Protocol(_)) => "protocol",
+        _ => "other",
+    }
+}
+
+#[cfg(test)]
+mod tests;

crates/buzz-admin/src/storage_snapshot_startup/tests.rs

@@ -0,0 +1,117 @@
+use std::cell::Cell;
+
+use super::*;
+
+#[test]
+fn worker_pool_overrides_do_not_change_relay_defaults_or_session_policy() {
+    let defaults = DbConfig::default();
+    let worker = worker_config(DbConfig {
+        lock_timeout_ms: 123,
+        idle_txn_timeout_ms: 456,
+        statement_timeout_ms: 789,
+        ..defaults.clone()
+    });
+    assert_eq!(worker.max_connections, 1);
+    assert_eq!(worker.min_connections, 0);
+    assert_eq!(worker.acquire_timeout_secs, 30);
+    assert_eq!(worker.lock_timeout_ms, 123);
+    assert_eq!(worker.idle_txn_timeout_ms, 456);
+    assert_eq!(worker.statement_timeout_ms, 789);
+    assert_eq!(defaults.max_connections, 20);
+    assert_eq!(defaults.min_connections, 2);
+    assert_eq!(defaults.acquire_timeout_secs, 3);
+}
+
+#[tokio::test(start_paused = true)]
+async fn cold_connection_can_exceed_the_old_three_second_budget() {
+    let started = Instant::now();
+    connect_with_retry(|| async {
+        sleep(Duration::from_secs(6)).await;
+        Ok(())
+    })
+    .await
+    .expect("cold startup succeeds");
+    assert_eq!(started.elapsed(), Duration::from_secs(6));
+}
+
+#[tokio::test(start_paused = true)]
+async fn transient_failure_backs_off_before_a_successful_attempt() {
+    let attempts = Cell::new(0);
+    let started = Instant::now();
+    let result = connect_with_retry(|| {
+        attempts.set(attempts.get() + 1);
+        let attempt = attempts.get();
+        async move {
+            if attempt == 1 {
+                Err(sqlx::Error::Io(ErrorKind::ConnectionReset.into()).into())
+            } else {
+                Ok(42)
+            }
+        }
+    })
+    .await;
+    assert_eq!(result.expect("second attempt succeeds"), 42);
+    assert_eq!(attempts.get(), 2);
+    assert_eq!(started.elapsed(), Duration::from_secs(2));
+}
+
+#[tokio::test(start_paused = true)]
+async fn retry_exhaustion_propagates_the_last_error_after_three_attempts() {
+    let attempts = Cell::new(0);
+    let started = Instant::now();
+    let error = connect_with_retry(|| {
+        attempts.set(attempts.get() + 1);
+        async { Err::<(), _>(sqlx::Error::Io(ErrorKind::ConnectionReset.into()).into()) }
+    })
+    .await
+    .expect_err("exhausted startup must fail");
+    assert_eq!(attempts.get(), 3);
+    assert_eq!(started.elapsed(), Duration::from_secs(7));
+    assert!(error.to_string().contains("3 attempt(s)"));
+    assert!(matches!(
+        error.downcast_ref::<DbError>(),
+        Some(DbError::Sqlx(sqlx::Error::Io(_)))
+    ));
+}
+
+#[tokio::test(start_paused = true)]
+async fn hung_initialization_is_bounded_to_97_seconds_including_backoff() {
+    let attempts = Cell::new(0);
+    let started = Instant::now();
+    let error = connect_with_retry(|| {
+        attempts.set(attempts.get() + 1);
+        std::future::pending::<buzz_db::Result<()>>()
+    })
+    .await
+    .expect_err("hung connection setup must time out");
+    assert_eq!(attempts.get(), 3);
+    assert_eq!(started.elapsed(), Duration::from_secs(97));
+    assert!(matches!(
+        error.downcast_ref::<DbError>(),
+        Some(DbError::Sqlx(sqlx::Error::PoolTimedOut))
+    ));
+}
+
+#[tokio::test(start_paused = true)]
+async fn permanent_errors_fail_immediately_and_do_not_expose_connection_details() {
+    let attempts = Cell::new(0);
+    let started = Instant::now();
+    for driver_error in [
+        sqlx::Error::Configuration("sensitive-url".into()),
+        sqlx::Error::Tls("sensitive-certificate".into()),
+        sqlx::Error::Protocol("sensitive-server-reply".into()),
+        sqlx::Error::Io(ErrorKind::PermissionDenied.into()),
+    ] {
+        let mut error = Some(driver_error);
+        let result = connect_with_retry(|| {
+            attempts.set(attempts.get() + 1);
+            let error = error.take().expect("permanent errors must not retry");
+            async move { Err::<(), _>(error.into()) }
+        })
+        .await
+        .expect_err("permanent error");
+        assert!(!result.to_string().contains("sensitive"));
+    }
+    assert_eq!(attempts.get(), 4);
+    assert_eq!(started.elapsed(), Duration::ZERO);
+}