diff --git a/Justfile b/Justfile index 3f5bec95a64..29814653674 100644 --- a/Justfile +++ b/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 diff --git a/crates/buzz-admin/Cargo.toml b/crates/buzz-admin/Cargo.toml index 263ba4eb319..6b3701bd817 100644 --- a/crates/buzz-admin/Cargo.toml +++ b/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"] } diff --git a/crates/buzz-admin/src/main.rs b/crates/buzz-admin/src/main.rs index bfb256e311c..49ca21204d3 100644 --- a/crates/buzz-admin/src/main.rs +++ b/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 { 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, Keys)> { } async fn connect_db() -> Result { + 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. diff --git a/crates/buzz-admin/src/storage_snapshot_startup.rs b/crates/buzz-admin/src/storage_snapshot_startup.rs new file mode 100644 index 00000000000..1a9955ed1b1 --- /dev/null +++ b/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 { + 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(mut connect: Connect) -> Result +where + Connect: FnMut() -> Attempt, + Attempt: Future>, +{ + 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; diff --git a/crates/buzz-admin/src/storage_snapshot_startup/tests.rs b/crates/buzz-admin/src/storage_snapshot_startup/tests.rs new file mode 100644 index 00000000000..e69f2c8a615 --- /dev/null +++ b/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::(), + 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::>() + }) + .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::(), + 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); +} diff --git a/crates/buzz-admin/tests/postgres_storage_snapshot.rs b/crates/buzz-admin/tests/postgres_storage_snapshot.rs new file mode 100644 index 00000000000..79b960983ec --- /dev/null +++ b/crates/buzz-admin/tests/postgres_storage_snapshot.rs @@ -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, + 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!( + "{key}{size}\ + 2026-09-20T00:00:00.000Z\ + \"test\"STANDARD" + ) + }) + .collect(); + let body = format!( + "\ + test-bucket1000\ + false{contents}\ + " + ); + 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 { + 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 { + 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"); +} diff --git a/scripts/run-tests.sh b/scripts/run-tests.sh index 8395e9d0580..f4a06635806 100755 --- a/scripts/run-tests.sh +++ b/scripts/run-tests.sh @@ -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.