From a8fe2b612c1c83de8f0983fa7a5cc1cb81ac7be8 Mon Sep 17 00:00:00 2001 From: Ian Oberst Date: Thu, 8 Oct 2026 12:59:10 -0700 Subject: [PATCH] feat(db): advance an empty partition catch-all under bounded locks (#7993) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds bounded, serialized advancement of an empty right-edge partition catch-all for `events` and `delivery_log`. The repaired catalogs currently end in an empty catch-all from 2027-01-01, and the merged manager (#6515) only audits them. Without this change, January writes land in the catch-all, after which only an operator can move them. When `BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED=true`, each relay keeps six dedicated monthlies after the current month. It does this at startup and on every audit interval by replacing the empty catch-all with canonical monthlies and a new canonical `_p_future` after the horizon. From today's catalogs the first run creates January–March 2027 and moves the catch-all to April 2027. ## Safety design - **Refuse before locking.** The following return `operator_required` without taking any table lock: - a populated catch-all, or one whose occupancy is unknown; - a catch-all with child-only triggers, which the replacements would not inherit; - a misaligned bound; - a `DEFAULT`, anomalous or pending-detach child; - a name collision. - **One winner.** A schema-scoped `pg_try_advisory_xact_lock` is taken before any DDL. Losers report `skipped_locked` and never issue DDL. The winner re-plans from a fresh audit before taking any table lock and checks for collisions. A runner that lost a race therefore sees the committed layout as a no-op, not as colliding names. - **Writer lock order.** First `ONLY` the parent is locked `ACCESS EXCLUSIVE` under a 2 s `lock_timeout`. Then each foreign-key counterpart (today only `communities`) is locked in OID order `SHARE ROW EXCLUSIVE NOWAIT`, then the catch-all `NOWAIT`. - The counterparts are read again once the parent is locked, and that set is the one locked. Any FK DDL on or to the parent conflicts with the parent lock, so the set can't change after that. An FK committed while maintenance waited is therefore locked `NOWAIT` too, so attaching never waits on it under the parent lock. The read before the lock remains as a cheap early refusal. - This is the order an insert takes locks (parent, then the FK check on `communities`), so maintenance cannot deadlock with ingest. The internal manual runbook used to lock `communities` first; it now uses this order too. - `SHARE ROW EXCLUSIVE` is the strongest lock the DDL takes on `communities`; measured on PG 16 and 17, the catch-all `DROP` takes none, as long as the catch-all's only foreign keys are those inherited from the parent. A catch-all with its own foreign key, or one that another table's foreign key references, is refused: once before any lock, and again after the catch-all lock, where the check is authoritative. Readers and other writers' `FOR KEY SHARE` checks proceed. A concurrent `communities` writer makes the run fail fast with `lock_timeout`. - A `40P01` deadlock victim is also reported as `lock_timeout`. - Parents that are referenced by a foreign key, or that reference an inherited table, are refused, including when such a key commits while maintenance waits on the parent. - **Re-plan under lock.** The catalog is re-audited under those locks. The transaction proceeds only if the plan is unchanged. - **Bounded.** `statement_timeout` 5 s, `idle_in_transaction_session_timeout` 5 s, and a 10 s client deadline per table. A deadline-abandoned connection is closed, never returned to the pool, and the server rolls back anything uncommitted. If the deadline fires after `COMMIT` was sent, the change may have landed; the next audit reports the truth. The orphaned backend keeps its locks until its in-flight statement ends, at most `statement_timeout`. - **Verify before commit.** The changed catalog is re-audited inside the transaction for bounds, canonical kinds, trigger parity and catch-all emptiness, and serving safety must not regress. It also checks that every new child has valid, ready clones of each parent index and validated clones of each PK/unique/FK constraint. - **Audits tolerate a concurrent advance.** Catalog rows come from the statement snapshot, but `pg_get_expr` renders from current catalog state. So an audit that overlaps another relay's `DROP` of a catch-all can see a partition that no longer exists. The read-only audit detects this (a NULL bound, or `42P01` from the emptiness probe) and retries up to twice with a fresh snapshot. - **Fail fast on audit probes.** If another relay's audit is probing the catch-all's emptiness, the catch-all `NOWAIT` lock fails and the run reports `lock_timeout`. The parent lock is then released immediately, and the next runner or interval completes the advance. Uncovered-month creation (`BUZZ_PARTITION_MANAGER_CREATE_ENABLED`, default `true`) now runs through the same locked transaction, replacing the unlocked per-month `CREATE`. That flag remains the kill switch for all automatic partition DDL: with it off, advancement is disabled too. With advancement disabled (the default) the periodic loop stays read-only and startup behaves as before, apart from the locking. A runner that planned DDL always re-audits before returning, even when a peer did the work, so it never reports the pre-maintenance layout. A serialized re-plan that finds only disabled work reports `skipped_disabled`. ## Behavior changes - The relay audit horizon and the `buzz-admin partition-audit` default are now 6 months (`PARTITION_MANAGER_MONTHS_AHEAD`), up from 3. Until advancement is enabled, months already covered by the catch-all make `degraded` true. That correctly reports a runway under six months. - New metric `buzz_partition_maintenance_runs_total{table,outcome}` with outcomes `noop | skipped_disabled | created | advanced | skipped_locked | operator_required | lock_timeout | deadline | error`. `buzz_partition_create_attempts_total` keeps its per-month `created | skipped_covered | error` labels. `error` there still means a DDL or collision failure; `lock_timeout` and `deadline` runs add nothing to it, and only `buzz_partition_maintenance_runs_total` records them. - `ensure_future_partitions` (free function and `Db` method) is removed; `maintain_partitions` is the one entry point. - New advisory-lock label `lock_type="partition_maintenance"`. ## Rollout Enable `BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED` per environment in BPCI: `bb-block-staging` first, then production. The staging canary proves a real advance because the six-month horizon already extends past January. Fallback: rerun the internal manual runbook, which uses the same parent-first lock order as of 2026-10-05. It keeps `ACCESS EXCLUSIVE` on the counterparts, because its detach/attach can drop FK triggers on them. ## Testing - `cargo test -p buzz-db --lib partition`: 17 unit tests. The 10 new ones cover: - plans: repaired layout, policy gating, no-op, refusal states including child-only triggers, December→January; - the create kill switch also stopping advancement; - the month cap at exactly 120; - per-month metric outcomes for transient failures versus DDL failures; - the mid-audit-drop retry classifier. - `cargo test -p buzz-db --lib partition -- --ignored` (PostgreSQL 17): 48/48 on 9 consecutive runs, including 17 new tests: - advance to the canonical layout, with routing and idempotent rerun; - replacing a canonically named catch-all in one transaction; - populated catch-all refused in under 1 s while the parent is locked; - a `communities` writer making the run fail fast and roll back; - a `communities` reader not blocking the advance; - an in-flight insert's FK check, and its write to `communities`, behind queued maintenance completing without a deadlock; - a foreign key committed while maintenance waits on the parent, with a writer on its new counterpart, failing fast as `lock_timeout` on that table instead of waiting under the parent lock; - a foreign key referencing `events`, committed while maintenance waits on the parent, being refused under the parent lock; - a catch-all's own foreign key, or a foreign key referencing it, being refused before the parent lock with the key intact, and a catch-all key committed while maintenance waits on the parent being refused after the catch-all lock; - a `40P01` deadlock victim being classified as `lock_timeout`; - parent contention timing out without blocking the other table; - held advisory lock skipping with no DDL; - concurrent runners converging, with at most a transient `lock_timeout`; - a stale audit after a concurrent advance being a no-op, and the losing runner returning the committed layout; - a re-plan that finds only disabled work reporting `skipped_disabled`; - an audit retrying when a listed partition is dropped mid-audit; - a catch-all with a child-only trigger being refused; - client deadline abandoning the attempt and closing the session (scoped to the test's own blocked backend); - catch-all name collision; - incoming foreign key refused. - Mutation checks: - restoring the old counterpart-first `ACCESS EXCLUSIVE` order fails the reader test (`lock_timeout`) and the writer test (a ~1 s deadlock wait); - dropping `NOWAIT` from the counterpart lock, or allowing a populated catch-all, fails the matching tests; - removing the parent-locked counterpart re-read fails the committed-FK tests (the parent is held 2.06 s waiting on the new counterpart), and ignoring its refusal fails the referencing-FK test; - locking the counterparts `SHARE ROW EXCLUSIVE` before the parent fails the writer test; - removing either catch-all foreign-key check fails its test. - Full PostgreSQL lane: 716/718. Both failures are timing-sensitive tests outside partition code, and each passes 3/3 when run alone: - `hanging_redis_peer_times_out_before_the_health_listener_binds`; - `readiness_check_cancellation_balances_waiter_and_inflight_connection`. - `cargo clippy -p buzz-db -p buzz-relay -p buzz-admin --all-targets -D warnings` and `cargo fmt --check` are clean. ## Related Refs #2396. This PR conflicts textually with #7994 (partition audit observability follow-ups) in `crates/buzz-db/src/store/partition.rs` and `crates/buzz-relay/src/main.rs`. Whichever PR merges second will be rebased onto the other. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Ian Oberst Co-authored-by: Ian Oberst Co-authored-by: Claude Opus 5.5 --- .env.example | 7 +- ARCHITECTURE.md | 2 +- crates/buzz-admin/src/main.rs | 2 +- crates/buzz-db/src/error.rs | 5 + crates/buzz-db/src/runtime/observability.rs | 6 +- .../src/store/artifact_postgres_tests.rs | 10 +- crates/buzz-db/src/store/partition.rs | 325 ++-- .../src/store/partition/maintenance.rs | 1117 ++++++++++++++ .../src/store/partition/maintenance_tests.rs | 1330 +++++++++++++++++ .../src/api/artifact_postgres_tests.rs | 12 +- crates/buzz-relay/src/config.rs | 39 +- crates/buzz-relay/src/main.rs | 31 +- 12 files changed, 2657 insertions(+), 229 deletions(-) create mode 100644 crates/buzz-db/src/store/partition/maintenance.rs create mode 100644 crates/buzz-db/src/store/partition/maintenance_tests.rs diff --git a/.env.example b/.env.example index d3deb6f536..f79eb23f73 100644 --- a/.env.example +++ b/.env.example @@ -49,8 +49,13 @@ REDIS_URL=redis://localhost:6379 # BUZZ_DB_POOL_SIZE=50 # Read-only partition catalog audit interval, clamped to 60s..24h (default 900s). # BUZZ_PARTITION_AUDIT_INTERVAL_SECS=900 -# Emergency kill switch for monthly partition CREATE; audits remain enabled. +# Emergency kill switch for all automatic partition DDL, including catch-all +# advancement; audits remain enabled. # BUZZ_PARTITION_MANAGER_CREATE_ENABLED=true +# Opt in to replacing an empty partition catch-all with six months of dedicated +# monthlies, at startup and on every audit interval (default false). Has no +# effect while the kill switch above is off. +# BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED=false # Writer-session Postgres timeouts for buzz-db-backed pools and the relay audit # pool, all in milliseconds; 0 disables. The separately deployed push gateway diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 55e3bbe969..82d7cac7bc 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -462,7 +462,7 @@ All database access. Uses `sqlx::query()` (runtime, not compile-time macros) — | `channel.rs` | Channel CRUD, membership management, role enforcement (transactional) | | `feed.rs` | `query_mentions` (INNER JOIN event_mentions), `query_needs_action`, `query_activity` | | `workflow.rs` | Full workflow/run/approval CRUD; SHA-256 hashed approval tokens | -| `partition.rs` | Monthly range partitioning for `events` and `delivery_log` tables | +| `partition.rs` | Monthly range partitioning for `events` and `delivery_log` tables: read-only catalog audit plus advisory-locked, time-bounded maintenance that creates uncovered months and advances an empty catch-all (`partition/maintenance.rs`) | | `dm.rs` | DM channel management | | `reaction.rs` | Reaction storage and retrieval | | `thread.rs` | Thread/reply tracking | diff --git a/crates/buzz-admin/src/main.rs b/crates/buzz-admin/src/main.rs index 312f01c1de..66780e2041 100644 --- a/crates/buzz-admin/src/main.rs +++ b/crates/buzz-admin/src/main.rs @@ -94,7 +94,7 @@ enum Command { /// Run the partition catalog audit using a read-only database session. PartitionAudit { /// Future months to include in the coverage check. - #[arg(long, default_value_t = 3)] + #[arg(long, default_value_t = buzz_db::partition::PARTITION_MANAGER_MONTHS_AHEAD)] months_ahead: u32, }, /// Inspect deployment-wide Buzz product feedback. diff --git a/crates/buzz-db/src/error.rs b/crates/buzz-db/src/error.rs index 4f15cc2b24..8c051a7a79 100644 --- a/crates/buzz-db/src/error.rs +++ b/crates/buzz-db/src/error.rs @@ -45,6 +45,11 @@ pub enum DbError { #[error("invalid data: {0}")] InvalidData(String), + /// A partition the catalog audit had already listed was dropped before the + /// audit finished. A fresh snapshot sees a consistent catalog. + #[error("partition catalog changed mid-audit: {0}")] + PartitionDroppedMidAudit(String), + /// A serving write admitted before the lifecycle transition is still live. /// This is an ordinary retryable drain condition, not a safety violation. #[error( diff --git a/crates/buzz-db/src/runtime/observability.rs b/crates/buzz-db/src/runtime/observability.rs index 99256f3c2f..d487b9a362 100644 --- a/crates/buzz-db/src/runtime/observability.rs +++ b/crates/buzz-db/src/runtime/observability.rs @@ -184,16 +184,18 @@ pub(crate) enum LockType { PushGate, Deletion, MigrationSchemaSafety, + PartitionMaintenance, } impl LockType { #[cfg(test)] - pub(crate) const ALL: [Self; 5] = [ + pub(crate) const ALL: [Self; 6] = [ Self::Replacement, Self::Membership, Self::PushGate, Self::Deletion, Self::MigrationSchemaSafety, + Self::PartitionMaintenance, ]; pub(crate) const fn as_str(self) -> &'static str { @@ -203,6 +205,7 @@ impl LockType { Self::PushGate => "push_gate", Self::Deletion => "deletion", Self::MigrationSchemaSafety => "migration_schema_safety", + Self::PartitionMaintenance => "partition_maintenance", } } } @@ -671,6 +674,7 @@ mod tests { "push_gate", "deletion", "migration_schema_safety", + "partition_maintenance", ] ); assert_eq!( diff --git a/crates/buzz-db/src/store/artifact_postgres_tests.rs b/crates/buzz-db/src/store/artifact_postgres_tests.rs index 078cefb8f6..4ac03a06f4 100644 --- a/crates/buzz-db/src/store/artifact_postgres_tests.rs +++ b/crates/buzz-db/src/store/artifact_postgres_tests.rs @@ -24,7 +24,15 @@ impl Fixture { crate::migration::run_migrations(&pool).await.unwrap(); } let db = Db::from_pool(pool); - db.ensure_future_partitions(1, true).await.unwrap(); + db.maintain_partitions( + 1, + crate::partition::PartitionMaintenancePolicy { + create_enabled: true, + advance_enabled: false, + }, + ) + .await + .unwrap(); let community = CommunityId::from_uuid(Uuid::new_v4()); sqlx::query("INSERT INTO communities(id,host) VALUES($1,$2)") .bind(community.as_uuid()) diff --git a/crates/buzz-db/src/store/partition.rs b/crates/buzz-db/src/store/partition.rs index 733ab73213..02dfa4e708 100644 --- a/crates/buzz-db/src/store/partition.rs +++ b/crates/buzz-db/src/store/partition.rs @@ -7,18 +7,24 @@ use std::time::{Duration, Instant}; use buzz_datastore_tracing::datastore_span; use chrono::{DateTime, Datelike, NaiveDate, TimeZone, Timelike, Utc}; use serde::Serialize; -use sqlx::{Connection, PgConnection, PgPool, Row}; -use tracing::info; +use sqlx::{PgConnection, PgPool, Row}; use crate::error::{DbError, Result}; use crate::Db; +mod maintenance; + +pub use maintenance::{ + maintain_partitions, PartitionMaintenanceOutcome, PartitionMaintenancePolicy, + PARTITION_MANAGER_MONTHS_AHEAD, +}; + /// Tables that may be partition-managed. The allowlist prevents DDL injection. const PARTITIONED_TABLES: &[&str] = &["events", "delivery_log"]; /// Maximum future-month horizon accepted by catalog audits and creation. /// -/// The relay uses three months. Keeping the public API bounded prevents an +/// The relay uses [`PARTITION_MANAGER_MONTHS_AHEAD`]. Keeping the public API bounded prevents an /// operator-supplied `buzz-admin partition-audit --months-ahead` value from /// allocating or iterating an effectively unbounded report. pub const MAX_PARTITION_MONTHS_AHEAD: u32 = 120; @@ -29,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)"), @@ -451,131 +461,24 @@ async fn audit_partition_catalog_report_at_with_timeout( } } -/// Audit first, then create only months proven to be uncovered. -/// -/// Covered ranges are never probed with DDL. Creation failures are collected -/// across all managed parents and months before an aggregate error is returned. -pub async fn ensure_future_partitions( - pool: &PgPool, - months_ahead: u32, - create_enabled: bool, -) -> Result { - ensure_future_partitions_at(pool, months_ahead, create_enabled, Utc::now()).await -} - +/// Create-only maintenance at a fixed clock, for tests of month creation. +#[cfg(test)] async fn ensure_future_partitions_at( pool: &PgPool, months_ahead: u32, create_enabled: bool, now: DateTime, ) -> Result { - let audit = audit_partition_catalog_at(pool, months_ahead, now).await?; - let mut errors = Vec::new(); - let mut created_any = false; - - for table in &audit.tables { - for month in &table.months { - match month.kind { - MonthCoverageKind::CoveredByMonthly - | MonthCoverageKind::CoveredByCatchAll - | MonthCoverageKind::CoveredByDefault => { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "skipped_covered" - ) - .increment(1); - } - MonthCoverageKind::Uncovered if !create_enabled => {} - MonthCoverageKind::Uncovered => { - let expected_name = partition_name(table.table, month.start); - let collision = match relation_name_exists(pool, &expected_name).await { - Ok(collision) => collision, - Err(error) => { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "error" - ) - .increment(1); - errors.push(format!( - "{} {}: failed to check canonical name {expected_name}: {error}", - table.table, - month.start.format("%Y-%m") - )); - continue; - } - }; - if collision { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "error" - ) - .increment(1); - errors.push(format!( - "{} {}: canonical name {expected_name} already exists without the expected attachment and bounds", - table.table, - month.start.format("%Y-%m") - )); - continue; - } - if !table.structurally_safe_for_creation() { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "error" - ) - .increment(1); - errors.push(format!( - "{} {}: catalog is not structurally safe for automatic partition creation", - table.table, - month.start.format("%Y-%m") - )); - continue; - } - match create_month_partition(pool, table.table, month.start).await { - Ok(name) => { - created_any = true; - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "created" - ) - .increment(1); - info!(table = table.table, partition = name, "added partition"); - } - Err(error) => { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "error" - ) - .increment(1); - errors.push(format!( - "{} {}: {error}", - table.table, - month.start.format("%Y-%m") - )); - } - } - } - } - } - } - - if errors.is_empty() { - if created_any { - audit_partition_catalog_at(pool, months_ahead, now).await - } else { - Ok(audit) - } - } else { - Err(DbError::InvalidData(format!( - "partition creation failed: {}", - errors.join("; ") - ))) - } + maintenance::maintain_partitions_at( + pool, + months_ahead, + PartitionMaintenancePolicy { + create_enabled, + advance_enabled: false, + }, + now, + ) + .await } async fn audit_table( @@ -585,6 +488,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) @@ -595,6 +518,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::PartitionDroppedMidAudit(_) => 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)") @@ -627,14 +564,14 @@ impl Db { audit_partition_catalog_report(&self.pool, months_ahead).await } - /// Ensures monthly partitions exist for the next N months when creation is enabled. - #[datastore_span(name = "ensure_future_partitions", system = "postgresql")] - pub async fn ensure_future_partitions( + /// Issues the bounded partition DDL `policy` permits, then returns a fresh audit. + #[datastore_span(name = "maintain_partitions", system = "postgresql")] + pub async fn maintain_partitions( &self, months_ahead: u32, - create_enabled: bool, + policy: PartitionMaintenancePolicy, ) -> Result { - ensure_future_partitions(&self.pool, months_ahead, create_enabled).await + maintain_partitions(&self.pool, months_ahead, policy).await } } @@ -758,14 +695,25 @@ 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::PartitionDroppedMidAudit(format!("a sibling of partition {name}")) + })?; 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::PartitionDroppedMidAudit(format!("partition {name}")))?; 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)| { @@ -1536,84 +1484,7 @@ fn parse_timestamp_literal(literal: &str) -> Option> { Some(Utc.from_utc_datetime(&date.and_hms_opt(0, 0, 0)?)) } -async fn create_month_partition( - pool: &PgPool, - table: &str, - start: DateTime, -) -> Result { - if !PARTITIONED_TABLES.contains(&table) { - return Err(DbError::InvalidData(format!( - "table not in partition allowlist: {table:?}" - ))); - } - let (end_year, end_month) = add_months(start.year(), start.month(), 1)?; - let end = month_start(end_year, end_month)?; - let partition_name = partition_name(table, start); - let start_date = start.format("%Y-%m-%d"); - let end_date = end.format("%Y-%m-%d"); - let sql = format!( - "CREATE TABLE {} PARTITION OF {} \ - FOR VALUES FROM ('{start_date}') TO ('{end_date}')", - quote_identifier(&partition_name), - quote_identifier(table) - ); - let mut connection = crate::observability::acquire_writer( - pool, - crate::observability::WriterOperation::Bootstrap, - ) - .await?; - let mut transaction = connection.begin().await?; - pin_catalog_rendering(&mut transaction).await?; - sqlx::query(sqlx::AssertSqlSafe(sql)) - .execute(&mut *transaction) - .await?; - let row = sqlx::query( - r#" - SELECT child.relkind::text AS relation_kind, - inherited.inhdetachpending AS detach_pending, - pg_catalog.pg_get_partkeydef(parent.oid) AS partition_key, - pg_catalog.pg_get_expr(child.relpartbound, child.oid) AS bound - FROM pg_catalog.pg_inherits inherited - JOIN pg_catalog.pg_class parent ON parent.oid = inherited.inhparent - JOIN pg_catalog.pg_namespace parent_ns ON parent_ns.oid = parent.relnamespace - JOIN pg_catalog.pg_class child ON child.oid = inherited.inhrelid - JOIN pg_catalog.pg_namespace child_ns ON child_ns.oid = child.relnamespace - WHERE parent_ns.nspname = current_schema() - AND child_ns.nspname = current_schema() - AND parent.relname = $1 - AND child.relname = $2 - "#, - ) - .bind(table) - .bind(&partition_name) - .fetch_optional(&mut *transaction) - .await?; - let exact = row.is_some_and(|row| { - let relation_kind = row.try_get::("relation_kind").ok(); - let detach_pending = row.try_get::("detach_pending").ok(); - let partition_key = row - .try_get::, _>("partition_key") - .ok() - .flatten(); - let bounds = row - .try_get::("bound") - .ok() - .and_then(|bound| parse_range_bounds(&bound)); - relation_kind.as_deref() == Some("r") - && detach_pending == Some(false) - && partition_key.as_deref() == expected_partition_key(table) - && bounds == Some((PartitionBound::Finite(start), PartitionBound::Finite(end))) - }); - if !exact { - transaction.rollback().await?; - return Err(DbError::InvalidData(format!( - "created partition {partition_name} failed exact catalog postcondition" - ))); - } - transaction.commit().await?; - Ok(partition_name) -} - +#[cfg(test)] async fn relation_name_exists(pool: &PgPool, relation: &str) -> Result { Ok(sqlx::query_scalar( r#" @@ -1687,6 +1558,20 @@ mod tests { .all(|error| error.error.contains("months_ahead must be at most 120"))); } + #[test] + fn only_a_mid_audit_drop_is_retried() { + assert!(is_concurrent_drop(&DbError::PartitionDroppedMidAudit( + "partition events_p2026_09".to_string() + ))); + assert!(!is_concurrent_drop(&DbError::NotFound("row".to_string()))); + assert!(!is_concurrent_drop(&DbError::InvalidData( + "bound".to_string() + ))); + assert!(!is_concurrent_drop(&DbError::Sqlx( + sqlx::Error::PoolTimedOut + ))); + } + #[test] fn parses_pg16_range_bound_formats() { let expected = Utc.with_ymd_and_hms(2026, 7, 1, 0, 0, 0).unwrap(); @@ -1903,17 +1788,17 @@ mod tests { assert_eq!(extra, Some(1.0)); } - mod postgres_tests { + pub(super) mod postgres_tests { use sqlx::postgres::PgPoolOptions; use uuid::Uuid; use super::*; - async fn scratch_pool() -> (PgPool, PgPool, String) { + pub(in crate::store::partition) async fn scratch_pool() -> (PgPool, PgPool, String) { scratch_pool_with_max_connections(4).await } - async fn scratch_pool_with_max_connections( + pub(in crate::store::partition) async fn scratch_pool_with_max_connections( max_connections: u32, ) -> (PgPool, PgPool, String) { let url = crate::test_support::database_url(); @@ -1941,7 +1826,7 @@ mod tests { (pool, admin, schema) } - async fn drop_schema(admin: &PgPool, schema: &str) { + pub(in crate::store::partition) async fn drop_schema(admin: &PgPool, schema: &str) { let _ = sqlx::query(sqlx::AssertSqlSafe(format!( "DROP SCHEMA IF EXISTS {schema} CASCADE" ))) @@ -1949,7 +1834,7 @@ mod tests { .await; } - async fn seed_parents(pool: &PgPool) { + pub(in crate::store::partition) async fn seed_parents(pool: &PgPool) { seed_parents_with_keys(pool, "created_at", "delivered_at").await; } @@ -1984,7 +1869,13 @@ mod tests { } } - async fn create_child(pool: &PgPool, table: &str, name: &str, lower: &str, upper: &str) { + pub(in crate::store::partition) async fn create_child( + pool: &PgPool, + table: &str, + name: &str, + lower: &str, + upper: &str, + ) { let sql = format!( "CREATE TABLE {name} PARTITION OF {table} FOR VALUES FROM ({lower}) TO ({upper})" ); @@ -2032,7 +1923,9 @@ mod tests { .expect("create nested child"); } - async fn catalog_snapshot(pool: &PgPool) -> Vec<(String, String, String, i64)> { + pub(in crate::store::partition) async fn catalog_snapshot( + pool: &PgPool, + ) -> Vec<(String, String, String, i64)> { let mut transaction = pool.begin().await.expect("begin catalog snapshot"); sqlx::query("SET TRANSACTION READ ONLY") .execute(&mut *transaction) @@ -2067,11 +1960,11 @@ mod tests { snapshot } - fn fixed_now() -> DateTime { + pub(in crate::store::partition) fn fixed_now() -> DateTime { Utc.with_ymd_and_hms(2026, 9, 15, 12, 0, 0).unwrap() } - async fn seed_fresh_layout(pool: &PgPool) { + pub(in crate::store::partition) async fn seed_fresh_layout(pool: &PgPool) { seed_parents(pool).await; for table in PARTITIONED_TABLES { create_child( diff --git a/crates/buzz-db/src/store/partition/maintenance.rs b/crates/buzz-db/src/store/partition/maintenance.rs new file mode 100644 index 0000000000..f19699181b --- /dev/null +++ b/crates/buzz-db/src/store/partition/maintenance.rs @@ -0,0 +1,1117 @@ +//! Serialized, bounded partition DDL for `events` and `delivery_log`. +//! +//! The manager issues only two kinds of change: creating a month that no +//! partition covers, and replacing an empty right-edge catch-all with +//! dedicated monthlies plus a new catch-all after the horizon. Each table's +//! change runs in one transaction that +//! +//! 1. holds a schema-scoped advisory lock, so concurrent relays produce one +//! winner and `skipped_locked` losers; +//! 2. locks the parent, then foreign-key counterparts read under the parent +//! lock, then the catch-all, in the order writers take them, waiting only on +//! the parent and failing fast on any other contention; +//! 3. re-audits under those locks and refuses if the plan changed; and +//! 4. re-audits the changed catalog and verifies index and constraint +//! inheritance before commit. +//! +//! A populated catch-all is always refused before any lock is taken: moving +//! rows requires an operator. + +use std::time::Duration; + +use chrono::{DateTime, Datelike, Utc}; +use serde::Serialize; +use sqlx::{Connection, PgConnection, PgPool, Row}; +use tracing::info; + +use super::{ + add_months, audit_partition_catalog_at, audit_table_on, month_start, partition_name, + pin_catalog_rendering, qualified_relation_name, validated_months_ahead, MonthCoverageKind, + PartitionAudit, PartitionBound, PartitionChildKind, PartitionTableAudit, + MAX_PARTITION_MONTHS_AHEAD, +}; +use crate::error::{DbError, Result}; + +/// Dedicated monthly runway, after the current month, that the relay maintains. +pub const PARTITION_MANAGER_MONTHS_AHEAD: u32 = 6; + +/// Maximum wait for the parent lock. Counterparts and children use `NOWAIT`. +const MAINTENANCE_LOCK_TIMEOUT: &str = "2s"; + +/// Server-side budget for each statement in the maintenance transaction. +const MAINTENANCE_STATEMENT_TIMEOUT: &str = "5s"; + +/// Server-side reaping of a maintenance session whose client stalls mid-transaction. +const MAINTENANCE_IDLE_IN_TRANSACTION_TIMEOUT: &str = "5s"; + +/// Client-side budget for one table's complete maintenance attempt. +const MAINTENANCE_TABLE_DEADLINE: Duration = Duration::from_secs(10); + +/// Namespace hashed with `current_schema()` into the maintenance advisory lock. +const MAINTENANCE_LOCK_NAMESPACE: &str = "buzz.partition_maintenance:"; + +/// Which automatic partition DDL the manager may issue. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct PartitionMaintenancePolicy { + /// Create months that no partition covers. Off stops all automatic DDL. + pub create_enabled: bool, + /// Replace an empty right-edge catch-all with dedicated monthlies. + /// Takes effect only while `create_enabled` is on. + pub advance_enabled: bool, +} + +/// Result of one table's maintenance attempt, used as the `outcome` metric label. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum PartitionMaintenanceOutcome { + /// Every desired month already has a dedicated leaf. + Noop, + /// DDL is needed but the policy forbids it. + SkippedDisabled, + /// Uncovered months were created without touching a catch-all. + Created, + /// An empty catch-all was replaced by monthlies and a later catch-all. + Advanced, + /// Another session holds the maintenance advisory lock. + SkippedLocked, + /// Automatic DDL would be unsafe; an operator must repair the catalog. + OperatorRequired, + /// A table lock was not available within the lock budget. + LockTimeout, + /// The server statement budget or the client deadline expired. + Deadline, + /// Any other failure. + Error, +} + +impl PartitionMaintenanceOutcome { + /// Stable metric label. + pub const fn as_str(self) -> &'static str { + match self { + Self::Noop => "noop", + Self::SkippedDisabled => "skipped_disabled", + Self::Created => "created", + Self::Advanced => "advanced", + Self::SkippedLocked => "skipped_locked", + Self::OperatorRequired => "operator_required", + Self::LockTimeout => "lock_timeout", + Self::Deadline => "deadline", + Self::Error => "error", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct CatchAllReplacement { + schema: String, + name: String, + new_lower: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct MaintenanceChange { + /// Ascending month starts whose canonical monthly is created. + create: Vec>, + /// The empty catch-all dropped and recreated after the horizon. + catch_all: Option, +} + +impl MaintenanceChange { + fn outcome(&self) -> PartitionMaintenanceOutcome { + if self.catch_all.is_some() { + PartitionMaintenanceOutcome::Advanced + } else { + PartitionMaintenanceOutcome::Created + } + } + + fn replaced(&self) -> Option<(String, String)> { + self.catch_all + .as_ref() + .map(|catch_all| (catch_all.schema.clone(), catch_all.name.clone())) + } + + fn target_names(&self, table: &str) -> Vec { + let mut names: Vec<_> = self + .create + .iter() + .map(|month| partition_name(table, *month)) + .collect(); + if self.catch_all.is_some() { + names.push(catch_all_name(table)); + } + names + } +} + +#[derive(Debug, PartialEq, Eq)] +enum MaintenancePlan { + Noop, + Disabled, + Refuse(String), + Apply(MaintenanceChange), +} + +enum LockedAttempt { + Committed(MaintenanceChange), + SkippedLocked, + Unchanged, + Disabled, + Refused(String), +} + +struct TableMaintenance { + outcome: PartitionMaintenanceOutcome, + change: Option, + error: Option, +} + +impl TableMaintenance { + fn done(outcome: PartitionMaintenanceOutcome) -> Self { + Self { + outcome, + change: None, + error: None, + } + } + + fn failed( + outcome: PartitionMaintenanceOutcome, + change: Option, + error: String, + ) -> Self { + Self { + outcome, + change, + error: Some(error), + } + } +} + +fn catch_all_name(table: &str) -> String { + format!("{table}_p_future") +} + +fn next_month(start: DateTime) -> Result> { + let (year, month) = add_months(start.year(), start.month(), 1)?; + month_start(year, month) +} + +/// Decide what DDL, if any, would bring `audit` to the desired layout. +/// +/// Pure so the same decision runs before locking, under locks, and in tests. +fn plan_maintenance( + audit: &PartitionTableAudit, + policy: PartitionMaintenancePolicy, +) -> MaintenancePlan { + let uncovered: Vec<_> = audit + .months + .iter() + .filter(|month| month.kind == MonthCoverageKind::Uncovered) + .map(|month| month.start) + .collect(); + let catch_all_covered = audit + .months + .iter() + .any(|month| month.kind == MonthCoverageKind::CoveredByCatchAll); + let create = policy.create_enabled && !uncovered.is_empty(); + // Advancement creates monthlies, so the create kill switch stops it too. + let advance = policy.create_enabled && policy.advance_enabled && catch_all_covered; + if !create && !advance { + return if uncovered.is_empty() && !catch_all_covered { + MaintenancePlan::Noop + } else { + MaintenancePlan::Disabled + }; + } + if !audit.structurally_safe_for_creation() { + return MaintenancePlan::Refuse( + "catalog is not structurally safe for automatic partition DDL".to_string(), + ); + } + + let mut months = if create { uncovered } else { Vec::new() }; + let mut catch_all = None; + if advance { + let catch_alls: Vec<_> = audit + .children + .iter() + .filter(|child| child.kind == PartitionChildKind::CatchAll) + .collect(); + let [child] = catch_alls.as_slice() else { + return MaintenancePlan::Refuse(format!( + "expected exactly one right-edge catch-all, found {}", + catch_alls.len() + )); + }; + match child.catch_all_nonempty { + Some(false) => {} + Some(true) => { + return MaintenancePlan::Refuse(format!( + "catch-all {} contains rows; advancing it requires an operator", + child.name + )) + } + None => { + return MaintenancePlan::Refuse(format!( + "catch-all {} occupancy is unknown", + child.name + )) + } + } + // 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", + child.name + )); + }; + if month_start(lower.year(), lower.month()).ok() != Some(lower) { + return MaintenancePlan::Refuse(format!( + "catch-all {} lower bound {lower} is not a UTC month start", + child.name + )); + } + let Some(new_lower) = audit + .months + .last() + .and_then(|month| next_month(month.start).ok()) + else { + return MaintenancePlan::Refuse("desired horizon is not representable".to_string()); + }; + let mut month = lower; + while month < new_lower { + if months.len() >= MAX_PARTITION_MONTHS_AHEAD as usize { + return MaintenancePlan::Refuse(format!( + "catch-all {} would require more than {MAX_PARTITION_MONTHS_AHEAD} monthlies", + child.name + )); + } + months.push(month); + let Ok(next) = next_month(month) else { + return MaintenancePlan::Refuse("desired horizon is not representable".to_string()); + }; + month = next; + } + catch_all = Some(CatchAllReplacement { + schema: child.schema.clone(), + name: child.name.clone(), + new_lower, + }); + } + months.sort(); + months.dedup(); + MaintenancePlan::Apply(MaintenanceChange { + create: months, + catch_all, + }) +} + +/// Audit, then issue only the DDL `policy` permits, one bounded transaction per table. +/// +/// A table that is already at the desired layout, or whose required DDL is +/// disabled, takes no lock. One table's failure does not prevent the other +/// table's attempt; failures are aggregated into the returned error. +pub async fn maintain_partitions( + pool: &PgPool, + months_ahead: u32, + policy: PartitionMaintenancePolicy, +) -> Result { + maintain_partitions_at(pool, months_ahead, policy, Utc::now()).await +} + +pub(super) async fn maintain_partitions_at( + pool: &PgPool, + months_ahead: u32, + policy: PartitionMaintenancePolicy, + now: DateTime, +) -> Result { + maintain_partitions_with_deadline(pool, months_ahead, policy, now, MAINTENANCE_TABLE_DEADLINE) + .await +} + +async fn maintain_partitions_with_deadline( + pool: &PgPool, + months_ahead: u32, + policy: PartitionMaintenancePolicy, + now: DateTime, + table_deadline: Duration, +) -> Result { + let audit = audit_partition_catalog_at(pool, months_ahead, now).await?; + maintain_audited(pool, audit, months_ahead, policy, now, table_deadline).await +} + +/// Maintain each table described by `audit`, which may be stale by now. +async fn maintain_audited( + pool: &PgPool, + audit: PartitionAudit, + months_ahead: u32, + policy: PartitionMaintenancePolicy, + now: DateTime, + table_deadline: Duration, +) -> Result { + let horizon = validated_months_ahead(months_ahead)?; + let mut errors = Vec::new(); + // Any table that planned DDL may have found the catalog changed under the + // lock, so the opening audit no longer describes it. + let mut stale = false; + for table_audit in &audit.tables { + let table = table_audit.table; + let plan = plan_maintenance(table_audit, policy); + stale |= matches!(plan, MaintenancePlan::Apply(_)); + let result = maintain_table( + pool, + table_audit, + plan, + horizon, + policy, + now, + table_deadline, + ) + .await; + emit_maintenance_metrics(table_audit, &result); + match (&result.error, result.outcome) { + (Some(error), outcome) => errors.push(format!("{table} {}: {error}", outcome.as_str())), + ( + None, + PartitionMaintenanceOutcome::Created | PartitionMaintenanceOutcome::Advanced, + ) => { + info!( + table, + outcome = result.outcome.as_str(), + partitions = ?result + .change + .as_ref() + .map(|change| change.target_names(table)), + "partition maintenance committed" + ); + } + (None, _) => {} + } + } + + if !errors.is_empty() { + Err(DbError::InvalidData(format!( + "partition maintenance failed: {}", + errors.join("; ") + ))) + } else if stale { + audit_partition_catalog_at(pool, months_ahead, now).await + } else { + Ok(audit) + } +} + +async fn maintain_table( + pool: &PgPool, + audit: &PartitionTableAudit, + plan: MaintenancePlan, + months_ahead: i32, + policy: PartitionMaintenancePolicy, + now: DateTime, + table_deadline: Duration, +) -> TableMaintenance { + let change = match plan { + MaintenancePlan::Noop => return TableMaintenance::done(PartitionMaintenanceOutcome::Noop), + MaintenancePlan::Disabled => { + return TableMaintenance::done(PartitionMaintenanceOutcome::SkippedDisabled) + } + 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, + error, + ); + } + MaintenancePlan::Apply(change) => change, + }; + + let deadline = tokio::time::Instant::now() + table_deadline; + let mut connection = match crate::observability::acquire_writer_until( + pool, + crate::observability::WriterOperation::Maintenance, + deadline, + ) + .await + { + Ok(connection) => connection, + Err(error) => { + let error = DbError::from(error); + return TableMaintenance::failed( + classify_error(&error), + Some(change), + error.to_string(), + ); + } + }; + let attempt = tokio::time::timeout_at( + deadline, + apply_change(&mut connection, audit.table, months_ahead, policy, now), + ) + .await; + let Ok(attempt) = attempt else { + // The abandoned transaction may still be open on the server. Close the + // session instead of returning it to the pool so PostgreSQL rolls back + // whatever did not commit. The deadline can fire after COMMIT was sent, + // so the change may have landed; the next audit reports the truth. The + // backend keeps its locks until its in-flight statement ends, which + // `statement_timeout` bounds. + connection.close_on_drop(); + drop(connection); + return TableMaintenance::failed( + PartitionMaintenanceOutcome::Deadline, + Some(change), + format!( + "exceeded the {}ms maintenance deadline; the attempt was abandoned and the \ + server rolls back any uncommitted work", + table_deadline.as_millis() + ), + ); + }; + match attempt { + Ok(LockedAttempt::Committed(committed)) => TableMaintenance { + outcome: committed.outcome(), + change: Some(committed), + error: None, + }, + Ok(LockedAttempt::SkippedLocked) => { + TableMaintenance::done(PartitionMaintenanceOutcome::SkippedLocked) + } + Ok(LockedAttempt::Unchanged) => TableMaintenance::done(PartitionMaintenanceOutcome::Noop), + Ok(LockedAttempt::Disabled) => { + TableMaintenance::done(PartitionMaintenanceOutcome::SkippedDisabled) + } + Ok(LockedAttempt::Refused(reason)) => TableMaintenance::failed( + PartitionMaintenanceOutcome::OperatorRequired, + Some(change), + reason, + ), + Err(error) => { + TableMaintenance::failed(classify_error(&error), Some(change), error.to_string()) + } + } +} + +async fn apply_change( + connection: &mut PgConnection, + table: &'static str, + months_ahead: i32, + policy: PartitionMaintenancePolicy, + now: DateTime, +) -> Result { + let mut transaction = connection.begin().await?; + configure_maintenance_transaction(&mut transaction).await?; + + let acquired: bool = crate::observability::observe_advisory_lock( + crate::observability::LockType::PartitionMaintenance, + sqlx::query_scalar( + "SELECT pg_catalog.pg_try_advisory_xact_lock(\ + pg_catalog.hashtextextended($1 || pg_catalog.current_schema()::text, 0))", + ) + .bind(MAINTENANCE_LOCK_NAMESPACE) + .fetch_one(&mut *transaction), + ) + .await?; + if !acquired { + transaction.rollback().await?; + 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 => { + transaction.rollback().await?; + return Ok(LockedAttempt::Unchanged); + } + MaintenancePlan::Disabled => { + transaction.rollback().await?; + return Ok(LockedAttempt::Disabled); + } + 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))); + } + + // Refuse an unproven lock set before taking any lock. + let parent = maintenance_parent(&mut transaction, table).await?; + if let Some(reason) = parent.refusal { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + if let Some(catch_all) = &expected.catch_all { + if let Some(reason) = catch_all_foreign_key_refusal(&mut transaction, catch_all).await? { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + } + // Parent first, in the order writers lock: an insert takes the parent and + // then the foreign-key counterparts its constraint check reads. Waiting on + // a counterpart while holding the parent could deadlock with such a writer. + // ONLY: without it PostgreSQL locks every historical child in turn. + lock( + &mut transaction, + &format!( + "LOCK TABLE ONLY {} IN ACCESS EXCLUSIVE MODE", + qualified_relation_name(&parent.schema, table) + ), + ) + .await?; + // Read the foreign keys again under the parent lock. The read above can + // miss one committed while this transaction waited on the parent, and + // attaching would then wait on its table while holding the parent. Adding + // or dropping a foreign key on or to the parent needs a lock that + // conflicts with this one, so the set cannot change from here on. + let parent = maintenance_parent(&mut transaction, table).await?; + if let Some(reason) = parent.refusal { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + // Attaching a partition adds its foreign key under SHARE ROW EXCLUSIVE on + // each referenced table. Taking that mode up front still admits readers and + // the `FOR KEY SHARE` checks of other writers, and fails at once rather + // than queue the counterpart's own writers behind this transaction. + // Dropping the catch-all removes only its inherited foreign key and check + // triggers, which lock the catch-all itself, never the referenced table + // (lock probes on PostgreSQL 16 and 17), so no stronger mode is needed. A + // catch-all with any other foreign key is refused before and after its lock. + for counterpart in &parent.counterparts { + lock( + &mut transaction, + &format!("LOCK TABLE ONLY {counterpart} IN SHARE ROW EXCLUSIVE MODE NOWAIT"), + ) + .await?; + } + // Blocks direct child inserts that bypass the parent. Never wait on a + // child while holding the parent. + if let Some(catch_all) = &expected.catch_all { + lock( + &mut transaction, + &format!( + "LOCK TABLE ONLY {} IN ACCESS EXCLUSIVE MODE NOWAIT", + qualified_relation_name(&catch_all.schema, &catch_all.name) + ), + ) + .await?; + // Foreign-key DDL on or to the catch-all conflicts with this lock but + // not with the parent's, so only this check is authoritative. + if let Some(reason) = catch_all_foreign_key_refusal(&mut transaction, catch_all).await? { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + } + + // 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::Noop => { + transaction.rollback().await?; + return Ok(LockedAttempt::Unchanged); + } + MaintenancePlan::Disabled => { + transaction.rollback().await?; + return Ok(LockedAttempt::Disabled); + } + MaintenancePlan::Refuse(reason) => { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + MaintenancePlan::Apply(_) => { + transaction.rollback().await?; + return Err(DbError::InvalidData( + "partition catalog changed between the audit and the maintenance locks".to_string(), + )); + } + } + + let parent_relation = qualified_relation_name(&parent.schema, table); + if let Some(catch_all) = &expected.catch_all { + execute_ddl( + &mut transaction, + format!( + "DROP TABLE {}", + qualified_relation_name(&catch_all.schema, &catch_all.name) + ), + ) + .await?; + } + for month in &expected.create { + let end = next_month(*month)?; + execute_ddl( + &mut transaction, + format!( + "CREATE TABLE {} PARTITION OF {parent_relation} \ + FOR VALUES FROM ('{}') TO ('{}')", + qualified_relation_name(&parent.schema, &partition_name(table, *month)), + month.format("%Y-%m-%d"), + end.format("%Y-%m-%d"), + ), + ) + .await?; + } + if let Some(catch_all) = &expected.catch_all { + execute_ddl( + &mut transaction, + format!( + "CREATE TABLE {} PARTITION OF {parent_relation} \ + FOR VALUES FROM ('{}') TO (MAXVALUE)", + qualified_relation_name(&parent.schema, &catch_all_name(table)), + catch_all.new_lower.format("%Y-%m-%d"), + ), + ) + .await?; + } + + let after = audit_table_on(&mut transaction, table, months_ahead, now).await?; + let verified = verify_applied(&locked, &after, &expected).and( + verify_inheritance(&mut transaction, parent.oid, &expected.target_names(table)).await?, + ); + if let Err(reason) = verified { + transaction.rollback().await?; + return Err(DbError::InvalidData(format!( + "partition maintenance failed its postcondition and rolled back: {reason}" + ))); + } + transaction.commit().await?; + Ok(LockedAttempt::Committed(expected)) +} + +async fn configure_maintenance_transaction(connection: &mut PgConnection) -> Result<()> { + pin_catalog_rendering(connection).await?; + for (setting, value) in [ + ("lock_timeout", MAINTENANCE_LOCK_TIMEOUT), + ("statement_timeout", MAINTENANCE_STATEMENT_TIMEOUT), + ( + "idle_in_transaction_session_timeout", + MAINTENANCE_IDLE_IN_TRANSACTION_TIMEOUT, + ), + ("application_name", "buzz_partition_maintenance"), + ] { + sqlx::query("SELECT pg_catalog.set_config($1, $2, true)") + .bind(setting) + .bind(value) + .execute(&mut *connection) + .await?; + } + Ok(()) +} + +async fn lock(connection: &mut PgConnection, sql: &str) -> Result<()> { + sqlx::query(sqlx::AssertSqlSafe(sql.to_string())) + .execute(&mut *connection) + .await?; + Ok(()) +} + +async fn execute_ddl(connection: &mut PgConnection, sql: String) -> Result<()> { + sqlx::query(sqlx::AssertSqlSafe(sql)) + .execute(&mut *connection) + .await?; + Ok(()) +} + +struct MaintenanceParent { + oid: i64, + schema: String, + /// Quoted, schema-qualified referenced tables in ascending OID order. + counterparts: Vec, + refusal: Option, +} + +async fn maintenance_parent( + connection: &mut PgConnection, + table: &str, +) -> Result { + let row = sqlx::query( + r#" + SELECT parent.oid::bigint AS parent_oid, + parent_ns.nspname AS parent_schema, + ARRAY( + SELECT pg_catalog.format('%I.%I', referenced_ns.nspname, referenced.relname) + FROM pg_catalog.pg_class referenced + JOIN pg_catalog.pg_namespace referenced_ns + ON referenced_ns.oid = referenced.relnamespace + WHERE referenced.oid IN ( + SELECT foreign_key.confrelid + FROM pg_catalog.pg_constraint foreign_key + WHERE foreign_key.conrelid = parent.oid + AND foreign_key.contype = 'f' + AND foreign_key.confrelid <> parent.oid + ) + ORDER BY referenced.oid + )::text[] AS counterparts, + EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint foreign_key + JOIN pg_catalog.pg_class referenced ON referenced.oid = foreign_key.confrelid + WHERE foreign_key.conrelid = parent.oid + AND foreign_key.contype = 'f' + AND (referenced.relkind <> 'r' OR referenced.relhassubclass) + ) AS counterpart_has_children, + EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint foreign_key + WHERE foreign_key.confrelid = parent.oid + AND foreign_key.contype = 'f' + ) AS referenced_by_foreign_key + FROM pg_catalog.pg_class parent + JOIN pg_catalog.pg_namespace parent_ns ON parent_ns.oid = parent.relnamespace + WHERE parent_ns.nspname = pg_catalog.current_schema() + AND parent.relname = $1 + "#, + ) + .bind(table) + .fetch_one(&mut *connection) + .await?; + let refusal = if row.try_get::("counterpart_has_children")? { + Some(format!( + "{table} references a partitioned or inherited table; the counterpart lock set is not proven" + )) + } else if row.try_get::("referenced_by_foreign_key")? { + Some(format!( + "{table} is referenced by a foreign key; the counterpart lock set is not proven" + )) + } else { + None + }; + Ok(MaintenanceParent { + oid: row.try_get("parent_oid")?, + schema: row.try_get("parent_schema")?, + counterparts: row.try_get("counterparts")?, + refusal, + }) +} + +/// Names the policy would create, and the catch-all it would replace, even +/// when the plan is refused. Used to report collisions ahead of other refusals. +fn wanted_names( + audit: &PartitionTableAudit, + policy: PartitionMaintenancePolicy, +) -> (Vec, Option<(String, String)>) { + let mut names: Vec<_> = audit + .months + .iter() + .filter(|month| { + (policy.create_enabled && month.kind == MonthCoverageKind::Uncovered) + || (policy.advance_enabled && month.kind == MonthCoverageKind::CoveredByCatchAll) + }) + .map(|month| partition_name(audit.table, month.start)) + .collect(); + let catch_alls: Vec<_> = audit + .children + .iter() + .filter(|child| child.kind == PartitionChildKind::CatchAll) + .collect(); + let advancing = policy.advance_enabled + && audit + .months + .iter() + .any(|month| month.kind == MonthCoverageKind::CoveredByCatchAll); + if advancing { + names.push(catch_all_name(audit.table)); + } + let replaced = match catch_alls.as_slice() { + [child] if advancing => Some((child.schema.clone(), child.name.clone())), + _ => None, + }; + (names, replaced) +} + +/// Return a target name already used by a relation other than the catch-all being replaced. +async fn colliding_relation<'e, E>( + executor: E, + names: Vec, + replaced: Option<(String, String)>, +) -> Result> +where + E: sqlx::PgExecutor<'e>, +{ + let (replaced_schema, replaced_name) = replaced.unwrap_or_default(); + Ok(sqlx::query_scalar( + r#" + SELECT relation.relname::text + FROM pg_catalog.pg_class relation + JOIN pg_catalog.pg_namespace namespace ON namespace.oid = relation.relnamespace + WHERE namespace.nspname = pg_catalog.current_schema() + AND relation.relname = ANY($1) + AND NOT (namespace.nspname = $2 AND relation.relname = $3) + ORDER BY relation.relname + LIMIT 1 + "#, + ) + .bind(names) + .bind(replaced_schema) + .bind(replaced_name) + .fetch_optional(executor) + .await?) +} + +fn collision_message(name: &str) -> String { + format!("canonical name {name} already exists without the expected attachment and bounds") +} + +/// Refuse a catch-all with a foreign key of its own, or one that references it. +/// +/// Such a key is invisible to the parent's counterpart set. Dropping the +/// catch-all would silently drop it and lock its other table without `NOWAIT` +/// while holding the parent. The referenced arm also matches the catch-all's +/// clone of a key referencing the parent; the parent refusal runs first, so +/// that arm is a backstop. +async fn catch_all_foreign_key_refusal( + connection: &mut PgConnection, + catch_all: &CatchAllReplacement, +) -> Result> { + let names: Option = sqlx::query_scalar( + r#" + SELECT pg_catalog.string_agg(foreign_key.conname::text, ', ' ORDER BY foreign_key.conname) + FROM pg_catalog.pg_constraint foreign_key + WHERE foreign_key.contype = 'f' + AND ((foreign_key.conrelid = $1::regclass AND foreign_key.conparentid = 0) + OR foreign_key.confrelid = $1::regclass) + "#, + ) + .bind(qualified_relation_name(&catch_all.schema, &catch_all.name)) + .fetch_one(&mut *connection) + .await?; + Ok(names.map(|names| { + format!( + "catch-all {} has foreign keys outside the parent's counterpart set ({names}); \ + the counterpart lock set is not proven", + catch_all.name + ) + })) +} + +/// Verify the changed catalog, audited inside the transaction, before commit. +fn verify_applied( + before: &PartitionTableAudit, + after: &PartitionTableAudit, + change: &MaintenanceChange, +) -> std::result::Result<(), String> { + let table = after.table; + if !after.structurally_safe_for_creation() { + return Err("catalog is not structurally safe after DDL".to_string()); + } + if before.serving_safe && !after.serving_safe { + return Err("current month lost its covering partition".to_string()); + } + let parity = |name: &str| { + after + .children + .iter() + .find(|child| child.name == name) + .filter(|child| child.missing_triggers.is_empty() && child.extra_triggers.is_empty()) + }; + for month in &change.create { + let name = partition_name(table, *month); + let end = next_month(*month).map_err(|error| error.to_string())?; + if !parity(&name).is_some_and(|child| { + child.kind == PartitionChildKind::CanonicalMonthly + && child.lower == Some(PartitionBound::Finite(*month)) + && child.upper == Some(PartitionBound::Finite(end)) + }) { + return Err(format!( + "{name} is missing, misbounded, or lacks trigger parity" + )); + } + } + if let Some(catch_all) = &change.catch_all { + let name = catch_all_name(table); + if !parity(&name).is_some_and(|child| { + child.kind == PartitionChildKind::CatchAll + && child.lower == Some(PartitionBound::Finite(catch_all.new_lower)) + && child.catch_all_nonempty == Some(false) + }) { + return Err(format!( + "{name} is missing, misbounded, or lacks trigger parity" + )); + } + if after + .children + .iter() + .filter(|child| child.kind == PartitionChildKind::CatchAll) + .count() + != 1 + { + return Err("expected exactly one catch-all after DDL".to_string()); + } + } + if let Some(month) = after.months.iter().find(|month| { + change.create.contains(&month.start) && month.kind != MonthCoverageKind::CoveredByMonthly + }) { + return Err(format!( + "{} is not covered by its dedicated monthly", + month.start.format("%Y-%m") + )); + } + Ok(()) +} + +/// Verify every new child carries a valid, ready clone of each parent index and +/// a validated clone of each parent primary-key, unique, and foreign-key constraint. +async fn verify_inheritance( + connection: &mut PgConnection, + parent_oid: i64, + names: &[String], +) -> Result> { + let rows = sqlx::query( + r#" + SELECT child.relname::text AS name, + ( + SELECT count(*) + FROM pg_catalog.pg_index parent_index + WHERE parent_index.indrelid = $1::bigint::oid + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_inherits index_edge + JOIN pg_catalog.pg_index child_index + ON child_index.indexrelid = index_edge.inhrelid + WHERE index_edge.inhparent = parent_index.indexrelid + AND child_index.indrelid = child.oid + AND child_index.indisvalid + AND child_index.indisready + ) + ) AS missing_indexes, + ( + SELECT count(*) + FROM pg_catalog.pg_constraint parent_constraint + WHERE parent_constraint.conrelid = $1::bigint::oid + AND parent_constraint.contype IN ('f', 'p', 'u') + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint child_constraint + WHERE child_constraint.conrelid = child.oid + AND child_constraint.conparentid = parent_constraint.oid + AND child_constraint.convalidated + ) + ) AS missing_constraints + FROM pg_catalog.pg_class child + JOIN pg_catalog.pg_namespace child_ns ON child_ns.oid = child.relnamespace + JOIN pg_catalog.pg_inherits edge + ON edge.inhrelid = child.oid AND edge.inhparent = $1::bigint::oid + WHERE child_ns.nspname = pg_catalog.current_schema() + AND child.relname = ANY($2) + "#, + ) + .bind(parent_oid) + .bind(names) + .fetch_all(&mut *connection) + .await?; + if rows.len() != names.len() { + return Ok(Err(format!( + "expected {} new children attached to the parent, found {}", + names.len(), + rows.len() + ))); + } + for row in rows { + let name: String = row.try_get("name")?; + let missing_indexes: i64 = row.try_get("missing_indexes")?; + let missing_constraints: i64 = row.try_get("missing_constraints")?; + if missing_indexes != 0 || missing_constraints != 0 { + return Ok(Err(format!( + "{name} is missing {missing_indexes} inherited indexes and \ + {missing_constraints} inherited constraints" + ))); + } + } + Ok(Ok(())) +} + +fn classify_error(error: &DbError) -> PartitionMaintenanceOutcome { + match error { + DbError::Sqlx(sqlx::Error::Database(database)) => match database.code().as_deref() { + // A deadlock victim lost a lock race; the next pass retries. + Some("55P03" | "40P01") => PartitionMaintenanceOutcome::LockTimeout, + Some("57014" | "25P03") => PartitionMaintenanceOutcome::Deadline, + _ => PartitionMaintenanceOutcome::Error, + }, + DbError::Sqlx(sqlx::Error::PoolTimedOut) => PartitionMaintenanceOutcome::Deadline, + _ => PartitionMaintenanceOutcome::Error, + } +} + +fn emit_maintenance_metrics(audit: &PartitionTableAudit, result: &TableMaintenance) { + let table = audit.table; + metrics::counter!( + "buzz_partition_maintenance_runs_total", + "table" => table, + "outcome" => result.outcome.as_str() + ) + .increment(1); + // Per-month creation accounting, continuous with the pre-advancement manager. + let changed = result + .change + .as_ref() + .map(|change| change.create.as_slice()) + .unwrap_or_default(); + // `error` means DDL or collision failure. Lock and deadline outcomes are + // transient and only `buzz_partition_maintenance_runs_total` records them. + let month_outcome = match result.outcome { + PartitionMaintenanceOutcome::Error | PartitionMaintenanceOutcome::OperatorRequired => { + Some("error") + } + PartitionMaintenanceOutcome::LockTimeout | PartitionMaintenanceOutcome::Deadline => None, + _ => Some("created"), + }; + for month in &audit.months { + let outcome = if changed.contains(&month.start) { + let Some(outcome) = month_outcome else { + continue; + }; + outcome + } else if month.kind == MonthCoverageKind::Uncovered { + continue; + } else { + "skipped_covered" + }; + metrics::counter!( + "buzz_partition_create_attempts_total", + "table" => table, + "outcome" => outcome + ) + .increment(1); + } +} + +#[cfg(test)] +#[path = "maintenance_tests.rs"] +mod tests; diff --git a/crates/buzz-db/src/store/partition/maintenance_tests.rs b/crates/buzz-db/src/store/partition/maintenance_tests.rs new file mode 100644 index 0000000000..99a1684814 --- /dev/null +++ b/crates/buzz-db/src/store/partition/maintenance_tests.rs @@ -0,0 +1,1330 @@ +use chrono::TimeZone; + +use super::*; +use crate::store::partition::{MonthCoverage, PartitionChildAudit, PARTITIONED_TABLES}; + +fn month(year: i32, month: u32) -> DateTime { + Utc.with_ymd_and_hms(year, month, 1, 0, 0, 0).unwrap() +} + +fn child( + name: &str, + kind: PartitionChildKind, + lower: PartitionBound, + upper: PartitionBound, + catch_all_nonempty: Option, +) -> PartitionChildAudit { + PartitionChildAudit { + schema: "public".to_string(), + name: name.to_string(), + relation_kind: "r".to_string(), + lower: Some(lower), + upper: Some(upper), + kind, + pending_detach: false, + catch_all_nonempty, + default_nonempty: None, + missing_triggers: Vec::new(), + extra_triggers: Vec::new(), + } +} + +/// The repaired production layout: a July/August legacy leaf, September through +/// December monthlies, and a catch-all from January 2027. +fn repaired_audit(catch_all_nonempty: bool) -> PartitionTableAudit { + let mut children = vec![ + child( + "events_p2026_07_08_legacy", + PartitionChildKind::LegacyLeaf, + PartitionBound::Finite(month(2026, 7)), + PartitionBound::Finite(month(2026, 9)), + None, + ), + child( + "events_p_future_next", + PartitionChildKind::CatchAll, + PartitionBound::Finite(month(2027, 1)), + PartitionBound::MaxValue, + Some(catch_all_nonempty), + ), + ]; + for m in 9..=12 { + children.push(child( + &format!("events_p2026_{m:02}"), + PartitionChildKind::CanonicalMonthly, + PartitionBound::Finite(month(2026, m)), + PartitionBound::Finite({ + let (year, next) = add_months(2026, m, 1).unwrap(); + month(year, next) + }), + None, + )); + } + let months = (0..=PARTITION_MANAGER_MONTHS_AHEAD as i32) + .map(|offset| { + let (year, m) = add_months(2026, 9, offset).unwrap(); + let start = month(year, m); + MonthCoverage { + start, + kind: if start < month(2027, 1) { + MonthCoverageKind::CoveredByMonthly + } else { + MonthCoverageKind::CoveredByCatchAll + }, + } + }) + .collect(); + PartitionTableAudit { + table: "events", + partition_key: Some("RANGE (created_at)".to_string()), + expected_partition_key: "RANGE (created_at)", + partition_key_valid: true, + children, + coverage_leaves: Vec::new(), + months, + serving_safe: true, + } +} + +const ADVANCE: PartitionMaintenancePolicy = PartitionMaintenancePolicy { + create_enabled: true, + advance_enabled: true, +}; + +const CREATE_ONLY: PartitionMaintenancePolicy = PartitionMaintenancePolicy { + create_enabled: true, + advance_enabled: false, +}; + +#[test] +fn repaired_layout_plans_three_monthlies_and_an_april_catch_all() { + assert_eq!( + plan_maintenance(&repaired_audit(false), ADVANCE), + MaintenancePlan::Apply(MaintenanceChange { + create: vec![month(2027, 1), month(2027, 2), month(2027, 3)], + catch_all: Some(CatchAllReplacement { + schema: "public".to_string(), + name: "events_p_future_next".to_string(), + new_lower: month(2027, 4), + }), + }) + ); +} + +#[test] +fn advancement_requires_its_own_policy_flag() { + assert_eq!( + plan_maintenance(&repaired_audit(false), CREATE_ONLY), + MaintenancePlan::Disabled + ); + assert_eq!( + plan_maintenance( + &repaired_audit(false), + PartitionMaintenancePolicy::default() + ), + MaintenancePlan::Disabled + ); +} + +#[test] +fn full_monthly_runway_is_a_noop_even_when_disabled() { + let mut audit = repaired_audit(false); + for coverage in &mut audit.months { + coverage.kind = MonthCoverageKind::CoveredByMonthly; + } + assert_eq!(plan_maintenance(&audit, ADVANCE), MaintenancePlan::Noop); + assert_eq!( + plan_maintenance(&audit, PartitionMaintenancePolicy::default()), + MaintenancePlan::Noop + ); +} + +#[test] +fn unsafe_catch_all_states_require_an_operator() { + let refused = |audit: &PartitionTableAudit, needle: &str| match plan_maintenance(audit, ADVANCE) + { + MaintenancePlan::Refuse(reason) => assert!(reason.contains(needle), "{reason}"), + plan => panic!("expected refusal containing {needle:?}, got {plan:?}"), + }; + + refused(&repaired_audit(true), "contains rows"); + + let mut unknown = repaired_audit(false); + unknown.children[1].catch_all_nonempty = None; + refused(&unknown, "occupancy is unknown"); + + let mut misaligned = repaired_audit(false); + misaligned.children[1].lower = Some(PartitionBound::Finite( + Utc.with_ymd_and_hms(2027, 1, 15, 0, 0, 0).unwrap(), + )); + refused(&misaligned, "not a UTC month start"); + + 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] +fn advancement_crosses_the_year_boundary() { + // December: the horizon runs through June and the catch-all moves to July. + let mut audit = repaired_audit(false); + audit.months = (0..=PARTITION_MANAGER_MONTHS_AHEAD as i32) + .map(|offset| { + let (year, m) = add_months(2026, 12, offset).unwrap(); + MonthCoverage { + start: month(year, m), + kind: if offset == 0 { + MonthCoverageKind::CoveredByMonthly + } else { + MonthCoverageKind::CoveredByCatchAll + }, + } + }) + .collect(); + let MaintenancePlan::Apply(change) = plan_maintenance(&audit, ADVANCE) else { + panic!("expected advancement"); + }; + assert_eq!(change.create.first(), Some(&month(2027, 1))); + assert_eq!(change.create.last(), Some(&month(2027, 6))); + assert_eq!(change.create.len(), 6); + assert_eq!(change.catch_all.map(|c| c.new_lower), Some(month(2027, 7))); +} + +#[test] +fn create_kill_switch_also_stops_advancement() { + let advance_only = PartitionMaintenancePolicy { + create_enabled: false, + advance_enabled: true, + }; + assert_eq!( + plan_maintenance(&repaired_audit(false), advance_only), + MaintenancePlan::Disabled + ); +} + +#[test] +fn advancement_is_capped_at_the_supported_horizon() { + // The horizon ends in March 2027, so a catch-all from April 2017 needs + // exactly the maximum number of monthlies and one from March 2017 needs one more. + let from = |lower| { + let mut audit = repaired_audit(false); + audit.children[1].lower = Some(PartitionBound::Finite(lower)); + plan_maintenance(&audit, ADVANCE) + }; + let MaintenancePlan::Apply(change) = from(month(2017, 4)) else { + panic!("expected advancement at the cap"); + }; + assert_eq!(change.create.len(), MAX_PARTITION_MONTHS_AHEAD as usize); + match from(month(2017, 3)) { + MaintenancePlan::Refuse(reason) => assert!(reason.contains("more than"), "{reason}"), + plan => panic!("expected refusal past the cap, got {plan:?}"), + } +} + +/// `buzz_partition_create_attempts_total` counts keyed by their `outcome` label. +fn create_attempts(result: &TableMaintenance) -> std::collections::BTreeMap { + let recorder = metrics_util::debugging::DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + metrics::with_local_recorder(&recorder, || { + emit_maintenance_metrics(&repaired_audit(false), result) + }); + snapshotter + .snapshot() + .into_vec() + .into_iter() + .filter(|(key, ..)| key.key().name() == "buzz_partition_create_attempts_total") + .map(|(key, _, _, value)| { + let outcome = key + .key() + .labels() + .find(|label| label.key() == "outcome") + .map(|label| label.value().to_string()) + .expect("outcome label"); + let metrics_util::debugging::DebugValue::Counter(count) = value else { + panic!("create attempts must be a counter"); + }; + (outcome, count) + }) + .collect() +} + +#[test] +fn only_ddl_failures_count_as_create_errors() { + let MaintenancePlan::Apply(change) = plan_maintenance(&repaired_audit(false), ADVANCE) else { + panic!("expected advancement"); + }; + let failed = |outcome| TableMaintenance::failed(outcome, Some(change.clone()), String::new()); + let covered = ("skipped_covered".to_string(), 4); + + for outcome in [ + PartitionMaintenanceOutcome::LockTimeout, + PartitionMaintenanceOutcome::Deadline, + ] { + assert_eq!( + create_attempts(&failed(outcome)), + [covered.clone()].into(), + "{outcome:?}" + ); + } + for outcome in [ + PartitionMaintenanceOutcome::Error, + PartitionMaintenanceOutcome::OperatorRequired, + ] { + assert_eq!( + create_attempts(&failed(outcome)), + [covered.clone(), ("error".to_string(), 3)].into(), + "{outcome:?}" + ); + } + let advanced = TableMaintenance { + outcome: PartitionMaintenanceOutcome::Advanced, + change: Some(change.clone()), + error: None, + }; + assert_eq!( + create_attempts(&advanced), + [covered, ("created".to_string(), 3)].into() + ); +} + +#[test] +fn outcome_labels_are_stable() { + use PartitionMaintenanceOutcome::*; + assert_eq!( + [ + Noop, + SkippedDisabled, + Created, + Advanced, + SkippedLocked, + OperatorRequired, + LockTimeout, + Deadline, + Error + ] + .map(PartitionMaintenanceOutcome::as_str), + [ + "noop", + "skipped_disabled", + "created", + "advanced", + "skipped_locked", + "operator_required", + "lock_timeout", + "deadline", + "error" + ] + ); +} + +mod postgres_tests { + use std::time::Instant; + + use sqlx::PgPool; + + use super::*; + use crate::store::partition::tests::postgres_tests::{ + catalog_snapshot, create_child, drop_schema, fixed_now, scratch_pool, + scratch_pool_with_max_connections, seed_parents, + }; + + /// The repaired incident layout for both managed parents, optionally with + /// outgoing foreign keys to a `communities` counterpart. + async fn seed_repaired_layout(pool: &PgPool, with_foreign_keys: bool) { + seed_parents(pool).await; + if with_foreign_keys { + sqlx::query("CREATE TABLE communities (id UUID PRIMARY KEY)") + .execute(pool) + .await + .expect("create communities"); + for table in PARTITIONED_TABLES { + sqlx::query(sqlx::AssertSqlSafe(format!( + "ALTER TABLE {table} ADD COLUMN community_id UUID REFERENCES communities(id)" + ))) + .execute(pool) + .await + .expect("add counterpart foreign key"); + } + } + for table in PARTITIONED_TABLES { + create_child( + pool, + table, + &format!("{table}_p_past"), + "MINVALUE", + "'2026-07-01'", + ) + .await; + create_child( + pool, + table, + &format!("{table}_p2026_07_08_legacy"), + "'2026-07-01'", + "'2026-09-01'", + ) + .await; + for m in 9..=12 { + let upper = if m == 12 { + "'2027-01-01'".to_string() + } else { + format!("'2026-{:02}-01'", m + 1) + }; + create_child( + pool, + table, + &format!("{table}_p2026_{m:02}"), + &format!("'2026-{m:02}-01'"), + &upper, + ) + .await; + } + create_child( + pool, + table, + &format!("{table}_p_future_next"), + "'2027-01-01'", + "MAXVALUE", + ) + .await; + } + } + + async fn maintain(pool: &PgPool) -> Result { + maintain_partitions_at(pool, PARTITION_MANAGER_MONTHS_AHEAD, ADVANCE, fixed_now()).await + } + + fn assert_advanced(audit: &PartitionAudit) { + for table in &audit.tables { + assert!( + table + .months + .iter() + .all(|month| month.kind == MonthCoverageKind::CoveredByMonthly), + "{}: {:?}", + table.table, + table.months + ); + let catch_all = table + .children + .iter() + .find(|child| child.kind == PartitionChildKind::CatchAll) + .expect("catch-all"); + assert_eq!(catch_all.name, format!("{}_p_future", table.table)); + assert_eq!( + catch_all.lower, + Some(PartitionBound::Finite(month(2027, 4))) + ); + assert_eq!(catch_all.catch_all_nonempty, Some(false)); + assert!(!table + .children + .iter() + .any(|child| child.name.ends_with("_p_future_next"))); + assert_eq!(table.missing_trigger_count(), 0); + assert_eq!(table.extra_trigger_count(), 0); + } + } + + /// Open a transaction on a separate session that holds `sql`'s lock until dropped. + async fn hold_lock(pool: &PgPool, sql: &str) -> sqlx::Transaction<'static, sqlx::Postgres> { + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(sqlx::AssertSqlSafe(sql.to_string())) + .execute(&mut *holder) + .await + .expect("take conflicting lock"); + holder + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn empty_catch_all_advances_to_canonical_layout_and_is_idempotent() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + + let audit = maintain(&pool).await.expect("advance"); + assert_advanced(&audit); + let fresh = audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()) + .await + .expect("fresh audit"); + assert_advanced(&fresh); + + // Rows route to the new leaves, including past the new horizon. + for (created_at, leaf) in [ + ("2027-02-10", "events_p2027_02"), + ("2031-01-01", "events_p_future"), + ] { + let routed: String = sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "WITH inserted AS (INSERT INTO events (created_at, alternate_at) \ + VALUES ('{created_at}', now()) RETURNING tableoid) \ + SELECT tableoid::regclass::text FROM inserted" + ))) + .fetch_one(&pool) + .await + .expect("route row"); + assert_eq!(routed, leaf); + } + sqlx::query("DELETE FROM events WHERE created_at >= '2027-01-01'") + .execute(&pool) + .await + .expect("clear routed rows"); + + let before = catalog_snapshot(&pool).await; + maintain(&pool).await.expect("idempotent rerun"); + assert_eq!(catalog_snapshot(&pool).await, before); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn canonical_catch_all_name_is_replaced_in_one_transaction() { + let (pool, admin, schema) = scratch_pool().await; + seed_parents(&pool).await; + for table in PARTITIONED_TABLES { + create_child( + &pool, + table, + &format!("{table}_p_past"), + "MINVALUE", + "'2026-09-01'", + ) + .await; + create_child( + &pool, + table, + &format!("{table}_p_future"), + "'2026-09-01'", + "MAXVALUE", + ) + .await; + } + let audit = maintain(&pool).await.expect("advance fresh layout"); + assert_advanced(&audit); + assert!(audit.tables.iter().all(|table| table + .children + .iter() + .any(|child| child.name == format!("{}_p2026_09", table.table)))); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn populated_catch_all_is_refused_before_any_lock() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query("INSERT INTO events (created_at, alternate_at) VALUES ('2027-02-01', now())") + .execute(&pool) + .await + .expect("populate catch-all"); + let before = catalog_snapshot(&pool).await; + // A lock-taking attempt would wait out the 2s parent lock timeout. + let holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + + let started = Instant::now(); + let result = maintain(&pool).await; + let elapsed = started.elapsed(); + drop(holder); + + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events operator_required") + && message.contains("events_p_future_next contains rows") + && !message.contains("delivery_log")), + "{result:?}" + ); + assert!(elapsed < std::time::Duration::from_secs(1), "{elapsed:?}"); + 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 counterpart_contention_fails_fast_and_rolls_back() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + let before = catalog_snapshot(&pool).await; + let holder = hold_lock( + &pool, + "INSERT INTO communities(id) VALUES (gen_random_uuid())", + ) + .await; + + let started = Instant::now(); + let result = maintain(&pool).await; + let elapsed = started.elapsed(); + drop(holder); + + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events lock_timeout") + && message.contains("delivery_log lock_timeout")), + "{result:?}" + ); + // NOWAIT: no 2s wait per table while a counterpart writer is active. + assert!(elapsed < std::time::Duration::from_secs(1), "{elapsed:?}"); + assert_eq!(catalog_snapshot(&pool).await, before); + + assert_advanced(&maintain(&pool).await.expect("advance after release")); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn counterpart_reader_does_not_block_advancement() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + let holder = hold_lock(&pool, "SELECT count(*) FROM communities").await; + + let result = maintain(&pool).await; + drop(holder); + + assert_advanced(&result.expect("advance beside a counterpart reader")); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn writer_checking_a_foreign_key_behind_maintenance_does_not_deadlock() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + // An in-flight insert: the parent first, then the foreign-key check. + let mut writer = hold_lock(&pool, "LOCK TABLE ONLY events IN ROW EXCLUSIVE MODE").await; + let writer_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *writer) + .await + .expect("writer backend pid"); + + let maintenance = tokio::spawn({ + let pool = pool.clone(); + async move { maintain(&pool).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(writer_pid) + .fetch_one(&admin) + .await + .expect("check queued maintenance"); + if queued { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + assert!(queued, "maintenance never queued behind the writer"); + + // Maintenance waits on the parent holding no counterpart lock, so the + // writer's key check, and a write to the counterpart itself, proceed + // instead of entering a lock cycle in any counterpart lock mode. + let started = Instant::now(); + sqlx::query("SELECT id FROM communities FOR KEY SHARE") + .fetch_all(&mut *writer) + .await + .expect("foreign-key check behind queued maintenance"); + sqlx::query("INSERT INTO communities(id) VALUES (gen_random_uuid())") + .execute(&mut *writer) + .await + .expect("counterpart write behind queued maintenance"); + assert!( + started.elapsed() < std::time::Duration::from_millis(500), + "{:?}", + started.elapsed() + ); + writer.commit().await.expect("commit writer"); + + assert_advanced( + &maintenance + .await + .expect("maintenance task") + .expect("advance after the writer commits"), + ); + drop_schema(&admin, &schema).await; + } + + async fn wait_until_blocked_by(admin: &PgPool, pid: i32, waiters: i64) { + for _ in 0..100 { + let blocked: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_stat_activity WHERE $1 = ANY(pg_blocking_pids(pid))", + ) + .bind(pid) + .fetch_one(admin) + .await + .expect("count blocked backends"); + if blocked >= waiters { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + panic!("{waiters} backends never queued behind {pid}"); + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn foreign_key_committed_while_waiting_on_the_parent_is_locked_nowait() { + let (pool, admin, schema) = scratch_pool_with_max_connections(6).await; + seed_repaired_layout(&pool, true).await; + sqlx::query("CREATE TABLE reviewers (id UUID PRIMARY KEY)") + .execute(&pool) + .await + .expect("create new counterpart"); + // An in-flight migration adds a foreign key that maintenance's + // pre-lock read cannot see yet. + let mut migration = hold_lock( + &pool, + "ALTER TABLE events ADD FOREIGN KEY (community_id) REFERENCES reviewers(id)", + ) + .await; + let migration_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *migration) + .await + .expect("migration backend pid"); + // A writer on the new counterpart queues behind the migration and is + // granted its lock when the migration commits. + let (release_writer, writer_released) = tokio::sync::oneshot::channel::<()>(); + let writer = tokio::spawn({ + let pool = pool.clone(); + async move { + let mut writer = pool.begin().await.expect("begin counterpart writer"); + sqlx::query("INSERT INTO reviewers(id) VALUES (gen_random_uuid())") + .execute(&mut *writer) + .await + .expect("counterpart writer insert"); + let _ = writer_released.await; + writer + .rollback() + .await + .expect("roll back counterpart writer"); + } + }); + wait_until_blocked_by(&admin, migration_pid, 1).await; + let maintenance = tokio::spawn({ + let pool = pool.clone(); + async move { maintain(&pool).await } + }); + wait_until_blocked_by(&admin, migration_pid, 2).await; + + migration.commit().await.expect("commit migration"); + let result = maintenance.await.expect("maintenance task"); + + // The parent-locked read finds the new counterpart and its NOWAIT lock + // fails naming it, instead of attaching against it under the parent + // lock until a lock timeout that names no relation. + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events lock_timeout") + && message.contains("could not obtain lock on relation") + && message.contains(r#".reviewers""#) + && !message.contains("delivery_log")), + "{result:?}" + ); + + release_writer.send(()).expect("release counterpart writer"); + writer.await.expect("counterpart writer task"); + assert_advanced(&maintain(&pool).await.expect("advance after release")); + drop_schema(&admin, &schema).await; + } + + async fn backend_pid(connection: &mut PgConnection) -> i32 { + sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(connection) + .await + .expect("backend pid") + } + + fn refused_events_only(result: &Result, reason: &str) -> bool { + matches!(result, Err(DbError::InvalidData(message)) + if message.contains("events operator_required") + && message.contains(reason) + && !message.contains("delivery_log")) + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn foreign_key_referencing_the_parent_committed_while_waiting_is_refused() { + let (pool, admin, schema) = scratch_pool_with_max_connections(6).await; + seed_repaired_layout(&pool, true).await; + sqlx::query( + "CREATE TABLE event_refs \ + (created_at TIMESTAMPTZ, alternate_at TIMESTAMPTZ, id BIGINT)", + ) + .execute(&pool) + .await + .expect("create referencing table"); + // An in-flight migration makes events a referenced table, which the + // pre-lock read cannot see yet. + let mut migration = hold_lock( + &pool, + "ALTER TABLE event_refs ADD FOREIGN KEY (created_at, alternate_at, id) \ + REFERENCES events (created_at, alternate_at, id)", + ) + .await; + let migration_pid = backend_pid(&mut migration).await; + let maintenance = tokio::spawn({ + let pool = pool.clone(); + async move { maintain(&pool).await } + }); + wait_until_blocked_by(&admin, migration_pid, 1).await; + migration.commit().await.expect("commit migration"); + + let result = maintenance.await.expect("maintenance task"); + assert!( + refused_events_only(&result, "events is referenced by a foreign key"), + "{result:?}" + ); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn catch_all_foreign_keys_outside_the_parent_are_refused() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + sqlx::query("CREATE TABLE reviewers (id UUID PRIMARY KEY)") + .execute(&pool) + .await + .expect("create counterpart"); + let keys = |pool: &PgPool| { + let pool = pool.clone(); + async move { + sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_constraint WHERE contype = 'f' \ + AND (conrelid = 'events_p_future_next'::regclass AND conparentid = 0 \ + OR confrelid = 'events_p_future_next'::regclass)", + ) + .fetch_one(&pool) + .await + .expect("count catch-all keys") + } + }; + // Refused before any lock: with the parent held, a run that reached + // the parent lock would report lock_timeout instead. + let holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + for (add, remove) in [ + ( + "ALTER TABLE ONLY events_p_future_next ADD CONSTRAINT catch_all_reviewer \ + FOREIGN KEY (community_id) REFERENCES reviewers(id)", + "ALTER TABLE events_p_future_next DROP CONSTRAINT catch_all_reviewer", + ), + ( + "CREATE TABLE catch_all_refs (created_at TIMESTAMPTZ, alternate_at TIMESTAMPTZ, \ + id BIGINT, CONSTRAINT catch_all_reviewer FOREIGN KEY (created_at, alternate_at, id) \ + REFERENCES events_p_future_next (created_at, alternate_at, id))", + "DROP TABLE catch_all_refs", + ), + ] { + sqlx::query(sqlx::AssertSqlSafe(add.to_string())) + .execute(&pool) + .await + .expect("add catch-all key"); + let result = maintain(&pool).await; + assert!( + refused_events_only(&result, "catch_all_reviewer"), + "{add}: {result:?}" + ); + // Refused, not dropped along with the catch-all. + assert_eq!(keys(&pool).await, 1, "{add}"); + sqlx::query(sqlx::AssertSqlSafe(remove.to_string())) + .execute(&pool) + .await + .expect("remove catch-all key"); + } + drop(holder); + assert_advanced(&maintain(&pool).await.expect("advance without the keys")); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn catch_all_foreign_key_committed_while_waiting_on_the_parent_is_refused() { + let (pool, admin, schema) = scratch_pool_with_max_connections(6).await; + seed_repaired_layout(&pool, true).await; + sqlx::query("CREATE TABLE reviewers (id UUID PRIMARY KEY)") + .execute(&pool) + .await + .expect("create counterpart"); + // The catch-all's own key conflicts with no parent lock, so a reader + // holds the parent to keep maintenance waiting while it commits. + let mut reader = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + let reader_pid = backend_pid(&mut reader).await; + let migration = hold_lock( + &pool, + "ALTER TABLE ONLY events_p_future_next ADD CONSTRAINT catch_all_reviewer \ + FOREIGN KEY (community_id) REFERENCES reviewers(id)", + ) + .await; + let maintenance = tokio::spawn({ + let pool = pool.clone(); + async move { maintain(&pool).await } + }); + wait_until_blocked_by(&admin, reader_pid, 1).await; + migration.commit().await.expect("commit migration"); + reader.commit().await.expect("release parent"); + + let result = maintenance.await.expect("maintenance task"); + assert!( + refused_events_only(&result, "catch_all_reviewer"), + "{result:?}" + ); + let kept: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_constraint WHERE conname = 'catch_all_reviewer' \ + AND connamespace = current_schema()::regnamespace", + ) + .fetch_one(&pool) + .await + .expect("count catch-all key"); + assert_eq!(kept, 1); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn deadlock_victim_is_classified_as_a_lock_timeout() { + let (pool, admin, schema) = scratch_pool().await; + seed_parents(&pool).await; + let mut first = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS EXCLUSIVE MODE").await; + let mut second = hold_lock( + &pool, + "LOCK TABLE ONLY delivery_log IN ACCESS EXCLUSIVE MODE", + ) + .await; + let first_waits = sqlx::query("LOCK TABLE ONLY delivery_log IN ACCESS EXCLUSIVE MODE") + .execute(&mut *first); + let second_waits = async { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + sqlx::query("LOCK TABLE ONLY events IN ACCESS EXCLUSIVE MODE") + .execute(&mut *second) + .await + }; + let (first_result, second_result) = tokio::join!(first_waits, second_waits); + let victim = first_result + .err() + .or(second_result.err()) + .expect("one session is the deadlock victim"); + assert_eq!( + victim + .as_database_error() + .and_then(|error| error.code()) + .as_deref(), + Some("40P01"), + "{victim:?}" + ); + assert_eq!( + classify_error(&DbError::from(victim)), + PartitionMaintenanceOutcome::LockTimeout + ); + drop((first, second)); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn parent_contention_times_out_without_blocking_the_other_table() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + let holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + + let result = maintain(&pool).await; + drop(holder); + + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events lock_timeout") && !message.contains("delivery_log")), + "{result:?}" + ); + let audit = audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()) + .await + .expect("audit"); + let events = audit.tables.iter().find(|t| t.table == "events").unwrap(); + assert!(events + .children + .iter() + .any(|child| child.name == "events_p_future_next")); + let delivery_log = audit + .tables + .iter() + .filter(|t| t.table == "delivery_log") + .cloned() + .collect::>(); + assert_advanced(&PartitionAudit { + audited_at: audit.audited_at, + tables: delivery_log, + }); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn held_maintenance_lock_skips_without_ddl() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + let before = catalog_snapshot(&pool).await; + let holder = hold_lock( + &pool, + "SELECT pg_advisory_xact_lock(hashtextextended(\ + 'buzz.partition_maintenance:' || current_schema()::text, 0))", + ) + .await; + + maintain(&pool) + .await + .expect("skipped_locked is not a failure"); + assert_eq!(catalog_snapshot(&pool).await, before); + drop(holder); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn concurrent_runners_converge_on_one_layout() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + + let (first, second) = tokio::join!(maintain(&pool), maintain(&pool)); + // 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; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn client_deadline_rolls_back_and_discards_the_session() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + let before = catalog_snapshot(&pool).await; + // The parent lock wait (2s) outlasts the 300ms client deadline. + 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, + PARTITION_MANAGER_MONTHS_AHEAD, + ADVANCE, + fixed_now(), + std::time::Duration::from_millis(300), + ) + .await; + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events deadline")), + "{result:?}" + ); + + // 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 { + 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!( + 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 + .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 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, + plan_maintenance(table_audit, ADVANCE), + 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); + + // The losing runner reports the layout the winner committed. + let reported = maintain_audited( + &pool, + stale, + PARTITION_MANAGER_MONTHS_AHEAD, + ADVANCE, + fixed_now(), + std::time::Duration::from_secs(10), + ) + .await + .expect("losing runner"); + assert_advanced(&reported); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn replan_that_finds_only_disabled_work_is_skipped_disabled() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query("DROP TABLE events_p_future_next") + .execute(&pool) + .await + .expect("drop catch-all"); + // Creation is planned for the uncovered months, but a catch-all appears + // before the lock and advancing it is disabled. + let stale = audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()) + .await + .expect("stale audit"); + let events = stale.tables.iter().find(|t| t.table == "events").unwrap(); + create_child( + &pool, + "events", + "events_p_future_next", + "'2027-01-01'", + "MAXVALUE", + ) + .await; + let before = catalog_snapshot(&pool).await; + + let result = maintain_table( + &pool, + events, + plan_maintenance(events, CREATE_ONLY), + PARTITION_MANAGER_MONTHS_AHEAD as i32, + CREATE_ONLY, + fixed_now(), + std::time::Duration::from_secs(10), + ) + .await; + assert_eq!( + (result.outcome, result.error.as_deref()), + (PartitionMaintenanceOutcome::SkippedDisabled, None) + ); + assert_eq!(catalog_snapshot(&pool).await, before); + 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() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query("CREATE TABLE events_p_future (unrelated BOOLEAN)") + .execute(&pool) + .await + .expect("create colliding relation"); + let before = catalog_snapshot(&pool).await; + + let result = maintain(&pool).await; + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("canonical name events_p_future already exists")), + "{result:?}" + ); + let after = catalog_snapshot(&pool).await; + assert!(after + .iter() + .any(|(name, ..)| name == "events_p_future_next")); + assert_eq!( + after + .iter() + .filter(|(name, ..)| name.starts_with("events_")) + .count(), + before + .iter() + .filter(|(name, ..)| name.starts_with("events_")) + .count() + ); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn incoming_foreign_key_is_refused_under_lock() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query( + "CREATE TABLE event_refs (created_at TIMESTAMPTZ, alternate_at TIMESTAMPTZ, id BIGINT, \ + FOREIGN KEY (created_at, alternate_at, id) REFERENCES events)", + ) + .execute(&pool) + .await + .expect("create referencing table"); + let before = catalog_snapshot(&pool).await; + + let result = maintain(&pool).await; + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events is referenced by a foreign key")), + "{result:?}" + ); + let after = catalog_snapshot(&pool).await; + assert!(after + .iter() + .any(|(name, ..)| name == "events_p_future_next")); + assert_ne!(after, before, "delivery_log still advances"); + drop_schema(&admin, &schema).await; + } +} diff --git a/crates/buzz-relay/src/api/artifact_postgres_tests.rs b/crates/buzz-relay/src/api/artifact_postgres_tests.rs index bb0a071e8f..e60010cc8a 100644 --- a/crates/buzz-relay/src/api/artifact_postgres_tests.rs +++ b/crates/buzz-relay/src/api/artifact_postgres_tests.rs @@ -25,7 +25,17 @@ impl Fixture { .unwrap(); let host = format!("artifact-review-{}.local", Uuid::new_v4()); state.db.ensure_configured_community(&host).await.unwrap(); - state.db.ensure_future_partitions(1, true).await.unwrap(); + state + .db + .maintain_partitions( + 1, + buzz_db::partition::PartitionMaintenancePolicy { + create_enabled: true, + advance_enabled: false, + }, + ) + .await + .unwrap(); let tenant = crate::tenant::bind_community(&state.db, &host) .await .unwrap(); diff --git a/crates/buzz-relay/src/config.rs b/crates/buzz-relay/src/config.rs index 589dc25719..ca895b989e 100644 --- a/crates/buzz-relay/src/config.rs +++ b/crates/buzz-relay/src/config.rs @@ -206,8 +206,13 @@ pub struct Config { pub metrics_port: u16, /// Interval between read-only partition catalog audits. pub partition_audit_interval: Duration, - /// Whether the partition manager may create uncovered monthly partitions. + /// Whether the partition manager may issue any partition DDL. Off also + /// disables catch-all advancement. pub partition_manager_create_enabled: bool, + /// Whether the partition manager may replace an empty right-edge catch-all + /// with dedicated monthlies. Also enables periodic (not only startup) DDL. + /// Requires `partition_manager_create_enabled`. + pub partition_manager_advance_enabled: bool, /// When true, NIP-42 pubkey-only authentication (no API token) is /// restricted to pubkeys in the `pubkey_allowlist` table. Users with valid @@ -938,6 +943,9 @@ impl Config { ); let partition_manager_create_enabled = parse_bool("BUZZ_PARTITION_MANAGER_CREATE_ENABLED", true)?; + // Off by default: each environment opts in after its canary. + let partition_manager_advance_enabled = + parse_bool("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED", false)?; let s3_addressing_style = match std::env::var("BUZZ_S3_ADDRESSING_STYLE") { Ok(value) => value.parse().map_err(ConfigError::InvalidValue)?, @@ -1392,6 +1400,7 @@ impl Config { metrics_port, partition_audit_interval, partition_manager_create_enabled, + partition_manager_advance_enabled, pubkey_allowlist_enabled, require_relay_membership, huddle_audio_available, @@ -1565,6 +1574,7 @@ mod tests { assert_eq!(config.max_frame_bytes, DEFAULT_MAX_FRAME_BYTES); assert_eq!(config.partition_audit_interval, Duration::from_secs(900)); assert!(config.partition_manager_create_enabled); + assert!(!config.partition_manager_advance_enabled); assert!(config.slow_client_grace_limit > 0); assert!( !config.pubkey_allowlist_enabled, @@ -2069,6 +2079,33 @@ mod tests { )); } + #[test] + fn partition_manager_advance_is_opt_in_and_parses_strictly() { + let _guards = env_guards(); + let previous = std::env::var_os("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED"); + + std::env::remove_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED"); + let default = Config::from_env().expect("default config"); + std::env::set_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED", " TRUE "); + let enabled = Config::from_env().expect("enabled config"); + std::env::set_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED", "enable-maybe"); + let invalid = Config::from_env(); + + if let Some(value) = previous { + std::env::set_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED", value); + } else { + std::env::remove_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED"); + } + + assert!(!default.partition_manager_advance_enabled); + assert!(enabled.partition_manager_advance_enabled); + assert!(matches!( + invalid, + Err(ConfigError::InvalidValue(ref message)) + if message.contains("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED") + )); + } + #[test] fn s3_addressing_style_env_accepts_virtual_and_rejects_invalid_values() { let _guards = env_guards(); diff --git a/crates/buzz-relay/src/main.rs b/crates/buzz-relay/src/main.rs index 9021d6644f..3df5c428ee 100644 --- a/crates/buzz-relay/src/main.rs +++ b/crates/buzz-relay/src/main.rs @@ -12,7 +12,7 @@ use uuid::Uuid; use buzz_audit::AuditService; use buzz_auth::AuthService; use buzz_core::CommunityId; -use buzz_db::{Db, DbConfig}; +use buzz_db::{partition::PARTITION_MANAGER_MONTHS_AHEAD, Db, DbConfig}; use buzz_pubsub::PubSubManager; use buzz_search::SearchService; @@ -480,16 +480,20 @@ async fn run_relay_main(boot: BootTracker) -> anyhow::Result<()> { info!("Skipping database migrations because BUZZ_AUTO_MIGRATE is not enabled"); } + let partition_policy = buzz_db::partition::PartitionMaintenancePolicy { + create_enabled: config.partition_manager_create_enabled, + advance_enabled: config.partition_manager_advance_enabled, + }; let mut step = StepTimer::start(StartupStep::PartitionEnsure); let startup_partition_audit = match db - .ensure_future_partitions(3, config.partition_manager_create_enabled) + .maintain_partitions(PARTITION_MANAGER_MONTHS_AHEAD, partition_policy) .await { Ok(audit) => Some(audit), Err(error) => { step.degrade(); error!(%error, "Failed to ensure partitions"); - match db.audit_partitions(3).await { + match db.audit_partitions(PARTITION_MANAGER_MONTHS_AHEAD).await { Ok(audit) => Some(audit), Err(error) => { error!(%error, "Initial partition catalog audit failed"); @@ -812,8 +816,9 @@ async fn run_relay_main(boot: BootTracker) -> anyhow::Result<()> { state.record_partition_audit(audit); } - // The periodic path is deliberately read-only. Partition creation only - // occurs during the bounded startup pass. + // The periodic path is read-only unless catch-all advancement is enabled. + // Maintenance takes no lock when the runway is already complete, so a + // periodic attempt usually costs one catalog audit. { let partition_state = Arc::clone(&state); let audit_interval = state.config.partition_audit_interval; @@ -822,7 +827,21 @@ async fn run_relay_main(boot: BootTracker) -> anyhow::Result<()> { PartitionAuditSchedule::new(audit_interval, has_startup_partition_audit); loop { schedule.wait().await; - let result = partition_state.db.audit_partitions(3).await; + let db = &partition_state.db; + let result = if partition_policy.advance_enabled { + match db + .maintain_partitions(PARTITION_MANAGER_MONTHS_AHEAD, partition_policy) + .await + { + Ok(audit) => Ok(audit), + Err(error) => { + warn!(%error, "Periodic partition maintenance failed"); + db.audit_partitions(PARTITION_MANAGER_MONTHS_AHEAD).await + } + } + } else { + db.audit_partitions(PARTITION_MANAGER_MONTHS_AHEAD).await + }; schedule.record_attempt(result.is_ok()); match result { Ok(audit) => partition_state.record_partition_audit(audit),