Use worker snapshots for relay storage metrics (#7845)

## Summary

Make the relay read completed storage-accounting snapshots from
PostgreSQL by default. An environment with no snapshot emits no storage
metrics. When a worker publishes its first result, the relay picks it up
on the next successful leader usage tick, without a restart or
configuration change.

This changes how storage usage is measured and reported. **S3 remains
the object store. PostgreSQL stores the completed counts, byte totals,
community breakdown, and calculation metadata—not the objects
themselves.**

### Why move relay storage accounting entirely to database reads?

Today the relay supports two sources: its own S3 scan (`inline`, the
default) and the worker's saved snapshot (`external`). That requires
operators to deploy a worker and then coordinate a separate relay-mode
change in each environment.

The worker already owns the S3 listing and calculation. It runs in a
separate process with its own object cap, memory limit, and deadline.
Keeping a second scan inside the serving relay retains the resource
pressure and failure modes that this isolation is meant to remove.

PostgreSQL provides a durable handoff. The worker replaces one completed
snapshot atomically after a successful calculation. The relay reads that
row and exposes its values through the existing metrics endpoint. A
relay restart or leadership change can load the last completed result
without scanning the bucket again.

There is deliberately no S3 fallback when a snapshot is absent or a read
fails. Such a fallback would put the expensive scan back in the relay
precisely when a worker is absent or unhealthy. It would also leave two
calculation paths with different resource limits and freshness behavior.

### What happens if no worker exists?

**A worker that has never published is a normal, successful inactive
state.** The snapshot table must exist through the normal database
schema, but it may contain zero rows. The reader returns success, emits
no storage totals or failed-load gauge, and checks again on the next
usage tick. It does not exit the relay, start a worker, or list S3.

The relay checks for a completed database row; it does not inspect
Kubernetes Jobs or CronJobs. Worker presence and snapshot presence are
therefore different:

| Situation | Storage-reader behavior |
|---|---|
| No worker and no snapshot row | Return success, emit no storage
metrics, and query again next tick. |
| Worker starts later and publishes its first row | Read and emit it on
the next successful leader usage tick; no restart or mode change. |
| Worker is removed or fails, but its last valid row remains | Continue
reporting the last completed totals with their increasing age. |
| Query succeeds but a previously present row has been removed | Clear
the cache and stop refreshing storage series; existing exporter series
expire through the idle timeout. |
| Database read fails, times out, or returns an invalid snapshot | Log a
warning, set `buzz_storage_snapshot_load_ok=0`, retain any last-good
cached totals, and retry next tick. |
| Snapshot table is missing | Treat this as a schema error, not normal
worker absence; follow the failed-read behavior above. |
| `BUZZ_STORAGE_METRICS=off` | Skip the snapshot query and all storage
metric emission. |

No row means “not measured,” not zero bytes. A valid completed snapshot
of an empty bucket can report zero. If a read fails before any good
snapshot has been cached, only failed-load health is emitted.

These reads happen in the existing background metrics task, not the
relay startup path. A storage-read failure does not request process exit
or change readiness directly. A wider database outage can still affect
the relay's existing database-dependent operations and readiness.

### Read, parse, and emit path

```text
Worker: S3 listing -> BucketSnapshot -> JSON -> PostgreSQL
Relay: PostgreSQL -> JSON -> BucketSnapshot -> existing metric gauges
Collector: relay /metrics endpoint -> Datadog
```

1. `main.rs::run_storage_sweep_tick` calls the reader from the existing
leader-only usage task. The default interval remains 300 seconds.
2. `storage_sweep.rs::run_storage_metrics_tick` coordinates the read,
cache update, failure health, and metric emission.
3. `refresh_persisted_snapshot` calls the existing
`Db::load_storage_accounting_snapshot()` method. A five-second timeout
covers pool acquisition and the query. It uses the relay's existing
pool.
4. The existing query in `buzz-db/src/store/storage_accounting.rs` reads
`snapshot`, `completed_at`, `duration_ms`, `max_objects`, and `code_sha`
from the singleton row.
5. `serde_json::from_value(stored.snapshot)` decodes the JSON into the
existing `buzz_media::BucketSnapshot` type. Rust infers that type from
`cache_persisted_snapshot`'s argument. The generated `Deserialize`
implementation handles all fields, including the UUID-keyed
`per_community` map.
6. `emit_cached_storage_metrics` sets the existing physical/logical
totals and community byte/object gauges. The existing Prometheus
exporter exposes them for Datadog collection; this code does not send
snapshot JSON to Datadog.

The worker's `BucketSnapshot`/`CommunityStorage` structures, JSON field
names, SQL publication method, and database schema are unchanged. The
old external-mode JSON conversion moves out of `main.rs` into the reader
helper. No new format or hand-written field parser is needed. Duration
and object-cap metadata are validated before replacing the cache.

The five-second relay read timeout is separate from the worker's
30-second startup acquisition budget introduced in #7770. This PR does
not change the worker's pool or startup logic.

### Configuration and code changes

| Configuration | Before | After |
|---|---|---|
| Unset, `inline`, or `on` | Relay-local S3 scan | Read completed
database snapshots. `inline` logs one migration warning. |
| `external` or `snapshot` | Read completed database snapshots |
Continue reading snapshots. |
| `off` | Disabled | Disabled. |
| Unknown or empty value | Disabled with a configuration error | Same. |

Existing `inline` settings intentionally become reader aliases. Merely
changing the default would leave older deployments on their explicit
`inline` setting and preserve the need for coordinated configuration
PRs.

Remove relay-local scan scheduling, in-flight scan state, scan
configuration, and obsolete scan-attempt tests and health metrics.
Preserve the physical/logical total and per-community metric names,
leader-only emission, community scope filtering, and cleanup of old
community labels.

The shared relay-state field remains in place; its comment now describes
a cached worker result. The Helm chart defaults to `relayMode: external`
while leaving `storageAccounting.enabled: false`. Its rendering test
verifies that the reader is configured even when no worker is created.
The reader regression tests cover missing rows, later publication,
updates, invalid data, timeout, removal, reactivation, and the explicit
off switch. `docs/storage-accounting.md` describes this operating model.

### Freshness and rollout impact

`buzz_storage_snapshot_age_seconds` uses the worker's original
completion timestamp. Re-reading an old row never makes it fresh.
`buzz_storage_snapshot_load_ok=1` means the row was read and decoded; it
does not prove that the worker's latest attempt succeeded.

Monitor snapshot age against the worker's schedule and allowed runtime,
and monitor Job failures separately. The old relay scan-attempt health
gauges are retired because the relay no longer runs those scans.

Deploy this relay version through the usual release process. Once it
reaches an environment, workers can be enabled independently.
Already-released charts that still set `inline` work with the new
binary; future chart releases default explicitly to `external`. An
intentional `off` override still requires an operator to enable the
reader.

**Environments with no completed snapshot lose relay-calculated storage
totals after this upgrade.** That is the intended tradeoff for removing
scans from the serving process. Environments with a prior snapshot
continue showing its last completed totals and age. This PR does not
install workers in additional environments or deploy the new relay
image.

### Related issue

Related to #4601. Follows the worker startup fix in #7770. The separate
fleet-usage collector work in #7176 overlaps this code and may need
reconciliation when it merges.

### Testing

Built and ran the relay locally against isolated PostgreSQL and Redis
with the legacy `inline` setting, five-second usage ticks, and a
15-second metric idle timeout. The same process was exercised through
these steps:

1. Start with an empty snapshot table: relay becomes ready and exposes
no storage series.
2. Publish a snapshot containing 123 bytes: the community gauge becomes
123.
3. Replace it with 456 bytes: the gauge becomes 456.
4. Remove the row: storage series expire without fabricating a zero-byte
measurement.
5. Publish 789 bytes: metrics reactivate without restarting the relay;
readiness remains healthy.

The local S3 endpoint recorded zero requests. The unrelated Git
conformance startup probe was disabled for this isolated reader check;
the test covers storage accounting, not every other use of S3 in the
relay. No new relay image was deployed to staging or production during
this test.

The broad local integration run also hit an existing
`p0_pool_acquisitions_use_typed_operation_pairs_without_other` failure
in `buzz-db/tests/observability_source.rs`. The same source-only test
fails on unmodified `main` at `0ef7a2222`; this PR changes neither the
test nor its failing source.

The full local `just ci` run stopped at mobile native-asset setup: the
`objective_c` build hook received no SDK path from `xcrun`. Mobile tests
did not run.

Generated with Codex

Signed-off-by: Ravneet Arora <rarora@squareup.com>
This commit is contained in:
ravarora2
2026-09-23 14:39:39 -04:00
committed by GitHub
parent 3b2e50b15c
commit 4142d2b84e
7 changed files with 411 additions and 866 deletions

View File

@@ -1745,84 +1745,26 @@ async fn run_usage_metrics_tick(
Ok(())
}
/// Storage-sweep half of the leader-only tick: harvest/spawn (never awaits
/// the sweep itself) then re-emit whatever snapshot is cached. Split out of
/// [`run_usage_metrics_tick`] because it has its own always-on config
/// (independent of `EmissionScope`) and a hard kill switch — a disabled
/// sweep must never touch a single storage-family gauge, including the
/// health ones, so a relay without `s3:ListBucket` can turn the whole
/// feature off cleanly.
/// Read worker storage snapshots only on the existing leader metrics tick.
async fn run_storage_sweep_tick(
state: &AppState,
emission_scope: &EmissionScope,
host_map: &HashMap<Uuid, String>,
) {
static SWEEP_CONFIG: std::sync::OnceLock<storage_sweep::StorageSweepConfig> =
static MODE: std::sync::OnceLock<storage_sweep::StorageMetricsMode> =
std::sync::OnceLock::new();
// SWEEP_CONFIG is a function-local OnceLock by design: it is localized
// feature config consumed only by this code path, read once on the first
// leader tick, and stable for the process lifetime (env is immutable).
// Keeping it here avoids widening Config/AppState for a single consumer.
let config = *SWEEP_CONFIG.get_or_init(storage_sweep::StorageSweepConfig::from_env);
match config.mode {
storage_sweep::StorageMetricsMode::Disabled => return,
storage_sweep::StorageMetricsMode::Inline => {
let media_storage = Arc::clone(&state.media_storage);
let max_objects = config.max_objects;
storage_sweep::maybe_spawn_sweep(
&state.storage_sweep,
config.interval,
config.timeout,
async move {
buzz_media::fold_bucket_listing(max_objects, move |token| {
let media_storage = Arc::clone(&media_storage);
async move { media_storage.list_page(token, 1000).await }
})
.await
},
)
.await;
}
storage_sweep::StorageMetricsMode::Snapshot => {
match state.db.load_storage_accounting_snapshot().await {
Ok(Some(stored)) => match serde_json::from_value(stored.snapshot) {
Ok(snapshot) => {
let duration = std::time::Duration::from_millis(
u64::try_from(stored.duration_ms).unwrap_or_default(),
);
let max_objects = u64::try_from(stored.max_objects).unwrap_or_default();
storage_sweep::cache_persisted_snapshot(
&state.storage_sweep,
snapshot,
stored.completed_at,
duration,
max_objects,
)
.await;
}
Err(error) => {
warn!(error = %error, "stored storage snapshot is invalid");
storage_sweep::record_persisted_snapshot_load_failure(&state.storage_sweep)
.await;
}
},
Ok(None) => {
storage_sweep::record_persisted_snapshot_load_failure(&state.storage_sweep)
.await;
}
Err(error) => {
warn!(error = %error, "failed to load stored storage snapshot");
storage_sweep::record_persisted_snapshot_load_failure(&state.storage_sweep)
.await;
}
}
}
let mode = *MODE.get_or_init(storage_sweep::StorageMetricsMode::from_env);
if let Err(error) = storage_sweep::run_storage_metrics_tick(
&state.db,
&state.storage_sweep,
mode,
host_map,
|id| emission_scope.allows(id),
)
.await
{
warn!(error = %error, "failed to load stored storage snapshot; retrying next usage tick");
}
storage_sweep::emit_storage_metrics(&state.storage_sweep, config.mode, host_map, |id| {
emission_scope.allows(id)
})
.await;
}
/// Emit the database-derived usage snapshot from the stable leader only.

View File

@@ -703,7 +703,7 @@ pub struct AppState {
pub audit_tx: Option<mpsc::Sender<buzz_audit::NewAuditEntry>>,
/// Media storage client (S3/MinIO).
pub media_storage: Arc<MediaStorage>,
/// Single-flight + cache state for the hourly S3 storage sweep. See
/// Cached worker snapshot and storage metric emission bookkeeping. See
/// `storage_sweep` module docs; shared with the usage-metrics tick via
/// `Arc` the same way other cross-tick poller state lives on `AppState`.
pub storage_sweep: Arc<tokio::sync::Mutex<crate::storage_sweep::StorageSweepState>>,

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,151 @@
//! Exercise the same database-to-metrics tick called by the relay leader.
use super::*;
use buzz_media::CommunityStorage;
use metrics_util::debugging::{DebugValue, DebuggingRecorder};
fn observe_tick(
db: &Db,
state: &Mutex<StorageSweepState>,
mode: StorageMetricsMode,
hosts: &HashMap<Uuid, String>,
) -> (anyhow::Result<()>, HashMap<String, f64>) {
let recorder = DebuggingRecorder::new();
let result = metrics::with_local_recorder(&recorder, || {
futures::executor::block_on(run_storage_metrics_tick(db, state, mode, hosts, |_| true))
});
let values = recorder
.snapshotter()
.snapshot()
.into_vec()
.into_iter()
.filter(|(key, _, _, _)| key.key().name().contains("_storage_"))
.filter_map(|(key, _, _, value)| match value {
DebugValue::Gauge(value) => Some((key.key().name().to_owned(), value.into_inner())),
_ => None,
})
.collect();
(result, values)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[ignore = "requires Postgres"]
async fn worker_publication_activates_reader_without_restart_or_mode_change() {
// The PostgreSQL lane supplies a separate desired-state database per test.
// Local callers must likewise supply an explicitly isolated database.
let url = std::env::var("BUZZ_TEST_DATABASE_URL").expect("isolated BUZZ_TEST_DATABASE_URL");
let pool = sqlx::postgres::PgPoolOptions::new()
.max_connections(4)
.connect(&url)
.await
.expect("connect isolated database");
let db = Db::from_pool(pool.clone());
let community = Uuid::new_v4();
let hosts = HashMap::from([(community, "storage.example.test".to_owned())]);
for configured in [None, Some("inline"), Some("external")] {
sqlx::query("DELETE FROM storage_accounting_snapshots")
.execute(&pool)
.await
.expect("clear isolated snapshot");
let state = Mutex::new(StorageSweepState::default());
let mode = parse_storage_metrics_mode(configured);
let (result, values) = observe_tick(&db, &state, mode, &hosts);
result.expect("no worker is normal");
assert!(
values.is_empty(),
"no row must not emit zeros or failed-load health"
);
let mut worker = db.try_lock_storage_accounting().await.unwrap().unwrap();
for bytes in [100, 250] {
let snapshot = BucketSnapshot {
physical_bytes: bytes,
physical_objects: 1,
logical_bytes: bytes,
logical_objects: 1,
per_community: HashMap::from([(community, CommunityStorage { bytes, objects: 1 })]),
..Default::default()
};
worker
.save_snapshot(
&serde_json::to_value(snapshot).unwrap(),
12,
1_000,
"test-worker",
)
.await
.expect("publish through production worker session");
let (result, values) = observe_tick(&db, &state, mode, &hosts);
result.expect("same running reader picks up worker publication");
assert_eq!(values["buzz_community_storage_bytes"], bytes as f64);
assert_eq!(values["buzz_total_storage_bytes"], bytes as f64);
assert_eq!(values["buzz_storage_snapshot_load_ok"], 1.0);
assert!(!values
.keys()
.any(|key| key.starts_with("buzz_storage_sweep_")));
}
// Corrupt data must not be mistaken for a worker that is not installed.
worker
.save_snapshot(
&serde_json::json!({"invalid": true}),
12,
1_000,
"bad-worker",
)
.await
.unwrap();
let (result, values) = observe_tick(&db, &state, mode, &hosts);
assert!(result.is_err());
assert_eq!(values["buzz_storage_snapshot_load_ok"], 0.0);
assert_eq!(values["buzz_community_storage_bytes"], 250.0);
// A blocked table read is bounded, preserves the cache, and can recover.
let mut blocker = pool.begin().await.unwrap();
sqlx::query("LOCK TABLE storage_accounting_snapshots IN ACCESS EXCLUSIVE MODE")
.execute(&mut *blocker)
.await
.unwrap();
let started = std::time::Instant::now();
let (result, values) = observe_tick(&db, &state, mode, &hosts);
assert!(result.is_err());
assert!(started.elapsed() < Duration::from_secs(8));
assert_eq!(values["buzz_storage_snapshot_load_ok"], 0.0);
assert_eq!(values["buzz_community_storage_bytes"], 250.0);
blocker.rollback().await.unwrap();
sqlx::query("DELETE FROM storage_accounting_snapshots")
.execute(&pool)
.await
.unwrap();
let (result, values) = observe_tick(&db, &state, mode, &hosts);
result.expect("removed worker snapshot is inactive");
assert!(
values.is_empty(),
"do not refresh stale totals or report zero storage"
);
worker
.save_snapshot(
&serde_json::to_value(BucketSnapshot::default()).unwrap(),
1,
100,
"restarted-worker",
)
.await
.unwrap();
let (result, values) = observe_tick(&db, &state, mode, &hosts);
result.expect("reader reactivates after removal and transient failure");
assert_eq!(values["buzz_storage_snapshot_load_ok"], 1.0);
assert_eq!(
values["buzz_total_storage_bytes"], 0.0,
"a real empty snapshot may report zero"
);
assert_eq!(
values["buzz_community_storage_bytes"], 0.0,
"clear the previous community when publication resumes"
);
drop(worker);
}
pool.close().await;
}

View File

@@ -4,7 +4,7 @@ templates:
- templates/deployment.yaml
- templates/secret-chart.yaml
tests:
- it: renders no worker by default
- it: enables the snapshot reader without requiring a worker
set:
relayUrl: wss://buzz.example.com
ownerPubkey: "0000000000000000000000000000000000000000000000000000000000000000"
@@ -17,6 +17,12 @@ tests:
- hasDocuments:
count: 0
template: templates/storage-accounting-cronjob.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: BUZZ_STORAGE_METRICS
value: external
template: templates/deployment.yaml
- it: renders a non-overlapping worker with an isolated cap
set:
relayUrl: wss://buzz.example.com

View File

@@ -329,19 +329,22 @@ externalRedis:
# the endpoint + autogenerated credentials automatically.
#
# Storage metrics are controlled by storageAccounting.relayMode below:
# - inline (default) keeps the relay-local hourly bucket sweep.
# - external reads the latest complete snapshot written by the CronJob.
# - external (default) reads the latest complete snapshot written by a worker.
# No row means no storage gauges; a later worker publication is picked up on
# the next usage tick without restarting or reconfiguring the relay.
# - inline is a legacy alias for external in current relay images.
# A directly configured BUZZ_STORAGE_METRICS=on also selects the reader.
# - off disables storage metrics in the relay.
# storageAccounting.enabled controls the CronJob independently, so operators
# can stage the worker before switching the relay to external mode. Inline and
# worker scans require `s3:ListBucket` on the bucket ARN itself (bucket-level;
# storageAccounting.enabled controls the CronJob independently of the reader.
# Only worker scans require `s3:ListBucket` on the bucket ARN itself (bucket-level;
# distinct from the object-level GetObject/PutObject/DeleteObject permissions
# already required for media). Without it the scan fails AccessDenied and no
# other media functionality is affected. External mode emits
# buzz_storage_snapshot_{age_seconds,duration_seconds,max_objects,
# cap_utilization,load_ok}; it does not infer worker-attempt health from the
# last successful snapshot. Monitor CronJob/Job status for worker failures,
# including deadline expiry and OOM termination.
# last successful snapshot. Old totals retain their original age; a missing
# row stops refreshing all storage series. Monitor snapshot age and CronJob/Job
# status for worker failures, including deadline expiry and OOM termination.
#
# Whole-community deletion (`buzz-admin deletions ...`) permanently removes
# tenant-owned object versions through the v5 deletion path. In addition to the
@@ -351,12 +354,8 @@ externalRedis:
# never-versioned buckets, which S3 exposes through `ListObjectVersions` with the
# `null` version id.
#
# Note: inline-only buzz_storage_sweep_failures is a process-local gauge — on
# leader failover it resets to the new leader's local count, not a global total.
# Note: on a failed sweep attempt, the next retry fires on the next usage tick
# (default 300 s BUZZ_USAGE_METRICS_INTERVAL_SECS), not at sweep-interval
# cadence — so a permanently missing s3:ListBucket yields one cheap LIST call
# per tick until the permission is added.
# Snapshot reads run on the usage tick (default 300 s), only on its elected
# leader. A read is bounded to five seconds and never falls back to S3.
s3:
endpoint: ""
bucket: "buzz-media"
@@ -370,8 +369,8 @@ s3:
accessKey: ""
secretKey: ""
# Isolated S3 accounting worker. Disabled by default, so chart upgrades retain
# the existing relay-local sweep until an operator enables and sizes the job.
# Isolated S3 accounting worker. Disabled by default. The relay is ready to
# consume snapshots whenever an operator enables a worker in this environment.
storageAccounting:
enabled: false
schedule: "0 3 * * *"
@@ -379,7 +378,7 @@ storageAccounting:
activeDeadlineSeconds: 3600
successfulJobsHistoryLimit: 1
failedJobsHistoryLimit: 3
relayMode: inline # inline | external | off
relayMode: external # external | off; inline is a legacy alias
serviceAccountName: "" # defaults to the main Buzz service account
podLabels: {}
podAnnotations:

View File

@@ -0,0 +1,56 @@
# Storage accounting
`buzz-admin storage-snapshot` lists S3, calculates physical and per-community
logical usage, and atomically replaces the completed snapshot in PostgreSQL.
The relay reads that snapshot and exposes the existing storage gauges on its
metrics endpoint for collection by Datadog or another metrics collector.
## Independent rollout
Deploy the snapshot-reading relay version through the usual release process.
It is ready before a worker exists: no database row means no storage gauges.
Add a scheduled worker to each environment when ready. Its first successful
publication becomes visible on the next leader usage tick, without restarting
the relay or changing its mode. The default tick interval is 300 seconds.
The snapshot table is part of the normal database schema, independent of worker
installation. A missing table or inaccessible database is an error, not evidence
that the worker is absent. Snapshot reads have a five-second budget, including
connection acquisition, and retry on the next usage tick. They do not run on
the relay startup path and never fall back to an S3 scan.
`BUZZ_STORAGE_METRICS` accepts:
| Value | Behavior |
|---|---|
| Unset, `external`, `snapshot`, `on` | Read completed worker snapshots. |
| `inline` | Legacy alias for the reader; logs a migration warning once. |
| `off` | Skip storage reads and metrics. |
| Anything else | Log a configuration error and disable storage metrics. |
Relay-local scans have been removed. Existing chart values that set `inline`
need no coordinated configuration change. Environments without a worker will
stop producing storage totals when this relay version arrives. The former
`BUZZ_STORAGE_SWEEP_*` settings no longer apply; size and bound the worker
instead. An intentional `off` override still requires an operator to enable it.
## Freshness and failures
The reader publishes the last completed totals together with
`buzz_storage_snapshot_age_seconds`. Re-reading an old row never resets its
completion time. Alert on age using a threshold appropriate for the worker's
schedule and allowed run time; for a daily worker, allow more than 24 hours.
`buzz_storage_snapshot_load_ok=1` means the row was read and decoded, not that
the worker's most recent attempt succeeded. Check Job failures separately.
A failed or invalid read logs an error, sets load health to zero, and retains
the last good totals with their original age. With no prior good snapshot,
only failed-load health is emitted. A successful read that finds no row clears
the cache and stops refreshing storage series. Existing exported series expire
through the metrics exporter's idle timeout; absence is never reported as zero
bytes. A valid empty snapshot can report zero bytes.
The reader preserves leader-only emission, community scope filtering, and
cleanup of series for renamed or removed communities. It does not need S3
listing permission. The worker still needs that permission and its database
configuration.