diff --git a/crates/buzz-db/src/store/partition.rs b/crates/buzz-db/src/store/partition.rs index 8ca8d14429..f7ef039091 100644 --- a/crates/buzz-db/src/store/partition.rs +++ b/crates/buzz-db/src/store/partition.rs @@ -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, ) -> Result { 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, +) -> Result { 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 = row.try_get("bound_partition_key")?; let detach_pending: bool = row.try_get("detach_pending")?; - let sibling_bounds: Vec = 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::>, _>("sibling_bounds")? + .into_iter() + .collect::>>() + .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::, _>("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)| { diff --git a/crates/buzz-db/src/store/partition/maintenance.rs b/crates/buzz-db/src/store/partition/maintenance.rs index 34cdfa56d2..c0a8d5beb2 100644 --- a/crates/buzz-db/src/store/partition/maintenance.rs +++ b/crates/buzz-db/src/store/partition/maintenance.rs @@ -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, 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, - expected: &MaintenanceChange, ) -> Result { 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<()> { diff --git a/crates/buzz-db/src/store/partition/maintenance_tests.rs b/crates/buzz-db/src/store/partition/maintenance_tests.rs index 4a437a88b0..6ce243193c 100644 --- a/crates/buzz-db/src/store/partition/maintenance_tests.rs +++ b/crates/buzz-db/src/store/partition/maintenance_tests.rs @@ -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 = 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::>() + }; + 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() {