jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-10-28 13:01:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Harden catch-all advancement against races and child-only triggers

Browse Source 
- Refuse to replace a catch-all that carries child-only triggers; the
  replacement partitions inherit only parent triggers.
- Re-plan from a fresh audit once the maintenance advisory lock is held,
  before any table lock, so a runner that lost a race sees the winner's
  layout as a no-op instead of reporting its monthlies as collisions.
  Name collisions on the refusal path are still reported before any lock.
- Retry a read-only table audit (up to twice) when a partition it listed
  was dropped concurrently: pg_get_expr renders from current catalog
  state and yields NULL for a relation missing since the snapshot, and
  the emptiness probe then fails with 42P01.
- Scope the deadline test to the backend blocked by its own holder, and
  let concurrent-runner convergence tolerate a fail-fast lock_timeout
  from the other runner's audit probing the catch-all.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Ian Oberst <ioberst@block.xyz>
This commit is contained in:
Ian Oberst
2026-09-30 09:07:41 -07:00
parent d603b1f197
commit ee1fb7e4ad
3 changed files with 294 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
crates/buzz-db/src/store/partition.rs
View File

@@ -35,6 +35,10 @@ const PARTITION_AUDIT_STATEMENT_TIMEOUT: &str = "5s";
/// Maximum wall-clock duration for one complete multi-table catalog audit.
const PARTITION_AUDIT_TOTAL_TIMEOUT: Duration = Duration::from_secs(10);
/// Fresh-snapshot retries after a partition vanishes mid-audit, such as when
/// another relay's maintenance replaces a catch-all.
const PARTITION_AUDIT_CONCURRENT_DROP_RETRIES: usize = 2;
fn expected_partition_key(table: &str) -> Option<&'static str> {
match table {
"events" => Some("RANGE (created_at)"),
@@ -495,6 +499,26 @@ async fn audit_table(
now: DateTime<Utc>,
) -> Result<PartitionTableAudit> {
let months_ahead = validated_months_ahead(months_ahead)?;
let mut retries = 0;
loop {
match audit_table_once(pool, table, months_ahead, now).await {
Err(error)
if retries < PARTITION_AUDIT_CONCURRENT_DROP_RETRIES
&& is_concurrent_drop(&error) =>
{
retries += 1;
}
result => return result,
}
}
}
async fn audit_table_once(
pool: &PgPool,
table: &'static str,
months_ahead: i32,
now: DateTime<Utc>,
) -> Result<PartitionTableAudit> {
let mut transaction = pool.begin().await?;
sqlx::query("SET TRANSACTION READ ONLY")
.execute(&mut *transaction)
@@ -505,6 +529,20 @@ async fn audit_table(
Ok(audit)
}
/// Whether an audit failed only because a partition it had already listed was
/// dropped by a concurrent transaction. A fresh snapshot sees a consistent
/// catalog, so the audit can be retried.
fn is_concurrent_drop(error: &DbError) -> bool {
match error {
DbError::NotFound(_) => true,
// `relation_contains_rows` names a leaf that no longer exists.
DbError::Sqlx(sqlx::Error::Database(database)) => {
database.code().as_deref() == Some("42P01")
}
_ => false,
}
}
async fn configure_audit_transaction(connection: &mut PgConnection) -> Result<()> {
pin_catalog_rendering(connection).await?;
sqlx::query("SELECT pg_catalog.set_config('statement_timeout', $1, true)")
@@ -678,14 +716,27 @@ async fn audit_table_on(
let is_leaf: bool = row.try_get("is_leaf")?;
let bound_partition_key: Option<String> = row.try_get("bound_partition_key")?;
let detach_pending: bool = row.try_get("detach_pending")?;
let sibling_bounds: Vec<String> = row.try_get("sibling_bounds")?;
// Catalog rows come from the statement snapshot, but `pg_get_expr`
// renders from current catalog state and yields NULL for a relation
// dropped since the snapshot was taken.
let sibling_bounds = row
.try_get::<Vec<Option<String>>, _>("sibling_bounds")?
.into_iter()
.collect::<Option<Vec<String>>>()
.ok_or_else(|| {
DbError::NotFound(format!(
"a sibling of partition {name} was dropped mid-audit"
))
})?;
if detach_pending {
pending_roots.insert(root_child_relation_oid);
}
let partition_key_compatible = partition_key_valid
&& bound_partition_key.as_deref() == Some(expected_partition_key)
&& !detach_pending;
let expression: String = row.try_get("bound")?;
let expression: String = row
.try_get::<Option<String>, _>("bound")?
.ok_or_else(|| DbError::NotFound(format!("partition {name} was dropped mid-audit")))?;
let is_default = expression.trim() == "DEFAULT";
let own_range = parse_range_bounds(&expression);
let is_catch_all = own_range.as_ref().is_some_and(|(lower, upper)| {

112
crates/buzz-db/src/store/partition/maintenance.rs
View File

@@ -159,7 +159,7 @@ enum MaintenancePlan {
}
enum LockedAttempt {
Committed,
Committed(MaintenanceChange),
SkippedLocked,
Unchanged,
Refused(String),
@@ -263,6 +263,15 @@ fn plan_maintenance(
))
}
}
// Replacements inherit only parent triggers; dropping the catch-all
// would silently discard child-specific behavior.
if !child.extra_triggers.is_empty() {
return MaintenancePlan::Refuse(format!(
"catch-all {} has child-only triggers ({}); replacing it would drop them",
child.name,
child.extra_triggers.join(", ")
));
}
let Some(PartitionBound::Finite(lower)) = child.lower else {
return MaintenancePlan::Refuse(format!(
"catch-all {} has no finite lower bound",
@@ -389,44 +398,33 @@ async fn maintain_table(
now: DateTime<Utc>,
table_deadline: Duration,
) -> TableMaintenance {
let table = audit.table;
let plan = plan_maintenance(audit, policy);
let (names, replaced) = match &plan {
let change = match plan_maintenance(audit, policy) {
MaintenancePlan::Noop => return TableMaintenance::done(PartitionMaintenanceOutcome::Noop),
MaintenancePlan::Disabled => {
return TableMaintenance::done(PartitionMaintenanceOutcome::SkippedDisabled)
}
MaintenancePlan::Refuse(_) => wanted_names(audit, policy),
MaintenancePlan::Apply(change) => (change.target_names(table), change.replaced()),
};
// Reject name collisions first, and before any lock, so a misnamed
// relation is reported precisely and never costs a parent outage.
match colliding_relation(pool, names, replaced).await {
Ok(None) => {}
Ok(Some(name)) => {
return TableMaintenance::failed(
PartitionMaintenanceOutcome::OperatorRequired,
None,
collision_message(&name),
)
}
Err(error) => {
return TableMaintenance::failed(classify_error(&error), None, error.to_string())
}
}
let change = match plan {
MaintenancePlan::Apply(change) => change,
MaintenancePlan::Refuse(reason) => {
// Report a name collision ahead of the structural reason it
// usually causes. A refusal takes no lock, so no race applies.
let (names, replaced) = wanted_names(audit, policy);
let error = match colliding_relation(pool, names, replaced).await {
Ok(Some(name)) => collision_message(&name),
Ok(None) => reason,
Err(error) => {
return TableMaintenance::failed(
classify_error(&error),
None,
error.to_string(),
)
}
};
return TableMaintenance::failed(
PartitionMaintenanceOutcome::OperatorRequired,
None,
reason,
)
}
MaintenancePlan::Noop | MaintenancePlan::Disabled => {
return TableMaintenance::done(PartitionMaintenanceOutcome::Noop)
error,
);
}
MaintenancePlan::Apply(change) => change,
};
let deadline = tokio::time::Instant::now() + table_deadline;
@@ -449,7 +447,7 @@ async fn maintain_table(
};
let attempt = tokio::time::timeout_at(
deadline,
apply_change(&mut connection, table, months_ahead, policy, now, &change),
apply_change(&mut connection, audit.table, months_ahead, policy, now),
)
.await;
let Ok(attempt) = attempt else {
@@ -467,9 +465,9 @@ async fn maintain_table(
);
};
match attempt {
Ok(LockedAttempt::Committed) => TableMaintenance {
outcome: change.outcome(),
change: Some(change),
Ok(LockedAttempt::Committed(committed)) => TableMaintenance {
outcome: committed.outcome(),
change: Some(committed),
error: None,
},
Ok(LockedAttempt::SkippedLocked) => {
@@ -493,7 +491,6 @@ async fn apply_change(
months_ahead: i32,
policy: PartitionMaintenancePolicy,
now: DateTime<Utc>,
expected: &MaintenanceChange,
) -> Result<LockedAttempt> {
let mut transaction = connection.begin().await?;
configure_maintenance_transaction(&mut transaction).await?;
@@ -513,6 +510,33 @@ async fn apply_change(
return Ok(LockedAttempt::SkippedLocked);
}
// Re-plan while serialized against other managers but before any table
// lock. A runner that lost a race sees the winner's committed layout here
// and stops, instead of mistaking the winner's monthlies for collisions.
let serialized = audit_table_on(&mut transaction, table, months_ahead, now).await?;
let expected = match plan_maintenance(&serialized, policy) {
MaintenancePlan::Apply(change) => change,
MaintenancePlan::Noop | MaintenancePlan::Disabled => {
transaction.rollback().await?;
return Ok(LockedAttempt::Unchanged);
}
MaintenancePlan::Refuse(reason) => {
transaction.rollback().await?;
return Ok(LockedAttempt::Refused(reason));
}
};
// Reject misnamed relations before any lock can cost a parent outage.
if let Some(name) = colliding_relation(
&mut *transaction,
expected.target_names(table),
expected.replaced(),
)
.await?
{
transaction.rollback().await?;
return Ok(LockedAttempt::Refused(collision_message(&name)));
}
let parent = maintenance_parent(&mut transaction, table).await?;
if let Some(reason) = parent.refusal {
transaction.rollback().await?;
@@ -551,10 +575,10 @@ async fn apply_change(
.await?;
}
// The pre-lock audit may be stale; re-plan from the locked catalog.
// Close the window between the serialized audit and the table locks.
let locked = audit_table_on(&mut transaction, table, months_ahead, now).await?;
match plan_maintenance(&locked, policy) {
MaintenancePlan::Apply(change) if change == *expected => {}
MaintenancePlan::Apply(change) if change == expected => {}
MaintenancePlan::Noop | MaintenancePlan::Disabled => {
transaction.rollback().await?;
return Ok(LockedAttempt::Unchanged);
@@ -570,16 +594,6 @@ async fn apply_change(
));
}
}
if let Some(name) = colliding_relation(
&mut *transaction,
expected.target_names(table),
expected.replaced(),
)
.await?
{
transaction.rollback().await?;
return Ok(LockedAttempt::Refused(collision_message(&name)));
}
let parent_relation = qualified_relation_name(&parent.schema, table);
if let Some(catch_all) = &expected.catch_all {
@@ -620,7 +634,7 @@ async fn apply_change(
}
let after = audit_table_on(&mut transaction, table, months_ahead, now).await?;
let verified = verify_applied(&locked, &after, expected).and(
let verified = verify_applied(&locked, &after, &expected).and(
verify_inheritance(&mut transaction, parent.oid, &expected.target_names(table)).await?,
);
if let Err(reason) = verified {
@@ -630,7 +644,7 @@ async fn apply_change(
)));
}
transaction.commit().await?;
Ok(LockedAttempt::Committed)
Ok(LockedAttempt::Committed(expected))
}
async fn configure_maintenance_transaction(connection: &mut PgConnection) -> Result<()> {

199
crates/buzz-db/src/store/partition/maintenance_tests.rs
View File

@@ -162,6 +162,10 @@ fn unsafe_catch_all_states_require_an_operator() {
let mut pending = repaired_audit(false);
pending.children[1].pending_detach = true;
refused(&pending, "not structurally safe");
let mut child_triggers = repaired_audit(false);
child_triggers.children[1].extra_triggers = vec!["audit_catch_all".to_string()];
refused(&child_triggers, "child-only triggers (audit_catch_all)");
}
#[test]
@@ -536,13 +540,26 @@ mod postgres_tests {
seed_repaired_layout(&pool, true).await;
let (first, second) = tokio::join!(maintain(&pool), maintain(&pool));
first.expect("first runner");
second.expect("second runner");
assert_advanced(
&audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now())
.await
.expect("audit"),
);
// A runner's catch-all NOWAIT fails fast if the other runner's audit is
// probing that catch-all; nothing else may go wrong.
for (runner, result) in [("first", first), ("second", second)] {
let Err(error) = result else { continue };
let message = error.to_string();
assert!(
message.contains("lock_timeout"),
"{runner} runner: {message}"
);
for table in PARTITIONED_TABLES {
for outcome in ["operator_required", "deadline", "error"] {
assert!(
!message.contains(&format!("{table} {outcome}:")),
"{runner} runner: {message}"
);
}
}
}
// Whatever a fail-fast loser left undone, the next pass converges.
assert_advanced(&maintain(&pool).await.expect("follow-up pass"));
drop_schema(&admin, &schema).await;
}
@@ -553,7 +570,11 @@ mod postgres_tests {
seed_repaired_layout(&pool, false).await;
let before = catalog_snapshot(&pool).await;
// The parent lock wait (2s) outlasts the 300ms client deadline.
let holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await;
let mut holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await;
let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()")
.fetch_one(&mut *holder)
.await
.expect("holder backend pid");
let result = maintain_partitions_with_deadline(
&pool,
@@ -569,25 +590,41 @@ mod postgres_tests {
"{result:?}"
);
// The abandoned session is closed, not pooled: its backend disappears
// and releases every lock before the holder commits.
let mut gone = false;
// The abandoned session is closed, not pooled: the backend that was
// queued behind the holder disappears and releases every lock before
// the holder commits. Only sessions blocked by this test's holder
// count, so maintenance elsewhere on the cluster cannot interfere.
let blocked: Vec<i32> = sqlx::query_scalar(
"SELECT pid FROM pg_stat_activity WHERE $1 = ANY(pg_blocking_pids(pid))",
)
.bind(holder_pid)
.fetch_all(&admin)
.await
.expect("sessions blocked by the holder");
let mut gone = blocked.is_empty();
for _ in 0..50 {
let sessions: i64 = sqlx::query_scalar(
"SELECT count(*) FROM pg_stat_activity \
WHERE application_name = 'buzz_partition_maintenance'",
)
.fetch_one(&admin)
.await
.expect("count maintenance sessions");
if sessions == 0 {
gone = true;
if gone {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let alive: bool = sqlx::query_scalar(
"SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid = ANY($1))",
)
.bind(&blocked)
.fetch_one(&admin)
.await
.expect("check abandoned backend");
gone = !alive;
}
drop(holder);
assert!(gone, "abandoned maintenance session was not closed");
assert!(
blocked.len() <= 1,
"unexpected blocked sessions {blocked:?}"
);
assert!(
gone,
"abandoned maintenance session {blocked:?} was not closed"
);
let after = catalog_snapshot(&pool).await;
let events_only = |snapshot: &[(String, String, String, i64)]| {
snapshot
@@ -600,6 +637,126 @@ mod postgres_tests {
drop_schema(&admin, &schema).await;
}
#[tokio::test]
#[ignore = "requires Postgres"]
async fn catch_all_with_child_only_trigger_is_refused() {
let (pool, admin, schema) = scratch_pool().await;
seed_repaired_layout(&pool, false).await;
sqlx::query(
"CREATE TRIGGER catch_all_only BEFORE INSERT ON events_p_future_next \
FOR EACH ROW EXECUTE FUNCTION partition_test_trigger()",
)
.execute(&pool)
.await
.expect("create child-only trigger");
let before = catalog_snapshot(&pool).await;
let result = maintain(&pool).await;
assert!(
matches!(result, Err(DbError::InvalidData(ref message))
if message.contains("events operator_required")
&& message.contains("child-only triggers (catch_all_only)")
&& !message.contains("delivery_log")),
"{result:?}"
);
let after = catalog_snapshot(&pool).await;
let events_only = |snapshot: &[(String, String, String, i64)]| {
snapshot
.iter()
.filter(|(name, ..)| name.starts_with("events_"))
.cloned()
.collect::<Vec<_>>()
};
assert_eq!(events_only(&after), events_only(&before));
drop_schema(&admin, &schema).await;
}
#[tokio::test]
#[ignore = "requires Postgres"]
async fn stale_audit_after_a_concurrent_advance_is_a_noop() {
let (pool, admin, schema) = scratch_pool().await;
seed_repaired_layout(&pool, false).await;
// A runner audits, then loses the race to a runner that commits first.
let stale = audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now())
.await
.expect("stale audit");
assert_advanced(&maintain(&pool).await.expect("winning runner"));
let advanced = catalog_snapshot(&pool).await;
for table_audit in &stale.tables {
let result = maintain_table(
&pool,
table_audit,
PARTITION_MANAGER_MONTHS_AHEAD as i32,
ADVANCE,
fixed_now(),
std::time::Duration::from_secs(10),
)
.await;
assert_eq!(
(result.outcome, result.error.as_deref()),
(PartitionMaintenanceOutcome::Noop, None),
"{}",
table_audit.table
);
}
assert_eq!(catalog_snapshot(&pool).await, advanced);
drop_schema(&admin, &schema).await;
}
#[tokio::test]
#[ignore = "requires Postgres"]
async fn audit_retries_when_a_listed_partition_is_dropped() {
let (pool, admin, schema) = scratch_pool().await;
seed_repaired_layout(&pool, false).await;
let mut dropper = pool.begin().await.expect("begin dropper");
let dropper_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()")
.fetch_one(&mut *dropper)
.await
.expect("dropper backend pid");
sqlx::query("DROP TABLE events_p_future_next")
.execute(&mut *dropper)
.await
.expect("drop catch-all");
// The audit lists the catch-all, then queues behind the drop's lock.
let audit = tokio::spawn({
let pool = pool.clone();
async move {
audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()).await
}
});
let mut queued = false;
for _ in 0..50 {
queued = sqlx::query_scalar(
"SELECT EXISTS (SELECT 1 FROM pg_stat_activity \
WHERE $1 = ANY(pg_blocking_pids(pid)))",
)
.bind(dropper_pid)
.fetch_one(&admin)
.await
.expect("check audit queue");
if queued {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
}
assert!(queued, "audit never queued behind the drop");
dropper.commit().await.expect("commit drop");
let audit = audit.await.expect("audit task").expect("audit retried");
let events = audit
.tables
.iter()
.find(|table| table.table == "events")
.expect("events audit");
assert!(events
.children
.iter()
.all(|child| child.name != "events_p_future_next"));
drop_schema(&admin, &schema).await;
}
#[tokio::test]
#[ignore = "requires Postgres"]
async fn unrelated_catch_all_name_is_an_operator_error() {