fix(db): give the replica-fence probe its own pool and instrument it (#9297)

## Problem

The replica-fence probe (`replica_fence::run_probe`) commits a heartbeat
token every 500 ms. It borrowed its connection from the **serving writer
pool**, so under load each beat queued behind request traffic.

In a production deployment during a reconnect/replay storm:

-
`buzz_db_pool_acquire_duration_seconds{pool_role="writer",operation="maintenance"}`
went from ≤50 ms max to 2.7–3.0 s, which is the writer pool's acquire
timeout.
- `buzz_db_replica_heartbeat_age_seconds` went from about 0.47 s to
1.25–3.02 s, in the same 10 s buckets.
- Bounded reads then failed the freshness check
(`buzz_db_route_decision{decision="writer",reason="stale"}` went from
0–6 to 1.3k–6.4k per 10 s), so read traffic spilled onto the writer
exactly when it was saturated. Replica lag stayed at tens of
milliseconds the whole time.
- No `replica fence probe failed` warning fired, because each beat
eventually succeeded, just late. The probe had no span or metric of its
own (`maintenance` is shared by 15+ call sites), so this was only
visible indirectly.

## Change

- **Dedicated pool** (`replica_fence::connect_probe_pool`):
  - lazy, one connection, `application_name=buzz-replica-fence-probe`;
- same credentials and `max_lifetime` as the writer, so it follows
failover;
  - modelled on the push producer's pool.
- It doesn't run the writer pool's `after_connect`: the probe never
writes `events`, so no floor-guard arming is needed. Instead it installs
its own session timeouts:
- `lock_timeout` 1 s. While the probe waits for the exclusive floor
lock, new channel writes queue behind it, so failing a beat fast is
cheaper than holding the queue for the serving writer's 5 s budget.
Retained proofs stand.
- `statement_timeout` and `idle_in_transaction_session_timeout` 2 s. The
idle timeout reaps a stalled client that still holds the exclusive lock.
  - A 5 s client-side deadline bounds the whole probe.
  - Cost: one extra writer connection per relay process.
- The `Db::connect_writer_pool` doc now names this pool and the
push-enqueue pool as deliberate exceptions to its "every writer pool
uses this constructor" rule, and says why.
- **Spans** (via `datastore_span`, which also records
`buzz_db_operation_duration_seconds{operation=…}`):
- `replica_fence_probe` carries `acquire_ms`, `lock_ms`, `scan_ms`,
`heartbeat_ms`, `commit_ms`, `beat_gap_ms`, `token` and `error_kind`.
- Child spans `replica_fence_probe_acquire` and
`replica_fence_probe_lock` cover the two waits.
- **Logs:**
- Failure: `warn` with `error_kind` (`acquire_timeout`, `lock_timeout`,
`statement_timeout`, `deadline`, `masked_activity`, …),
`consecutive_failures`, `heartbeat_age_ms` and per-phase timings.
  - Recovery after failures: `info`.

- **Reader spill metric:** new `buzz_db_reader_spill_total{path,
reason}`. It counts only the routed reads the replica could have served
but the writer did. `buzz_db_route_decision{decision="writer"}` already
existed but merges three different causes under `reason="stale"` and
also counts configuration (`disabled`). Reasons:

  | Reason | Meaning |
  |---|---|
| `heartbeat_stale` | The newest beat is older than the bounded-read
budget. The probe is late, not the replica. |
| `heartbeat_missing` | No live proof: before the first beat, or nothing
within `FENCE_STALENESS`. |
| `fence_wall_behind` | The read's upper bound is newer than the fence
wall. |
| `replica_lag` | The session proved only an older entry than the read
needs. |
  | `reader_acquire_timeout` | Couldn't check out a reader connection. |
| `reader_proof_timeout`, `reader_token_behind`,
`reader_validation_error` | The reader couldn't prove the fence. |

`buzz_db_route_decision` keeps its labels and values unchanged for
existing dashboards. History reads rejected on a saturated reader stay
`decision="rejected"` and aren't counted as spills.
- **Heartbeat interval versus budget:**
- The problem: at 500 ms beats, a 1000 ms `BUZZ_REPLICA_READ_MAX_AGE_MS`
tolerates no missed beat. Just before each commit the newest entry is
already about one interval old, so a single late or failed beat puts
every bounded read over budget.
- The default beat interval stays at **500 ms**. Every beat takes the
deployment-wide exclusive floor lock that channel writes queue behind,
so beating faster costs more lock drains. The budget is what should
absorb a slow beat.
- New `min_read_budget(interval)` = interval + `PROBE_LOCK_TIMEOUT` (1
s) + 200 ms, which is **1700 ms** at the default. That covers one beat
that waits out the lock timeout, plus commit latency. `Db::new` warns at
startup when the configured budget is below it. A deployment running a
1000 ms budget will see this warning until the budget is raised.
- The interval is now configurable with
`BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS`. Unset means the 500 ms default.
Values outside 300–10000 ms fail startup. The 300 ms floor keeps the
120-entry proof ring covering more than the 30 s staleness gate, and the
ring assertion checks against that floor.

The proof ordering is unchanged: exclusive lock, then `S`, then activity
scan, then token commit last. `committed_at` is still captured before
the `UPDATE` round trip.

## Tests

- New `fence_probe_beats_while_serving_writer_pool_is_exhausted`: goes
through `Db::new` + `spawn_fence_probe`, holds the serving pool's only
connection, and requires two fresh beats and exactly one probe session.
It fails on `main`'s wiring, which I checked by temporarily reverting
the `spawn_fence_probe` change.
- New `cluster_global_probe_pool_bounds_the_floor_lock_wait`: a held
shared floor lock makes the probe fail with `lock_timeout` after about 1
s, the retained proof survives, and the next beat succeeds after
release.
- New `reader_spills_are_counted_by_cause`:
  - fresh heartbeat → no spill;
- late heartbeat → `heartbeat_stale`, with the legacy `stale` still
emitted;
  - past the staleness gate → `heartbeat_missing`;
  - routing disabled → no spill.
- Unit tests cover the log throttle and the interval/budget helpers. The
epoch-rotation test now runs through the one-connection probe pool.
- The `observability_source` source-policy test now asserts the probe
never takes a serving-writer `Maintenance` checkout and is spawned
through `connect_probe_pool`.
- Existing replica-fence and `spawn_fence_probe` tests pass unchanged,
apart from signature updates.

## Not done / follow-ups

- Locally: `buzz-db` unit tests and clippy (`-D warnings`) pass, and 38
replica-fence and routing Postgres tests pass on a fresh database. I
didn't run the full Postgres suite, because it needs the nextest
per-test-database wrapper, so CI's Postgres lane is the gate there.
- No local relay run yet. This needs a reader-configured relay, then
pool exhaustion to see the spans and logs end to end.
- The deployed `BUZZ_REPLICA_READ_MAX_AGE_MS` lives in deployment
config. Check it against the new startup warning.
- The `PROBE_DEADLINE` path has no test. It's only reachable through a
stall that no server-side timeout catches.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Signed-off-by: Jemiah Westerman <jemiah@squareup.com>
Co-authored-by: Son of Anton <8e0ef2ab02d69a7ac089ef5eca90d1a3fbbdc546f6322dc34d0c4f3deecea7f4@buzz.block.builderlab.xyz>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Jemiah Westerman
2026-11-09 14:06:53 -08:00
committed by GitHub
parent dd9da3ac6d
commit 5b1698ec24
7 changed files with 892 additions and 115 deletions

View File

@@ -299,6 +299,9 @@ pub struct Db {
/// bounded-stale read semantics are a product decision, not an
/// invariant, so the gate ships off.
pub(crate) replica_read_max_age: Option<Duration>,
/// Interval at which [`Db::spawn_fence_probe`]'s probe commits
/// heartbeat tokens ([`DbConfig::replica_fence_probe_interval_ms`]).
pub(crate) fence_probe_interval: Duration,
/// Whether the reader endpoint supports the Aurora PostgreSQL identity
/// function ([`replica_fence::AURORA_IDENTITY_FN`]) — probed
/// once per process on the first routed read (on a plain autocommit
@@ -700,6 +703,14 @@ pub struct DbConfig {
/// than the staleness gate never routes anyway, so a larger budget
/// would only misrepresent the config.
pub replica_read_max_age_ms: u64,
/// Replica fence heartbeat interval in milliseconds (env
/// `BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS`). `0` selects
/// [`replica_fence::DEFAULT_PROBE_INTERVAL`]; other values are clamped
/// to [`replica_fence::MIN_PROBE_INTERVAL`]..=
/// [`replica_fence::MAX_PROBE_INTERVAL`]. Only used when a read pool is
/// configured. The bounded-read budget should be at least
/// [`replica_fence::min_read_budget`] of this interval.
pub replica_fence_probe_interval_ms: u64,
/// Session `lock_timeout` in milliseconds for writer connections (env
/// `BUZZ_DB_LOCK_TIMEOUT_MS`). `0` disables the timeout.
pub lock_timeout_ms: u64,
@@ -735,6 +746,7 @@ impl Default for DbConfig {
max_lifetime_secs: 300,
idle_timeout_secs: 600,
replica_read_max_age_ms: 0,
replica_fence_probe_interval_ms: 0,
lock_timeout_ms: DEFAULT_LOCK_TIMEOUT_MS,
idle_txn_timeout_ms: DEFAULT_IDLE_TXN_TIMEOUT_MS,
statement_timeout_ms: DEFAULT_STATEMENT_TIMEOUT_MS,
@@ -818,6 +830,19 @@ impl Db {
None => None,
};
let replica_read_max_age = read_budget_from_ms(config.replica_read_max_age_ms);
let fence_probe_interval =
replica_fence::probe_interval_from_ms(config.replica_fence_probe_interval_ms);
let min_read_budget = replica_fence::min_read_budget(fence_probe_interval);
if let Some(budget) = replica_read_max_age.filter(|b| *b < min_read_budget) {
tracing::warn!(
budget_ms = budget.as_millis() as u64,
probe_interval_ms = fence_probe_interval.as_millis() as u64,
recommended_min_ms = min_read_budget.as_millis() as u64,
"BUZZ_REPLICA_READ_MAX_AGE_MS cannot absorb one fence heartbeat that waits out \
its lock timeout; bounded replica reads will spill to the writer whenever \
a beat is late"
);
}
let push_enqueue = config
.push_enabled
.then(|| push_enqueue::PushEnqueue::new(&pool));
@@ -830,6 +855,7 @@ impl Db {
read_max_connections,
fence: std::sync::Arc::new(replica_fence::ReplicaFence::new()),
replica_read_max_age,
fence_probe_interval,
reader_aurora_identity: std::sync::Arc::new(std::sync::OnceLock::new()),
history_reads: history_reads::gate_for_read_pool(read_max_connections),
})
@@ -843,6 +869,19 @@ impl Db {
/// Additional writer pools, including the relay audit pool, must use this
/// constructor so they inherit the timeout, floor-guard, and isolation
/// policy installed by [`Db::new`].
///
/// Two single-connection side pools deliberately do not, so serving-pool
/// saturation cannot delay their work. Neither inserts into `events`, so
/// neither needs the `buzz.created_at_floor` guard, and each installs
/// its own shorter timeouts so it fails fast instead of queueing:
/// - [`replica_fence::connect_probe_pool`]: the fence probe takes the
/// floor lock itself (exclusive) and writes only `replica_heartbeat`.
/// - the push-enqueue pool (`push_enqueue::PushEnqueue::new`): it writes
/// only `push_match_queue`, and is lazy so relay startup never depends
/// on push.
///
/// New pools follow the rule above unless they have a reason as specific
/// as these.
pub async fn connect_writer_pool(config: &DbConfig) -> Result<PgPool> {
use connection_observability::{
classify_pool_outcome, record_milestone, DbConnectionStep, DbConnectionStepAttempt,
@@ -1382,6 +1421,7 @@ impl Db {
read_pool: None,
fence: std::sync::Arc::new(replica_fence::ReplicaFence::new()),
replica_read_max_age: None,
fence_probe_interval: replica_fence::DEFAULT_PROBE_INTERVAL,
reader_aurora_identity: std::sync::Arc::new(std::sync::OnceLock::new()),
}
}
@@ -1406,6 +1446,7 @@ impl Db {
read_pool: Some(read_pool),
fence: std::sync::Arc::new(replica_fence::ReplicaFence::new()),
replica_read_max_age: None,
fence_probe_interval: replica_fence::DEFAULT_PROBE_INTERVAL,
reader_aurora_identity: std::sync::Arc::new(std::sync::OnceLock::new()),
}
}
@@ -1442,9 +1483,12 @@ impl Db {
return Ok(false);
}
self.verify_replica_fence_at_boot().await?;
// The probe gets its own connection so serving-pool saturation
// cannot delay the beat (see `replica_fence::connect_probe_pool`).
tokio::spawn(replica_fence::run_probe(
self.pool.clone(),
replica_fence::connect_probe_pool(&self.pool),
std::sync::Arc::clone(&self.fence),
self.fence_probe_interval,
));
Ok(true)
}
@@ -1637,6 +1681,32 @@ impl Db {
}
}
/// Count one routed read the replica could have served that the writer
/// served instead (`buzz_db_reader_spill_total{path, reason}`).
///
/// A subset of `buzz_db_route_decision{decision="writer"}` that leaves
/// out configuration (`disabled`) and splits its ambiguous `stale`
/// reason by cause. Reasons:
/// - `heartbeat_stale`: the newest heartbeat is older than the
/// bounded-read budget (`BUZZ_REPLICA_READ_MAX_AGE_MS`) — the probe
/// is late, not the replica;
/// - `heartbeat_missing`: no live heartbeat proof at all (the probe has
/// not committed one, or none in [`replica_fence::FENCE_STALENESS`]);
/// - `fence_wall_behind`: the read's upper bound is newer than the
/// newest fence wall;
/// - `replica_lag`: the serving session proved only an older entry than
/// the read needs;
/// - `reader_acquire_timeout`, `reader_proof_timeout`,
/// `reader_token_behind`, `reader_validation_error`: the reader
/// session could not be checked out or could not prove the fence.
///
/// History reads rejected on a saturated reader are not spills; they
/// stay `buzz_db_route_decision{decision="rejected"}`.
fn record_spill(path: &'static str, reason: &'static str) {
metrics::counter!("buzz_db_reader_spill_total", "path" => path, "reason" => reason)
.increment(1);
}
/// Record one route decision (Rev 2 observability): which path, where it
/// went, and why.
pub(crate) fn record_route(path: &'static str, decision: &'static str, reason: &'static str) {
@@ -1890,23 +1960,27 @@ impl Db {
// checkout; the connection-local observation still has to prove it.
let Some(newest) = self.fence.newest() else {
Self::record_route(path, "writer", "uninitialized");
Self::record_spill(path, "heartbeat_missing");
return Ok(RouteDecision::Writer);
};
// Precheck helpers against the newest shared entry: if the newest
// cannot satisfy an arm, no proved (older-or-equal) entry can.
let bounded_precheck =
|budget: &Option<Duration>| -> std::result::Result<(), &'static str> {
match budget {
Some(budget) if newest.committed_at.elapsed() <= *budget => Ok(()),
Some(_) => Err("stale"),
None => Err("disabled"),
}
};
let covered_precheck = |upper: &DateTime<Utc>| -> std::result::Result<(), &'static str> {
// cannot satisfy an arm, no proved (older-or-equal) entry can. Each
// failure carries its `buzz_db_route_decision` reason and, when the
// read is a spill rather than a configuration choice, its
// `buzz_db_reader_spill_total` reason.
type Precheck = std::result::Result<(), (&'static str, Option<&'static str>)>;
let bounded_precheck = |budget: &Option<Duration>| -> Precheck {
match budget {
Some(budget) if newest.committed_at.elapsed() <= *budget => Ok(()),
Some(_) => Err(("stale", Some("heartbeat_stale"))),
None => Err(("disabled", None)),
}
};
let covered_precheck = |upper: &DateTime<Utc>| -> Precheck {
if *upper <= newest.fence_wall {
Ok(())
} else {
Err("stale")
Err(("stale", Some("fence_wall_behind")))
}
};
let precheck = match &predicate {
@@ -1914,13 +1988,20 @@ impl Db {
RoutePredicate::Covered { upper, .. } => covered_precheck(upper),
// No upper bound: the caller post-verifies served rows.
RoutePredicate::CoveredPostVerified { .. } => Ok(()),
// Covered first (no budget dependence), else bounded.
// Covered first (no budget dependence), else bounded. With the
// bounded arm disabled the read still spilled for its wall.
RoutePredicate::BoundedOrCovered { upper, .. } => {
covered_precheck(upper).or_else(|_| bounded_precheck(&self.replica_read_max_age))
covered_precheck(upper).or_else(|(_, covered_spill)| {
bounded_precheck(&self.replica_read_max_age)
.map_err(|(reason, spill)| (reason, spill.or(covered_spill)))
})
}
};
if let Err(reason) = precheck {
if let Err((reason, spill)) = precheck {
Self::record_route(path, "writer", reason);
if let Some(spill) = spill {
Self::record_spill(path, spill);
}
return Ok(RouteDecision::Writer);
}
let history = operation == observability::ReaderOperation::SubscriptionHistory;
@@ -1966,6 +2047,7 @@ impl Db {
// The session proves an older entry than the
// predicate needs (replication lag) — fail closed.
Self::record_route(path, "writer", "stale");
Self::record_spill(path, "replica_lag");
Ok(RouteDecision::Writer)
}
}
@@ -1977,6 +2059,7 @@ impl Db {
}
Err(reason) => {
Self::record_route(path, "writer", reason);
Self::record_spill(path, reason);
Ok(RouteDecision::Writer)
}
}

View File

@@ -66,7 +66,9 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use chrono::{DateTime, Utc};
use sqlx::{PgConnection, PgPool, Row};
use sqlx::pool::PoolConnection;
use sqlx::postgres::PgPoolOptions;
use sqlx::{PgConnection, PgPool, Postgres, Row};
use uuid::Uuid;
use buzz_datastore_tracing::datastore_span;
@@ -95,14 +97,52 @@ pub const REPLICA_FLOOR_LOCK_KEY: i64 = 0x62757a7a666c6f72;
/// between machines.
pub const FENCE_CLOCK_MARGIN_SECS: i64 = 5;
/// How often the probe samples the writer and commits a heartbeat token.
/// Default interval at which the probe samples the writer and commits a
/// heartbeat token (env `BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS`).
///
/// 500ms keeps the cadence at least 2x under the smallest sensible bounded
/// budget (`BUZZ_REPLICA_READ_MAX_AGE_MS`, deploy plan 1000ms) so
/// eligibility doesn't flap between beats. Cost is one single-row UPDATE
/// tuple of WAL per beat per pod — ~20 beats/s fleet-wide, <0.1% of the
/// writer.
pub const PROBE_INTERVAL: Duration = Duration::from_millis(500);
/// Every beat takes the deployment-wide exclusive floor lock, and channel
/// writes queue behind it while it is held, so beating faster costs more
/// lock drains per pod, not just WAL. Tolerating a slow beat is the read
/// budget's job (see [`min_read_budget`]), not the interval's.
pub const DEFAULT_PROBE_INTERVAL: Duration = Duration::from_millis(500);
/// Shortest accepted probe interval. Bounded by the retained ring, which
/// must hold more history than [`FENCE_STALENESS`] (see the assertion at
/// [`RING_CAPACITY`]), and by the writer cost of beating faster.
pub const MIN_PROBE_INTERVAL: Duration = Duration::from_millis(300);
/// Longest accepted probe interval: a probe slower than this leaves only a
/// few beats inside [`FENCE_STALENESS`] and needs a multi-second read
/// budget to route at all.
pub const MAX_PROBE_INTERVAL: Duration = Duration::from_secs(10);
/// Map a configured probe interval in milliseconds to the effective one:
/// `0` selects [`DEFAULT_PROBE_INTERVAL`]; other values are clamped to
/// [`MIN_PROBE_INTERVAL`]..=[`MAX_PROBE_INTERVAL`].
pub fn probe_interval_from_ms(ms: u64) -> Duration {
if ms == 0 {
return DEFAULT_PROBE_INTERVAL;
}
Duration::from_millis(ms).clamp(MIN_PROBE_INTERVAL, MAX_PROBE_INTERVAL)
}
/// Commit and scheduling latency a bounded-read budget should absorb on
/// top of one lock-timed-out beat.
const READ_BUDGET_LATENCY_MARGIN: Duration = Duration::from_millis(200);
/// Smallest bounded-read budget (`BUZZ_REPLICA_READ_MAX_AGE_MS`) that
/// survives one beat failing on [`PROBE_LOCK_TIMEOUT`] at `probe_interval`:
/// the newest entry is one interval old when that beat starts, ages by the
/// lock wait, and the next beat (fired at once, since missed ticks are
/// delayed rather than skipped) commits after
/// [`READ_BUDGET_LATENCY_MARGIN`]. 1700ms at the default interval. A
/// smaller budget still works, but bounded reads spill to the writer
/// whenever a beat waits out the floor lock; the relay warns at startup.
pub const fn min_read_budget(probe_interval: Duration) -> Duration {
probe_interval
.saturating_add(PROBE_LOCK_TIMEOUT)
.saturating_add(READ_BUDGET_LATENCY_MARGIN)
}
/// A fence whose newest entry is older than this is stale: the probe has
/// stopped committing tokens and routing eligibility closes until a new
@@ -114,8 +154,43 @@ pub const PROBE_INTERVAL: Duration = Duration::from_millis(500);
/// just stops spending reader checkouts once the probe is evidently dead.
pub const FENCE_STALENESS: Duration = Duration::from_secs(30);
/// `application_name` of the probe's dedicated connection, so it is
/// identifiable in `pg_stat_activity` and database-side wait reports.
pub const PROBE_APPLICATION_NAME: &str = "buzz-replica-fence-probe";
/// Bound on waiting for the probe's connection. The probe loop is the
/// pool's only, strictly sequential user, so this only ever covers a
/// (re)dial — never queueing behind serving traffic.
const PROBE_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(2);
/// Session `lock_timeout` on the probe connection.
///
/// While the probe waits for the exclusive floor lock, new shared
/// requesters (channel-event writes) queue behind it, so a long wait
/// stalls channel writes deployment-wide. A beat that gives up costs
/// nothing — retained proofs stand — so fail fast instead of inheriting
/// the serving writer's longer lock budget.
///
/// Every pod's probe takes the same exclusive lock and the same heartbeat
/// row lock, so probes also queue behind each other. Each holds them for a
/// few milliseconds, so this budget assumes well under a hundred relay
/// processes per writer.
const PROBE_LOCK_TIMEOUT: Duration = Duration::from_secs(1);
/// Session `statement_timeout` and `idle_in_transaction_session_timeout`
/// on the probe connection. Every probe statement is a sub-millisecond
/// point operation; the idle bound reaps a stalled client that still
/// holds the exclusive floor lock.
const PROBE_STATEMENT_TIMEOUT: Duration = Duration::from_secs(2);
/// Client-side bound on one whole probe, covering stalls no server-side
/// timeout observes (for example a dead TCP peer). Comfortably above the
/// acquire timeout plus one lock wait plus the remaining statements.
const PROBE_DEADLINE: Duration = Duration::from_secs(5);
/// How many `(token, fence_wall)` entries the fence retains. At one probe
/// per [`PROBE_INTERVAL`] (500ms) this is ~60 seconds of history — a reader
/// per [`DEFAULT_PROBE_INTERVAL`] (500ms) this is ~60 seconds of history,
/// and at least 36 seconds at [`MIN_PROBE_INTERVAL`] — a reader
/// session lagging further than that behind the newest token fails closed
/// (routes to the writer) rather than proving from thin air. Aurora reader
/// lag is typically tens of milliseconds; a reader minutes behind is a
@@ -128,7 +203,8 @@ const RING_CAPACITY: usize = 120;
// for capacity rather than lag. Compile-checked so a future cadence or
// capacity tweak can't silently shrink the window below the gate.
const _: () = assert!(
RING_CAPACITY as u64 * PROBE_INTERVAL.as_millis() as u64 > FENCE_STALENESS.as_millis() as u64,
RING_CAPACITY as u64 * MIN_PROBE_INTERVAL.as_millis() as u64
> FENCE_STALENESS.as_millis() as u64,
"fence ring must retain more history than the staleness gate"
);
@@ -281,8 +357,13 @@ impl ReplicaFence {
/// Age of the newest retained entry, ungated (observability: how long
/// since the probe last committed a token).
pub fn heartbeat_age(&self) -> Option<Duration> {
self.last_committed_at().map(|at| at.elapsed())
}
/// Commit instant of the newest retained entry, ungated.
fn last_committed_at(&self) -> Option<Instant> {
let inner = self.inner.lock().expect("fence lock poisoned");
inner.ring.back().map(|entry| entry.committed_at.elapsed())
inner.ring.back().map(|entry| entry.committed_at)
}
/// The newest fence wall, or `None` when closed or stale.
@@ -563,6 +644,146 @@ pub enum ProbeError {
/// The single heartbeat row (migration 0026) is missing on the writer.
#[error("replica_heartbeat row missing on the writer — migration 0026 not applied?")]
HeartbeatRowMissing,
/// The whole probe overran [`PROBE_DEADLINE`] without any server-side
/// timeout firing (for example a stalled connection).
#[error("replica fence probe exceeded its {0:?} deadline")]
Deadline(Duration),
}
impl ProbeError {
/// Stable, low-cardinality failure class for spans and logs.
pub fn kind(&self) -> &'static str {
match self {
Self::Writer(sqlx::Error::PoolTimedOut) => "acquire_timeout",
Self::Writer(sqlx::Error::PoolClosed) => "pool_closed",
Self::Writer(sqlx::Error::Database(e)) => match e.code().as_deref() {
// lock_not_available: the probe's `lock_timeout` fired.
Some("55P03") => "lock_timeout",
// query_canceled: the probe's `statement_timeout` fired.
Some("57014") => "statement_timeout",
_ => "database",
},
Self::Writer(sqlx::Error::Io(_)) => "io",
Self::Writer(_) => "writer",
Self::MaskedActivity { .. } => "masked_activity",
Self::HeartbeatRowMissing => "heartbeat_row_missing",
Self::Deadline(_) => "deadline",
}
}
}
/// Connect the probe's dedicated one-connection pool.
///
/// The probe must not borrow the serving writer pool: under load its beat
/// then queues behind request traffic for up to that pool's acquire
/// timeout, every pod's newest token ages past the bounded-read budget at
/// once, and reads spill onto the writer that is already saturated. A pool
/// of its own keeps the beat independent of serving load, the way the push
/// producer's pool keeps push work off it. It costs one extra writer
/// connection per relay process.
///
/// It reuses the writer's credentials and connection lifetime, so it
/// follows writer failover the same way. It does not run the writer pool's
/// `after_connect` policy: the probe never writes `events`, so it needs no
/// floor-guard arming, and it installs its own tighter timeouts (see
/// [`PROBE_LOCK_TIMEOUT`], [`PROBE_STATEMENT_TIMEOUT`]). The pool is lazy,
/// so a writer outage cannot make startup depend on it.
pub fn connect_probe_pool(writer: &PgPool) -> PgPool {
PgPoolOptions::new()
.max_connections(1)
.min_connections(1)
.acquire_timeout(PROBE_ACQUIRE_TIMEOUT)
.max_lifetime(writer.options().get_max_lifetime())
.after_connect(|conn, _meta| {
Box::pin(async move {
sqlx::query(
"SELECT set_config('lock_timeout', $1, false), \
set_config('statement_timeout', $2, false), \
set_config('idle_in_transaction_session_timeout', $2, false)",
)
.bind(millis(PROBE_LOCK_TIMEOUT).to_string())
.bind(millis(PROBE_STATEMENT_TIMEOUT).to_string())
.execute(&mut *conn)
.await?;
Ok(())
})
})
.connect_lazy_with(
(*writer.connect_options())
.clone()
.application_name(PROBE_APPLICATION_NAME),
)
}
/// Wall time of each probe phase, recorded on the probe span and in the
/// failure log. Phases a failed probe never reached stay zero.
#[derive(Debug, Default, Clone, Copy)]
struct ProbeTimings {
/// Waiting for the probe connection (a dial, at most).
acquire: Duration,
/// `BEGIN` plus waiting for the exclusive floor lock.
lock: Duration,
/// Sampling `S` and scanning `pg_stat_activity`.
scan: Duration,
/// The heartbeat `UPDATE ... RETURNING`.
heartbeat: Duration,
/// `COMMIT`.
commit: Duration,
}
/// What one probe observed about itself, beyond its result.
#[derive(Debug, Default, Clone, Copy)]
struct ProbeTrace {
timings: ProbeTimings,
/// Time between the previously retained beat and this one's commit;
/// `None` before the first beat or when this probe failed.
beat_gap: Option<Duration>,
}
fn millis(elapsed: Duration) -> u64 {
u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX)
}
/// Await `phase`, storing its wall time in `slot` whatever its outcome.
///
/// The time is stored on drop, so a phase cut off by [`PROBE_DEADLINE`]
/// still reports how long it had been running — the stalled phase is the
/// one the failure log most needs to name.
async fn timed<T>(slot: &mut Duration, phase: impl std::future::Future<Output = T>) -> T {
struct Stamp<'a> {
slot: &'a mut Duration,
started: Instant,
}
impl Drop for Stamp<'_> {
fn drop(&mut self) {
*self.slot = self.started.elapsed();
}
}
let _stamp = Stamp {
slot,
started: Instant::now(),
};
phase.await
}
/// Check out the probe connection. Its own span (and the
/// `buzz_db_operation_duration_seconds` series the span macro records)
/// makes the probe's connection wait directly visible.
#[datastore_span(name = "replica_fence_probe_acquire", system = "postgresql")]
async fn acquire_probe_connection(pool: &PgPool) -> Result<PoolConnection<Postgres>, sqlx::Error> {
pool.acquire().await
}
/// Take the exclusive replica-floor lock, draining compliant shared-lock
/// channel writers. Spanned separately: this is the probe's only wait on
/// other database sessions.
#[datastore_span(name = "replica_fence_probe_lock", system = "postgresql")]
async fn take_floor_lock(conn: &mut PgConnection) -> Result<(), sqlx::Error> {
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(REPLICA_FLOOR_LOCK_KEY)
.execute(conn)
.await?;
Ok(())
}
/// Take one ordered writer sample: under the exclusive replica-floor lock,
@@ -571,75 +792,79 @@ pub enum ProbeError {
/// The statements are separately awaited on one pinned writer transaction;
/// a single SELECT would not guarantee evaluation order across the
/// subexpressions, reopening the race this ordering exists to close.
async fn sample_writer(writer: &PgPool) -> Result<WriterSample, ProbeError> {
let connection = crate::observability::acquire_writer(
writer,
crate::observability::WriterOperation::Maintenance,
)
async fn sample_writer(
writer: &PgPool,
timings: &mut ProbeTimings,
) -> Result<WriterSample, ProbeError> {
let connection = timed(&mut timings.acquire, acquire_probe_connection(writer)).await?;
let mut tx = timed(&mut timings.lock, async {
let mut tx = sqlx::Transaction::begin(connection, None).await?;
take_floor_lock(&mut tx).await?;
Ok::<_, sqlx::Error>(tx)
})
.await?;
let mut tx = sqlx::Transaction::begin(connection, None).await?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(REPLICA_FLOOR_LOCK_KEY)
.execute(&mut *tx)
.await?;
let (sampled_at, oldest_xact_start) = timed(&mut timings.scan, async {
// 1. S first.
let sampled_at: DateTime<Utc> = sqlx::query_scalar("SELECT clock_timestamp()")
.fetch_one(&mut *tx)
.await?;
// 1. S first.
let sampled_at: DateTime<Utc> = sqlx::query_scalar("SELECT clock_timestamp()")
// 2. Activity scan. Classification (fail closed on anything unknown):
// - `backend_type IS NULL` → masked. CRITICAL: an unprivileged view
// masks `backend_type` itself along with `state`/`xact_start`
// (verified on PG 16/17), so filtering on
// `backend_type = 'client backend'` would silently EXCLUDE masked
// rows and fail open. Masked rows must be detected before any
// backend-type filter;
// - client backend with `state IS NULL`, or a transactional/unknown
// state with NULL `xact_start` → masked → error;
// - client backend, `state = 'idle'`, NULL `xact_start`
// → no transaction → safely ignore;
// - any row with non-NULL `xact_start` (any backend type)
// → include in the minimum. Background
// workers cannot insert events rows, but counting them is the
// conservative direction (a long autovacuum merely holds the fence
// back), and it keeps the classification simple.
// Scope is every other backend regardless of role or application
// name: an admin psql transaction writes under the same trigger and
// must be representable in the oldest-xact term.
//
// Prepared transactions (2PC) are a bucket of their own: while
// prepared they have left `pg_stat_activity` but can still commit
// after the token. Their deferred floor guard already ran at PREPARE,
// so `pg_prepared_xacts.prepared` bounds their rows exactly like
// `xact_start`; fold it into the same minimum.
let row = sqlx::query(
r#"
SELECT
least(
(SELECT min(xact_start)
FROM pg_stat_activity
WHERE pid <> pg_backend_pid()),
(SELECT min(prepared) FROM pg_prepared_xacts)
) AS oldest_xact_start,
(SELECT count(*)
FROM pg_stat_activity
WHERE pid <> pg_backend_pid()
AND (backend_type IS NULL
OR (backend_type = 'client backend'
AND (state IS NULL
OR (state <> 'idle' AND xact_start IS NULL))))
) AS masked
"#,
)
.fetch_one(&mut *tx)
.await?;
// 2. Activity scan. Classification (fail closed on anything unknown):
// - `backend_type IS NULL` → masked. CRITICAL: an unprivileged view
// masks `backend_type` itself along with `state`/`xact_start`
// (verified on PG 16/17), so filtering on
// `backend_type = 'client backend'` would silently EXCLUDE masked
// rows and fail open. Masked rows must be detected before any
// backend-type filter;
// - client backend with `state IS NULL`, or a transactional/unknown
// state with NULL `xact_start` → masked → error;
// - client backend, `state = 'idle'`, NULL `xact_start`
// → no transaction → safely ignore;
// - any row with non-NULL `xact_start` (any backend type)
// → include in the minimum. Background
// workers cannot insert events rows, but counting them is the
// conservative direction (a long autovacuum merely holds the fence
// back), and it keeps the classification simple.
// Scope is every other backend regardless of role or application
// name: an admin psql transaction writes under the same trigger and
// must be representable in the oldest-xact term.
//
// Prepared transactions (2PC) are a bucket of their own: while
// prepared they have left `pg_stat_activity` but can still commit
// after the token. Their deferred floor guard already ran at PREPARE,
// so `pg_prepared_xacts.prepared` bounds their rows exactly like
// `xact_start`; fold it into the same minimum.
let row = sqlx::query(
r#"
SELECT
least(
(SELECT min(xact_start)
FROM pg_stat_activity
WHERE pid <> pg_backend_pid()),
(SELECT min(prepared) FROM pg_prepared_xacts)
) AS oldest_xact_start,
(SELECT count(*)
FROM pg_stat_activity
WHERE pid <> pg_backend_pid()
AND (backend_type IS NULL
OR (backend_type = 'client backend'
AND (state IS NULL
OR (state <> 'idle' AND xact_start IS NULL))))
) AS masked
"#,
)
.fetch_one(&mut *tx)
let masked: i64 = row.get("masked");
if masked > 0 {
return Err(ProbeError::MaskedActivity { masked });
}
let oldest_xact_start: Option<DateTime<Utc>> = row.get("oldest_xact_start");
Ok((sampled_at, oldest_xact_start))
})
.await?;
let masked: i64 = row.get("masked");
if masked > 0 {
return Err(ProbeError::MaskedActivity { masked });
}
let oldest_xact_start: Option<DateTime<Utc>> = row.get("oldest_xact_start");
// 3. Token commit LAST, on the same pinned connection. The single-row
// UPDATE serializes concurrent pods' probes, so RETURNING token is
@@ -647,10 +872,13 @@ async fn sample_writer(writer: &PgPool) -> Result<WriterSample, ProbeError> {
// round trip so `elapsed()` over-estimates the observation's age —
// the conservative direction for the head-freshness bound.
let committed_at = Instant::now();
let row = sqlx::query(
"UPDATE replica_heartbeat SET token = token + 1 WHERE id = 1 RETURNING token, epoch",
let row = timed(
&mut timings.heartbeat,
sqlx::query(
"UPDATE replica_heartbeat SET token = token + 1 WHERE id = 1 RETURNING token, epoch",
)
.fetch_optional(&mut *tx),
)
.fetch_optional(&mut *tx)
.await?
.ok_or(ProbeError::HeartbeatRowMissing)?;
@@ -661,7 +889,7 @@ async fn sample_writer(writer: &PgPool) -> Result<WriterSample, ProbeError> {
epoch: row.get("epoch"),
committed_at,
};
tx.commit().await?;
timed(&mut timings.commit, tx.commit()).await?;
Ok(sample)
}
@@ -685,8 +913,75 @@ fn fence_wall(sample_s: DateTime<Utc>, oldest_xact_start: Option<DateTime<Utc>>)
/// writer and records the rotated token, so a reader still serving the
/// pre-rewind timeline (whose old, higher token would otherwise satisfy
/// `token >= M`) fails the epoch check instead of proving stale coverage.
///
/// `writer` is the probe's own pool in production ([`connect_probe_pool`]).
pub async fn probe_once(writer: &PgPool, fence: &ReplicaFence) -> Result<TokenEntry, ProbeError> {
let sample = sample_writer(writer).await?;
probe_traced(writer, fence, &mut ProbeTrace::default()).await
}
/// [`probe_once`] under one span carrying every phase's wall time, the gap
/// since the previous beat, and the failure class, bounded by
/// [`PROBE_DEADLINE`].
#[datastore_span(
name = "replica_fence_probe",
system = "postgresql",
fields(
acquire_ms = ::tracing::field::Empty,
lock_ms = ::tracing::field::Empty,
scan_ms = ::tracing::field::Empty,
heartbeat_ms = ::tracing::field::Empty,
commit_ms = ::tracing::field::Empty,
beat_gap_ms = ::tracing::field::Empty,
token = ::tracing::field::Empty,
error_kind = ::tracing::field::Empty,
)
)]
async fn probe_traced(
writer: &PgPool,
fence: &ReplicaFence,
trace: &mut ProbeTrace,
) -> Result<TokenEntry, ProbeError> {
let previous = fence.last_committed_at();
let result = match tokio::time::timeout(
PROBE_DEADLINE,
probe_untraced(writer, fence, &mut trace.timings),
)
.await
{
Ok(result) => result,
Err(_) => Err(ProbeError::Deadline(PROBE_DEADLINE)),
};
if let Ok(entry) = &result {
trace.beat_gap = previous.map(|at| entry.committed_at.saturating_duration_since(at));
}
let span = tracing::Span::current();
let timings = trace.timings;
span.record("acquire_ms", millis(timings.acquire));
span.record("lock_ms", millis(timings.lock));
span.record("scan_ms", millis(timings.scan));
span.record("heartbeat_ms", millis(timings.heartbeat));
span.record("commit_ms", millis(timings.commit));
match &result {
Ok(entry) => {
span.record("token", entry.token);
if let Some(gap) = trace.beat_gap {
span.record("beat_gap_ms", millis(gap));
}
}
Err(error) => {
span.record("error_kind", error.kind());
}
}
result
}
async fn probe_untraced(
writer: &PgPool,
fence: &ReplicaFence,
timings: &mut ProbeTimings,
) -> Result<TokenEntry, ProbeError> {
let sample = sample_writer(writer, timings).await?;
let wall = fence_wall(sample.sampled_at, sample.oldest_xact_start);
match fence.record(sample.token, sample.epoch, sample.committed_at, wall) {
RecordOutcome::Recorded => Ok(TokenEntry {
@@ -703,11 +998,7 @@ pub async fn probe_once(writer: &PgPool, fence: &ReplicaFence) -> Result<TokenEn
// so the same three-bucket argument (and the same wall) holds
// for the rotated token.
let committed_at = Instant::now();
let mut connection = crate::observability::acquire_writer(
writer,
crate::observability::WriterOperation::Maintenance,
)
.await?;
let mut connection = acquire_probe_connection(writer).await?;
let row = sqlx::query(
"UPDATE replica_heartbeat SET epoch = gen_random_uuid(), token = token + 1 \
WHERE id = 1 RETURNING token, epoch",
@@ -808,24 +1099,99 @@ pub struct HeartbeatObservation {
pub backend: String,
}
/// Background probe loop: commit a heartbeat token every `PROBE_INTERVAL`.
/// Runs for the life of the process.
pub async fn run_probe(writer: PgPool, fence: Arc<ReplicaFence>) {
let mut interval = tokio::time::interval(PROBE_INTERVAL);
/// Background probe loop: commit a heartbeat token every `probe_interval`
/// (see [`probe_interval_from_ms`]). Runs for the life of the process.
/// `writer` is the probe's own pool ([`connect_probe_pool`]), never the
/// serving writer pool.
pub async fn run_probe(writer: PgPool, fence: Arc<ReplicaFence>, probe_interval: Duration) {
let mut interval = tokio::time::interval(probe_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut state = ProbeLoopState::default();
loop {
interval.tick().await;
probe_tick(&writer, &fence).await;
probe_tick(&writer, &fence, &mut state).await;
}
}
/// How often a run of failed, slow or late beats is logged. Every beat is
/// still on its span; the log only marks the
/// episode, so a degraded writer cannot turn the probe cadence into a log
/// flood.
const PROBE_LOG_INTERVAL: Duration = Duration::from_secs(10);
/// Admits at most one log line per [`PROBE_LOG_INTERVAL`], counting the
/// rest.
#[derive(Debug, Default)]
struct LogThrottle {
last: Option<Instant>,
suppressed: u64,
}
impl LogThrottle {
/// `Some(suppressed since the last admitted line)` when a line may be
/// logged now.
fn admit(&mut self, now: Instant) -> Option<u64> {
if self
.last
.is_some_and(|at| now.duration_since(at) < PROBE_LOG_INTERVAL)
{
self.suppressed += 1;
return None;
}
self.last = Some(now);
Some(std::mem::take(&mut self.suppressed))
}
}
/// Log state carried across probe ticks.
#[derive(Debug, Default)]
struct ProbeLoopState {
/// Failed probes since the last success.
consecutive_failures: u64,
/// Failure lines; reset on recovery so a new run's first failure
/// always logs.
failure_log: LogThrottle,
}
/// One probe-loop tick. A failure is logged and keeps the retained entries:
/// they were proved by earlier handshakes and stay valid, so one failed
/// probe must not route every read to the writer. Only staleness
/// ([`FENCE_STALENESS`]) closes a fence whose probe keeps failing.
async fn probe_tick(writer: &PgPool, fence: &ReplicaFence) {
if let Err(e) = probe_once(writer, fence).await {
tracing::warn!(error = %e, "replica fence probe failed; retained proofs stand");
async fn probe_tick(writer: &PgPool, fence: &ReplicaFence, state: &mut ProbeLoopState) {
let mut trace = ProbeTrace::default();
let result = probe_traced(writer, fence, &mut trace).await;
let timings = trace.timings;
match result {
Err(e) => {
state.consecutive_failures += 1;
let Some(suppressed) = state.failure_log.admit(Instant::now()) else {
return;
};
tracing::warn!(
error = %e,
error_kind = e.kind(),
consecutive_failures = state.consecutive_failures,
suppressed,
heartbeat_age_ms = fence.heartbeat_age().map(millis),
acquire_ms = millis(timings.acquire),
lock_ms = millis(timings.lock),
scan_ms = millis(timings.scan),
heartbeat_ms = millis(timings.heartbeat),
commit_ms = millis(timings.commit),
"replica fence probe failed; retained proofs stand"
);
}
Ok(entry) if state.consecutive_failures > 0 => {
tracing::info!(
token = entry.token,
failed_probes = state.consecutive_failures,
beat_gap_ms = trace.beat_gap.map(millis),
"replica fence probe recovered"
);
state.consecutive_failures = 0;
state.failure_log = LogThrottle::default();
}
Ok(_) => {}
}
}
@@ -897,7 +1263,7 @@ mod postgres_tests {
.await
.expect_err("a closed writer pool fails the probe");
probe_tick(&writer, &fence).await;
probe_tick(&writer, &fence, &mut ProbeLoopState::default()).await;
let newest = fence
.newest()
@@ -917,6 +1283,41 @@ mod postgres_tests {
);
}
/// An unset interval selects the default, a set one is clamped to the
/// accepted range, and the budget floor covers one lock-timed-out beat.
#[test]
fn probe_interval_defaults_and_clamps() {
assert_eq!(probe_interval_from_ms(0), DEFAULT_PROBE_INTERVAL);
assert_eq!(probe_interval_from_ms(400), Duration::from_millis(400));
assert_eq!(probe_interval_from_ms(1), MIN_PROBE_INTERVAL);
assert_eq!(probe_interval_from_ms(u64::MAX), MAX_PROBE_INTERVAL);
assert_eq!(
min_read_budget(DEFAULT_PROBE_INTERVAL),
Duration::from_millis(1700)
);
assert_eq!(
min_read_budget(Duration::from_millis(400)),
Duration::from_millis(1600)
);
}
/// The probe's log throttle admits one line per interval and reports
/// how many it held back.
#[test]
fn probe_log_throttle_admits_one_line_per_interval() {
let mut throttle = LogThrottle::default();
let start = Instant::now();
assert_eq!(throttle.admit(start), Some(0), "first line always logs");
assert_eq!(throttle.admit(start + Duration::from_secs(1)), None);
assert_eq!(throttle.admit(start + Duration::from_secs(2)), None);
assert_eq!(
throttle.admit(start + PROBE_LOG_INTERVAL),
Some(2),
"next line reports what was suppressed"
);
assert_eq!(throttle.admit(start + PROBE_LOG_INTERVAL), None);
}
/// Resolve picks the greatest retained entry <= the observed token —
/// a lagged reader proves from an older wall, never from thin air.
#[test]
@@ -1038,7 +1439,9 @@ mod postgres_tests {
.expect("connect idle");
let _idle_conn = idle_pool.acquire().await.expect("idle conn");
let before = sample_writer(&pool).await.expect("sample without tx");
let before = sample_writer(&pool, &mut ProbeTimings::default())
.await
.expect("sample without tx");
// Now hold a transaction open on a second connection.
let tx_pool = PgPool::connect(&crate::test_support::database_url())
@@ -1050,7 +1453,9 @@ mod postgres_tests {
.await
.expect("touch tx");
let during = sample_writer(&pool).await.expect("sample with open tx");
let during = sample_writer(&pool, &mut ProbeTimings::default())
.await
.expect("sample with open tx");
let oldest = during
.oldest_xact_start
.expect("open transaction must appear in the oldest-xact term");
@@ -1070,6 +1475,67 @@ mod postgres_tests {
tx.rollback().await.expect("rollback");
}
/// The probe pool installs its own lock budget: while a shared floor
/// lock outlives it, the probe fails fast with `lock_timeout` rather
/// than holding an exclusive request (which queues every new channel
/// write behind it) for the serving writer's longer budget. Retained
/// proofs are untouched and the next beat succeeds once the lock clears.
#[tokio::test]
#[ignore = "requires Postgres"]
async fn cluster_global_probe_pool_bounds_the_floor_lock_wait() {
let (pool, _scratch) = scratch_db().await;
let probe_pool = connect_probe_pool(&pool);
let fence = ReplicaFence::new();
let first = probe_once(&probe_pool, &fence)
.await
.expect("probe through its own pool");
let (statement, idle): (String, String) = sqlx::query_as(
"SELECT current_setting('statement_timeout'), \
current_setting('idle_in_transaction_session_timeout')",
)
.fetch_one(&probe_pool)
.await
.expect("read probe session timeouts");
// PROBE_STATEMENT_TIMEOUT, as Postgres renders it.
assert_eq!(
(statement.as_str(), idle.as_str()),
("2s", "2s"),
"probe session must bound statements and idle transactions"
);
let mut blocker = pool.begin().await.expect("begin shared-lock blocker");
sqlx::query("SELECT pg_advisory_xact_lock_shared($1)")
.bind(REPLICA_FLOOR_LOCK_KEY)
.execute(&mut *blocker)
.await
.expect("hold shared floor lock");
let started = Instant::now();
let err = probe_once(&probe_pool, &fence)
.await
.expect_err("a held shared floor lock must time the probe out");
let waited = started.elapsed();
assert_eq!(err.kind(), "lock_timeout", "unexpected error: {err:?}");
assert!(
waited >= PROBE_LOCK_TIMEOUT && waited < PROBE_LOCK_TIMEOUT + Duration::from_secs(1),
"probe waited {waited:?}, expected about {PROBE_LOCK_TIMEOUT:?}"
);
assert_eq!(
fence.newest().map(|entry| entry.token),
Some(first.token),
"a timed-out beat keeps the retained proof"
);
blocker
.rollback()
.await
.expect("release shared-lock blocker");
let next = probe_once(&probe_pool, &fence)
.await
.expect("probe recovers once the lock clears");
assert!(next.token > first.token);
}
/// The sample must not capture `S` while a compliant shared writer lock is
/// still held. It must block until release, then sample.
#[tokio::test]
@@ -1085,7 +1551,10 @@ mod postgres_tests {
.expect("hold shared floor lock");
let sample_pool = pool.clone();
let mut sampling = tokio::spawn(async move { sample_writer(&sample_pool).await });
let mut sampling =
tokio::spawn(
async move { sample_writer(&sample_pool, &mut ProbeTimings::default()).await },
);
assert!(
tokio::time::timeout(Duration::from_millis(100), &mut sampling)
.await
@@ -1155,7 +1624,7 @@ mod postgres_tests {
};
let unpriv = PgPool::connect(&unpriv_url).await.expect("connect unpriv");
let err = sample_writer(&unpriv)
let err = sample_writer(&unpriv, &mut ProbeTimings::default())
.await
.expect_err("masked pg_stat_activity must fail closed");
assert!(
@@ -1263,9 +1732,12 @@ mod postgres_tests {
#[ignore = "requires Postgres"]
async fn cluster_global_probe_rotates_epoch_on_same_epoch_token_regression() {
let (pool, _scratch) = scratch_db().await;
// Through the production one-connection pool: rotation re-acquires
// after the sample has released its connection.
let probe_pool = connect_probe_pool(&pool);
let fence = ReplicaFence::new();
let before = probe_once(&pool, &fence).await.expect("probe");
let before = probe_once(&probe_pool, &fence).await.expect("probe");
let mut conn = pool.acquire().await.expect("conn");
let old_epoch = observe_heartbeat(&mut conn, false)
.await
@@ -1279,7 +1751,9 @@ mod postgres_tests {
.await
.expect("rewind token");
let after = probe_once(&pool, &fence).await.expect("recovery probe");
let after = probe_once(&probe_pool, &fence)
.await
.expect("recovery probe");
// The pre-rewind observation must no longer prove anything.
assert_eq!(
fence.resolve(before.token, old_epoch),

View File

@@ -2592,6 +2592,66 @@ async fn lazy_reader_pool_still_spawns_fence_probe() {
assert!(spawned, "a configured (lazy) reader must spawn the probe");
}
/// The fence probe must keep beating while every serving writer connection
/// is checked out. On the serving pool its beat queued behind request
/// traffic, so under load every pod's newest token aged past the
/// bounded-read budget and reads spilled onto the saturated writer.
#[tokio::test]
#[ignore = "requires Postgres"]
async fn fence_probe_beats_while_serving_writer_pool_is_exhausted() {
let (seed, scratch) = ScratchDb::migrated(&admin_url().await, "fence_own_pool").await;
seed.close().await;
let url = scratch.url().to_owned();
let db = Db::new(&DbConfig {
database_url: url.clone(),
read_database_url: Some(url),
max_connections: 1,
min_connections: 0,
acquire_timeout_secs: 1,
..DbConfig::default()
})
.await
.expect("connect armed Db with one serving writer connection");
assert!(
db.spawn_fence_probe().await.expect("verification passes"),
"a configured reader must spawn the probe"
);
// Hold the serving writer pool's only connection for the whole test.
let mut held = db.pool.acquire().await.expect("hold serving connection");
let baseline = db.fence().newest().map(|entry| entry.token);
// Two beats strictly after the hold: one in flight when the hold began
// cannot satisfy this on its own.
let mut beats = 0;
let mut last = baseline;
let deadline = std::time::Instant::now() + Duration::from_secs(5);
while beats < 2 {
assert!(
std::time::Instant::now() < deadline,
"probe stopped beating while the serving writer pool was exhausted"
);
tokio::time::sleep(Duration::from_millis(50)).await;
let newest = db.fence().newest().map(|entry| entry.token);
if newest > last {
beats += 1;
last = newest;
}
}
// The beats come from the probe's own, identifiable session.
let probe_sessions: i64 = sqlx::query_scalar(
"SELECT count(*) FROM pg_stat_activity \
WHERE datname = current_database() AND application_name = $1",
)
.bind(replica_fence::PROBE_APPLICATION_NAME)
.fetch_one(&mut *held)
.await
.expect("count probe sessions");
assert_eq!(probe_sessions, 1, "the probe holds exactly one connection");
}
/// Thread replies: head fetch reads the writer; a FULL cursor page is
/// served by the replica; an UNDER-limit cursor page (candidate terminal
/// page) is re-run on the writer so a lagged replica can never truncate

View File

@@ -516,3 +516,91 @@ async fn stalled_reader_sessions_release_their_slots_for_new_connections() {
drop(_guard);
read_pool.close().await;
}
/// Reads that spill to the writer are counted by cause, so a late fence
/// heartbeat is distinguishable from a busy or lagging reader, while the
/// legacy `buzz_db_route_decision` reasons stay as they were.
#[tokio::test]
#[ignore = "requires Postgres"]
async fn reader_spills_are_counted_by_cause() {
let (mut db, _scratch) = same_db_reader("spill_cause", 2).await;
db.set_replica_read_max_age_for_tests(Some(Duration::from_secs(1)));
let query = EventQuery::for_community(CommunityId::from_uuid(Uuid::new_v4()));
let recorder = DebuggingRecorder::new();
let snapshotter = recorder.snapshotter();
let _guard = metrics::set_default_local_recorder(&recorder);
// Counters drained since the last step: (route decisions, spills).
let step = || {
let mut routes = std::collections::HashMap::new();
let mut spills = std::collections::HashMap::new();
for ((name, decision, reason), n) in drain_counters(&snapshotter) {
// Drained counters linger at zero.
if n == 0 {
continue;
}
match name.as_str() {
"buzz_db_route_decision" => {
routes.insert((decision, reason), n);
}
"buzz_db_reader_spill_total" => {
spills.insert(reason, n);
}
_ => {}
}
}
(routes, spills)
};
let writer = |reason: &str| ("writer".to_owned(), reason.to_owned());
// A fresh heartbeat: the replica serves the read, nothing spills.
db.fence().force_open_for_tests(chrono::Utc::now());
db.count_events_routed("spill_cause", &query)
.await
.expect("replica read");
let (_, spills) = step();
assert!(
spills.is_empty(),
"a served read is not a spill: {spills:?}"
);
// Older than the 1s budget but inside the staleness gate: the probe is
// late, the replica is fine.
db.fence()
.force_open_for_tests_at(chrono::Utc::now(), Instant::now() - Duration::from_secs(2));
db.count_events_routed("spill_cause", &query)
.await
.expect("writer answers a late-heartbeat read");
let (routes, spills) = step();
assert_eq!(routes.get(&writer("stale")), Some(&1), "{routes:?}");
assert_eq!(
spills,
[("heartbeat_stale".to_owned(), 1)].into(),
"a late heartbeat must be its own spill reason"
);
// Past the staleness gate: no live heartbeat proof at all.
db.fence().force_open_for_tests_at(
chrono::Utc::now(),
Instant::now() - replica_fence::FENCE_STALENESS - Duration::from_secs(1),
);
db.count_events_routed("spill_cause", &query)
.await
.expect("writer answers a missing-heartbeat read");
let (routes, spills) = step();
assert_eq!(routes.get(&writer("uninitialized")), Some(&1), "{routes:?}");
assert_eq!(spills, [("heartbeat_missing".to_owned(), 1)].into());
// A disabled bounded arm is configuration, not a spill.
db.set_replica_read_max_age_for_tests(None);
db.fence().force_open_for_tests(chrono::Utc::now());
db.count_events_routed("spill_cause", &query)
.await
.expect("writer answers with routing disabled");
let (routes, spills) = step();
assert_eq!(routes.get(&writer("disabled")), Some(&1), "{routes:?}");
assert!(
spills.is_empty(),
"disabled routing is not a spill: {spills:?}"
);
}

View File

@@ -454,7 +454,12 @@ fn p0_pool_acquisitions_use_typed_operation_pairs_without_other() {
.next()
.expect("replica-fence production source");
assert!(replica_fence_production.contains("WriterOperation::Bootstrap"));
assert!(replica_fence_production.contains("WriterOperation::Maintenance"));
// The heartbeat probe runs on its own pool, never a serving writer
// checkout: queued behind request traffic, its beat goes stale exactly
// when the writer is saturated.
assert!(!replica_fence_production.contains("WriterOperation::Maintenance"));
assert!(replica_fence_production.contains("pub fn connect_probe_pool("));
assert!(runtime.contains("replica_fence::connect_probe_pool(&self.pool)"));
assert!(!replica_fence_production.contains("pool.begin().await"));
assert!(!replica_fence_production.contains("writer.acquire().await"));
assert!(!replica_fence_production.contains("fetch_optional(writer)"));

View File

@@ -127,6 +127,11 @@ pub struct Config {
/// `0` (the default) disables bounded-staleness replica routing; see
/// [`buzz_db::DbConfig::replica_read_max_age_ms`].
pub replica_read_max_age_ms: u64,
/// Replica fence heartbeat interval in milliseconds
/// (`BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS`). Unset (`0`) selects the
/// default; set values must lie within the accepted range. See
/// [`buzz_db::DbConfig::replica_fence_probe_interval_ms`].
pub replica_fence_probe_interval_ms: u64,
/// Upper bound, in milliseconds, of the per-connection random delay applied
/// when sending the `1012 Service Restart` close frame during graceful
@@ -709,6 +714,28 @@ impl Config {
Err(_) => 0,
};
// Fence heartbeat interval: unset or empty selects the default. An
// out-of-range value fails startup rather than being clamped, so the
// deployed cadence is always the configured one.
let replica_fence_probe_interval_ms =
match std::env::var("BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS") {
Ok(raw) if !raw.trim().is_empty() => {
let min = buzz_db::replica_fence::MIN_PROBE_INTERVAL.as_millis() as u64;
let max = buzz_db::replica_fence::MAX_PROBE_INTERVAL.as_millis() as u64;
raw.trim()
.parse::<u64>()
.ok()
.filter(|ms| (min..=max).contains(ms))
.ok_or_else(|| {
ConfigError::InvalidValue(format!(
"BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS must be an integer \
between {min} and {max}"
))
})?
}
_ => 0,
};
// Drain jitter: 0 = off (default). Clamp oversized values so every
// delayed close is initiated with ten seconds left in the relay's
// hard-drain budget. An empty/whitespace-only value is treated as unset
@@ -1453,6 +1480,7 @@ impl Config {
database_url,
read_database_url,
replica_read_max_age_ms,
replica_fence_probe_interval_ms,
drain_jitter_ms,
redis_url,
db_pool_size,
@@ -2485,6 +2513,44 @@ mod tests {
}
}
#[test]
fn replica_fence_probe_interval_defaults_and_rejects_out_of_range() {
let _guards = env_guards();
let previous = std::env::var_os("BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS");
let read = |value: Option<&str>| {
match value {
Some(value) => std::env::set_var("BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS", value),
None => std::env::remove_var("BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS"),
}
Config::from_env().map(|config| config.replica_fence_probe_interval_ms)
};
let unset = read(None);
let empty = read(Some(" "));
let set = read(Some("500"));
let rejected =
["soon", "0", "299", "10001", "-400"].map(|value| (value, read(Some(value))));
if let Some(value) = previous {
std::env::set_var("BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS", value);
} else {
std::env::remove_var("BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS");
}
assert_eq!(unset.expect("unset"), 0, "unset selects the default");
assert_eq!(empty.expect("empty"), 0, "empty selects the default");
assert_eq!(set.expect("set"), 500);
for (value, result) in rejected {
match result {
Err(ConfigError::InvalidValue(message)) => assert!(
message.contains("BUZZ_REPLICA_FENCE_PROBE_INTERVAL_MS"),
"{value}: {message}"
),
other => panic!("{value} must fail startup, got {other:?}"),
}
}
}
#[test]
fn drain_jitter_defaults_off_and_rejects_junk() {
let _guards = env_guards();

View File

@@ -349,6 +349,7 @@ async fn run_relay_main(boot: BootTracker) -> anyhow::Result<()> {
database_url: config.database_url.clone(),
read_database_url: config.read_database_url.clone(),
replica_read_max_age_ms: config.replica_read_max_age_ms,
replica_fence_probe_interval_ms: config.replica_fence_probe_interval_ms,
max_connections: config.db_pool_size,
read_max_connections: config.db_read_pool_size,
..DbConfig::default()