fix(admin): allow cold storage worker DB startup (#7770)

## Summary

The standalone S3 accounting worker can exit before its scan begins
because database startup inherits the relay’s three-second connection
budget. A cold DNS probe took about five seconds.

Give `buzz-admin storage-snapshot` its own pool: one connection, no
minimum idle connections, and a 30-second acquisition budget. Retry
startup timeouts and transport failures at most twice, after two and
five seconds. Total startup wait is bounded to 97 seconds.
Configuration, authentication, TLS, and protocol errors fail
immediately.

Log each startup attempt’s elapsed time and error category without
connection credentials. Relay defaults stay unchanged. Retries stop
before lock acquisition; the same lock-owning database session still
publishes the completed snapshot.

### Related issue

Follow-up to #7543. No duplicate startup fix found.

### Testing

Archived-command check: built a nextest archive, temporarily moved
`target/debug/buzz-admin` out of the checkout, and ran
`scripts/postgres-test-run.sh --archive-file <archive>`. Both tests
failed before the path fix and passed afterward. The tests use nextest’s
relocated executable path, with the original Cargo path retained for
local runs.

Manually restored the old startup call and ran
`scripts/postgres-test-run.sh -p buzz-admin --test
postgres_storage_snapshot -E 'test(storage_snapshot_retries_then)'`. It
failed before reaching S3; restoring the fix passed. This confirms the
test exercises the actual worker command.

Local `just ci` passed all stages before mobile tests, which stopped at
SDK discovery in the `objective_c` native hook. The default Xcode tools
report an unaccepted license. Two existing PostgreSQL tests hit
connection timeouts in the parallel run and passed individually.

After merge, build a worker image and repeat the cold-pod staging smoke
test through snapshot publication before changing relay metrics mode.

Generated with Codex

---------

Signed-off-by: Ravneet Arora <rarora@squareup.com>
This commit is contained in:
ravarora2
2026-09-22 13:32:20 -04:00
committed by GitHub
parent 77729abfb6
commit a929532a04
7 changed files with 581 additions and 13 deletions

View File

@@ -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

View File

@@ -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"] }

View File

@@ -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.

View File

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

View File

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

View File

@@ -0,0 +1,292 @@
//! Run the real command through delayed/broken PostgreSQL connections.
//! The PostgreSQL nextest wrapper supplies a separate database per test.
use std::process::{ExitStatus, Stdio};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use serde_json::{json, Value};
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::process::Command;
use tokio::task::{JoinHandle, JoinSet};
use tokio::time::{sleep, timeout};
use url::Url;
struct Server {
address: std::net::SocketAddr,
requests: Arc<AtomicUsize>,
task: JoinHandle<()>,
}
impl Drop for Server {
fn drop(&mut self) {
self.task.abort();
}
}
async fn postgres_proxy(database_url: &str, reject: usize, delay: Duration) -> Server {
let target = Url::parse(database_url).expect("test database URL");
let target = (
target.host_str().expect("database host").to_owned(),
target.port().unwrap_or(5432),
);
let listener = TcpListener::bind("127.0.0.1:0").await.expect("proxy");
let address = listener.local_addr().expect("proxy address");
let requests = Arc::new(AtomicUsize::new(0));
let connections = Arc::clone(&requests);
let task = tokio::spawn(async move {
let mut sessions = JoinSet::new();
loop {
tokio::select! {
accepted = listener.accept() => {
let (mut client, _) = accepted.expect("accept database connection");
let index = connections.fetch_add(1, Ordering::SeqCst);
let target = target.clone();
sessions.spawn(async move {
if index < reject {
// Wait for the client handshake before dropping the
// connection, so this is an I/O failure, not a refused
// connection SQLx would retry inside the same attempt.
let mut header = [0; 4];
let _ = client.read_exact(&mut header).await;
return;
}
sleep(delay).await;
let mut upstream = TcpStream::connect(target).await.expect("real Postgres");
let _ = tokio::io::copy_bidirectional(&mut client, &mut upstream).await;
});
}
result = sessions.join_next(), if !sessions.is_empty() => {
result.expect("session").expect("proxy session completed");
}
}
}
});
Server {
address,
requests,
task,
}
}
async fn s3_listing() -> Server {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("S3 listener");
let address = listener.local_addr().expect("S3 address");
let requests = Arc::new(AtomicUsize::new(0));
let observed = Arc::clone(&requests);
let task = tokio::spawn(async move {
let sha = "a".repeat(64);
let contents: String = [
(format!("{sha}.bin"), 100),
(format!("{sha}.thumb.jpg"), 10),
(
format!("_meta/00000000-0000-0000-0000-000000000001/{sha}.json"),
30,
),
]
.into_iter()
.map(|(key, size)| {
format!(
"<Contents><Key>{key}</Key><Size>{size}</Size>\
<LastModified>2026-09-20T00:00:00.000Z</LastModified>\
<ETag>\"test\"</ETag><StorageClass>STANDARD</StorageClass></Contents>"
)
})
.collect();
let body = format!(
"<ListBucketResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\">\
<Name>test-bucket</Name><Prefix></Prefix><MaxKeys>1000</MaxKeys>\
<IsTruncated>false</IsTruncated>{contents}\
</ListBucketResult>"
);
loop {
let (mut stream, _) = listener.accept().await.expect("S3 connection");
observed.fetch_add(1, Ordering::SeqCst);
let mut request = Vec::new();
let mut chunk = [0; 1024];
while !request.windows(4).any(|bytes| bytes == b"\r\n\r\n") {
let read = stream.read(&mut chunk).await.expect("S3 request");
assert!(read > 0 && request.len() + read <= 16 * 1024);
request.extend_from_slice(&chunk[..read]);
}
assert!(String::from_utf8_lossy(&request).contains("list-type=2"));
let response = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/xml\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
);
stream
.write_all(response.as_bytes())
.await
.expect("S3 response");
}
});
Server {
address,
requests,
task,
}
}
async fn database() -> (String, PgPool) {
let url = std::env::var("BUZZ_TEST_DATABASE_URL")
.expect("run through scripts/postgres-test-run.sh for an isolated database");
let pool = PgPoolOptions::new()
.max_connections(1)
.connect(&url)
.await
.expect("test DB");
sqlx::query(
"INSERT INTO storage_accounting_snapshots \
(singleton, snapshot, completed_at, duration_ms, max_objects, code_sha) \
VALUES (TRUE, $1, NOW(), 1, 100, 'before')",
)
.bind(json!({"previous": "complete"}))
.execute(&pool)
.await
.expect("seed last-good snapshot");
(url, pool)
}
async fn capture(reader: impl AsyncRead + Unpin) -> std::io::Result<String> {
let mut bytes = Vec::new();
reader.take(64 * 1024).read_to_end(&mut bytes).await?;
Ok(String::from_utf8_lossy(&bytes).into_owned())
}
struct WorkerOutput {
status: ExitStatus,
stdout: String,
stderr: String,
}
async fn run_worker(database_url: &str, proxy: &Server, s3: &Server) -> WorkerOutput {
let mut url = Url::parse(database_url).expect("database URL");
url.set_host(Some("127.0.0.1")).expect("proxy host");
url.set_port(Some(proxy.address.port()))
.expect("proxy port");
// The proxy deliberately interrupts the PostgreSQL handshake. Keep TLS
// out of this local transport test so it exercises sqlx::Error::Io.
url.query_pairs_mut().append_pair("sslmode", "disable");
// nextest relocates executables when extracting a test archive. Keep the
// compile-time Cargo path only as the fallback for local cargo test runs.
let worker_binary = std::env::var_os("NEXTEST_BIN_EXE_buzz_admin")
.unwrap_or_else(|| env!("CARGO_BIN_EXE_buzz-admin").into());
let mut child = Command::new(worker_binary)
.args(["storage-snapshot", "--max-objects", "100"])
.env_clear()
.env("DATABASE_URL", url.as_str())
.env("BUZZ_S3_ENDPOINT", format!("http://{}", s3.address))
.env("BUZZ_S3_BUCKET", "test-bucket")
.env("BUZZ_S3_REGION", "us-east-1")
.env("BUZZ_S3_ACCESS_KEY", "test-access-key")
.env("BUZZ_S3_SECRET_KEY", "test-secret-key")
.env("BUZZ_STORAGE_SNAPSHOT_CODE_SHA", "startup-test")
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true)
.spawn()
.expect("start real worker");
let stdout = child.stdout.take().expect("stdout");
let stderr = child.stderr.take().expect("stderr");
let completed = timeout(Duration::from_secs(25), async {
tokio::try_join!(child.wait(), capture(stdout), capture(stderr))
})
.await;
let (status, stdout, stderr) = match completed {
Ok(output) => output.expect("worker output"),
Err(_) => {
child.kill().await.expect("kill hung worker");
panic!("worker did not finish within 25 seconds");
}
};
WorkerOutput {
status,
stdout,
stderr,
}
}
fn events(output: &str) -> Vec<Value> {
output
.lines()
.filter_map(|line| serde_json::from_str(line).ok())
.collect()
}
#[tokio::test]
#[ignore = "requires Postgres"]
async fn storage_snapshot_retries_then_connects_after_the_old_three_second_budget() {
let (url, pool) = database().await;
let proxy = postgres_proxy(&url, 1, Duration::from_secs(5)).await;
let s3 = s3_listing().await;
let output = run_worker(&url, &proxy, &s3).await;
assert!(
output.status.success(),
"{}\n{}",
output.stdout,
output.stderr
);
assert_eq!(
proxy.requests.load(Ordering::SeqCst),
2,
"one rejected connection, one lock-owning session, no spare pool connections"
);
assert_eq!(s3.requests.load(Ordering::SeqCst), 1);
let startup = events(&output.stderr);
assert!(startup.iter().any(
|event| event["event"] == "storage_snapshot_db_connect_failed"
&& event["attempt"] == 1
&& event["retry_in_ms"] == 2000
));
assert!(startup.iter().any(
|event| event["event"] == "storage_snapshot_db_connect_completed"
&& event["attempt"] == 2
&& event["attempt_elapsed_ms"].as_u64().unwrap() >= 5000
));
let scan = events(&output.stdout);
assert_eq!(scan.first().unwrap()["event"], "storage_snapshot_started");
assert_eq!(scan.last().unwrap()["event"], "storage_snapshot_completed");
let (snapshot, revision): (Value, String) = sqlx::query_as(
"SELECT snapshot, code_sha FROM storage_accounting_snapshots WHERE singleton = TRUE",
)
.fetch_one(&pool)
.await
.expect("saved snapshot");
assert_eq!(revision, "startup-test");
assert_eq!(snapshot["physical_objects"], 3);
assert_eq!(snapshot["physical_bytes"], 140);
assert_eq!(snapshot["logical_bytes"], 110);
assert_eq!(snapshot["logical_objects"], 1);
}
#[tokio::test]
#[ignore = "requires Postgres"]
async fn storage_snapshot_exhaustion_never_lists_s3_or_replaces_the_last_good_snapshot() {
let (url, pool) = database().await;
let proxy = postgres_proxy(&url, usize::MAX, Duration::ZERO).await;
let s3 = s3_listing().await;
let output = run_worker(&url, &proxy, &s3).await;
assert_eq!(output.status.code(), Some(5), "{}", output.stderr);
assert_eq!(proxy.requests.load(Ordering::SeqCst), 3);
assert_eq!(s3.requests.load(Ordering::SeqCst), 0);
assert!(!output.stdout.contains("storage_snapshot_started"));
assert!(!output.stdout.contains("storage_snapshot_completed"));
let failures: Vec<_> = events(&output.stderr)
.into_iter()
.filter(|event| event["event"] == "storage_snapshot_db_connect_failed")
.collect();
assert_eq!(failures.len(), 3);
assert!(failures.last().unwrap()["retry_in_ms"].is_null());
let (snapshot, revision): (Value, String) = sqlx::query_as(
"SELECT snapshot, code_sha FROM storage_accounting_snapshots WHERE singleton = TRUE",
)
.fetch_one(&pool)
.await
.expect("last good snapshot");
assert_eq!(snapshot, json!({"previous": "complete"}));
assert_eq!(revision, "before");
}

View File

@@ -107,8 +107,8 @@ run_unit_tests() {
run_test_step "buzz-media storage snapshot serialization test" \
cargo test -p buzz-media --lib bucket_index::tests::bucket_snapshot_json_round_trip_preserves_community_keys -- --exact --nocapture
run_test_step "buzz-admin completed snapshot persistence test" \
cargo test -p buzz-admin storage_snapshot_tests::failed_fold_never_invokes_snapshot_persistence -- --exact --nocapture
run_test_step "buzz-admin storage snapshot tests" \
cargo test -p buzz-admin storage_snapshot -- --nocapture
# Multi-tenant conformance gate: independent replay checker + golden
# fixtures (buzz-conformance). Pure in-process trace replay, no infra.