From a4499b2fdf640734f7410790fc586b77cbc71afb Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 08:54:05 +1300 Subject: [PATCH 01/18] fix(worker): keep planner statistics fresh on tables autovacuum never analyses Autovacuum does not analyse partitioned tables, so the thirteen ring parents have no inherited statistics, and the two head tables (queue_claim_heads, queue_enqueue_heads) never reach the autoanalyze threshold, so the planner prices lane joins against them at its defaults. The maintenance leader now runs ANALYZE on them when it is elected and every five minutes (ClientBuilder::statistics_refresh_interval, Python statistics_refresh_interval_ms; zero disables). One relation per transaction, with a bounded lock_timeout and statement_timeout, partitioned parents with SKIP_LOCKED, in a task beside the leader loop so a slow ANALYZE never delays rotate, prune or promote ticks. Migration v053 adds awa.analyze_queue_storage_relation(), a SECURITY DEFINER function so a runtime role that does not own the tables can refresh statistics without a MAINTAIN grant. A missing function or an owner that may not analyse the tables is logged once and recorded in awa.maintenance.analyze.attempts without affecting job processing. --- CHANGELOG.md | 34 ++ awa-metrics/src/lib.rs | 49 +- .../v053_analyze_queue_storage_statistics.sql | 109 ++++ awa-model/src/migrations.rs | 8 +- awa-python/python/awa/_awa.pyi | 1 + awa-python/python/awa/client.py | 7 + awa-python/src/client.rs | 6 +- awa-python/tests/test_start_config.py | 14 + awa-worker/src/client.rs | 21 + awa-worker/src/lib.rs | 1 + awa-worker/src/maintenance.rs | 34 ++ awa-worker/src/statistics.rs | 347 +++++++++++ awa/tests/statistics_refresh_test.rs | 561 ++++++++++++++++++ docs/architecture.md | 2 + docs/configuration.md | 11 + docs/grafana/README.md | 2 + docs/security/database-roles.md | 6 + docs/upgrade-0.6-to-0.7.md | 16 + skills/awa-operations/SKILL.md | 13 +- 19 files changed, 1238 insertions(+), 4 deletions(-) create mode 100644 awa-model/migrations/v053_analyze_queue_storage_statistics.sql create mode 100644 awa-worker/src/statistics.rs create mode 100644 awa/tests/statistics_refresh_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c50c51a9..0d1f27174 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,40 @@ Notable changes between releases. Detailed migration notes for storage transitio ## [Unreleased] +### Added + +- **The maintenance leader keeps planner statistics fresh on the queue-storage + tables autovacuum never analyses.** Autovacuum does not analyse partitioned + tables, so the thirteen ring parents (`ready_entries`, `done_entries`, + `leases`, ...) had no inherited statistics, and the two small head tables + (`queue_claim_heads`, `queue_enqueue_heads`) never reach the autoanalyze + threshold, so the planner priced lane joins against them at its defaults + while they drove joins over hundreds of thousands of retained rows. The + leader now runs `ANALYZE` on them when it is elected and then every five + minutes. `ClientBuilder::statistics_refresh_interval` (Python + `statistics_refresh_interval_ms`) changes the cadence; zero disables it. One + pass takes under a second on a small ring and a few seconds on a large one. + It analyses one relation per transaction with a 2 s `lock_timeout`, skips + children that prune holds, and never makes prune wait: a prune attempt that + lands while ANALYZE samples the oldest slot's tables returns `blocked` and + retries on its normal backoff. Results are in + `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration` + (scope `heads` or `ring_parent`; outcome `ok`, `skipped_privilege`, + `skipped_lock`, `unavailable` or `failed`). +- **Migration v053 adds `awa.analyze_queue_storage_relation(schema, relation)`, + so a runtime role that does not own the tables can refresh statistics.** + `ANALYZE` skips a table its caller may not analyse, and the documented + runtime grants do not include `MAINTAIN`, so the refresh runs through this + `SECURITY DEFINER` function owned by the migration role. It analyses only a + head table or partitioned table of a schema that + `awa.queue_storage_schemas()` reports as trusted, and `EXECUTE` is granted + to `PUBLIC` like the other runtime helpers. Run `awa migrate` as the schema + owner (see [database roles](docs/security/database-roles.md)); no grant + changes are needed. A binary running ahead of the migration, or against a + function whose owner may not analyse the tables, logs one warning naming what + is missing and skips the refresh without affecting job processing. Earlier + binaries never call the function. + ### Fixed - **Batch enqueues take their lane locks in one order, so they no longer diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 1a1c9097b..14e3cb540 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -77,6 +77,8 @@ pub mod names { pub const MAINTENANCE_PRUNE_ATTEMPTS: &str = "awa.maintenance.prune.attempts"; pub const MAINTENANCE_PRUNE_SKIPPED_ROWS: &str = "awa.maintenance.prune.skipped_rows"; pub const MAINTENANCE_PRUNE_DURATION: &str = "awa.maintenance.prune.duration"; + pub const MAINTENANCE_ANALYZE_ATTEMPTS: &str = "awa.maintenance.analyze.attempts"; + pub const MAINTENANCE_ANALYZE_DURATION: &str = "awa.maintenance.analyze.duration"; pub const STORAGE_TRANSITION_READY: &str = "awa.storage.transition_ready"; pub const STORAGE_CANONICAL_LIVE_BACKLOG: &str = "awa.storage.canonical_live_backlog"; pub const STORAGE_LIVE_RUNTIME_CAPABILITY: &str = "awa.storage.live_runtime_capability"; @@ -277,6 +279,13 @@ pub struct AwaMetrics { /// rather than one per tick. Fleets alert on this rather than /// scraping the matching `tracing::warn!` line. Issue #242. pub maintenance_branch_overrun_total: Counter, + /// Planner-statistics `ANALYZE` calls by scope (`heads`, `ring_parent`) + /// and outcome (`ok`, `skipped_privilege`, `skipped_lock`, + /// `unavailable`, `failed`). One increment per relation per refresh pass. + pub maintenance_analyze_attempts: Counter, + /// Wall-clock duration of one planner-statistics `ANALYZE` call, by scope + /// and outcome. + pub maintenance_analyze_duration_seconds: Histogram, /// Current ring `current_slot` per ring, sampled from each rotate call. /// The slot number itself isn't meaningful but the rate of advance is — /// dashboards plot `rate(slot_changes)` to see whether rotation is @@ -620,6 +629,21 @@ impl AwaMetrics { ) .with_unit("{episode}") .build(), + maintenance_analyze_attempts: meter + .u64_counter(names::MAINTENANCE_ANALYZE_ATTEMPTS) + .with_description( + "Planner-statistics ANALYZE calls by scope (heads/ring_parent) and outcome", + ) + .with_unit("{attempt}") + .build(), + maintenance_analyze_duration_seconds: meter + .f64_histogram(names::MAINTENANCE_ANALYZE_DURATION) + .with_description( + "Duration of one planner-statistics ANALYZE call by scope and outcome", + ) + .with_unit("s") + .with_boundaries(WAIT_DURATION_BUCKETS_SECONDS.to_vec()) + .build(), ring_current_slot: meter .i64_gauge(names::RING_CURRENT_SLOT) .with_description("Current slot index per ring (queue/lease/claim)") @@ -1280,6 +1304,24 @@ impl AwaMetrics { self.maintenance_branch_overrun_total.add(1, &attrs); } + /// Record one planner-statistics `ANALYZE` call. `scope` is `"heads"` or + /// `"ring_parent"`; `outcome` is `"ok"`, `"skipped_privilege"`, + /// `"skipped_lock"`, `"unavailable"` or `"failed"`. + pub fn record_statistics_analyze( + &self, + scope: &'static str, + outcome: &'static str, + duration: Duration, + ) { + let attrs = [ + opentelemetry::KeyValue::new("awa.analyze.scope", scope), + opentelemetry::KeyValue::new("awa.analyze.outcome", outcome), + ]; + self.maintenance_analyze_attempts.add(1, &attrs); + self.maintenance_analyze_duration_seconds + .record(duration.as_secs_f64(), &attrs); + } + /// Record a ring prune outcome. pub fn record_prune_outcome(&self, ring: &'static str, outcome: &awa_model::PruneOutcome) { let (label, reason, count) = match outcome { @@ -1457,7 +1499,7 @@ mod tests { .build(); let metrics = AwaMetrics::new(&provider.meter("awa-metrics-test")); - let seconds: [(&str, &Histogram, &[f64]); 9] = [ + let seconds: [(&str, &Histogram, &[f64]); 10] = [ ( names::ENQUEUE_DURATION, &metrics.enqueue_duration_seconds, @@ -1478,6 +1520,11 @@ mod tests { &metrics.maintenance_branch_duration_seconds, &WAIT_DURATION_BUCKETS_SECONDS, ), + ( + names::MAINTENANCE_ANALYZE_DURATION, + &metrics.maintenance_analyze_duration_seconds, + &WAIT_DURATION_BUCKETS_SECONDS, + ), ( names::DISPATCH_WAKE_TO_CLAIM_DURATION, &metrics.dispatch_wake_to_claim_seconds, diff --git a/awa-model/migrations/v053_analyze_queue_storage_statistics.sql b/awa-model/migrations/v053_analyze_queue_storage_statistics.sql new file mode 100644 index 000000000..a82a52f8a --- /dev/null +++ b/awa-model/migrations/v053_analyze_queue_storage_statistics.sql @@ -0,0 +1,109 @@ +-- v053: Planner statistics for the queue-storage tables autovacuum never analyses +-- +-- Autovacuum does not analyse partitioned tables, so the ring parents +-- (ready_entries, done_entries, leases, ...) have no inherited statistics. It +-- also never reaches the autoanalyze threshold on the two small head tables +-- (queue_claim_heads, queue_enqueue_heads), so the planner prices lane joins +-- against them at its defaults (320 rows, n_distinct 200) while they drive +-- joins over hundreds of thousands of retained rows. +-- +-- The maintenance leader refreshes both on an interval. Runtimes usually +-- connect as a role that does not own the tables, and ANALYZE skips with a +-- warning a table its caller may not analyse, so the refresh runs through +-- this SECURITY DEFINER function owned by the migration role, which owns the +-- tables. +-- +-- The function analyses one relation per call and the caller runs each call as +-- its own transaction, so the locks ANALYZE takes on a partitioned table's +-- children are held for that table rather than for the whole refresh. +-- SKIP_LOCKED skips a child that ring prune holds ACCESS EXCLUSIVE instead of +-- waiting for it and aborting the statistics gathered so far. +-- +-- Definer boundary: the function only analyses a head table or a partitioned +-- table of a schema that awa.queue_storage_schemas() reports, which requires +-- the schema's ownership, grants and triggers to be trusted, and builds SQL +-- from nothing but quoted identifiers. ANALYZE writes statistics and nothing +-- else. +-- +-- Returns 'analyzed'; 'skipped_privilege' when the function owner may not +-- analyse the relation; or 'missing' when the relation does not exist. +-- +-- N-1 binaries never call it. A binary newer than this schema logs once that +-- the function is missing and skips the refresh. Completes in milliseconds and +-- takes no lock on any queue-storage table. + +CREATE OR REPLACE FUNCTION awa.analyze_queue_storage_relation( + p_schema TEXT, + p_relation TEXT +) +RETURNS TEXT +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + v_relid OID; + v_relkind "char"; + v_may_analyze BOOLEAN; +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM awa.queue_storage_schemas() AS substrate + WHERE substrate.schema_name = p_schema + ) THEN + RAISE EXCEPTION '% is not a queue-storage schema', p_schema + USING ERRCODE = '22023'; + END IF; + + v_relid := to_regclass(format('%I.%I', p_schema, p_relation)); + IF v_relid IS NULL THEN + RETURN 'missing'; + END IF; + + SELECT class.relkind INTO v_relkind FROM pg_class AS class WHERE class.oid = v_relid; + IF NOT ( + v_relkind = 'p' + OR p_relation IN ('queue_claim_heads', 'queue_enqueue_heads') + ) THEN + RAISE EXCEPTION '%.% is not a partitioned or head table of a queue-storage schema', + p_schema, p_relation + USING ERRCODE = '22023'; + END IF; + + -- ANALYZE skips, with a warning rather than an error, a table the caller + -- may not analyse: the owner, a MAINTAIN holder (PostgreSQL 17+), a + -- member of pg_maintain, or the database owner may analyse it. + SELECT + pg_has_role(current_user, class.relowner, 'USAGE') + OR pg_has_role( + current_user, + (SELECT database.datdba FROM pg_database AS database + WHERE database.datname = current_database()), + 'USAGE' + ) + OR CASE + WHEN current_setting('server_version_num')::int >= 170000 + THEN has_table_privilege(class.oid, 'MAINTAIN') + ELSE FALSE + END + INTO v_may_analyze + FROM pg_class AS class + WHERE class.oid = v_relid; + + IF NOT v_may_analyze THEN + RETURN 'skipped_privilege'; + END IF; + + IF v_relkind = 'p' THEN + EXECUTE format('ANALYZE (SKIP_LOCKED) %I.%I', p_schema, p_relation); + ELSE + EXECUTE format('ANALYZE %I.%I', p_schema, p_relation); + END IF; + RETURN 'analyzed'; +END +$$; + +GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(TEXT, TEXT) TO PUBLIC; + +INSERT INTO awa.schema_version (version, description) +VALUES (53, 'Planner statistics refresh for queue-storage tables autovacuum never analyses') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 385035b65..13d3beb49 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -5,7 +5,7 @@ use sqlx::{Connection, PgPool}; use tracing::{info, warn}; /// Current schema version. -pub const CURRENT_VERSION: i32 = 52; +pub const CURRENT_VERSION: i32 = 53; /// Migrations that require an exclusive (no-live-runtime) upgrade window. /// @@ -292,6 +292,11 @@ const MIGRATIONS: &[(i32, &str, &[&str])] = &[ ), (51, "Per-slot claim functions", &[V23_UP, V51_UP]), (52, "Retained job output", &[V23_UP, V52_UP]), + ( + 53, + "Planner statistics refresh for queue-storage tables autovacuum never analyses", + &[V23_UP, V53_UP], + ), ]; const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql"); @@ -352,6 +357,7 @@ const V49_UP: &str = include_str!("../migrations/v049_prepared_lane_sequences.sq const V50_UP: &str = include_str!("../migrations/v050_drop_ready_lane_shard_index.sql"); const V51_UP: &str = include_str!("../migrations/v051_claim_slot_functions.sql"); const V52_UP: &str = include_str!("../migrations/v052_retained_job_output.sql"); +const V53_UP: &str = include_str!("../migrations/v053_analyze_queue_storage_statistics.sql"); /// Old version numbers from pre-0.4 releases that used V3/V4/V5 numbering. /// Also tolerates the unreleased inline-V6 branch numbering used during review. diff --git a/awa-python/python/awa/_awa.pyi b/awa-python/python/awa/_awa.pyi index 2219bfc3b..34f5edbab 100644 --- a/awa-python/python/awa/_awa.pyi +++ b/awa-python/python/awa/_awa.pyi @@ -605,6 +605,7 @@ class Client: deadline_rescue_interval_ms: int | None = None, callback_rescue_interval_ms: int | None = None, terminal_count_rollup_interval_ms: int | None = None, + statistics_refresh_interval_ms: int | None = None, queue_storage_schema: str | None = None, queue_storage_queue_slot_count: int = 16, queue_storage_lease_slot_count: int = 8, diff --git a/awa-python/python/awa/client.py b/awa-python/python/awa/client.py index f734ca679..131d3391b 100644 --- a/awa-python/python/awa/client.py +++ b/awa-python/python/awa/client.py @@ -922,6 +922,7 @@ async def start( deadline_rescue_interval_ms: int | None = None, callback_rescue_interval_ms: int | None = None, terminal_count_rollup_interval_ms: int | None = None, + statistics_refresh_interval_ms: int | None = None, queue_storage_schema: str | None = None, queue_storage_queue_slot_count: int = DEFAULT_QUEUE_STORAGE_QUEUE_SLOT_COUNT, queue_storage_lease_slot_count: int = DEFAULT_QUEUE_STORAGE_LEASE_SLOT_COUNT, @@ -967,6 +968,11 @@ async def start( counters. Exact reads include pending deltas, so the default favors fewer hot counter updates over immediate folding. + ``statistics_refresh_interval_ms`` controls how often the maintenance + leader runs ``ANALYZE`` on the queue-storage head tables and + partitioned ring parents, which autovacuum never analyses (default 5 + minutes; ``0`` disables the refresh). + Queue dicts also accept ``claimers`` and ``claim_batch_size`` for hot queue-storage queues. ``claimers`` adds dispatcher/claimer loops that share the queue's worker permits; benchmark ``2`` or ``4`` before @@ -988,6 +994,7 @@ async def start( deadline_rescue_interval_ms=deadline_rescue_interval_ms, callback_rescue_interval_ms=callback_rescue_interval_ms, terminal_count_rollup_interval_ms=terminal_count_rollup_interval_ms, + statistics_refresh_interval_ms=statistics_refresh_interval_ms, queue_storage_schema=queue_storage_schema, queue_storage_queue_slot_count=queue_storage_queue_slot_count, queue_storage_lease_slot_count=queue_storage_lease_slot_count, diff --git a/awa-python/src/client.rs b/awa-python/src/client.rs index f141508a7..e17b0c15e 100644 --- a/awa-python/src/client.rs +++ b/awa-python/src/client.rs @@ -1796,7 +1796,7 @@ impl PyClient { }) } - #[pyo3(signature = (queues=None, *, poll_interval_ms=200, global_max_workers=None, completed_retention_hours=None, failed_retention_hours=None, descriptor_retention_days=None, cleanup_batch_size=None, leader_election_interval_ms=None, heartbeat_interval_ms=None, promote_interval_ms=None, heartbeat_rescue_interval_ms=None, heartbeat_staleness_ms=None, deadline_rescue_interval_ms=None, callback_rescue_interval_ms=None, terminal_count_rollup_interval_ms=None, queue_storage_schema=None, queue_storage_queue_slot_count=16, queue_storage_lease_slot_count=8, queue_storage_claim_slot_count=8, queue_storage_queue_stripe_count=1, queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=None, queue_storage_rotate_min_slot_age_ms=None, queue_storage_rotate_target_slot_bytes=None, queue_storage_rotate_min_slot_rows=None, queue_storage_rotate_max_slot_age_ms=None, storage_transition_role=None))] + #[pyo3(signature = (queues=None, *, poll_interval_ms=200, global_max_workers=None, completed_retention_hours=None, failed_retention_hours=None, descriptor_retention_days=None, cleanup_batch_size=None, leader_election_interval_ms=None, heartbeat_interval_ms=None, promote_interval_ms=None, heartbeat_rescue_interval_ms=None, heartbeat_staleness_ms=None, deadline_rescue_interval_ms=None, callback_rescue_interval_ms=None, terminal_count_rollup_interval_ms=None, statistics_refresh_interval_ms=None, queue_storage_schema=None, queue_storage_queue_slot_count=16, queue_storage_lease_slot_count=8, queue_storage_claim_slot_count=8, queue_storage_queue_stripe_count=1, queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=None, queue_storage_rotate_min_slot_age_ms=None, queue_storage_rotate_target_slot_bytes=None, queue_storage_rotate_min_slot_rows=None, queue_storage_rotate_max_slot_age_ms=None, storage_transition_role=None))] #[allow(clippy::too_many_arguments)] fn start<'py>( &self, @@ -1816,6 +1816,7 @@ impl PyClient { deadline_rescue_interval_ms: Option, callback_rescue_interval_ms: Option, terminal_count_rollup_interval_ms: Option, + statistics_refresh_interval_ms: Option, queue_storage_schema: Option, queue_storage_queue_slot_count: u32, queue_storage_lease_slot_count: u32, @@ -1950,6 +1951,9 @@ impl PyClient { } builder = builder.terminal_count_rollup_interval(Duration::from_millis(ms)); } + if let Some(ms) = statistics_refresh_interval_ms { + builder = builder.statistics_refresh_interval(Duration::from_millis(ms)); + } if queue_storage_queue_slot_count == 0 { return Err(pyo3::exceptions::PyValueError::new_err( "queue_storage_queue_slot_count must be > 0", diff --git a/awa-python/tests/test_start_config.py b/awa-python/tests/test_start_config.py index 302df0fc2..4d19775d0 100644 --- a/awa-python/tests/test_start_config.py +++ b/awa-python/tests/test_start_config.py @@ -316,12 +316,26 @@ async def handle(job): deadline_rescue_interval_ms=1_000, callback_rescue_interval_ms=1_000, terminal_count_rollup_interval_ms=1_000, + statistics_refresh_interval_ms=1_000, leader_election_interval_ms=200, promote_interval_ms=500, ) await client.shutdown() +@pytest.mark.asyncio +async def test_statistics_refresh_can_be_disabled(client): + """statistics_refresh_interval_ms=0 disables the refresh; start() accepts it.""" + queue = "cfg_statistics_refresh_off" + + @client.task(ConfigTestJob, queue=queue) + async def handle(job): + return None + + await client.start([(queue, 1)], statistics_refresh_interval_ms=0) + await client.shutdown() + + @pytest.mark.asyncio async def test_per_queue_retention_in_dict_config(client): """Dict form with per-queue retention config starts successfully.""" diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 253bbde7a..d32e8b977 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -163,6 +163,7 @@ pub struct ClientBuilder { leader_check_interval: Option, priority_aging_interval: Option, terminal_count_rollup_interval: Option, + statistics_refresh_interval: Option, completed_retention: Option, failed_retention: Option, descriptor_retention: Option, @@ -226,6 +227,7 @@ impl ClientBuilder { leader_check_interval: None, priority_aging_interval: None, terminal_count_rollup_interval: None, + statistics_refresh_interval: None, completed_retention: None, failed_retention: None, descriptor_retention: None, @@ -905,6 +907,20 @@ impl ClientBuilder { self } + /// Set how often the maintenance leader refreshes planner statistics for + /// the queue-storage tables autovacuum never analyses (default: 5 minutes). + /// + /// Autovacuum does not analyse partitioned tables, and never reaches its + /// threshold on the two small head tables, so the leader runs `ANALYZE` + /// on both when it is elected and then every interval, through + /// `awa.analyze_queue_storage_relation()` (schema v053) so that a runtime + /// role that does not own the tables still works. `Duration::ZERO` + /// disables the refresh. Queue storage only. + pub fn statistics_refresh_interval(mut self, interval: Duration) -> Self { + self.statistics_refresh_interval = Some(interval); + self + } + /// Set how often queue depth/lag metrics are published (default: 30s). pub fn queue_stats_interval(mut self, interval: Duration) -> Self { self.queue_stats_interval = Some(interval); @@ -1256,6 +1272,7 @@ impl ClientBuilder { leader_check_interval: self.leader_check_interval, priority_aging_interval: self.priority_aging_interval, terminal_count_rollup_interval: self.terminal_count_rollup_interval, + statistics_refresh_interval: self.statistics_refresh_interval, completed_retention: self.completed_retention, failed_retention: self.failed_retention, descriptor_retention: self.descriptor_retention, @@ -1362,6 +1379,7 @@ pub struct Client { leader_check_interval: Option, priority_aging_interval: Option, terminal_count_rollup_interval: Option, + statistics_refresh_interval: Option, completed_retention: Option, failed_retention: Option, descriptor_retention: Option, @@ -1932,6 +1950,9 @@ impl Client { if let Some(interval) = self.terminal_count_rollup_interval { maintenance = maintenance.terminal_count_rollup_interval(interval); } + if let Some(interval) = self.statistics_refresh_interval { + maintenance = maintenance.statistics_refresh_interval(interval); + } if let Some(retention) = self.completed_retention { maintenance = maintenance.completed_retention(retention); } diff --git a/awa-worker/src/lib.rs b/awa-worker/src/lib.rs index 4768e59bb..f9a171104 100644 --- a/awa-worker/src/lib.rs +++ b/awa-worker/src/lib.rs @@ -17,6 +17,7 @@ pub mod metrics; mod notify_hub; mod output; mod runtime; +mod statistics; mod storage; // Re-exports diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index d84960d37..042936b9a 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -1,5 +1,6 @@ use crate::executor::DlqPolicy; use crate::runtime::InFlightMap; +use crate::statistics::{RefreshGuard, StatisticsRefresher}; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; use awa_model::audited_sql; use awa_model::cron::{ @@ -738,6 +739,9 @@ pub struct MaintenanceService { priority_aging_interval: Duration, batch_operations_interval: Duration, terminal_count_rollup_interval: Duration, + /// How often the leader refreshes planner statistics for the queue-storage + /// tables autovacuum never analyses. Zero disables the refresh. + statistics_refresh_interval: Duration, /// How long a descriptor catalog row can sit without being refreshed /// before the maintenance leader deletes it. Zero disables cleanup. /// Default: 30 days. @@ -757,6 +761,8 @@ pub struct MaintenanceService { ring_authority_flip_aware_since: std::sync::Mutex>, } +/// Default cadence of the planner-statistics refresh on the maintenance leader. +pub(crate) const DEFAULT_STATISTICS_REFRESH_INTERVAL: Duration = Duration::from_secs(300); const PROMOTE_BATCH_SIZE: i64 = 4_096; const PROMOTE_MAX_BATCHES_PER_TICK: usize = 32; const CRON_CATCH_UP_LIMIT: usize = 1_000; @@ -817,6 +823,7 @@ impl MaintenanceService { priority_aging_interval: Duration::from_secs(60), batch_operations_interval: Duration::from_secs(1), terminal_count_rollup_interval: Duration::from_secs(30), + statistics_refresh_interval: DEFAULT_STATISTICS_REFRESH_INTERVAL, descriptor_retention: Duration::from_secs(30 * 86400), // 30d ring_authority_auto_flip: true, ring_authority_auto_flip_stable_period: Duration::from_secs(600), // 10min @@ -873,6 +880,14 @@ impl MaintenanceService { self } + /// Set how often the leader runs `ANALYZE` on the queue-storage tables + /// autovacuum never analyses (default: 5 minutes). `Duration::ZERO` + /// disables the refresh. + pub fn statistics_refresh_interval(mut self, interval: Duration) -> Self { + self.statistics_refresh_interval = interval; + self + } + /// How long a descriptor catalog row can go without being re-synced /// before the maintenance leader deletes it (default: 30 days). Set /// to `Duration::ZERO` to disable — useful if you maintain the catalog @@ -1121,6 +1136,7 @@ impl MaintenanceService { cron_eval_cancel.clone(), self.cron_eval_interval, )); + let _statistics_refresh = self.spawn_statistics_refresh(); loop { tokio::select! { @@ -1314,6 +1330,24 @@ impl MaintenanceService { } } + /// Start the planner-statistics refresh for this leadership term. The + /// returned guard stops it when dropped. + fn spawn_statistics_refresh(&self) -> Option { + if self.statistics_refresh_interval.is_zero() { + return None; + } + let runtime = self.storage.queue_storage()?; + Some(RefreshGuard::spawn( + StatisticsRefresher::new( + self.pool.clone(), + runtime.store.schema().to_string(), + self.metrics.clone(), + ), + self.statistics_refresh_interval, + &self.cancel, + )) + } + fn stop_cron_evaluator(cancel: &CancellationToken, task: &JoinHandle<()>) { cancel.cancel(); task.abort(); diff --git a/awa-worker/src/statistics.rs b/awa-worker/src/statistics.rs new file mode 100644 index 000000000..4fe0aab59 --- /dev/null +++ b/awa-worker/src/statistics.rs @@ -0,0 +1,347 @@ +//! Planner-statistics refresh for the queue-storage tables autovacuum never +//! analyses: the two head tables, which stay below the autoanalyze threshold, +//! and the partitioned ring parents, which autovacuum skips entirely. +//! +//! The maintenance leader runs one pass on election and one every interval. +//! Each relation is analysed through `awa.analyze_queue_storage_relation` +//! (migration v053, `SECURITY DEFINER`) in its own short transaction with a +//! bounded `lock_timeout` and `statement_timeout`, so the refresh never holds +//! locks across relations and cannot wait on ring prune for long. + +use crate::metrics::AwaMetrics; +use sqlx::PgPool; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; +use tokio_util::sync::CancellationToken; +use tracing::{debug, warn}; + +/// How long one relation's ANALYZE may wait for a lock before the pass moves +/// on. Ring prune holds `ACCESS EXCLUSIVE` on a slot's children only +/// briefly, so a wait this long means something else holds the table. +const LOCK_TIMEOUT: Duration = Duration::from_secs(2); + +/// Upper bound for one relation. A ring parent analyses in under a few +/// seconds even on a loaded host; this only stops a stuck statement. +const STATEMENT_TIMEOUT: Duration = Duration::from_secs(120); + +const HEAD_TABLES: [&str; 2] = ["queue_claim_heads", "queue_enqueue_heads"]; + +const SQLSTATE_INSUFFICIENT_PRIVILEGE: &str = "42501"; +const SQLSTATE_UNDEFINED_FUNCTION: &str = "42883"; +const SQLSTATE_LOCK_NOT_AVAILABLE: &str = "55P03"; + +/// Which kind of relation an analyse call covered; bounds metric cardinality. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RelationScope { + Heads, + RingParent, +} + +impl RelationScope { + pub(crate) fn as_str(self) -> &'static str { + match self { + RelationScope::Heads => "heads", + RelationScope::RingParent => "ring_parent", + } + } +} + +/// Result of analysing one relation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AnalyzeOutcome { + Ok, + /// The function owner (or the runtime role, if it cannot execute the + /// function) may not analyse the relation. + SkippedPrivilege, + /// `lock_timeout` expired waiting on a conflicting lock. + SkippedLock, + /// `awa.analyze_queue_storage_relation` is not installed (schema older + /// than v053). + Unavailable, + Failed, +} + +impl AnalyzeOutcome { + pub(crate) fn as_str(self) -> &'static str { + match self { + AnalyzeOutcome::Ok => "ok", + AnalyzeOutcome::SkippedPrivilege => "skipped_privilege", + AnalyzeOutcome::SkippedLock => "skipped_lock", + AnalyzeOutcome::Unavailable => "unavailable", + AnalyzeOutcome::Failed => "failed", + } + } +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(crate) struct PassSummary { + pub(crate) ok: usize, + pub(crate) skipped_privilege: usize, + pub(crate) skipped_lock: usize, + pub(crate) unavailable: usize, + pub(crate) failed: usize, +} + +impl PassSummary { + fn record(&mut self, outcome: AnalyzeOutcome) { + match outcome { + AnalyzeOutcome::Ok => self.ok += 1, + AnalyzeOutcome::SkippedPrivilege => self.skipped_privilege += 1, + AnalyzeOutcome::SkippedLock => self.skipped_lock += 1, + AnalyzeOutcome::Unavailable => self.unavailable += 1, + AnalyzeOutcome::Failed => self.failed += 1, + } + } +} + +/// Refreshes planner statistics for one queue-storage schema. +pub(crate) struct StatisticsRefresher { + pool: PgPool, + schema: String, + metrics: AwaMetrics, + warned_privilege: AtomicBool, + warned_unavailable: AtomicBool, +} + +impl StatisticsRefresher { + pub(crate) fn new(pool: PgPool, schema: String, metrics: AwaMetrics) -> Self { + Self { + pool, + schema, + metrics, + warned_privilege: AtomicBool::new(false), + warned_unavailable: AtomicBool::new(false), + } + } + + /// Run a pass now and then every `interval` until `cancel` fires. + pub(crate) async fn run(&self, interval: Duration, cancel: CancellationToken) { + let mut timer = tokio::time::interval(interval); + timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = timer.tick() => { + tokio::select! { + _ = cancel.cancelled() => return, + _ = self.run_pass() => {} + } + } + } + } + } + + /// Analyse the head tables, then every partitioned table of the schema. + pub(crate) async fn run_pass(&self) -> PassSummary { + let started = Instant::now(); + let mut summary = PassSummary::default(); + + match awa_model::storage::queue_storage_schema_ready(&self.pool, &self.schema).await { + Ok(true) => {} + Ok(false) => { + debug!(schema = %self.schema, "Queue-storage schema is not prepared yet; skipping statistics refresh"); + return summary; + } + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not check the queue-storage schema for statistics refresh"); + self.record( + RelationScope::RingParent, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + return summary; + } + } + + let ring_parents = match self.ring_parents().await { + Ok(parents) => parents, + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not list ring tables for statistics refresh"); + self.record( + RelationScope::RingParent, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + Vec::new() + } + }; + + let targets = HEAD_TABLES + .iter() + .map(|name| (RelationScope::Heads, (*name).to_string())) + .chain( + ring_parents + .into_iter() + .map(|name| (RelationScope::RingParent, name)), + ); + + for (scope, relation) in targets { + let relation_started = Instant::now(); + let outcome = self.analyze_relation(&relation).await; + self.record(scope, outcome, relation_started.elapsed()); + summary.record(outcome); + } + + debug!( + schema = %self.schema, + elapsed_ms = started.elapsed().as_millis() as u64, + ok = summary.ok, + skipped_privilege = summary.skipped_privilege, + skipped_lock = summary.skipped_lock, + unavailable = summary.unavailable, + failed = summary.failed, + "Planner statistics refresh finished" + ); + summary + } + + fn record(&self, scope: RelationScope, outcome: AnalyzeOutcome, elapsed: Duration) { + self.metrics + .record_statistics_analyze(scope.as_str(), outcome.as_str(), elapsed); + } + + async fn ring_parents(&self) -> Result, sqlx::Error> { + sqlx::query_scalar( + r#" + SELECT class.relname::text + FROM pg_class AS class + JOIN pg_namespace AS namespace ON namespace.oid = class.relnamespace + WHERE namespace.nspname = $1 AND class.relkind = 'p' + ORDER BY class.relname + "#, + ) + .bind(&self.schema) + .fetch_all(&self.pool) + .await + } + + async fn analyze_relation(&self, relation: &str) -> AnalyzeOutcome { + match self.call_analyze_function(relation).await { + Ok(status) => match status.as_str() { + "analyzed" => AnalyzeOutcome::Ok, + "skipped_privilege" => { + self.warn_privilege_once( + relation, + "the owner of awa.analyze_queue_storage_relation() may not ANALYZE this table; \ + run migrations as the role that owns the queue-storage tables, or grant the \ + function owner MAINTAIN on them (PostgreSQL 17+)", + ); + AnalyzeOutcome::SkippedPrivilege + } + // The relation is not part of this schema's shape; nothing to analyse. + "missing" => AnalyzeOutcome::Ok, + other => { + warn!( + status = other, + relation, "Unexpected statistics refresh status" + ); + AnalyzeOutcome::Failed + } + }, + Err(err) => self.classify_error(relation, &err), + } + } + + async fn call_analyze_function(&self, relation: &str) -> Result { + let mut tx = self.pool.begin().await?; + sqlx::query( + "SELECT set_config('lock_timeout', $1, true), set_config('statement_timeout', $2, true)", + ) + .bind(format!("{}ms", LOCK_TIMEOUT.as_millis())) + .bind(format!("{}ms", STATEMENT_TIMEOUT.as_millis())) + .execute(&mut *tx) + .await?; + let status: String = + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(&self.schema) + .bind(relation) + .fetch_one(&mut *tx) + .await?; + tx.commit().await?; + Ok(status) + } + + fn classify_error(&self, relation: &str, err: &sqlx::Error) -> AnalyzeOutcome { + let code = match err { + sqlx::Error::Database(db_err) => db_err.code().map(|code| code.into_owned()), + _ => None, + }; + match code.as_deref() { + Some(SQLSTATE_LOCK_NOT_AVAILABLE) => { + debug!( + relation, + "Statistics refresh timed out waiting for a lock; retrying next interval" + ); + AnalyzeOutcome::SkippedLock + } + Some(SQLSTATE_UNDEFINED_FUNCTION) => { + if !self.warned_unavailable.swap(true, Ordering::Relaxed) { + warn!( + schema = %self.schema, + "Planner statistics refresh is disabled: awa.analyze_queue_storage_relation() \ + is missing. Run `awa migrate` (schema v053 or later) to install it" + ); + } + AnalyzeOutcome::Unavailable + } + Some(SQLSTATE_INSUFFICIENT_PRIVILEGE) => { + self.warn_privilege_once( + relation, + "the runtime role may not execute awa.analyze_queue_storage_relation(); \ + GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) to it", + ); + AnalyzeOutcome::SkippedPrivilege + } + _ => { + warn!(relation, error = %err, "Planner statistics refresh failed for relation"); + AnalyzeOutcome::Failed + } + } + } + + fn warn_privilege_once(&self, relation: &str, missing: &str) { + if self.warned_privilege.swap(true, Ordering::Relaxed) { + return; + } + warn!( + schema = %self.schema, + relation, + "Planner statistics refresh skipped: {missing}. Set the statistics refresh \ + interval to zero to disable the refresh" + ); + } +} + +/// Stops the refresh task when the maintenance leader loses leadership or +/// shuts down. +pub(crate) struct RefreshGuard { + cancel: CancellationToken, +} + +impl RefreshGuard { + pub(crate) fn spawn( + refresher: StatisticsRefresher, + interval: Duration, + parent: &CancellationToken, + ) -> Self { + let cancel = parent.child_token(); + let task_cancel = cancel.clone(); + tokio::spawn(async move { + debug!( + schema = %refresher.schema, + interval_s = interval.as_secs(), + "Planner statistics refresh started" + ); + refresher.run(interval, task_cancel).await; + }); + Self { cancel } + } +} + +impl Drop for RefreshGuard { + fn drop(&mut self) { + self.cancel.cancel(); + } +} diff --git a/awa/tests/statistics_refresh_test.rs b/awa/tests/statistics_refresh_test.rs new file mode 100644 index 000000000..829aa3ea8 --- /dev/null +++ b/awa/tests/statistics_refresh_test.rs @@ -0,0 +1,561 @@ +//! The maintenance leader refreshes planner statistics for the queue-storage +//! tables autovacuum never analyses (the head tables and the partitioned ring +//! parents) on election and every interval, without delaying ring prune. + +use awa::model::{migrations, PruneOutcome, QueueStorage, QueueStorageConfig, RotateOutcome}; +use awa::{audited_sql, Client, JobContext, JobError, JobResult, QueueConfig, Worker}; +use opentelemetry::KeyValue; +use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData}; +use opentelemetry_sdk::metrics::{InMemoryMetricExporter, SdkMeterProvider}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::str::FromStr; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::{Duration, Instant}; + +const ROLE_PASSWORD: &str = "statistics_refresh_test"; +const RING_PARENTS: [&str; 13] = [ + "done_entries", + "lease_claim_batches", + "lease_claim_closure_batches", + "lease_claim_closures", + "lease_claims", + "leases", + "queue_terminal_count_deltas", + "ready_claim_attempt_batches", + "ready_entries", + "ready_segments", + "ready_tombstones", + "receipt_completion_batches", + "receipt_completion_tombstones", +]; + +static DATABASE_COUNTER: AtomicU32 = AtomicU32::new(0); + +fn base_database_url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) +} + +struct StatisticsProbeWorker; + +#[async_trait::async_trait] +impl Worker for StatisticsProbeWorker { + fn kind(&self) -> &'static str { + "statistics_probe_job" + } + + async fn perform(&self, _ctx: &JobContext) -> Result { + Ok(JobResult::Completed) + } +} + +/// A throwaway database with the awa schema migrated, optionally owned by a +/// non-superuser role the way a production schema owner is. +struct TestDatabase { + admin: PgPool, + name: String, + roles: Vec, + options: PgConnectOptions, + /// Connects as the schema owner (or the superuser when no owner role). + owner_pool: PgPool, +} + +impl TestDatabase { + async fn create(owner_role: Option<&str>, extra_roles: &[&str]) -> Self { + let admin = PgPoolOptions::new() + .max_connections(2) + .connect(&base_database_url()) + .await + .expect("connect admin pool"); + let name = format!( + "awa_stats_{}_{}", + std::process::id(), + DATABASE_COUNTER.fetch_add(1, Ordering::SeqCst) + ); + let mut roles = Vec::new(); + for role in owner_role.iter().chain(extra_roles.iter()) { + sqlx::raw_sql(audited_sql(format!( + "DROP ROLE IF EXISTS {role}; CREATE ROLE {role} LOGIN PASSWORD '{ROLE_PASSWORD}'" + ))) + .execute(&admin) + .await + .expect("create role"); + roles.push((*role).to_string()); + } + let owner_clause = owner_role + .map(|role| format!(" OWNER {role}")) + .unwrap_or_default(); + sqlx::raw_sql(audited_sql(format!("CREATE DATABASE {name}{owner_clause}"))) + .execute(&admin) + .await + .expect("create database"); + let options = PgConnectOptions::from_str(&base_database_url()) + .expect("parse DATABASE_URL") + .database(&name); + let owner_options = match owner_role { + Some(role) => options.clone().username(role).password(ROLE_PASSWORD), + None => options.clone(), + }; + let owner_pool = PgPoolOptions::new() + .max_connections(8) + .connect_with(owner_options) + .await + .expect("connect owner pool"); + migrations::run(&owner_pool).await.expect("run migrations"); + Self { + admin, + name, + roles, + options, + owner_pool, + } + } + + /// A pool for a role other than the schema owner, with the runtime grants + /// from the database-roles guide. + async fn runtime_pool(&self, role: &str) -> PgPool { + sqlx::raw_sql(audited_sql(format!( + r#" + GRANT USAGE ON SCHEMA awa TO {role}; + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {role}; + GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {role}; + GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {role}; + REVOKE EXECUTE ON FUNCTION + awa.install_queue_storage_substrate(TEXT, INT, INT, INT, BOOLEAN) FROM {role}; + "# + ))) + .execute(&self.owner_pool) + .await + .expect("grant runtime privileges"); + PgPoolOptions::new() + .max_connections(8) + .connect_with(self.options.clone().username(role).password(ROLE_PASSWORD)) + .await + .expect("connect runtime pool") + } + + async fn drop(self) { + self.owner_pool.close().await; + let _ = sqlx::raw_sql(audited_sql(format!( + "DROP DATABASE IF EXISTS {} WITH (FORCE)", + self.name + ))) + .execute(&self.admin) + .await; + for role in &self.roles { + let _ = sqlx::raw_sql(audited_sql(format!("DROP ROLE IF EXISTS {role}"))) + .execute(&self.admin) + .await; + } + self.admin.close().await; + } +} + +async fn wait_for(what: &str, timeout: Duration, mut check: F) +where + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if check().await { + return; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + panic!("timed out waiting for {what}"); +} + +async fn stats_rows(pool: &PgPool, schema: &str, table: &str, inherited_only: bool) -> i64 { + sqlx::query_scalar( + "SELECT count(*) FROM pg_stats + WHERE schemaname = $1 AND tablename = $2 AND (inherited OR NOT $3)", + ) + .bind(schema) + .bind(table) + .bind(inherited_only) + .fetch_one(pool) + .await + .expect("read pg_stats") +} + +async fn parent_reltuples(pool: &PgPool, schema: &str, table: &str) -> f32 { + sqlx::query_scalar("SELECT reltuples FROM pg_class WHERE oid = to_regclass($1)") + .bind(format!("{schema}.{table}")) + .fetch_one(pool) + .await + .expect("read reltuples") +} + +async fn analyze_count(pool: &PgPool, schema: &str, table: &str) -> i64 { + sqlx::query_scalar( + "SELECT COALESCE(analyze_count, 0)::bigint FROM pg_stat_user_tables + WHERE schemaname = $1 AND relname = $2", + ) + .bind(schema) + .bind(table) + .fetch_one(pool) + .await + .expect("read analyze_count") +} + +fn probe_client( + pool: PgPool, + statistics_interval: Duration, + storage: Option, +) -> Client { + let mut builder = Client::builder(pool) + .queue( + "statistics_probe", + QueueConfig { + max_workers: 1, + poll_interval: Duration::from_millis(100), + ..QueueConfig::default() + }, + ) + .register_worker(StatisticsProbeWorker) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(100)) + .statistics_refresh_interval(statistics_interval); + if let Some(config) = storage { + builder = builder + .queue_storage(config, Duration::from_secs(600), Duration::from_secs(600)) + .claim_rotate_interval(Duration::from_secs(600)); + } + builder.build().expect("build client") +} + +async fn seed_ready_rows(store: &QueueStorage, pool: &PgPool, rows_per_queue: i64) { + for queue in ["stats_a", "stats_b", "stats_c"] { + store + .enqueue_batch(pool, queue, 2, rows_per_queue) + .await + .expect("enqueue rows"); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn leader_refreshes_statistics_on_election_and_every_interval() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 200).await; + + assert_eq!(stats_rows(pool, "awa", "queue_claim_heads", false).await, 0); + assert_eq!( + stats_rows(pool, "awa", "queue_enqueue_heads", false).await, + 0 + ); + assert_eq!(stats_rows(pool, "awa", "ready_entries", true).await, 0); + assert!(parent_reltuples(pool, "awa", "ready_entries").await < 0.0); + + let client = probe_client(pool.clone(), Duration::from_secs(1), None); + client.start().await.expect("start client"); + + wait_for( + "first statistics pass on election", + Duration::from_secs(30), + || async { + stats_rows(pool, "awa", "queue_claim_heads", false).await > 0 + && stats_rows(pool, "awa", "queue_enqueue_heads", false).await > 0 + && stats_rows(pool, "awa", "ready_entries", true).await > 0 + }, + ) + .await; + let partitioned: Vec = sqlx::query_scalar( + "SELECT relname::text FROM pg_class + WHERE relnamespace = 'awa'::regnamespace AND relkind = 'p' ORDER BY relname", + ) + .fetch_all(pool) + .await + .expect("list partitioned tables"); + assert_eq!( + partitioned, RING_PARENTS, + "the refresh covers every ring parent" + ); + wait_for( + "every ring parent analysed", + Duration::from_secs(30), + || async { + for parent in RING_PARENTS { + if analyze_count(pool, "awa", parent).await < 1 { + return false; + } + } + true + }, + ) + .await; + let first_tuples = parent_reltuples(pool, "awa", "ready_entries").await; + assert!( + (550.0..=650.0).contains(&first_tuples), + "parent reltuples after the first pass: {first_tuples}" + ); + + seed_ready_rows(&store, pool, 2_000).await; + wait_for( + "a later pass to see the new rows", + Duration::from_secs(30), + || async { parent_reltuples(pool, "awa", "ready_entries").await > 5_000.0 }, + ) + .await; + assert!(analyze_count(pool, "awa", "ready_entries").await >= 2); + assert!(analyze_count(pool, "awa", "queue_claim_heads").await >= 2); + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn zero_interval_disables_the_refresh() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 100).await; + + let client = probe_client(pool.clone(), Duration::ZERO, None); + client.start().await.expect("start client"); + wait_for("leadership", Duration::from_secs(10), || async { + client.health_check().await.maintenance_alive + }) + .await; + tokio::time::sleep(Duration::from_secs(3)).await; + + assert_eq!(stats_rows(pool, "awa", "queue_claim_heads", false).await, 0); + assert_eq!(stats_rows(pool, "awa", "ready_entries", true).await, 0); + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn non_owner_runtime_role_refreshes_statistics() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime]).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, &database.owner_pool, 200).await; + let runtime_pool = database.runtime_pool(&runtime).await; + + let is_superuser_or_owner: bool = sqlx::query_scalar( + "SELECT rolsuper OR pg_has_role(current_user, 'pg_maintain', 'USAGE') + OR has_table_privilege('awa.queue_claim_heads', 'MAINTAIN') + OR (SELECT relowner = (SELECT oid FROM pg_roles WHERE rolname = current_user) + FROM pg_class WHERE oid = 'awa.queue_claim_heads'::regclass) + FROM pg_roles WHERE rolname = current_user", + ) + .fetch_one(&runtime_pool) + .await + .expect("check runtime role privileges"); + assert!( + !is_superuser_or_owner, + "the runtime role must not be able to analyse the tables itself" + ); + + let client = probe_client(runtime_pool.clone(), Duration::from_secs(1), None); + client + .start() + .await + .expect("start client as the runtime role"); + wait_for( + "statistics written through a non-owner role", + Duration::from_secs(30), + || async { + stats_rows(&database.owner_pool, "awa", "queue_claim_heads", false).await > 0 + && stats_rows(&database.owner_pool, "awa", "ready_entries", true).await > 0 + }, + ) + .await; + + client.shutdown(Duration::from_secs(5)).await; + runtime_pool.close().await; + database.drop().await; +} + +fn sum_by_outcome( + resource_metrics: &[opentelemetry_sdk::metrics::data::ResourceMetrics], + outcome: &str, +) -> u64 { + let mut total = 0; + for resource in resource_metrics { + for scope in resource.scope_metrics() { + for metric in scope.metrics() { + if metric.name() != "awa.maintenance.analyze.attempts" { + continue; + } + if let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() { + total += sum + .data_points() + .filter(|point| { + point.attributes().any(|KeyValue { key, value, .. }| { + key.as_str() == "awa.analyze.outcome" && value.as_str() == outcome + }) + }) + .map(|point| point.value()) + .sum::(); + } + } + } + } + total +} + +/// A refresh that cannot run is recorded and leaves the runtime healthy: the +/// helper's owner may not analyse the tables, or the schema predates the +/// helper. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unrunnable_refresh_is_recorded_and_does_not_fail_the_runtime() { + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + opentelemetry::global::set_meter_provider(provider.clone()); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let helper = format!("stats_helper_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&helper]).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, &database.owner_pool, 50).await; + + // The function now runs as a role that owns nothing and holds no MAINTAIN. + sqlx::raw_sql(audited_sql(format!("GRANT {helper} TO {owner}"))) + .execute(&database.admin) + .await + .expect("let the schema owner hand over the function"); + sqlx::raw_sql(audited_sql(format!( + "GRANT USAGE, CREATE ON SCHEMA awa TO {helper}; + ALTER FUNCTION awa.analyze_queue_storage_relation(text, text) OWNER TO {helper}" + ))) + .execute(&database.owner_pool) + .await + .expect("hand the helper to an unprivileged role"); + + let client = probe_client( + database.owner_pool.clone(), + Duration::from_millis(200), + None, + ); + client.start().await.expect("start client"); + wait_for("a skipped pass", Duration::from_secs(30), || async { + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + sum_by_outcome(&metrics, "skipped_privilege") > 0 + }) + .await; + assert_eq!( + stats_rows(&database.owner_pool, "awa", "queue_claim_heads", false).await, + 0, + "a skipped pass writes no statistics" + ); + assert!(client.health_check().await.healthy); + client.shutdown(Duration::from_secs(5)).await; + + // A schema that predates the helper skips the same way. + sqlx::raw_sql("DROP FUNCTION awa.analyze_queue_storage_relation(text, text)") + .execute(&database.owner_pool) + .await + .expect("drop the helper"); + let client = probe_client( + database.owner_pool.clone(), + Duration::from_millis(200), + None, + ); + client + .start() + .await + .expect("start client without the helper"); + wait_for("an unavailable pass", Duration::from_secs(30), || async { + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + sum_by_outcome(&metrics, "unavailable") > 0 + }) + .await; + assert!(client.health_check().await.healthy); + client.shutdown(Duration::from_secs(5)).await; + + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + assert_eq!(sum_by_outcome(&metrics, "failed"), 0); + database.drop().await; +} + +/// Back-to-back statistics passes (a far higher duty cycle than the default +/// five-minute interval) run while rotation and prune cycle the ring: prune +/// keeps reclaiming slots and only defers an attempt that lands while ANALYZE +/// is sampling the slot's tables. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn statistics_refresh_does_not_stop_prune_reclaiming_slots() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let schema = "awa_stats_ring"; + let config = QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }; + let store = QueueStorage::new(config.clone()).expect("queue storage config"); + store.prepare_schema(pool).await.expect("prepare schema"); + + let queue = "statistics_ring"; + let analyze_before = analyze_count(pool, schema, "ready_entries").await; + let client = probe_client(pool.clone(), Duration::from_millis(1), Some(config)); + client.start().await.expect("start client"); + + let (mut pruned, mut deferred) = (0_u32, 0_u32); + let started = Instant::now(); + for _ in 0..60 { + store + .enqueue_batch(pool, queue, 1, 600) + .await + .expect("enqueue"); + loop { + let claims = store.claim_batch(pool, queue, 600).await.expect("claim"); + if claims.is_empty() { + break; + } + store.complete_batch(pool, &claims).await.expect("complete"); + } + match store.rotate(pool).await.expect("rotate") { + RotateOutcome::Rotated { .. } | RotateOutcome::SkippedBusy { .. } => {} + other => panic!("unexpected rotate outcome {other:?}"), + } + match store + .prune_oldest(pool, Duration::ZERO) + .await + .expect("prune never errors under ANALYZE") + { + PruneOutcome::Pruned { .. } => pruned += 1, + PruneOutcome::Blocked { .. } => deferred += 1, + PruneOutcome::Noop | PruneOutcome::SkippedActive { .. } => {} + } + if started.elapsed() > Duration::from_secs(60) { + break; + } + } + let passes = analyze_count(pool, schema, "ready_entries").await - analyze_before; + eprintln!("prune under ANALYZE: pruned={pruned} deferred={deferred} analyze_passes={passes}"); + + client.shutdown(Duration::from_secs(5)).await; + assert!( + passes >= 3, + "the refresh ran {passes} times during the loop" + ); + assert!(pruned >= 20, "prune reclaimed {pruned} slots"); + + // With the refresh stopped, the oldest sealed slot reclaims immediately. + store.rotate(pool).await.expect("rotate"); + assert!(!matches!( + store + .prune_oldest(pool, Duration::ZERO) + .await + .expect("prune"), + PruneOutcome::Blocked { .. } + )); + database.drop().await; +} diff --git a/docs/architecture.md b/docs/architecture.md index d3c527f61..388a133a5 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -215,6 +215,8 @@ Queue storage has three independent rings, each advanced by the elected maintena The maintenance tick for each ring is deliberately small: attempt one rotate, then attempt one prune. If a partition is busy, blocked by a lock, current, or still live, the tick records a skipped/blocked outcome and tries again on a future interval. +The leader also refreshes planner statistics on the two head tables and the ring parents when it is elected and every five minutes ([Planner statistics](configuration.md#planner-statistics)): autovacuum never analyses partitioned tables or reaches its threshold on the head tables. It runs beside the ring ticks rather than in them, so a slow `ANALYZE` never delays a rotate or prune tick. + The maintenance tick decides whether to seal the current slot. `ClientBuilder::ring_rotation_cadence` configures `RingRotationCadence`: a non-empty slot becomes eligible after `min_slot_age` since the maintenance store first observed rows (default 5 s) once it holds at least `min_slot_rows` (default 256, sampled with a bounded probe), when total physical storage reaches `target_slot_bytes` (default 2 MiB), or after `max_slot_age` (default 60 s) regardless of rows. At a trickle of writes a slot therefore seals about once a minute instead of every window. The minimum age uses a local monotonic clock: an empty observation clears the timer, and a newly elected leader starts a fresh window. `max_slot_age` also counts from the slot's `rotated_at` in the rotation ledger, on the database clock, so frequent leader changes cannot keep a trickle slot open past it. Under compat (pre-flip) authority an older rotator can leave the current generation without a ledger row; the ceiling then relies on the local observation, which a leader change restarts. The size measurement includes heap, indexes, TOAST, free-space maps and visibility maps. `RingRotationCadence::EAGER` restores eligibility on every tick. Queue pruning queues a bounded exclusive-lock request on routing metadata first, then acquires every other child with `NOWAIT`. Contention rolls back the attempt and retries next tick; active-data skip gates retain their backoff. This lets segment-first claim readers and entry-first producers finish without a prune lock cycle. It does not guarantee progress under perpetual contention. diff --git a/docs/configuration.md b/docs/configuration.md index 930bc6bde..6db6cec59 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -508,6 +508,7 @@ DLQ rows are different: `dlq_entries` is a separate hold table and the maintenan | `cleanup_batch_size` | `1000` | Max rows deleted per pass. Raise for very high throughput; lower if you want gentler IO. | | `dlq_cleanup_batch_size` | `1000` | DLQ-specific batch size. | | `ClientBuilder::terminal_count_rollup_interval` (Rust) / `terminal_count_rollup_interval_ms` (Python kwarg) | `30s` | Queue storage only: how often maintenance folds pending `done_entries` terminal-count deltas for sealed queue slots into compact live counters. Exact reads include pending deltas and retained compact receipt batches, so this affects compaction/WAL pressure rather than correctness. | +| `ClientBuilder::statistics_refresh_interval` (Rust) / `statistics_refresh_interval_ms` (Python kwarg) | `5m` | Queue storage only: how often the maintenance leader runs `ANALYZE` on the two head tables and the partitioned ring parents, which autovacuum never analyses; also run once when a runtime becomes leader. `0` disables it. See [Planner statistics](#planner-statistics). | ## Queue storage tuning @@ -531,6 +532,16 @@ The maintenance leader folds pending `done_entries` deltas from sealed queue slo Raising the interval reduces folded-counter churn but leaves more pending delta rows for exact reads to sum. Lowering it folds counters sooner, which can help if exact-count reads are frequent and the `done_entries` delta ledger is growing faster than maintenance drains it. Under long reader transactions, lowering the interval does not force rollup; the useful fix is to release the pinned snapshot or let queue prune reclaim the whole sealed segment. Compact receipt successes are already retained as compact batch rows, so this interval does not change their hot-path write count. +### Planner statistics + +Autovacuum does not analyse partitioned tables, so the ring parents (`ready_entries`, `done_entries`, `leases`, and the other ten partitioned tables of the schema) have no inherited statistics. It also never reaches its analyze threshold on `queue_claim_heads` and `queue_enqueue_heads`, which hold a few rows per lane, so the planner prices joins against them at its defaults. Admin and metrics queries that join retained ready rows to the heads tables then pick plans for a few hundred rows while reading hundreds of thousands. + +The maintenance leader runs `ANALYZE` on both when it is elected and every `statistics_refresh_interval` (default `5m`). One pass takes well under a second on a small ring and a few seconds on a ring with millions of retained rows. It has no measurable effect on claim or completion throughput. + +- **Locks.** Each relation is analysed in its own transaction with a 2 s `lock_timeout` and a 120 s `statement_timeout`. Children that ring prune holds are skipped (`SKIP_LOCKED`). Prune never waits for the refresh: its own 10 ms lock timeout applies, so an attempt that lands while `ANALYZE` samples the oldest slot's tables returns `blocked` and retries on its usual backoff. `ANALYZE` of a partitioned table samples every child, so that window cannot be avoided. +- **Privileges.** The refresh calls `awa.analyze_queue_storage_relation(schema, relation)` (schema v053), a `SECURITY DEFINER` function owned by the role that ran the migration, so the runtime role needs no `MAINTAIN` grant and does not have to own the tables. If the function is missing (run `awa migrate`) or its owner may not analyse the tables, the leader logs one warning that names what is missing and skips the pass; job processing is unaffected. +- **Metrics.** `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration`, by `awa.analyze.scope` (`heads`, `ring_parent`) and `awa.analyze.outcome` (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`). A sustained `skipped_lock` rate means something else holds the tables; any `skipped_privilege` or `unavailable` means the refresh is not running. + ### Rust ```rust diff --git a/docs/grafana/README.md b/docs/grafana/README.md index 0ca4ed331..4b8fa278a 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -207,6 +207,8 @@ In Prometheus they appear under the OpenTelemetry [OTLP-to-Prometheus translatio | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | | `awa.maintenance.prune.attempts` | Counter | ring, outcome, reason | Ring prune outcomes, including destructive `pruned` and no-DDL `already_pruned` | | `awa.maintenance.prune.duration` | Histogram (s) | ring, phase | Successful destructive prune time split into `lock`, `truncate`, and `commit` | +| `awa.maintenance.analyze.attempts` | Counter | scope, outcome | Planner-statistics `ANALYZE` calls on the head tables and ring parents (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`) | +| `awa.maintenance.analyze.duration` | Histogram (s) | scope, outcome | Duration of one planner-statistics `ANALYZE` call | ### Heartbeat diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index b6c1557e9..2e2d19a60 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -63,6 +63,12 @@ ALTER DEFAULT PRIVILEGES FOR ROLE awa_owner IN SCHEMA awa If an existing installation created objects as `awa_migrator` without first setting the owner role, transfer every existing schema object to `awa_owner` before relying on owner-scoped defaults. New migrations should use the effective-role command above. +## Planner statistics + +The maintenance leader runs `ANALYZE` on the queue-storage head tables and partitioned ring parents (see [Planner statistics](../configuration.md#planner-statistics)). `ANALYZE` skips a table its caller may not analyse, so a runtime that does not own the tables cannot do this itself; the leader calls `awa.analyze_queue_storage_relation(schema, relation)` instead, a `SECURITY DEFINER` function from migration v053 owned by whichever role ran the migration. Nothing is needed beyond the grants above: the function is executable by `PUBLIC`, and the runtime does not need `MAINTAIN`. + +Run migrations with the schema owner as the effective role, as shown above, so the function's owner owns the tables. A function owned by a role that may not analyse them (it is not the table owner, the database owner, or a `MAINTAIN` holder) makes every pass report `skipped_privilege` and logs one warning. The function only analyses a head table or partitioned table of a schema that `awa.queue_storage_schemas()` reports as trusted, and `ANALYZE` changes statistics only. + ## Provision queues before starting restricted runtimes From schema v049, runtime calls use existing lane sequences without issuing diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 0afc48afb..50f6cc587 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -253,6 +253,22 @@ workers report not ready until the migration runs, as for any version bump. External migration runners apply the complete `awa migrate --sql --from 51 --to 52` bundle, which defines the helper before the tail calls it. +## v053: planner statistics helper + +v053 adds `awa.analyze_queue_storage_relation(schema, relation)`, a +`SECURITY DEFINER` function owned by the migration role that lets the +maintenance leader refresh planner statistics for the head tables and ring +parents when the runtime role does not own them (see +[Planner statistics](configuration.md#planner-statistics)). It creates one +function and grants `EXECUTE` to `PUBLIC`; it takes no lock on any table and +completes in milliseconds. Run the migration as the schema owner so the +function's owner owns the tables. + +Binaries before v053 never call it. A 0.7 binary running before the migration +logs once that the function is missing and skips the refresh; the +`awa.maintenance.analyze.attempts` counter reports `unavailable`. External +migration runners apply the `awa migrate --sql --from 52 --to 53` bundle. + ## The one gate `awa migrate` on a 0.7 binary refuses to apply migrations unless one of these holds: diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index 4cfb039b0..90ad16015 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -135,7 +135,8 @@ processes. Exactly one worker per database is the elected maintenance leader (session-scoped advisory lock — if it dies, another is elected on the next tick). The leader alone runs cluster-wide scans: deferred→ready promotion, the three rescue paths (stale-heartbeat, hard-deadline, callback-timeout), ring -rotation and prune, DLQ and descriptor cleanup, and cron evaluation. During an +rotation and prune, planner-statistics refresh, DLQ and descriptor cleanup, and +cron evaluation. During an unfinalized transition it promotes both planes' deferred backlogs, so drain and post-flip execution do not depend on which runtime holds leadership. Every worker dispatches and heartbeats its own attempts. There is no `pg_cron` or @@ -313,6 +314,16 @@ storage limit. Queue prune lock contention retries on the next decision tick; active-data gates retain exponential backoff. Persistent contention can still stall reclamation. See `docs/configuration.md` for the scheduling contract. +On queue storage the leader runs `ANALYZE` on `queue_claim_heads`, +`queue_enqueue_heads` and the partitioned ring parents when elected and every +`statistics_refresh_interval` (default 5 minutes; Python +`statistics_refresh_interval_ms`; zero disables), because autovacuum never +analyses them. It goes through the `SECURITY DEFINER` function +`awa.analyze_queue_storage_relation` from migration v053, so a non-owner runtime +role needs no `MAINTAIN` grant. Watch `awa.maintenance.analyze.attempts`: any +`skipped_privilege` or `unavailable` means it is not running (run `awa migrate` +as the schema owner). See `docs/configuration.md#planner-statistics`. + ## Managed Postgres Prefer PostgreSQL 18 where available; PG17 on some managed platforms caps From 9e4517caf8de13612c7c31938c56b1880a3ca1df Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 10:10:24 +1300 Subject: [PATCH 02/18] test(worker): wait for the last ANALYZE before asserting prune is unblocked --- awa/tests/statistics_refresh_test.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/awa/tests/statistics_refresh_test.rs b/awa/tests/statistics_refresh_test.rs index 829aa3ea8..7215fcf53 100644 --- a/awa/tests/statistics_refresh_test.rs +++ b/awa/tests/statistics_refresh_test.rs @@ -548,7 +548,24 @@ async fn statistics_refresh_does_not_stop_prune_reclaiming_slots() { ); assert!(pruned >= 20, "prune reclaimed {pruned} slots"); - // With the refresh stopped, the oldest sealed slot reclaims immediately. + // Shutdown can leave one ANALYZE statement finishing server-side; once it + // is gone the oldest sealed slot reclaims without being blocked. + wait_for( + "the last ANALYZE to finish", + Duration::from_secs(30), + || async { + let running: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_stat_activity + WHERE state = 'active' AND query LIKE '%analyze_queue_storage_relation%' + AND pid <> pg_backend_pid()", + ) + .fetch_one(pool) + .await + .expect("read pg_stat_activity"); + running == 0 + }, + ) + .await; store.rotate(pool).await.expect("rotate"); assert!(!matches!( store From 40f94f65fdd3362e2c88e54a0fb2a67e45b02ed6 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 11:07:07 +1300 Subject: [PATCH 03/18] test(bench): benchmark per-lane enqueue contention and its effect on maintenance Many concurrent producers batch-inserting into one queue lane serialise on the per-lane reservation lock, which a transaction holds until it commits. The new enqueue_lane_contention_test reports per-lane throughput and batch latency percentiles at high producer concurrency, with and without unique keys, across enqueue_shards, with a unique-conflict replay mix, and optionally drains every shard afterwards. Its nightly gate runs sixteen producers on one lane (and on four shards), checks lane integrity, and asserts contention-scaled floors; the rate and p99 regression thresholds live in benchmarks/baseline.json. enqueue_lane_maintenance_coupling_test measures a leader-shaped loop (promotion and heartbeat rescue on one task) with and without producers saturating a lane, and with a lane lock held for a fixed time: promotion batch latency, promotion lag on an unrelated queue, and how late rescue starts. --- .github/workflows/nightly-chaos.yml | 11 + awa/tests/enqueue_lane_contention_test.rs | 775 ++++++++++++++++++ .../enqueue_lane_maintenance_coupling_test.rs | 429 ++++++++++ benchmarks/baseline.json | 7 + 4 files changed, 1222 insertions(+) create mode 100644 awa/tests/enqueue_lane_contention_test.rs create mode 100644 awa/tests/enqueue_lane_maintenance_coupling_test.rs diff --git a/.github/workflows/nightly-chaos.yml b/.github/workflows/nightly-chaos.yml index 9c69c0889..a57f22aa2 100644 --- a/.github/workflows/nightly-chaos.yml +++ b/.github/workflows/nightly-chaos.yml @@ -197,6 +197,17 @@ jobs: cargo test --release --package awa --test concurrent_lifecycle_test \ test_concurrent_queue_sweep -- --exact --ignored --nocapture \ | tee artifacts/nightly/rust-queue-sweep.txt + - name: Run per-lane enqueue contention benchmark + if: always() + timeout-minutes: 5 + run: | + set -o pipefail + # Many producers batch-inserting into one queue lane. Gates the + # per-lane rate and p99 through benchmarks/baseline.json, and the + # lane-lock hold time relative to the per-row insert function. + cargo test --release --package awa --test enqueue_lane_contention_test \ + test_enqueue_lane_contention_gate -- --exact --ignored --nocapture \ + | tee artifacts/nightly/rust-enqueue-lane.txt - name: Run MVCC horizon overlap benchmark if: always() timeout-minutes: 10 diff --git a/awa/tests/enqueue_lane_contention_test.rs b/awa/tests/enqueue_lane_contention_test.rs new file mode 100644 index 000000000..46a89625b --- /dev/null +++ b/awa/tests/enqueue_lane_contention_test.rs @@ -0,0 +1,775 @@ +//! Per-lane enqueue contention: many concurrent producers batch-inserting +//! into one queue lane. +//! +//! Every enqueue on queue storage reserves its lane sequence under a +//! transaction-scoped advisory lock (`reserve_enqueue_seq`) that is held until +//! the enqueue transaction ends, so producers for one `(queue, priority, +//! shard)` lane serialise on it. This suite measures what that serialisation +//! costs: lane throughput and batch-insert latency at high producer +//! concurrency, with and without unique keys, with and without a unique +//! conflict fallback, and across `enqueue_shards`. +//! +//! * `test_enqueue_lane_contention_gate` runs a fixed set of scenarios, checks +//! lane integrity, asserts contention-scaled floors, and emits +//! `@@BENCH_JSON@@` records that `.github/scripts/check-bench-regression.py` +//! compares against `benchmarks/baseline.json`. +//! * `test_enqueue_lane_contention_matrix` is the exploration entry point: it +//! sweeps producer counts for one configuration chosen through environment +//! variables and prints every scenario. +//! +//! Matrix tunables (all optional): +//! AWA_LANE_BENCH_PRODUCERS comma list, default "1,4,8,16,32" +//! AWA_LANE_BENCH_BATCH jobs per batch, default 185 +//! AWA_LANE_BENCH_UNIQUE 0|1, default 1 +//! AWA_LANE_BENCH_SHARDS `enqueue_shards` for the queue, default 1 +//! AWA_LANE_BENCH_MODE insert_many | insert_many_copy | queue_storage | +//! legacy_row_function +//! AWA_LANE_BENCH_DUP_PCT % of batches carrying one duplicate key that +//! fall back to row-at-a-time inserts, default 0 +//! AWA_LANE_BENCH_SECS measured seconds per scenario, default 8 +//! AWA_LANE_BENCH_DRAIN 0|1: claim everything that landed afterwards and +//! check every shard drains +//! AWA_LANE_BENCH_SPLIT 0|1: run each batch in an explicit +//! transaction and time statement and commit +//! separately (insert_many mode only) +//! +//! Requires a running Postgres instance. +//! Run with: +//! cargo test --release --package awa --test enqueue_lane_contention_test \ +//! -- --ignored --nocapture --test-threads=1 + +mod bench_output; +mod ci_timing; + +use awa::model::insert_prepared; +use awa::model::{ + insert_many, insert_many_copy_from_pool, prepare_raw_job_insert, AwaError, QueueStorage, +}; +use awa::{InsertOpts, InsertParams, QueueStorageConfig, UniqueOpts}; +use awa_testing::setup; +use bench_output::{BenchLatency, BenchMetrics, BenchmarkResult, SCHEMA_VERSION}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::collections::HashMap; +use std::str::FromStr; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use uuid::Uuid; + +const APPLICATION_NAME: &str = "awa_lane_bench"; +const WARMUP: Duration = Duration::from_millis(1_500); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Mode { + InsertMany, + InsertManyCopy, + QueueStorage, + LegacyRowFunction, +} + +impl Mode { + fn parse(raw: &str) -> Self { + match raw { + "insert_many" => Mode::InsertMany, + "insert_many_copy" => Mode::InsertManyCopy, + "queue_storage" => Mode::QueueStorage, + "legacy_row_function" => Mode::LegacyRowFunction, + other => panic!("unknown AWA_LANE_BENCH_MODE {other}"), + } + } + + fn label(self) -> &'static str { + match self { + Mode::InsertMany => "insert_many", + Mode::InsertManyCopy => "insert_many_copy", + Mode::QueueStorage => "queue_storage", + Mode::LegacyRowFunction => "legacy_row_function", + } + } +} + +#[derive(Debug, Clone)] +struct Scenario { + name: String, + mode: Mode, + producers: usize, + batch_size: usize, + unique: bool, + shards: i16, + duplicate_batch_pct: u32, + measured: Duration, + split_commit: bool, + drain: bool, +} + +#[derive(Debug, Default, Clone)] +struct Samples { + /// Latency of the producer call, in milliseconds: one `insert_many` + /// statement, or for a duplicate batch the failed batch plus the + /// row-at-a-time fallback. + batch_ms: Vec, + /// Latency of each single-row statement issued by the fallback. + fallback_row_ms: Vec, + /// Statement and commit latency when the batch runs in an explicit + /// transaction. + statement_ms: Vec, + commit_ms: Vec, + jobs: u64, + batches: u64, + duplicate_batches: u64, + fallback_rows: u64, +} + +impl Samples { + fn merge(&mut self, other: Samples) { + self.batch_ms.extend(other.batch_ms); + self.fallback_row_ms.extend(other.fallback_row_ms); + self.statement_ms.extend(other.statement_ms); + self.commit_ms.extend(other.commit_ms); + self.jobs += other.jobs; + self.batches += other.batches; + self.duplicate_batches += other.duplicate_batches; + self.fallback_rows += other.fallback_rows; + } +} + +#[derive(Debug)] +struct ScenarioResult { + scenario: Scenario, + jobs_per_s: f64, + p50_ms: f64, + p95_ms: f64, + p99_ms: f64, + max_ms: f64, + mean_waiters: f64, + samples: Samples, + landed: i64, + distinct_lane_seqs: i64, + drained: Option<(i64, f64)>, +} + +fn percentile(sorted: &[f64], quantile: f64) -> f64 { + if sorted.is_empty() { + return 0.0; + } + let index = ((sorted.len() as f64 - 1.0) * quantile).round() as usize; + sorted[index.min(sorted.len() - 1)] +} + +fn sorted(values: &[f64]) -> Vec { + let mut copy = values.to_vec(); + copy.sort_by(|left, right| left.partial_cmp(right).expect("latencies are finite")); + copy +} + +fn env_string(name: &str, default: &str) -> String { + std::env::var(name).unwrap_or_else(|_| default.to_string()) +} + +fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|raw| raw.parse().ok()) + .unwrap_or(default) +} + +fn job_params(queue: &str, tag: &str, producer: usize, counter: u64, unique: bool) -> InsertParams { + InsertParams { + kind: "lane_bench_job".to_string(), + args: serde_json::json!({ + "tag": tag, + "producer": producer, + "n": counter, + "payload": "x".repeat(64), + }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + } +} + +fn is_unique_violation(err: &AwaError) -> bool { + match err { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(db_err)) => { + db_err.code().as_deref() == Some("23505") + } + _ => false, + } +} + +async fn bench_pool(max_connections: u32) -> PgPool { + let options = PgConnectOptions::from_str(&setup::database_url()) + .expect("parse DATABASE_URL") + .application_name(APPLICATION_NAME); + PgPoolOptions::new() + .max_connections(max_connections) + .connect_with(options) + .await + .expect("connect bench pool") +} + +async fn waiter_count(pool: &PgPool) -> i64 { + sqlx::query_scalar( + r#" + SELECT count(*)::bigint + FROM pg_stat_activity + WHERE application_name = $1 + AND wait_event_type = 'Lock' + AND wait_event = 'advisory' + "#, + ) + .bind(APPLICATION_NAME) + .fetch_one(pool) + .await + .unwrap_or(0) +} + +async fn seed_queue_shards(pool: &PgPool, queue: &str, shards: i16) { + sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(pool) + .await + .expect("seed queue_meta.enqueue_shards"); +} + +/// One statement that calls `awa.insert_job_compat` once per row, the shape +/// batch inserts had before `awa.insert_jobs_compat_batch`. It measures what a +/// batch costs the lane lock on this runner when every row goes through the +/// per-row function, which is the yardstick the gate compares `insert_many` +/// against. +const LEGACY_ROW_FUNCTION_SQL: &str = r#" + SELECT inserted.* + FROM unnest( + $1::text[], $2::text[], $3::jsonb[], $4::text[], $5::smallint[], $6::smallint[], + $7::timestamptz[], $8::jsonb[], $9::jsonb[], $10::bytea[], $11::text[] + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + CROSS JOIN LATERAL awa.insert_job_compat( + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL + ) AS inserted + ORDER BY input.ord +"#; + +async fn legacy_row_function_insert(pool: &PgPool, jobs: &[InsertParams]) -> Result<(), AwaError> { + let rows: Vec<_> = jobs + .iter() + .map(|job| { + prepare_raw_job_insert(job.kind.clone(), job.args.clone(), job.opts.clone()) + .expect("prepare row") + }) + .collect(); + sqlx::query(LEGACY_ROW_FUNCTION_SQL) + .bind(rows.iter().map(|row| row.kind()).collect::>()) + .bind(rows.iter().map(|row| row.queue()).collect::>()) + .bind(rows.iter().map(|row| row.args()).collect::>()) + .bind( + rows.iter() + .map(|row| row.state_db_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.priority()).collect::>()) + .bind( + rows.iter() + .map(|row| row.max_attempts()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.run_at()).collect::>()) + .bind(rows.iter().map(|row| row.metadata()).collect::>()) + .bind( + rows.iter() + .map(|row| serde_json::json!(row.tags())) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.unique_key()).collect::>()) + .bind( + rows.iter() + .map(|row| row.unique_states_bit_string()) + .collect::>(), + ) + .fetch_all(pool) + .await?; + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn producer_loop( + pool: PgPool, + store: Arc, + scenario: Scenario, + queue: String, + tag: String, + producer: usize, + measure_from: Instant, + stop_at: Instant, + counter_base: Arc, + stop: Arc, +) -> Samples { + let mut samples = Samples::default(); + let mut batch_index: u64 = 0; + let mut previous_batch: Vec = Vec::new(); + + while Instant::now() < stop_at && !stop.load(Ordering::Relaxed) { + let base = counter_base.fetch_add(scenario.batch_size as u64, Ordering::Relaxed); + let mut jobs: Vec = (0..scenario.batch_size as u64) + .map(|offset| job_params(&queue, &tag, producer, base + offset, scenario.unique)) + .collect(); + + // A replayed key: the first job repeats one this producer already + // inserted, so the batch hits exactly one unique conflict. + let duplicate = scenario.unique + && scenario.duplicate_batch_pct > 0 + && !previous_batch.is_empty() + && (batch_index + 1) * u64::from(scenario.duplicate_batch_pct) / 100 + > batch_index * u64::from(scenario.duplicate_batch_pct) / 100; + if duplicate { + jobs[0] = previous_batch[0].clone(); + } + batch_index += 1; + + let started = Instant::now(); + let mut statement_ms = None; + let mut commit_ms = None; + let mut fallback_rows_ms: Vec = Vec::new(); + + let outcome: Result<(), AwaError> = match scenario.mode { + Mode::InsertMany if scenario.split_commit => { + let mut tx = pool.begin().await.expect("begin"); + let statement_started = Instant::now(); + let result = insert_many(&mut *tx, &jobs).await; + let statement_elapsed = statement_started.elapsed(); + match result { + Ok(_) => { + let commit_started = Instant::now(); + tx.commit().await.expect("commit"); + statement_ms = Some(statement_elapsed.as_secs_f64() * 1e3); + commit_ms = Some(commit_started.elapsed().as_secs_f64() * 1e3); + Ok(()) + } + Err(err) => Err(err), + } + } + Mode::InsertMany => insert_many(&pool, &jobs).await.map(|_| ()), + Mode::InsertManyCopy => insert_many_copy_from_pool(&pool, &jobs).await.map(|_| ()), + Mode::QueueStorage => store.enqueue_params_batch(&pool, &jobs).await.map(|_| ()), + Mode::LegacyRowFunction => legacy_row_function_insert(&pool, &jobs).await, + }; + + let mut fallback_count = 0u64; + match outcome { + Ok(()) => {} + Err(err) if is_unique_violation(&err) => { + // What a replaying producer does on a unique conflict: + // insert the batch again one row at a time, skipping + // duplicates. + for job in &jobs { + let row = prepare_raw_job_insert( + job.kind.clone(), + job.args.clone(), + job.opts.clone(), + ) + .expect("prepare row"); + let row_started = Instant::now(); + match insert_prepared(&pool, &row).await { + Ok(_) => {} + Err(err) if is_unique_violation(&err) => {} + Err(err) => panic!("fallback insert failed: {err}"), + } + fallback_rows_ms.push(row_started.elapsed().as_secs_f64() * 1e3); + fallback_count += 1; + } + } + Err(err) => panic!("enqueue failed: {err}"), + } + let elapsed = started.elapsed(); + + if started >= measure_from { + samples.batch_ms.push(elapsed.as_secs_f64() * 1e3); + samples.jobs += scenario.batch_size as u64; + samples.batches += 1; + if let Some(ms) = statement_ms { + samples.statement_ms.push(ms); + } + if let Some(ms) = commit_ms { + samples.commit_ms.push(ms); + } + if fallback_count > 0 { + samples.duplicate_batches += 1; + samples.fallback_rows += fallback_count; + samples.fallback_row_ms.extend(fallback_rows_ms); + } + } + previous_batch = jobs; + } + samples +} + +async fn run_scenario(pool: &PgPool, scenario: &Scenario) -> ScenarioResult { + let queue = format!( + "lane_bench_{}", + Uuid::new_v4().simple().to_string().get(..12).unwrap() + ); + let tag = Uuid::new_v4().simple().to_string(); + seed_queue_shards(pool, &queue, scenario.shards).await; + + let producer_pool = bench_pool(scenario.producers as u32 + 4).await; + let store = Arc::new(QueueStorage::new(QueueStorageConfig::default()).expect("queue storage")); + + let started = Instant::now(); + let measure_from = started + WARMUP; + let stop_at = measure_from + scenario.measured; + let counter_base = Arc::new(AtomicU64::new(0)); + let stop = Arc::new(AtomicBool::new(false)); + + let mut tasks = Vec::new(); + for producer in 0..scenario.producers { + tasks.push(tokio::spawn(producer_loop( + producer_pool.clone(), + Arc::clone(&store), + scenario.clone(), + queue.clone(), + tag.clone(), + producer, + measure_from, + stop_at, + Arc::clone(&counter_base), + Arc::clone(&stop), + ))); + } + + // Sample advisory-lock waiters over the measured window only. + tokio::time::sleep_until(tokio::time::Instant::from_std(measure_from)).await; + let mut waiter_samples: Vec = Vec::new(); + while Instant::now() < stop_at { + waiter_samples.push(waiter_count(pool).await); + tokio::time::sleep(Duration::from_millis(25)).await; + } + + let mut samples = Samples::default(); + for task in tasks { + samples.merge(task.await.expect("producer task")); + } + producer_pool.close().await; + + let (landed, distinct_lane_seqs): (i64, i64) = sqlx::query_as( + r#" + SELECT count(*)::bigint, + count(DISTINCT (enqueue_shard, lane_seq))::bigint + FROM awa.ready_entries + WHERE queue = $1 + "#, + ) + .bind(&queue) + .fetch_one(pool) + .await + .expect("count landed rows"); + + // Claim everything that landed, from every shard, with the claim path the + // workers use. + let drained = if scenario.drain { + let drain_started = Instant::now(); + let mut claimed = 0i64; + loop { + let batch = store + .claim_runtime_batch(pool, &queue, 500, Duration::from_secs(3_600)) + .await + .expect("claim batch"); + if batch.is_empty() { + break; + } + claimed += batch.len() as i64; + } + Some((claimed, drain_started.elapsed().as_secs_f64())) + } else { + None + }; + + let latencies = sorted(&samples.batch_ms); + let window = scenario.measured.as_secs_f64(); + ScenarioResult { + scenario: scenario.clone(), + jobs_per_s: samples.jobs as f64 / window, + p50_ms: percentile(&latencies, 0.50), + p95_ms: percentile(&latencies, 0.95), + p99_ms: percentile(&latencies, 0.99), + max_ms: latencies.last().copied().unwrap_or(0.0), + mean_waiters: if waiter_samples.is_empty() { + 0.0 + } else { + waiter_samples.iter().sum::() as f64 / waiter_samples.len() as f64 + }, + samples, + landed, + distinct_lane_seqs, + drained, + } +} + +fn print_result(result: &ScenarioResult) { + let scenario = &result.scenario; + println!( + "[lane-bench] {name} mode={mode} producers={producers} batch={batch} unique={unique} \ + shards={shards} dup_pct={dup} | {rate:.0} jobs/s | batch ms p50={p50:.1} p95={p95:.1} \ + p99={p99:.1} max={max:.1} | waiters_mean={waiters:.1} | landed={landed} \ + distinct_seqs={distinct}", + name = scenario.name, + mode = scenario.mode.label(), + producers = scenario.producers, + batch = scenario.batch_size, + unique = scenario.unique, + shards = scenario.shards, + dup = scenario.duplicate_batch_pct, + rate = result.jobs_per_s, + p50 = result.p50_ms, + p95 = result.p95_ms, + p99 = result.p99_ms, + max = result.max_ms, + waiters = result.mean_waiters, + landed = result.landed, + distinct = result.distinct_lane_seqs, + ); + if !result.samples.statement_ms.is_empty() { + let statement = sorted(&result.samples.statement_ms); + let commit = sorted(&result.samples.commit_ms); + println!( + "[lane-bench] split: statement ms p50={:.1} p99={:.1} | commit ms p50={:.2} p99={:.2}", + percentile(&statement, 0.5), + percentile(&statement, 0.99), + percentile(&commit, 0.5), + percentile(&commit, 0.99), + ); + } + if let Some((claimed, seconds)) = result.drained { + println!( + "[lane-bench] drain: claimed {claimed} of {} landed in {seconds:.2}s ({:.0}/s)", + result.landed, + claimed as f64 / seconds.max(1e-9), + ); + } + if result.samples.duplicate_batches > 0 { + let rows = sorted(&result.samples.fallback_row_ms); + println!( + "[lane-bench] fallback: {} duplicate batches, {} single-row statements, \ + row ms p50={:.1} p99={:.1} max={:.1}", + result.samples.duplicate_batches, + result.samples.fallback_rows, + percentile(&rows, 0.5), + percentile(&rows, 0.99), + rows.last().copied().unwrap_or(0.0), + ); + } +} + +fn emit_json(result: &ScenarioResult) { + let mut outcomes = HashMap::new(); + outcomes.insert("jobs".to_string(), result.samples.jobs); + outcomes.insert("batches".to_string(), result.samples.batches); + BenchmarkResult { + schema_version: SCHEMA_VERSION, + scenario: result.scenario.name.clone(), + language: "rust".to_string(), + seeded: result.landed.max(0) as u64, + metrics: BenchMetrics { + throughput: None, + enqueue_per_s: Some(result.jobs_per_s), + drain_time_s: None, + latency_ms: Some(BenchLatency { + p50: Some(result.p50_ms), + p95: Some(result.p95_ms), + p99: Some(result.p99_ms), + }), + rescue: None, + }, + outcomes, + metadata: Some(serde_json::json!({ + "producers": result.scenario.producers, + "batch_size": result.scenario.batch_size, + "unique": result.scenario.unique, + "shards": result.scenario.shards, + "mean_lock_waiters": result.mean_waiters, + })), + } + .emit(); +} + +/// Lane integrity: every job that was acknowledged landed exactly once, under +/// a distinct `(shard, lane_seq)`, and every shard's sequence numbers are +/// dense (no producer lost or duplicated a reservation). +fn assert_lane_integrity(result: &ScenarioResult) { + if let Some((claimed, _)) = result.drained { + assert_eq!( + claimed, result.landed, + "{}: claim path did not drain every shard", + result.scenario.name + ); + } + assert_eq!( + result.landed, result.distinct_lane_seqs, + "{}: ready rows share a (shard, lane_seq)", + result.scenario.name + ); + assert!( + result.landed as u64 >= result.samples.jobs, + "{}: {} jobs acknowledged in the window but only {} rows landed", + result.scenario.name, + result.samples.jobs, + result.landed + ); +} + +async fn bench_setup() -> PgPool { + let pool = setup::setup(16).await; + setup::activate_queue_storage(&pool).await; + pool +} + +fn scenario( + name: &str, + producers: usize, + shards: i16, + unique: bool, + duplicate_batch_pct: u32, +) -> Scenario { + Scenario { + name: name.to_string(), + mode: Mode::InsertMany, + producers, + batch_size: 185, + unique, + shards, + duplicate_batch_pct, + measured: Duration::from_secs(8), + split_commit: false, + drain: false, + } +} + +/// Fixed scenarios whose results gate regressions in the nightly run. +/// +/// Absolute floors sit well below the rates a shared runner sustains, and the +/// latency ceiling well above, so they only catch a lane that has stopped +/// making progress; `benchmarks/baseline.json` carries the regression +/// thresholds on rate and p99. The one-producer scenarios measure what a batch +/// costs the lane lock, through `insert_many` and through the per-row +/// function, and are reported for comparison. Every bound goes through +/// `ci_timing` so a contended runner loosens it. +#[tokio::test(flavor = "multi_thread", worker_threads = 6)] +#[ignore] +async fn test_enqueue_lane_contention_gate() { + let pool = bench_setup().await; + + let mut row_function = scenario("enqueue_lane_1p_row_function_185", 1, 1, true, 0); + row_function.mode = Mode::LegacyRowFunction; + row_function.measured = Duration::from_secs(4); + let row_function = run_scenario(&pool, &row_function).await; + print_result(&row_function); + assert_lane_integrity(&row_function); + + let mut one_producer = scenario("enqueue_lane_1p_unique_185", 1, 1, true, 0); + one_producer.measured = Duration::from_secs(4); + let one_producer = run_scenario(&pool, &one_producer).await; + print_result(&one_producer); + assert_lane_integrity(&one_producer); + + let single_lane = run_scenario( + &pool, + &scenario("enqueue_lane_16p_unique_185", 16, 1, true, 0), + ) + .await; + print_result(&single_lane); + emit_json(&single_lane); + assert_lane_integrity(&single_lane); + + let sharded = run_scenario( + &pool, + &scenario("enqueue_lane_16p_unique_185_shards4", 16, 4, true, 0), + ) + .await; + print_result(&sharded); + emit_json(&sharded); + assert_lane_integrity(&sharded); + + let floor = ci_timing::contention_floor(GATE_MIN_JOBS_PER_S); + assert!( + single_lane.jobs_per_s >= floor as f64, + "single-lane enqueue rate {:.0}/s fell below the floor of {floor}/s", + single_lane.jobs_per_s + ); + let p99_ceiling = ci_timing::scaled_timeout(Duration::from_millis(GATE_MAX_P99_MS)); + assert!( + Duration::from_secs_f64(single_lane.p99_ms / 1e3) <= p99_ceiling, + "single-lane p99 batch latency {:.0}ms exceeded {:?}", + single_lane.p99_ms, + p99_ceiling + ); +} + +const GATE_MIN_JOBS_PER_S: i64 = 2_000; +const GATE_MAX_P99_MS: u64 = 15_000; + +/// Exploration sweep; see the module docs for the tunables. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_enqueue_lane_contention_matrix() { + let pool = bench_setup().await; + let producers: Vec = env_string("AWA_LANE_BENCH_PRODUCERS", "1,4,8,16,32") + .split(',') + .filter_map(|raw| raw.trim().parse().ok()) + .collect(); + let mode = Mode::parse(&env_string("AWA_LANE_BENCH_MODE", "insert_many")); + let batch_size = env_usize("AWA_LANE_BENCH_BATCH", 185); + let unique = env_usize("AWA_LANE_BENCH_UNIQUE", 1) != 0; + let shards = env_usize("AWA_LANE_BENCH_SHARDS", 1) as i16; + let duplicate_pct = env_usize("AWA_LANE_BENCH_DUP_PCT", 0) as u32; + let secs = env_usize("AWA_LANE_BENCH_SECS", 8) as u64; + let split = env_usize("AWA_LANE_BENCH_SPLIT", 0) != 0; + + for producer_count in producers { + let result = run_scenario( + &pool, + &Scenario { + name: format!("matrix_{producer_count}p"), + mode, + producers: producer_count, + batch_size, + unique, + shards, + duplicate_batch_pct: duplicate_pct, + measured: Duration::from_secs(secs), + split_commit: split, + drain: env_usize("AWA_LANE_BENCH_DRAIN", 0) != 0, + }, + ) + .await; + print_result(&result); + assert_lane_integrity(&result); + } +} diff --git a/awa/tests/enqueue_lane_maintenance_coupling_test.rs b/awa/tests/enqueue_lane_maintenance_coupling_test.rs new file mode 100644 index 000000000..67a55808a --- /dev/null +++ b/awa/tests/enqueue_lane_maintenance_coupling_test.rs @@ -0,0 +1,429 @@ +//! How a saturated enqueue lane affects the maintenance leader. +//! +//! Promotion of scheduled and retryable jobs re-enqueues them into their lane, +//! so it queues on the same per-lane advisory lock as producers. A promotion +//! batch is one transaction across every queue it touches, and the leader runs +//! its branches (promotion, rescue, rotation, ...) one after another on a +//! single task, so a batch parked behind a hot lane holds up the due jobs of +//! other queues in that batch and every other maintenance branch. +//! +//! `test_maintenance_under_hot_lane` measures, with and without concurrent +//! producers on one hot lane: +//! * promotion batch latency and how long a backlog on the hot queue takes +//! to drain, +//! * promotion lag on an unrelated queue (time from a job becoming due to +//! it being visible as ready), with the promotion loop shaped like the +//! leader's (a tick, then batches until one comes back short), +//! * heartbeat rescue of abandoned claims on the hot queue: how long each run +//! takes and how late it starts once its tick is due (1s here, so a +//! blocked promotion shows up as a late start). +//! +//! Tunables (optional): +//! AWA_LANE_MAINT_PRODUCERS producers on the hot lane, default 16 +//! AWA_LANE_MAINT_BACKLOG due scheduled jobs on the hot queue, default 24576 +//! AWA_LANE_MAINT_RESCUE abandoned claims on the hot queue, default 6000 +//! +//! Requires a running Postgres instance. +//! Run with: +//! cargo test --release --package awa --test enqueue_lane_maintenance_coupling_test \ +//! -- --ignored --nocapture --test-threads=1 + +use awa::model::{insert_many, QueueStorage}; +use awa::{InsertOpts, InsertParams, JobState, QueueStorageConfig}; +use awa_testing::setup; +use sqlx::PgPool; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use uuid::Uuid; + +const PROMOTE_BATCH: i64 = 4_096; +const PROMOTE_TICK: Duration = Duration::from_millis(250); +const PROBE_PERIOD: Duration = Duration::from_millis(100); +const RESCUE_TICK: Duration = Duration::from_secs(1); + +fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|raw| raw.parse().ok()) + .unwrap_or(default) +} + +fn percentile(values: &[f64], quantile: f64) -> f64 { + if values.is_empty() { + return 0.0; + } + let mut sorted = values.to_vec(); + sorted.sort_by(|left, right| left.partial_cmp(right).expect("finite")); + sorted[((sorted.len() as f64 - 1.0) * quantile).round() as usize] +} + +fn params( + queue: &str, + tag: &str, + n: u64, + run_at: Option>, +) -> InsertParams { + InsertParams { + kind: "lane_maint_job".to_string(), + args: serde_json::json!({ "tag": tag, "n": n }), + opts: InsertOpts { + queue: queue.to_string(), + run_at, + ..Default::default() + }, + } +} + +async fn seed_due_backlog(pool: &PgPool, queue: &str, count: usize) { + let due = chrono::Utc::now() - chrono::Duration::seconds(60); + let mut inserted = 0usize; + while inserted < count { + let chunk = (count - inserted).min(2_000); + let jobs: Vec = (0..chunk) + .map(|offset| params(queue, "backlog", (inserted + offset) as u64, Some(due))) + .collect(); + insert_many(pool, &jobs).await.expect("seed backlog"); + inserted += chunk; + } +} + +async fn deferred_due(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.deferred_jobs WHERE queue = $1 AND state = 'scheduled'", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("count deferred") +} + +struct Producers { + stop: Arc, + tasks: Vec>, +} + +impl Producers { + fn start(pool: &PgPool, queue: &str, count: usize) -> Self { + let stop = Arc::new(AtomicBool::new(false)); + let tag = Uuid::new_v4().simple().to_string(); + let tasks = (0..count) + .map(|producer| { + let pool = pool.clone(); + let stop = Arc::clone(&stop); + let queue = queue.to_string(); + let tag = tag.clone(); + tokio::spawn(async move { + let mut jobs_inserted = 0u64; + let mut counter = (producer as u64) << 40; + while !stop.load(Ordering::Relaxed) { + let jobs: Vec = (0..185) + .map(|offset| params(&queue, &tag, counter + offset, None)) + .collect(); + counter += 185; + insert_many(&pool, &jobs).await.expect("producer insert"); + jobs_inserted += 185; + } + jobs_inserted + }) + }) + .collect(); + Self { stop, tasks } + } + + async fn stop(self) -> u64 { + self.stop.store(true, Ordering::Relaxed); + let mut total = 0; + for task in self.tasks { + total += task.await.expect("producer task"); + } + total + } +} + +#[derive(Default)] +struct LeaderReport { + /// Promotion batches that moved a full share of the backlog. + batch_ms: Vec, + drain_s: f64, + probe_lag_ms: Vec, + /// How long after its tick each heartbeat-rescue run started, and how long + /// the run itself took. + rescue_start_delay_ms: Vec, + rescue_run_ms: Vec, +} + +/// A leader loop shaped like the maintenance leader's: one task selecting over +/// a promotion timer (250ms; on each tick batches of up to 4096 until one is +/// short) and a heartbeat-rescue timer, running each branch to completion +/// before taking the next tick. Runs until the hot queue's backlog is gone and +/// `min_window` has passed, while a probe feeds one due job per 100ms into an +/// unrelated queue and times how long each takes to become ready. +async fn run_leader( + pool: &PgPool, + store: &QueueStorage, + hot: &str, + other: &str, + min_window: Duration, +) -> LeaderReport { + let mut report = LeaderReport::default(); + let started = Instant::now(); + let stop = Arc::new(AtomicBool::new(false)); + + let probe = { + let pool = pool.clone(); + let other = other.to_string(); + let stop = Arc::clone(&stop); + tokio::spawn(async move { + let mut lags = Vec::new(); + let mut n = 0u64; + while !stop.load(Ordering::Relaxed) { + n += 1; + // Older than the hot backlog, so the probe is always in the + // first promotion batch rather than behind the backlog's + // run_at order. + let due_at = chrono::Utc::now() - chrono::Duration::seconds(120); + insert_many(&pool, &[params(&other, "probe", n, Some(due_at))]) + .await + .expect("insert probe"); + let inserted = Instant::now(); + loop { + let visible: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM awa.ready_entries + WHERE queue = $1 AND args->>'tag' = 'probe' AND (args->>'n')::bigint = $2)", + ) + .bind(&other) + .bind(n as i64) + .fetch_one(&pool) + .await + .expect("probe visibility"); + if visible { + lags.push(inserted.elapsed().as_secs_f64() * 1e3); + break; + } + if stop.load(Ordering::Relaxed) { + return lags; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + tokio::time::sleep(PROBE_PERIOD).await; + } + lags + }) + }; + + let mut promote_timer = tokio::time::interval(PROMOTE_TICK); + promote_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut rescue_timer = tokio::time::interval(RESCUE_TICK); + rescue_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + promote_timer.tick().await; + rescue_timer.tick().await; + + let mut backlog_gone_at: Option = None; + loop { + tokio::select! { + tick = promote_timer.tick() => { + let _ = tick; + loop { + let batch_started = Instant::now(); + let promoted = store + .promote_due(pool, JobState::Scheduled, PROMOTE_BATCH) + .await + .expect("promote_due"); + if promoted >= 1_000 { + report.batch_ms.push(batch_started.elapsed().as_secs_f64() * 1e3); + } + if (promoted as i64) < PROMOTE_BATCH { + break; + } + } + if backlog_gone_at.is_none() && deferred_due(pool, hot).await == 0 { + backlog_gone_at = Some(Instant::now()); + } + if backlog_gone_at.is_some() && started.elapsed() >= min_window { + break; + } + } + tick = rescue_timer.tick() => { + let delay = tokio::time::Instant::now().saturating_duration_since(tick); + let run_started = Instant::now(); + let rescued = store + .rescue_stale_heartbeats(pool, Duration::ZERO) + .await + .expect("rescue_stale_heartbeats"); + if !rescued.is_empty() { + report.rescue_start_delay_ms.push(delay.as_secs_f64() * 1e3); + report.rescue_run_ms.push(run_started.elapsed().as_secs_f64() * 1e3); + } + } + } + } + report.drain_s = backlog_gone_at + .expect("backlog drained") + .duration_since(started) + .as_secs_f64(); + stop.store(true, Ordering::Relaxed); + report.probe_lag_ms = probe.await.expect("probe task"); + report +} + +fn max_of(values: &[f64]) -> f64 { + values.iter().copied().fold(0.0, f64::max) +} + +fn print_leader(label: &str, report: &LeaderReport) { + println!( + "[lane-maint] {label} promote batch ms p50={:.0} max={:.0} ({} batches) | hot backlog \ + drained in {:.2}s | unrelated-queue promotion lag ms p50={:.0} p99={:.0} max={:.0} ({} probes) \ + | rescue ran {} times, run ms p50={:.0} max={:.0}, started late by ms p50={:.0} max={:.0}", + percentile(&report.batch_ms, 0.5), + max_of(&report.batch_ms), + report.batch_ms.len(), + report.drain_s, + percentile(&report.probe_lag_ms, 0.5), + percentile(&report.probe_lag_ms, 0.99), + max_of(&report.probe_lag_ms), + report.probe_lag_ms.len(), + report.rescue_run_ms.len(), + percentile(&report.rescue_run_ms, 0.5), + max_of(&report.rescue_run_ms), + percentile(&report.rescue_start_delay_ms, 0.5), + max_of(&report.rescue_start_delay_ms), + ); +} + +/// Enqueue `count` jobs on `queue` and claim every one of them without ever +/// completing or heartbeating, so each becomes an abandoned attempt. +async fn seed_abandoned_claims(pool: &PgPool, store: &QueueStorage, queue: &str, count: usize) { + let tag = Uuid::new_v4().simple().to_string(); + let mut enqueued = 0usize; + while enqueued < count { + let chunk = (count - enqueued).min(2_000); + let jobs: Vec = (0..chunk) + .map(|offset| params(queue, &tag, (enqueued + offset) as u64, None)) + .collect(); + store + .enqueue_params_batch(pool, &jobs) + .await + .expect("enqueue claimable jobs"); + enqueued += chunk; + } + let mut claimed = 0usize; + while claimed < count { + let batch = store + .claim_runtime_batch(pool, queue, 500, Duration::from_secs(3_600)) + .await + .expect("claim batch"); + assert!(!batch.is_empty(), "ran out of jobs to claim"); + claimed += batch.len(); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_maintenance_under_hot_lane() { + let producers = env_usize("AWA_LANE_MAINT_PRODUCERS", 16); + let backlog = env_usize("AWA_LANE_MAINT_BACKLOG", 24_576); + let rescue_claims = env_usize("AWA_LANE_MAINT_RESCUE", 6_000); + + let pool = setup::setup(producers as u32 + 12).await; + setup::activate_queue_storage(&pool).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + + for (label, producer_count) in [("idle lane ", 0usize), ("hot lane ", producers)] { + let suffix = Uuid::new_v4().simple().to_string(); + let hot = format!("lane_maint_hot_{}", &suffix[..10]); + let other = format!("lane_maint_other_{}", &suffix[..10]); + + // Warm both lanes so lane provisioning is not measured. + insert_many( + &pool, + &[ + params(&hot, "warm", 0, None), + params(&other, "warm", 0, None), + ], + ) + .await + .expect("warm lanes"); + + seed_due_backlog(&pool, &hot, backlog).await; + seed_abandoned_claims(&pool, &store, &hot, rescue_claims).await; + + let running = (producer_count > 0).then(|| Producers::start(&pool, &hot, producer_count)); + tokio::time::sleep(Duration::from_millis(1_500)).await; + + let report = run_leader(&pool, &store, &hot, &other, Duration::from_secs(12)).await; + print_leader(label, &report); + + if let Some(running) = running { + let produced = running.stop().await; + println!("[lane-maint] {label} producers inserted {produced} jobs meanwhile"); + } + } +} + +/// The mechanism without producer load: one transaction holds the hot lane's +/// reservation lock for a fixed time while the leader loop runs. Promotion of +/// a batch that contains hot-queue rows waits for it, and takes the unrelated +/// queue's due job in the same batch and the leader's other branches with it. +/// Heartbeat rescue called directly does not wait for the lane at all: it +/// writes retryable rows to the deferred table, which promotion later moves. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_leader_behind_a_held_lane_lock() { + let hold = Duration::from_millis(env_usize("AWA_LANE_MAINT_HOLD_MS", 3_000) as u64); + let pool = setup::setup(16).await; + setup::activate_queue_storage(&pool).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + + let suffix = Uuid::new_v4().simple().to_string(); + let hot = format!("lane_maint_hot_{}", &suffix[..10]); + let other = format!("lane_maint_other_{}", &suffix[..10]); + insert_many( + &pool, + &[ + params(&hot, "warm", 0, None), + params(&other, "warm", 0, None), + ], + ) + .await + .expect("warm lanes"); + seed_due_backlog(&pool, &hot, 8_192).await; + seed_abandoned_claims(&pool, &store, &hot, 3_000).await; + + let key: i64 = sqlx::query_scalar( + "SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, 2::smallint, 0::smallint)), + 0)", + ) + .bind(&hot) + .fetch_one(&pool) + .await + .expect("derive lane lock key"); + let mut gate = pool.begin().await.expect("begin gate"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(key) + .execute(&mut *gate) + .await + .expect("take lane lock"); + + let rescue_started = Instant::now(); + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue with the lane lock held"); + println!( + "[lane-maint] rescue of {} claims with the hot lane lock held took {:.0}ms", + rescued.len(), + rescue_started.elapsed().as_secs_f64() * 1e3 + ); + let remaining_hold = hold.saturating_sub(rescue_started.elapsed()); + + let release = tokio::spawn(async move { + tokio::time::sleep(remaining_hold).await; + gate.commit().await.expect("release lane lock"); + }); + let report = run_leader(&pool, &store, &hot, &other, Duration::from_secs(9)).await; + release.await.expect("release task"); + print_leader(&format!("lane held {}ms", hold.as_millis()), &report); +} diff --git a/benchmarks/baseline.json b/benchmarks/baseline.json index 9f7fee869..3fdb7bd08 100644 --- a/benchmarks/baseline.json +++ b/benchmarks/baseline.json @@ -32,6 +32,13 @@ "rust/stale_heartbeat_rescue": { "handler_per_s": 160.3 }, + "rust/enqueue_lane_16p_unique_185": { + "enqueue_per_s": 2000.0, + "p99_ms": 6000.0 + }, + "rust/enqueue_lane_16p_unique_185_shards4": { + "enqueue_per_s": 2000.0 + }, "rust/mvcc_horizon_overlap": { "min_overlap_vs_baseline": 0.75, "min_cooldown_vs_baseline": 0.75, From 91f3a137cb04749ba1bfd63543c7174a72e71922 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 11:07:07 +1300 Subject: [PATCH 04/18] perf(insert): hold a lane lock only for the lane write in batch inserts insert_many and insert_many_copy called awa.insert_job_compat once per row. The function takes the (queue, priority, shard) lane lock in reserve_enqueue_seq on the batch's first row, and the lock lasts to commit, so it was held while every remaining row was processed. A 185-row batch held it for about 40 ms on a local database, and producers on one lane queued behind whole batches: the lane's ceiling was roughly 4,000 jobs/s at any producer count, and a producer's call took (producers x 40 ms). Migration v053 adds awa.insert_jobs_compat_batch. It allocates job ids, claims unique keys in key order, registers new lanes and builds the rows before it reserves each lane's sequence range once, in (queue, priority, shard) order, then writes the ready rows and segments in two statements. It proposes head rows only for lanes that do not exist yet, because the head tables' row trigger takes the lane lock before the conflict check. insert_many and insert_many_copy call it. insert_many_copy now skips unique conflicts in one statement instead of one savepoint per row, and insert_many_skip_unique_conflicts exposes the same behaviour for callers that replay work: after a unique conflict they inserted the batch again one job at a time, taking the lane lock once per job. The lane-order test now parks a batch on one lane and checks which lane locks it holds, since job ids no longer follow lock order. --- CHANGELOG.md | 33 +- .../v053_insert_jobs_batch_compat.sql | 575 ++++++++++++++++++ awa-model/src/insert.rs | 357 ++++------- awa-model/src/lib.rs | 3 +- awa-model/src/migrations.rs | 8 +- awa/tests/enqueue_lane_contention_test.rs | 48 +- awa/tests/enqueue_lane_lock_order_test.rs | 188 +++++- awa/tests/insert_many_batch_test.rs | 426 +++++++++++++ awa/tests/integration_test.rs | 3 + benchmarks/baseline.json | 6 +- docs/troubleshooting.md | 20 + docs/upgrade-0.6-to-0.7.md | 19 + 12 files changed, 1396 insertions(+), 290 deletions(-) create mode 100644 awa-model/migrations/v053_insert_jobs_batch_compat.sql create mode 100644 awa/tests/insert_many_batch_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 065d50db1..a09350c9e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,34 @@ Notable changes between releases. Detailed migration notes for storage transitio ## [Unreleased] +### Added + +- **`insert_many_skip_unique_conflicts` inserts a batch and skips the jobs + whose unique key is already claimed.** `insert_many` fails the whole batch + with a unique violation, so a producer that replays work had to insert the + batch again one job at a time, taking the lane lock once per job. The new + function returns the inserted rows in input order, with skipped jobs absent. + +### Changed + +- **Batch enqueues hold a lane lock for the lane write only.** `insert_many` + and `insert_many_copy` called `awa.insert_job_compat` per row, which takes + the `(queue, priority, shard)` lane lock on the batch's first row and holds + it, because the lock lasts to commit, while every remaining row is + processed: concurrent producers on one lane queued behind whole batches. + Migration v053 adds `awa.insert_jobs_compat_batch`, which allocates job ids, + claims unique keys (in key order), registers new lanes and builds the rows + first, then reserves each lane's sequence range once, in lane order, and + writes the ready rows and segments in two statements. A 185-row batch into + one lane holds the lock for a fraction of the time it did, which raises the + per-lane ceiling several times over and cuts the time producers, promotion + and rescue spend waiting for the lane. `insert_many_copy` skips unique + conflicts in one pass instead of one savepoint per row. Run `awa migrate` + before upgrading producers: a binary that calls the new function against a + database without v053 fails with `function awa.insert_jobs_compat_batch ... + does not exist`. Binaries before this release keep calling + `awa.insert_job_compat` and share the lane lock with the new path. + ### Fixed - **Batch enqueues take their lane locks in one order, so they no longer @@ -15,10 +43,7 @@ Notable changes between releases. Detailed migration notes for storage transitio a `QueueStorage` enqueue could deadlock with it (`40P01`). The two insert paths now sort rows by `(queue, priority, shard)`, queues compared bytewise, before they reach `awa.insert_job_compat`, the order `QueueStorage` already - used, and still return rows in input order. Job ids within a batch that - spans several lanes are now allocated lane by lane rather than in input - order, and when an `insert_many_copy` batch carries the same unique key in - more than one lane, the row in the lowest lane is the one inserted. A + used, and still return rows in input order. A transaction that enqueues with separate statements still takes locks in statement order; see "`deadlock detected` while enqueueing" in the troubleshooting guide. diff --git a/awa-model/migrations/v053_insert_jobs_batch_compat.sql b/awa-model/migrations/v053_insert_jobs_batch_compat.sql new file mode 100644 index 000000000..a6f67912a --- /dev/null +++ b/awa-model/migrations/v053_insert_jobs_batch_compat.sql @@ -0,0 +1,575 @@ +-- v053: Batch producer insert that holds each enqueue lane lock only for the +-- lane write. +-- +-- awa.insert_job_compat takes the (queue, priority, shard) lane lock in +-- reserve_enqueue_seq on a batch's first row and, because the lock is +-- transaction-scoped, holds it while every remaining row is processed. Many +-- concurrent producers on one lane therefore queue behind whole batches. This +-- migration adds awa.insert_jobs_compat_batch, which does the lane-independent +-- work first (job ids, unique claims in key order, lane rows, payloads), then +-- reserves each lane's sequence range once, in (queue, priority, shard) order, +-- and writes the ready rows and segments in two statements. The lock is held +-- across the reservation, those two statements and the commit. +-- +-- The column arrays are parallel and must have equal length. Every element is +-- required except run_at, unique_key, unique_states and ordering_hash; +-- insert_job_compat's per-row defaults are the caller's to apply. +-- +-- p_skip_unique_conflicts inserts a batch while skipping rows whose unique key +-- is already claimed (the rows are absent from the result) instead of failing +-- the whole batch with a unique violation. +-- +-- N-1 binaries: unchanged. They keep calling awa.insert_job_compat, whose +-- definition this migration does not touch, and the two paths share the lane +-- lock, so producers on either version serialise exactly as before. +-- Binaries that call the new function need this migration first. +-- Live load: one CREATE OR REPLACE FUNCTION, no table DDL. + +CREATE OR REPLACE FUNCTION awa.insert_jobs_compat_batch( + p_kind TEXT[], + p_queue TEXT[], + p_args JSONB[], + p_state TEXT[], + p_priority SMALLINT[], + p_max_attempts SMALLINT[], + p_run_at TIMESTAMPTZ[], + p_metadata JSONB[], + p_tags JSONB[], + p_unique_key BYTEA[], + p_unique_states TEXT[], + p_ordering_hash TEXT[], + p_skip_unique_conflicts BOOLEAN DEFAULT FALSE +) +RETURNS SETOF awa.jobs +AS $$ +DECLARE + v_count INT := COALESCE(cardinality(p_kind), 0); + v_schema TEXT; + v_old_search_path TEXT; + v_now TIMESTAMPTZ := clock_timestamp(); + v_job_ids BIGINT[]; + v_skipped BIGINT[] := ARRAY[]::BIGINT[]; + v_lane RECORD; + v_lane_queues TEXT[] := ARRAY[]::TEXT[]; + v_lane_priorities SMALLINT[] := ARRAY[]::SMALLINT[]; + v_lane_shards SMALLINT[] := ARRAY[]::SMALLINT[]; + v_lane_starts BIGINT[] := ARRAY[]::BIGINT[]; + v_start BIGINT; + v_ready_slot INT; + v_ready_generation BIGINT; + v_row RECORD; + v_shards SMALLINT[]; +BEGIN + IF v_count = 0 THEN + RETURN; + END IF; + + IF cardinality(p_queue) IS DISTINCT FROM v_count + OR cardinality(p_args) IS DISTINCT FROM v_count + OR cardinality(p_state) IS DISTINCT FROM v_count + OR cardinality(p_priority) IS DISTINCT FROM v_count + OR cardinality(p_max_attempts) IS DISTINCT FROM v_count + OR cardinality(p_run_at) IS DISTINCT FROM v_count + OR cardinality(p_metadata) IS DISTINCT FROM v_count + OR cardinality(p_tags) IS DISTINCT FROM v_count + OR cardinality(p_unique_key) IS DISTINCT FROM v_count + OR cardinality(p_unique_states) IS DISTINCT FROM v_count + OR cardinality(p_ordering_hash) IS DISTINCT FROM v_count + THEN + RAISE EXCEPTION 'insert_jobs_compat_batch: column arrays must have equal length' + USING ERRCODE = '22023'; + END IF; + + IF EXISTS ( + SELECT 1 + FROM unnest(p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_metadata, p_tags) + AS input(kind, queue, args, state, priority, max_attempts, metadata, tags) + WHERE input.kind IS NULL OR input.queue IS NULL OR input.args IS NULL + OR input.state IS NULL OR input.priority IS NULL OR input.max_attempts IS NULL + OR input.metadata IS NULL OR input.tags IS NULL + ) THEN + RAISE EXCEPTION 'insert_jobs_compat_batch: kind, queue, args, state, priority, max_attempts, metadata and tags must not be NULL' + USING ERRCODE = '22004'; + END IF; + + IF EXISTS (SELECT 1 FROM unnest(p_kind) AS kind WHERE length(kind) > 200) THEN + RAISE EXCEPTION 'job kind length must be <= 200 characters' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS (SELECT 1 FROM unnest(p_queue) AS queue WHERE length(queue) > 200) THEN + RAISE EXCEPTION 'queue name length must be <= 200 characters' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_priority) AS priority + WHERE priority < 1 OR priority > 4 + ) THEN + RAISE EXCEPTION 'priority must be between 1 and 4' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_max_attempts) AS max_attempts + WHERE max_attempts < 1 OR max_attempts > 1000 + ) THEN + RAISE EXCEPTION 'max_attempts must be between 1 and 1000' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_tags) AS tags + WHERE jsonb_array_length(tags) > 20 + ) THEN + RAISE EXCEPTION 'job tags must contain at most 20 values' + USING ERRCODE = '23514'; + END IF; + + v_schema := awa.active_queue_storage_schema(); + + -- Canonical storage has no lane lock to shorten: insert row by row. + IF v_schema IS NULL THEN + IF NOT p_skip_unique_conflicts THEN + RETURN QUERY + SELECT inserted.* + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, + p_run_at, p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + CROSS JOIN LATERAL awa.insert_job_compat( + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL + ) AS inserted + ORDER BY input.ord; + RETURN; + END IF; + + -- Skipping needs a subtransaction per row. + FOR v_row IN + SELECT + input.ord, + input.kind, + input.queue, + input.args, + input.state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ) AS tags, + input.unique_key, + input.unique_states + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, + p_run_at, p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + ORDER BY input.ord + LOOP + BEGIN + RETURN NEXT awa.insert_job_compat( + v_row.kind, + v_row.queue, + v_row.args, + v_row.state::awa.job_state, + v_row.priority, + v_row.max_attempts, + v_row.run_at, + v_row.metadata, + v_row.tags, + v_row.unique_key, + v_row.unique_states::bit(8), + NULL + ); + EXCEPTION WHEN unique_violation THEN + NULL; + END; + END LOOP; + RETURN; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_state) AS state + WHERE state NOT IN ('available', 'scheduled', 'retryable') + ) THEN + RAISE EXCEPTION 'queue storage does not support initial state %', + (SELECT state FROM unnest(p_state) AS state + WHERE state NOT IN ('available', 'scheduled', 'retryable') LIMIT 1) + USING ERRCODE = '22023'; + END IF; + + v_old_search_path := current_setting('search_path'); + PERFORM set_config('search_path', format('%I,awa,public', v_schema), true); + + -- Everything up to the lane reservation is lane-independent and runs + -- without a lane lock. + EXECUTE format( + 'SELECT array_agg(nextval(%L::regclass)::bigint) FROM generate_series(1, $1)', + format('%I.job_id_seq', v_schema) + ) + INTO v_job_ids + USING v_count; + + -- Rows that enqueue as `available` go to a lane. Their shard is derived as + -- insert_job_compat derives it: ordering-key rows from the key hash, the + -- rest from the backend pid. + SELECT array_agg( + CASE WHEN input.state = 'available' THEN + MOD( + CASE WHEN input.ordering_hash IS NULL + THEN pg_backend_pid()::numeric + ELSE input.ordering_hash::numeric + END, + COALESCE(meta.enqueue_shards, 1)::numeric + )::smallint + END + ORDER BY input.ord + ) + INTO v_shards + FROM unnest(p_queue, p_state, p_ordering_hash) + WITH ORDINALITY AS input(queue, state, ordering_hash, ord) + LEFT JOIN awa.queue_meta AS meta ON meta.queue = input.queue; + + -- Claims go in unique-key order, so two producers with overlapping keys + -- cannot wait on each other in opposite orders. + IF p_skip_unique_conflicts THEN + WITH claimed AS ( + INSERT INTO awa.job_unique_claims (unique_key, job_id) + SELECT input.unique_key, v_job_ids[input.ord] + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ORDER BY input.unique_key, input.ord + ON CONFLICT (unique_key) DO NOTHING + RETURNING job_id + ) + SELECT COALESCE(array_agg(wanted.job_id), ARRAY[]::BIGINT[]) + INTO v_skipped + FROM ( + SELECT v_job_ids[input.ord] AS job_id + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ) AS wanted + WHERE NOT EXISTS (SELECT 1 FROM claimed WHERE claimed.job_id = wanted.job_id); + ELSE + INSERT INTO awa.job_unique_claims (unique_key, job_id) + SELECT input.unique_key, v_job_ids[input.ord] + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ORDER BY input.unique_key, input.ord; + END IF; + + -- Register lanes this batch is the first to use. The head tables carry a + -- row trigger that takes the lane lock before the conflict check, so only + -- missing lanes are proposed: a batch for an existing lane must not queue + -- behind that lane's producers before it has any work to show for it. + INSERT INTO queue_lanes (queue, priority) + SELECT lanes.queue, lanes.priority + FROM ( + SELECT DISTINCT input.queue, input.priority + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_lanes AS existing + WHERE existing.queue = lanes.queue AND existing.priority = lanes.priority + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority + ON CONFLICT ON CONSTRAINT queue_lanes_pkey DO NOTHING; + + INSERT INTO queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT lanes.queue, lanes.priority, lanes.shard + FROM ( + SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_enqueue_heads AS existing + WHERE existing.queue = lanes.queue + AND existing.priority = lanes.priority + AND existing.enqueue_shard = lanes.shard + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + INSERT INTO queue_claim_heads (queue, priority, enqueue_shard) + SELECT lanes.queue, lanes.priority, lanes.shard + FROM ( + SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_claim_heads AS existing + WHERE existing.queue = lanes.queue + AND existing.priority = lanes.priority + AND existing.enqueue_shard = lanes.shard + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + -- Re-enqueues that are not yet due carry no lane. + INSERT INTO deferred_jobs ( + job_id, kind, queue, args, state, priority, attempt, run_lease, + max_attempts, run_at, attempted_at, finalized_at, created_at, + unique_key, unique_states, payload + ) + SELECT + v_job_ids[input.ord], + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + 0, + 0, + input.max_attempts, + COALESCE(input.run_at, v_now), + NULL, + NULL, + v_now, + input.unique_key, + input.unique_states, + jsonb_build_object( + 'metadata', input.metadata, + 'tags', input.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state <> 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped); + + -- The result rows do not depend on the lanes, so build them before the + -- lock window opens. + RETURN QUERY + SELECT + v_job_ids[input.ord], + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + 0::smallint, + input.max_attempts, + COALESCE(input.run_at, v_now), + NULL::timestamptz, + NULL::timestamptz, + NULL::timestamptz, + NULL::timestamptz, + v_now, + ARRAY[]::jsonb[], + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL::uuid, + NULL::timestamptz, + NULL::text, + NULL::text, + NULL::text, + NULL::text, + 0::bigint, + NULL::jsonb + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE v_job_ids[input.ord] <> ALL (v_skipped) + ORDER BY input.ord; + + -- Lane lock window: reserve each lane's sequence range once, in + -- (queue, priority, shard) order, then write the ready rows. The locks are + -- held to the end of the transaction, which keeps the rows of a lane + -- visible in lane_seq order. + FOR v_lane IN + SELECT lanes.queue, lanes.priority, lanes.shard, lanes.row_count + FROM ( + SELECT input.queue, input.priority, v_shards[input.ord] AS shard, + count(*) AS row_count + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + GROUP BY input.queue, input.priority, v_shards[input.ord] + ) AS lanes + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + LOOP + EXECUTE format('SELECT %I.reserve_enqueue_seq($1, $2, $3, $4)', v_schema) + INTO v_start + USING v_lane.queue, v_lane.priority, v_lane.shard, v_lane.row_count::BIGINT; + + v_lane_queues := v_lane_queues || v_lane.queue; + v_lane_priorities := v_lane_priorities || v_lane.priority; + v_lane_shards := v_lane_shards || v_lane.shard; + v_lane_starts := v_lane_starts || v_start; + END LOOP; + + IF cardinality(v_lane_queues) > 0 THEN + SELECT rc.slot, rc.generation + INTO v_ready_slot, v_ready_generation + FROM ring_cursor('queue') AS rc; + + INSERT INTO ready_entries ( + ready_slot, ready_generation, job_id, kind, queue, args, priority, + attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, created_at, unique_key, unique_states, payload + ) + SELECT + v_ready_slot, + v_ready_generation, + v_job_ids[placed.ord], + placed.kind, + placed.queue, + placed.args, + placed.priority, + 0, + 0, + placed.max_attempts, + lane.start_seq + placed.lane_rank - 1, + lane.shard, + COALESCE(placed.run_at, v_now), + NULL, + v_now, + placed.unique_key, + placed.unique_states, + jsonb_build_object( + 'metadata', placed.metadata, + 'tags', placed.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM ( + SELECT + input.*, + v_shards[input.ord] AS shard, + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) AS lane_rank + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = placed.queue + AND lane.priority = placed.priority + AND lane.shard = placed.shard; + + -- One segment per run of consecutive lane_seq values sharing a run_at. + INSERT INTO ready_segments ( + ready_slot, ready_generation, queue, priority, enqueue_shard, + first_lane_seq, next_lane_seq, first_run_at + ) + SELECT + v_ready_slot, + v_ready_generation, + runs.queue, + runs.priority, + runs.shard, + min(runs.lane_seq), + max(runs.lane_seq) + 1, + runs.run_at + FROM ( + SELECT + placed.*, + placed.lane_seq - row_number() OVER ( + PARTITION BY placed.queue, placed.priority, placed.shard, placed.run_at + ORDER BY placed.lane_seq + ) AS run_id + FROM ( + SELECT + input.queue, + input.priority, + v_shards[input.ord] AS shard, + COALESCE(input.run_at, v_now) AS run_at, + lane.start_seq + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) - 1 AS lane_seq + FROM unnest(p_queue, p_priority, p_state, p_run_at) + WITH ORDINALITY AS input(queue, priority, state, run_at, ord) + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = input.queue + AND lane.priority = input.priority + AND lane.shard = v_shards[input.ord] + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + ) AS runs + GROUP BY runs.queue, runs.priority, runs.shard, runs.run_at, runs.run_id + ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) + DO NOTHING; + + PERFORM pg_notify('awa:' || notified.queue, '') + FROM (SELECT DISTINCT unnest(v_lane_queues) AS queue) AS notified; + END IF; + + PERFORM set_config('search_path', v_old_search_path, true); +END; +$$ LANGUAGE plpgsql VOLATILE +SET search_path = pg_catalog, awa, public; + +INSERT INTO awa.schema_version (version, description) +VALUES (53, 'Batch producer insert with a short enqueue lane lock window') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/src/insert.rs b/awa-model/src/insert.rs index e221b0fa7..15b5bf18b 100644 --- a/awa-model/src/insert.rs +++ b/awa-model/src/insert.rs @@ -11,13 +11,14 @@ use sqlx::{PgExecutor, PgPool}; // On queue storage every available-state insert holds a transaction-scoped // advisory lock on its `(queue, priority, enqueue_shard)` lane until the // enclosing transaction ends, so two transactions that enqueue into the same -// lanes in opposite orders can deadlock. Every multi-row enqueue therefore -// takes its lane locks in one canonical order: `(queue, priority, shard)` -// ascending, queues compared bytewise (`COLLATE "C"`, the order of Rust's -// `str`). The batch inserts here sort rows by that key, stably, before they -// reach `awa.insert_job_compat`, and put the rows back in input order for the -// caller. `QueueStorage` batch enqueues group rows in a `BTreeMap` keyed the -// same way. +// lanes in opposite orders can deadlock. The batch inserts here go through +// `awa.insert_jobs_compat_batch`, which does the lane-independent work first +// and then reserves its lanes in one canonical order: `(queue, priority, +// shard)` ascending, queues compared bytewise (`COLLATE "C"`, the order of +// Rust's `str`). `QueueStorage` batch enqueues group rows in a `BTreeMap` +// keyed the same way. Holding the lanes only across the reservation and the +// ready-row writes keeps a lane's lock window independent of the batch's +// unique-claim and payload work. // // The shard is derived as `insert_job_compat` derives it: rows with an // ordering key use `ordering_key_hash64 % enqueue_shards`, other rows @@ -211,65 +212,27 @@ impl PreparedJobInsert { /// job, because sqlx has no array encoding for the enum, the bit string, or a /// ragged array of arrays. `ordering_hash` carries each row's /// [`ordering_key_hash64`] as decimal text (the value exceeds `bigint`) so the -/// statement can place keyed rows on the shard `insert_job_compat` will pick. +/// function can place keyed rows on the shard `insert_job_compat` would pick. /// -/// Rows are inserted in canonical lane order (see "Lane lock order" above) and -/// returned in input order. -const INSERT_MANY_SQL: &str = r#" - WITH input AS MATERIALIZED ( - SELECT - raw.*, - MOD( - CASE WHEN raw.ordering_hash IS NULL - THEN pg_backend_pid()::numeric - ELSE raw.ordering_hash::numeric - END, - COALESCE( - (SELECT meta.enqueue_shards FROM awa.queue_meta AS meta WHERE meta.queue = raw.queue), - 1 - )::numeric - ) AS lane_shard - FROM unnest( - $1::text[], - $2::text[], - $3::jsonb[], - $4::text[], - $5::smallint[], - $6::smallint[], - $7::timestamptz[], - $8::jsonb[], - $9::jsonb[], - $10::bytea[], - $11::text[], - $12::bytea[], - $13::text[] - ) WITH ORDINALITY AS raw( - kind, queue, args, state, priority, max_attempts, run_at, - metadata, tags, unique_key, unique_states, ordering_key, ordering_hash, ord - ) - ORDER BY raw.queue COLLATE "C", raw.priority, lane_shard, raw.ord +/// `$13` skips rows whose unique key is already claimed instead of failing the +/// batch. Rows come back in input order; skipped rows are absent. +const INSERT_MANY_BATCH_SQL: &str = r#" + SELECT * + FROM awa.insert_jobs_compat_batch( + $1::text[], + $2::text[], + $3::jsonb[], + $4::text[], + $5::smallint[], + $6::smallint[], + $7::timestamptz[], + $8::jsonb[], + $9::jsonb[], + $10::bytea[], + $11::text[], + $12::text[], + $13 ) - SELECT inserted.* - FROM input - CROSS JOIN LATERAL awa.insert_job_compat( - input.kind, - input.queue, - input.args, - input.state::awa.job_state, - input.priority, - input.max_attempts, - input.run_at, - input.metadata, - ARRAY( - SELECT tag.value - FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, ord) - ORDER BY tag.ord - ), - input.unique_key, - input.unique_states::bit(8), - input.ordering_key - ) AS inserted - ORDER BY input.ord "#; /// Compute unique_key and unique_states from opts. @@ -476,7 +439,9 @@ fn ordering_hash_text(row: &PreparedRow) -> Option { /// Insert multiple jobs in a single statement. /// /// Supports uniqueness constraints — jobs with `unique` opts will have their -/// `unique_key` and `unique_states` computed and included. +/// `unique_key` and `unique_states` computed and included. A job whose unique +/// key is already claimed fails the whole batch with a unique violation; use +/// [`insert_many_skip_unique_conflicts`] to insert the rest instead. #[tracing::instrument( skip(executor, jobs), fields( @@ -489,6 +454,43 @@ fn ordering_hash_text(row: &PreparedRow) -> Option { ) )] pub async fn insert_many<'e, E>(executor: E, jobs: &[InsertParams]) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + insert_many_batch(executor, jobs, false).await +} + +/// Insert multiple jobs in a single statement, skipping those whose unique key +/// is already claimed. +/// +/// Returns the rows that were inserted, in input order; a skipped job has no +/// entry in the result. Jobs without `unique` opts are always inserted. +#[tracing::instrument( + skip(executor, jobs), + fields( + job.count = jobs.len(), + otel.name = "send", + otel.kind = "producer", + messaging.system = "awa", + messaging.operation.r#type = "send", + messaging.batch.message_count = jobs.len(), + ) +)] +pub async fn insert_many_skip_unique_conflicts<'e, E>( + executor: E, + jobs: &[InsertParams], +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + insert_many_batch(executor, jobs, true).await +} + +async fn insert_many_batch<'e, E>( + executor: E, + jobs: &[InsertParams], + skip_unique_conflicts: bool, +) -> Result, AwaError> where E: PgExecutor<'e>, { @@ -497,7 +499,7 @@ where } let rows = precompute_rows(jobs)?; - let results = sqlx::query_as::<_, JobRow>(INSERT_MANY_SQL) + let results = sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL) .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) .bind( rows.iter() @@ -529,48 +531,14 @@ where .map(|row| row.unique_states.as_deref()) .collect::>(), ) - .bind( - rows.iter() - .map(|row| row.ordering_key.as_deref()) - .collect::>(), - ) .bind(rows.iter().map(ordering_hash_text).collect::>()) + .bind(skip_unique_conflicts) .fetch_all(executor) .await?; Ok(results) } -/// CTEs that number the staged rows in staging order and re-sort them into -/// canonical lane order (see "Lane lock order" above); `staged_in_lane_order` -/// carries the staging-order position as `ord`. -macro_rules! staged_in_lane_order_ctes { - () => { - r#" - staged AS ( - SELECT *, row_number() OVER () AS ord - FROM pg_temp.awa_copy_staging - ), - staged_in_lane_order AS MATERIALIZED ( - SELECT - staged.*, - MOD( - CASE WHEN staged.ordering_hash IS NULL - THEN pg_backend_pid()::numeric - ELSE staged.ordering_hash::numeric - END, - COALESCE( - (SELECT meta.enqueue_shards FROM awa.queue_meta AS meta WHERE meta.queue = staged.queue), - 1 - )::numeric - ) AS lane_shard - FROM staged - ORDER BY staged.queue COLLATE "C", staged.priority, lane_shard, staged.ord - ) - "# - }; -} - /// Insert many jobs through the compatibility COPY surface. /// /// Uses a temp staging table with no constraints for fast COPY ingestion, @@ -646,150 +614,53 @@ pub async fn insert_many_copy( copy_in.finish().await?; } - // 3. INSERT...SELECT from staging into real table + // 3. Insert the staged rows through the batch function, which skips rows + // whose unique key is already claimed when the batch has any. let has_unique = rows.iter().any(|r| r.unique_key.is_some()); - - let results = if has_unique { - // The compatibility `awa.jobs` surface is now a view backed by hot and - // deferred tables, so the old `ON CONFLICT` path is no longer available - // here. Keep COPY for staging/parsing, then insert unique rows one at a - // time and skip duplicates explicitly. - let staged_rows = sqlx::query_as::< - _, - ( - i64, - String, - String, - serde_json::Value, - String, - i16, - i16, - Option>, - serde_json::Value, - Vec, - Option>, - Option, - Option>, - ), - >(concat!( - "WITH", - staged_in_lane_order_ctes!(), - r#" + let results = sqlx::query_as::<_, JobRow>( + r#" + WITH staged AS ( + SELECT *, row_number() OVER () AS ord + FROM pg_temp.awa_copy_staging + ), + columns AS ( SELECT - ord, - kind, - queue, - args, - state::text, - priority, - max_attempts, - run_at, - metadata, - tags, - unique_key, - unique_states::text, - ordering_key - FROM staged_in_lane_order - ORDER BY queue COLLATE "C", priority, lane_shard, ord - "# - )) - .fetch_all(&mut *conn) - .await?; - - let mut inserted = Vec::with_capacity(staged_rows.len()); - for ( - staged_position, - kind, - queue, - args, - state, - priority, - max_attempts, - run_at, - metadata, - tags, - unique_key, - unique_states, - ordering_key, - ) in staged_rows - { - sqlx::query("SAVEPOINT awa_copy_unique_row") - .execute(&mut *conn) - .await?; - - let result = sqlx::query_as::<_, JobRow>(POSTGRES_INSERT_JOB_SQL) - .bind(&kind) - .bind(&queue) - .bind(&args) - .bind(&state) - .bind(priority) - .bind(max_attempts) - .bind(run_at) - .bind(&metadata) - .bind(&tags) - .bind(&unique_key) - .bind(&unique_states) - .bind(&ordering_key) - .fetch_one(&mut *conn) - .await; - - match result { - Ok(row) => { - inserted.push((staged_position, row)); - sqlx::query("RELEASE SAVEPOINT awa_copy_unique_row") - .execute(&mut *conn) - .await?; - } - Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("23505") => { - sqlx::query("ROLLBACK TO SAVEPOINT awa_copy_unique_row") - .execute(&mut *conn) - .await?; - sqlx::query("RELEASE SAVEPOINT awa_copy_unique_row") - .execute(&mut *conn) - .await?; - continue; - } - Err(err) => { - sqlx::query("ROLLBACK TO SAVEPOINT awa_copy_unique_row") - .execute(&mut *conn) - .await?; - sqlx::query("RELEASE SAVEPOINT awa_copy_unique_row") - .execute(&mut *conn) - .await?; - return Err(AwaError::Database(err)); - } - } - } - - inserted.sort_unstable_by_key(|(staged_position, _)| *staged_position); - inserted.into_iter().map(|(_, row)| row).collect() - } else { - sqlx::query_as::<_, JobRow>(concat!( - "WITH", - staged_in_lane_order_ctes!(), - r#" - SELECT inserted.* - FROM staged_in_lane_order AS staged - CROSS JOIN LATERAL awa.insert_job_compat( - staged.kind, - staged.queue, - staged.args, - staged.state, - staged.priority, - staged.max_attempts, - staged.run_at, - staged.metadata, - staged.tags, - staged.unique_key, - staged.unique_states, - staged.ordering_key - ) AS inserted - ORDER BY staged.ord - "# - )) - .fetch_all(&mut *conn) - .await? - }; + array_agg(kind ORDER BY ord) AS kinds, + array_agg(queue ORDER BY ord) AS queues, + array_agg(args ORDER BY ord) AS args, + array_agg(state::text ORDER BY ord) AS states, + array_agg(priority ORDER BY ord) AS priorities, + array_agg(max_attempts ORDER BY ord) AS max_attempts, + array_agg(run_at ORDER BY ord) AS run_ats, + array_agg(metadata ORDER BY ord) AS metadata, + array_agg(to_jsonb(tags) ORDER BY ord) AS tags, + array_agg(unique_key ORDER BY ord) AS unique_keys, + array_agg(unique_states::text ORDER BY ord) AS unique_states, + array_agg(ordering_hash ORDER BY ord) AS ordering_hashes + FROM staged + ) + SELECT inserted.* + FROM columns + CROSS JOIN LATERAL awa.insert_jobs_compat_batch( + columns.kinds, + columns.queues, + columns.args, + columns.states, + columns.priorities, + columns.max_attempts, + columns.run_ats, + columns.metadata, + columns.tags, + columns.unique_keys, + columns.unique_states, + columns.ordering_hashes, + $1 + ) AS inserted + "#, + ) + .bind(has_unique) + .fetch_all(&mut *conn) + .await?; // Keep the session-local staging table reusable across multiple COPY calls // within the same outer transaction. diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index 66a5dd4ee..7f23ec69a 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -45,7 +45,8 @@ pub use cron::{CronJobRow, CronMissedFirePolicy, PeriodicJob, PeriodicJobBuilder pub use dlq::{DlqMetadata, DlqRow, ListDlqFilter, RetryFromDlqOpts}; pub use error::{map_sqlx_error, AwaError}; pub use insert::{ - insert, insert_many, insert_many_copy, insert_many_copy_from_pool, insert_prepared, insert_with, + insert, insert_many, insert_many_copy, insert_many_copy_from_pool, + insert_many_skip_unique_conflicts, insert_prepared, insert_with, }; pub use job::{InsertOpts, InsertParams, JobRow, JobState, UniqueOpts}; pub use output::{AttemptError, JobOutput, OutputDescriptor, OutputError, MAX_OUTPUT_BYTES}; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 385035b65..23afccc77 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -5,7 +5,7 @@ use sqlx::{Connection, PgPool}; use tracing::{info, warn}; /// Current schema version. -pub const CURRENT_VERSION: i32 = 52; +pub const CURRENT_VERSION: i32 = 53; /// Migrations that require an exclusive (no-live-runtime) upgrade window. /// @@ -292,6 +292,11 @@ const MIGRATIONS: &[(i32, &str, &[&str])] = &[ ), (51, "Per-slot claim functions", &[V23_UP, V51_UP]), (52, "Retained job output", &[V23_UP, V52_UP]), + ( + 53, + "Batch producer insert with a short enqueue lane lock window", + &[V53_UP], + ), ]; const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql"); @@ -352,6 +357,7 @@ const V49_UP: &str = include_str!("../migrations/v049_prepared_lane_sequences.sq const V50_UP: &str = include_str!("../migrations/v050_drop_ready_lane_shard_index.sql"); const V51_UP: &str = include_str!("../migrations/v051_claim_slot_functions.sql"); const V52_UP: &str = include_str!("../migrations/v052_retained_job_output.sql"); +const V53_UP: &str = include_str!("../migrations/v053_insert_jobs_batch_compat.sql"); /// Old version numbers from pre-0.4 releases that used V3/V4/V5 numbering. /// Also tolerates the unreleased inline-V6 branch numbering used during review. diff --git a/awa/tests/enqueue_lane_contention_test.rs b/awa/tests/enqueue_lane_contention_test.rs index 46a89625b..6495213a8 100644 --- a/awa/tests/enqueue_lane_contention_test.rs +++ b/awa/tests/enqueue_lane_contention_test.rs @@ -22,8 +22,8 @@ //! AWA_LANE_BENCH_BATCH jobs per batch, default 185 //! AWA_LANE_BENCH_UNIQUE 0|1, default 1 //! AWA_LANE_BENCH_SHARDS `enqueue_shards` for the queue, default 1 -//! AWA_LANE_BENCH_MODE insert_many | insert_many_copy | queue_storage | -//! legacy_row_function +//! AWA_LANE_BENCH_MODE insert_many | insert_many_copy | insert_many_skip | +//! queue_storage | legacy_row_function //! AWA_LANE_BENCH_DUP_PCT % of batches carrying one duplicate key that //! fall back to row-at-a-time inserts, default 0 //! AWA_LANE_BENCH_SECS measured seconds per scenario, default 8 @@ -43,7 +43,8 @@ mod ci_timing; use awa::model::insert_prepared; use awa::model::{ - insert_many, insert_many_copy_from_pool, prepare_raw_job_insert, AwaError, QueueStorage, + insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, + prepare_raw_job_insert, AwaError, QueueStorage, }; use awa::{InsertOpts, InsertParams, QueueStorageConfig, UniqueOpts}; use awa_testing::setup; @@ -64,6 +65,7 @@ const WARMUP: Duration = Duration::from_millis(1_500); enum Mode { InsertMany, InsertManyCopy, + InsertManySkip, QueueStorage, LegacyRowFunction, } @@ -73,6 +75,7 @@ impl Mode { match raw { "insert_many" => Mode::InsertMany, "insert_many_copy" => Mode::InsertManyCopy, + "insert_many_skip" => Mode::InsertManySkip, "queue_storage" => Mode::QueueStorage, "legacy_row_function" => Mode::LegacyRowFunction, other => panic!("unknown AWA_LANE_BENCH_MODE {other}"), @@ -83,6 +86,7 @@ impl Mode { match self { Mode::InsertMany => "insert_many", Mode::InsertManyCopy => "insert_many_copy", + Mode::InsertManySkip => "insert_many_skip", Mode::QueueStorage => "queue_storage", Mode::LegacyRowFunction => "legacy_row_function", } @@ -379,6 +383,9 @@ async fn producer_loop( } Mode::InsertMany => insert_many(&pool, &jobs).await.map(|_| ()), Mode::InsertManyCopy => insert_many_copy_from_pool(&pool, &jobs).await.map(|_| ()), + Mode::InsertManySkip => insert_many_skip_unique_conflicts(&pool, &jobs) + .await + .map(|_| ()), Mode::QueueStorage => store.enqueue_params_batch(&pool, &jobs).await.map(|_| ()), Mode::LegacyRowFunction => legacy_row_function_insert(&pool, &jobs).await, }; @@ -677,10 +684,12 @@ fn scenario( /// Absolute floors sit well below the rates a shared runner sustains, and the /// latency ceiling well above, so they only catch a lane that has stopped /// making progress; `benchmarks/baseline.json` carries the regression -/// thresholds on rate and p99. The one-producer scenarios measure what a batch -/// costs the lane lock, through `insert_many` and through the per-row -/// function, and are reported for comparison. Every bound goes through -/// `ci_timing` so a contended runner loosens it. +/// thresholds on rate and p99. What the gate asserts about lock hold time is +/// relative, because it must hold on any runner: one producer's 185-row batch +/// has to take well under the time the same rows take through the per-row +/// function, measured in the same run, and sixteen producers on one lane have +/// to out-run what the per-row function could ever sustain for it. Every +/// bound goes through `ci_timing` so a contended runner loosens it. #[tokio::test(flavor = "multi_thread", worker_threads = 6)] #[ignore] async fn test_enqueue_lane_contention_gate() { @@ -730,10 +739,35 @@ async fn test_enqueue_lane_contention_gate() { single_lane.p99_ms, p99_ceiling ); + + // The batch holds the lane lock for a fraction of the per-row function's + // time, and the lane's ceiling follows. Both bounds are ratios of numbers + // measured moments apart on the same machine. + let max_hold_ratio = (GATE_MAX_HOLD_RATIO * ci_timing::chaos_timeout_multiplier()).min(0.9); + let hold_ratio = one_producer.p50_ms / row_function.p50_ms; + assert!( + hold_ratio <= max_hold_ratio, + "a 185-row batch took {:.1}ms against {:.1}ms through the per-row function \ + (ratio {hold_ratio:.2}, limit {max_hold_ratio:.2})", + one_producer.p50_ms, + row_function.p50_ms + ); + let min_scaling = GATE_MIN_LANE_SCALING / ci_timing::chaos_timeout_multiplier(); + let lane_scaling = single_lane.jobs_per_s / row_function.jobs_per_s; + assert!( + lane_scaling >= min_scaling.max(1.0), + "16 producers on one lane sustain {:.0}/s, {lane_scaling:.1}x the per-row function's \ + single-producer {:.0}/s (need {:.1}x)", + single_lane.jobs_per_s, + row_function.jobs_per_s, + min_scaling.max(1.0) + ); } const GATE_MIN_JOBS_PER_S: i64 = 2_000; const GATE_MAX_P99_MS: u64 = 15_000; +const GATE_MAX_HOLD_RATIO: f64 = 0.4; +const GATE_MIN_LANE_SCALING: f64 = 2.0; /// Exploration sweep; see the module docs for the tunables. #[tokio::test(flavor = "multi_thread", worker_threads = 8)] diff --git a/awa/tests/enqueue_lane_lock_order_test.rs b/awa/tests/enqueue_lane_lock_order_test.rs index 924d212ae..51ec9a4ac 100644 --- a/awa/tests/enqueue_lane_lock_order_test.rs +++ b/awa/tests/enqueue_lane_lock_order_test.rs @@ -282,6 +282,74 @@ fn ordering_key_per_shard(shards: i16) -> Vec> { .collect() } +/// Key of the advisory lock `reserve_enqueue_seq` holds for one lane. +async fn lane_reservation_lock_key(pool: &PgPool, queue: &str, priority: i16, shard: i16) -> i64 { + sqlx::query_scalar( + r#" + SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, $2, $3)), + 0 + ) + "#, + ) + .bind(queue) + .bind(priority) + .bind(shard) + .fetch_one(pool) + .await + .expect("derive lane reservation lock key") +} + +/// The backend waiting for the advisory lock `key`, once there is one. +async fn advisory_waiter_pid(pool: &PgPool, key: i64) -> i32 { + let started = Instant::now(); + loop { + let pid: Option = sqlx::query_scalar( + r#" + SELECT pid + FROM pg_locks + WHERE locktype = 'advisory' + AND NOT granted + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) + AND ((classid::bigint << 32) | objid::bigint) = $1 + LIMIT 1 + "#, + ) + .bind(key) + .fetch_optional(pool) + .await + .expect("poll pg_locks"); + if let Some(pid) = pid { + return pid; + } + assert!( + started.elapsed() < WAIT_DEADLINE, + "no backend started waiting for advisory lock {key}" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +async fn granted_advisory_keys(pool: &PgPool, pid: i32) -> std::collections::BTreeSet { + sqlx::query_scalar::<_, i64>( + r#" + SELECT (classid::bigint << 32) | objid::bigint + FROM pg_locks + WHERE locktype = 'advisory' AND granted AND pid = $1 + "#, + ) + .bind(pid) + .fetch_all(pool) + .await + .expect("read granted advisory locks") + .into_iter() + .collect() +} + +/// A batch over sharded lanes reserves them in `(queue, priority, shard)` +/// order, whatever order its rows arrive in. A gate holds one lane in the +/// middle of that order; the producer must be parked on it holding exactly the +/// lanes before it. async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { let _guard = LANE_LOCK_TEST_GUARD.lock().await; let (pool, _config) = setup_queue_storage().await; @@ -309,6 +377,7 @@ async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { let keys = ordering_key_per_shard(shards); let mut jobs = Vec::new(); + let mut lanes = std::collections::BTreeSet::new(); let rounds = match producer { Producer::CopyLarge => 700, _ => 12, @@ -317,55 +386,112 @@ async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { for queue in &queues { let mut params = job(queue, round, false); params.opts.priority = if round % 5 == 0 { 1 } else { 2 }; - // Keyed rows spread over every shard; every third row has no key. - if round % 3 != 0 { - params.opts.ordering_key = Some(keys[(round + queue.len()) % keys.len()].clone()); - } + // Every row is keyed, so its shard, and so its lane, is known here. + let shard = (round + queue.len()) % keys.len(); + params.opts.ordering_key = Some(keys[shard].clone()); + lanes.insert(( + queue.as_bytes().to_vec(), + params.opts.priority, + shard as i16, + queue.clone(), + )); jobs.push(params); } } + assert!(lanes.len() >= 6, "the batch must span many lanes"); - let inserted = producer.run(&pool, &jobs).await.expect("insert batch"); + let canonical: Vec<(String, i16, i16)> = lanes + .into_iter() + .map(|(_, priority, shard, queue)| (queue, priority, shard)) + .collect(); + let mut canonical_keys = Vec::new(); + for (queue, priority, shard) in &canonical { + canonical_keys.push(lane_reservation_lock_key(&pool, queue, *priority, *shard).await); + } + + // Provision the lanes so the gate lock is the only thing the producer waits for. + for (index, queue) in queues.iter().enumerate() { + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + for (priority, shard) in [ + (1, 0), + (1, 1), + (1, 2), + (1, 3), + (2, 0), + (2, 1), + (2, 2), + (2, 3), + ] { + let mut params = job(queue, 10_000 + index, false); + params.opts.priority = priority; + params.opts.ordering_key = Some(keys[shard as usize].clone()); + store + .enqueue_params_batch(&pool, &[params]) + .await + .expect("provision lane"); + } + } + + let gated = canonical.len() / 2; + let mut gate = pool.begin().await.expect("begin gate transaction"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(canonical_keys[gated]) + .execute(&mut *gate) + .await + .expect("take gate lock"); + + let competing = { + let pool = pool.clone(); + let jobs = jobs.clone(); + tokio::spawn(async move { producer.run(&pool, &jobs).await }) + }; + let waiter = advisory_waiter_pid(&pool, canonical_keys[gated]).await; + let held = granted_advisory_keys(&pool, waiter).await; + let expected: std::collections::BTreeSet = + canonical_keys[..gated].iter().copied().collect(); + assert_eq!( + held, + expected, + "a batch parked on lane {gated} of {} must hold the lanes before it and no later one", + canonical.len() + ); + + gate.commit().await.expect("release gate"); + let inserted = tokio::time::timeout(WAIT_DEADLINE, competing) + .await + .expect("producer finished") + .expect("producer task") + .expect("insert batch"); + + // Lane ordering must not leak into the result: rows come back in input order. assert_eq!(inserted.len(), jobs.len()); for (row, params) in inserted.iter().zip(&jobs) { assert_eq!(row.queue, params.opts.queue); assert_eq!(row.args, params.args); } - // Job ids are allocated as each row reaches `insert_job_compat`, so id - // order is lock order. - let lane_by_job: std::collections::HashMap = - sqlx::query_as::<_, (i64, String, i16, i16)>( - "SELECT job_id, queue, priority, enqueue_shard FROM awa.ready_entries WHERE job_id = ANY($1)", + // Within a lane, lane_seq follows input order. + let lane_by_job: std::collections::HashMap = + sqlx::query_as::<_, (i64, String, i16, i16, i64)>( + "SELECT job_id, queue, priority, enqueue_shard, lane_seq FROM awa.ready_entries WHERE job_id = ANY($1)", ) .bind(inserted.iter().map(|row| row.id).collect::>()) .fetch_all(&pool) .await .expect("read lanes") .into_iter() - .map(|(job_id, queue, priority, shard)| (job_id, (queue, priority, shard))) + .map(|(job_id, queue, priority, shard, lane_seq)| { + (job_id, (queue, priority, shard, lane_seq)) + }) .collect(); assert_eq!(lane_by_job.len(), jobs.len()); - - let mut by_lock_order: Vec<(i64, usize)> = inserted - .iter() - .enumerate() - .map(|(position, row)| (row.id, position)) - .collect(); - by_lock_order.sort_unstable(); - let lock_order: Vec<(&(String, i16, i16), usize)> = by_lock_order - .iter() - .map(|(job_id, position)| (&lane_by_job[job_id], *position)) - .collect(); - for pair in lock_order.windows(2) { - let (lane, position) = pair[0]; - let (next_lane, next_position) = pair[1]; - assert!( - lane.0.as_bytes() < next_lane.0.as_bytes() - || (lane.0 == next_lane.0 && (lane.1, lane.2) < (next_lane.1, next_lane.2)) - || (lane == next_lane && position < next_position), - "lane {lane:?} (input {position}) was locked before lane {next_lane:?} (input {next_position})" - ); + let mut last_seq_by_lane: std::collections::HashMap<(String, i16, i16), i64> = + std::collections::HashMap::new(); + for row in &inserted { + let (queue, priority, shard, lane_seq) = lane_by_job[&row.id].clone(); + if let Some(previous) = last_seq_by_lane.insert((queue, priority, shard), lane_seq) { + assert_eq!(lane_seq, previous + 1, "lane_seq must follow input order"); + } } } diff --git a/awa/tests/insert_many_batch_test.rs b/awa/tests/insert_many_batch_test.rs new file mode 100644 index 000000000..aafdf9202 --- /dev/null +++ b/awa/tests/insert_many_batch_test.rs @@ -0,0 +1,426 @@ +//! `insert_many` and `insert_many_copy` on the batch insert function. +//! +//! A batch does its lane-independent work (job ids, unique claims, lane +//! registration) before it reserves its lanes, and holds a lane lock only for +//! the reservation, the ready-row writes and the commit. These tests cover that +//! ordering, skip-on-conflict inserts, and the shape of what lands in the +//! lanes. +//! +//! Requires a running Postgres instance. +//! Run with: `cargo test --package awa --test insert_many_batch_test -- --nocapture` + +use awa::model::{ + insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, AwaError, + QueueStorage, +}; +use awa::{InsertOpts, InsertParams, JobState, QueueStorageConfig, UniqueOpts}; +use awa_testing::setup; +use sqlx::PgPool; +use std::sync::LazyLock; +use std::time::{Duration, Instant}; +use tokio::sync::Mutex; +use uuid::Uuid; + +/// The tests switch the storage engine of a shared database. +static ENGINE_GUARD: LazyLock> = LazyLock::new(|| Mutex::new(())); + +const WAIT_DEADLINE: Duration = Duration::from_secs(20); + +async fn queue_storage_pool() -> PgPool { + let pool = setup::setup(8).await; + setup::activate_queue_storage(&pool).await; + pool +} + +fn unique_queue(label: &str) -> String { + format!( + "batch_{label}_{}", + &Uuid::new_v4().simple().to_string()[..12] + ) +} + +fn job(queue: &str, seq: i64, unique: bool) -> InsertParams { + InsertParams { + kind: "insert_many_batch_job".to_string(), + args: serde_json::json!({ "seq": seq }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + } +} + +async fn ready_lane_seqs(pool: &PgPool, queue: &str) -> Vec<(i64, i64)> { + sqlx::query_as( + "SELECT (args->>'seq')::bigint, lane_seq FROM awa.ready_entries + WHERE queue = $1 ORDER BY lane_seq", + ) + .bind(queue) + .fetch_all(pool) + .await + .expect("read ready rows") +} + +async fn claims_for_queue(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.job_unique_claims AS claims + WHERE claims.job_id IN (SELECT job_id FROM awa.ready_entries WHERE queue = $1)", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("count claims") +} + +fn is_unique_violation(err: &AwaError) -> bool { + match err { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(db_err)) => { + db_err.code().as_deref() == Some("23505") + } + _ => false, + } +} + +async fn lane_lock_key(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, 2::smallint, 0::smallint)), + 0)", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("derive lane lock key") +} + +#[tokio::test] +async fn batch_lands_in_one_lane_in_input_order_and_is_claimable() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("order"); + + let jobs: Vec = (0..300).map(|seq| job(&queue, seq, seq % 2 == 0)).collect(); + let inserted = insert_many(&pool, &jobs).await.expect("insert batch"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + (0..300).collect::>(), + "rows come back in input order" + ); + assert!(inserted.iter().all(|row| row.state == JobState::Available)); + + let lane = ready_lane_seqs(&pool, &queue).await; + assert_eq!(lane.len(), 300); + let first = lane[0].1; + for (offset, (seq, lane_seq)) in lane.iter().enumerate() { + assert_eq!(*seq, offset as i64, "lane_seq follows input order"); + assert_eq!(*lane_seq, first + offset as i64, "lane_seq is dense"); + } + + // The lane's segment metadata must let the claim path find every row. + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let mut claimed = Vec::new(); + while claimed.len() < 300 { + let batch = store + .claim_runtime_batch(&pool, &queue, 100, Duration::from_secs(60)) + .await + .expect("claim"); + assert!(!batch.is_empty(), "claim path lost rows of the batch"); + claimed.extend( + batch + .into_iter() + .map(|job| job.job.args["seq"].as_i64().unwrap()), + ); + } + assert_eq!( + claimed, + (0..300).collect::>(), + "claimed in lane order" + ); +} + +#[tokio::test] +async fn batch_with_mixed_states_routes_each_row() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("mixed"); + + let later = chrono::Utc::now() + chrono::Duration::hours(1); + let mut jobs: Vec = (0..6).map(|seq| job(&queue, seq, false)).collect(); + jobs[1].opts.run_at = Some(later); + jobs[4].opts.run_at = Some(later); + + let inserted = insert_many(&pool, &jobs).await.expect("insert mixed batch"); + let states: Vec = inserted.iter().map(|row| row.state).collect(); + assert_eq!( + states, + vec![ + JobState::Available, + JobState::Scheduled, + JobState::Available, + JobState::Available, + JobState::Scheduled, + JobState::Available, + ] + ); + let ready: Vec = ready_lane_seqs(&pool, &queue) + .await + .into_iter() + .map(|(seq, _)| seq) + .collect(); + assert_eq!(ready, vec![0, 2, 3, 5]); + let deferred: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.deferred_jobs WHERE queue = $1 AND state = 'scheduled'", + ) + .bind(&queue) + .fetch_one(&pool) + .await + .expect("count deferred"); + assert_eq!(deferred, 2); +} + +#[tokio::test] +async fn unique_conflict_fails_the_batch_atomically() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("atomic"); + + insert_many(&pool, &[job(&queue, 1, true)]) + .await + .expect("first insert"); + + let batch = vec![ + job(&queue, 0, true), + job(&queue, 1, true), + job(&queue, 2, true), + ]; + let err = insert_many(&pool, &batch) + .await + .expect_err("duplicate key must fail the batch"); + assert!(is_unique_violation(&err), "{err:?}"); + + let landed: Vec = ready_lane_seqs(&pool, &queue) + .await + .into_iter() + .map(|(seq, _)| seq) + .collect(); + assert_eq!(landed, vec![1], "nothing from the failed batch landed"); + assert_eq!(claims_for_queue(&pool, &queue).await, 1); +} + +#[tokio::test] +async fn skip_unique_conflicts_inserts_the_rest_in_order() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("skip"); + + insert_many(&pool, &[job(&queue, 1, true), job(&queue, 3, true)]) + .await + .expect("seed"); + + // Keys 1 and 3 are already claimed, 4 repeats inside the batch, and the + // last row carries no unique key at all. + let batch = vec![ + job(&queue, 0, true), + job(&queue, 1, true), + job(&queue, 2, true), + job(&queue, 3, true), + job(&queue, 4, true), + job(&queue, 4, true), + job(&queue, 5, false), + ]; + let inserted = insert_many_skip_unique_conflicts(&pool, &batch) + .await + .expect("skip insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 2, 4, 5] + ); + + let lane = ready_lane_seqs(&pool, &queue).await; + assert_eq!( + lane.iter().map(|(seq, _)| *seq).collect::>(), + vec![1, 3, 0, 2, 4, 5], + "skipped rows take no lane position" + ); + for pair in lane.windows(2) { + assert_eq!(pair[1].1, pair[0].1 + 1, "lane_seq stays dense"); + } + + // Everything skipped: an empty result, not an error. + let inserted = insert_many_skip_unique_conflicts(&pool, &[job(&queue, 0, true)]) + .await + .expect("all skipped"); + assert!(inserted.is_empty()); +} + +#[tokio::test] +async fn copy_skips_unique_conflicts_in_one_pass() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("copy_skip"); + + insert_many(&pool, &[job(&queue, 2, true)]) + .await + .expect("seed"); + let batch: Vec = (0..6).map(|seq| job(&queue, seq, true)).collect(); + let inserted = insert_many_copy_from_pool(&pool, &batch) + .await + .expect("copy insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 1, 3, 4, 5] + ); + assert_eq!(ready_lane_seqs(&pool, &queue).await.len(), 6); +} + +#[tokio::test] +async fn canonical_storage_inserts_and_skips_row_by_row() { + let _guard = ENGINE_GUARD.lock().await; + let pool = setup::setup(8).await; + setup::reset_to_canonical(&pool).await; + let queue = unique_queue("canonical"); + + insert_many(&pool, &[job(&queue, 1, true)]) + .await + .expect("seed"); + let batch: Vec = (0..4).map(|seq| job(&queue, seq, true)).collect(); + let err = insert_many(&pool, &batch).await.expect_err("duplicate"); + assert!(is_unique_violation(&err), "{err:?}"); + + let inserted = insert_many_skip_unique_conflicts(&pool, &batch) + .await + .expect("skip insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 2, 3] + ); +} + +/// While another transaction holds a lane lock, a batch bound for that lane +/// has already done its lane-independent work (it holds its unique claims) and +/// has not yet written a ready row. Every row's unique claim, not just the +/// first's, is held by then. +#[tokio::test] +async fn batch_waits_for_the_lane_after_its_lane_independent_work() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("window"); + + // Provision the lane so the gate lock is the only thing to wait for. + insert_many(&pool, &[job(&queue, -1, false)]) + .await + .expect("provision lane"); + + let key = lane_lock_key(&pool, &queue).await; + let mut gate = pool.begin().await.expect("begin gate"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(key) + .execute(&mut *gate) + .await + .expect("take lane lock"); + + let jobs_for_probe: Vec = (0..50).map(|seq| job(&queue, seq, true)).collect(); + let producer = { + let pool = pool.clone(); + let jobs = jobs_for_probe.clone(); + tokio::spawn(async move { insert_many(&pool, &jobs).await }) + }; + + let started = Instant::now(); + let waiter = loop { + let pid: Option = sqlx::query_scalar( + "SELECT pid FROM pg_locks + WHERE locktype = 'advisory' AND NOT granted + AND ((classid::bigint << 32) | objid::bigint) = $1 + LIMIT 1", + ) + .bind(key) + .fetch_optional(&pool) + .await + .expect("poll pg_locks"); + if let Some(pid) = pid { + break pid; + } + assert!(started.elapsed() < WAIT_DEADLINE, "producer never waited"); + tokio::time::sleep(Duration::from_millis(10)).await; + }; + + let holds_relation_lock = |relation: &'static str| { + let pool = pool.clone(); + async move { + sqlx::query_scalar::<_, bool>( + "SELECT EXISTS ( + SELECT 1 FROM pg_locks + WHERE pid = $1 AND locktype = 'relation' AND granted + AND relation = $2::regclass + )", + ) + .bind(waiter) + .bind(relation) + .fetch_one(&pool) + .await + .expect("read relation locks") + } + }; + // The claim of the batch's last row is already held, uncommitted: an + // insert of the same key waits on the producer's transaction. + let last = jobs_for_probe.last().expect("probe job"); + let last_key = + awa::model::prepare_raw_job_insert(last.kind.clone(), last.args.clone(), last.opts.clone()) + .expect("prepare last job") + .unique_key() + .expect("unique key") + .to_vec(); + let mut probe = pool.begin().await.expect("begin probe"); + sqlx::query("SET LOCAL lock_timeout = '300ms'") + .execute(&mut *probe) + .await + .expect("set lock timeout"); + let probe_result = sqlx::query( + "INSERT INTO awa.job_unique_claims (unique_key, job_id) VALUES ($1, 0) + ON CONFLICT (unique_key) DO NOTHING", + ) + .bind(&last_key) + .execute(&mut *probe) + .await; + match probe_result { + Err(sqlx::Error::Database(db_err)) => assert_eq!( + db_err.code().as_deref(), + Some("55P03"), + "the last row's claim must already be held while the lane lock is awaited" + ), + other => panic!("claim of the last row was not held while waiting for the lane: {other:?}"), + } + drop(probe); + assert!( + !holds_relation_lock("awa.ready_entries").await, + "no ready row is written while the lane lock is awaited" + ); + + gate.commit().await.expect("release lane lock"); + let inserted = tokio::time::timeout(WAIT_DEADLINE, producer) + .await + .expect("producer finished") + .expect("producer task") + .expect("insert batch"); + assert_eq!(inserted.len(), 50); + assert_eq!(ready_lane_seqs(&pool, &queue).await.len(), 51); +} diff --git a/awa/tests/integration_test.rs b/awa/tests/integration_test.rs index c816f0675..a4b9c6567 100644 --- a/awa/tests/integration_test.rs +++ b/awa/tests/integration_test.rs @@ -552,6 +552,9 @@ async fn test_insert_many() { #[tokio::test] async fn test_insert_many_exceeding_bind_parameter_limit() { const JOB_COUNT: usize = 6_000; + // A concurrent test that activates queue storage would route part of the + // batch away from the canonical tables this test reads back. + let _guard = test_lock().lock().await; let client = setup_with_connections(4).await; let queue = "integ_insert_many_bind_limit"; clean_queue(client.pool(), queue).await; diff --git a/benchmarks/baseline.json b/benchmarks/baseline.json index 3fdb7bd08..e4ec128e0 100644 --- a/benchmarks/baseline.json +++ b/benchmarks/baseline.json @@ -33,11 +33,11 @@ "handler_per_s": 160.3 }, "rust/enqueue_lane_16p_unique_185": { - "enqueue_per_s": 2000.0, - "p99_ms": 6000.0 + "enqueue_per_s": 10000.0, + "p99_ms": 2500.0 }, "rust/enqueue_lane_16p_unique_185_shards4": { - "enqueue_per_s": 2000.0 + "enqueue_per_s": 10000.0 }, "rust/mvcc_horizon_overlap": { "min_overlap_vs_baseline": 0.75, diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 1672c5865..ff6085d91 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -434,6 +434,26 @@ Fix: awa --database-url "$DATABASE_URL" migrate ``` +### Producers wait on one queue's enqueue lane + +Symptoms: + +- many connections wait with `wait_event = 'advisory'` while the database is otherwise idle +- `insert_many` takes seconds although the work per batch is milliseconds, and single-job inserts into the same queue take much longer than usual +- scheduled-job promotion and rescue of the same queue run slowly while producers saturate it + +Cause: + +- on queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until it commits, so concurrent producers for one lane run one at a time; a lane's ceiling is a batch's rows divided by the time the lock is held +- from migration v053, `insert_many` and `insert_many_copy` hold the lock across the lane write only, a few milliseconds for a 185-row batch; before it, the lock was held while every row of the batch was processed +- maintenance promotion moves due scheduled and retryable jobs into the lane under the same lock; a promotion batch is one transaction across every queue it touches, and the leader runs its branches one at a time, so a batch parked behind a hot lane delays the due jobs of other queues in that batch and the leader's other branches. Rescue itself does not take the lane lock; it writes retryable rows that promotion moves later + +Fix: + +- run `awa migrate` so producers use the batch function +- when `insert_many` fails on a unique conflict, retry with `insert_many_skip_unique_conflicts` rather than inserting the batch again row by row: each single-row insert queues for the lane separately +- raise `awa.queue_meta.enqueue_shards` for a queue whose producers still saturate one lane, accepting partitioned rather than strict FIFO ([ADR-025](adr/025-sharded-enqueue-heads.md)) + ### `deadlock detected` while enqueueing Cause: diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 0afc48afb..0a114d970 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -253,6 +253,25 @@ workers report not ready until the migration runs, as for any version bump. External migration runners apply the complete `awa migrate --sql --from 51 --to 52` bundle, which defines the helper before the tail calls it. +## v053: batch producer insert + +v053 adds `awa.insert_jobs_compat_batch`, the function `insert_many` and +`insert_many_copy` call. It is one `CREATE OR REPLACE FUNCTION`: no table DDL, +no scan, no lock beyond the catalog, so it applies instantly under live load. +`awa.insert_job_compat` is unchanged, and binaries before v053 keep calling it; +both paths take the same per-lane reservation lock, so producers on either +version serialise on a lane as before. + +The runtime role needs `EXECUTE` on the new function; the +[default privileges](security/database-roles.md#runtime-grants) cover functions +the migrator creates. + +A 0.7 binary that calls the new function before the migration has run fails +its `insert_many` and `insert_many_copy` calls with `function +awa.insert_jobs_compat_batch(...) does not exist`. Migrate before rolling +producers. External migration runners apply the `awa migrate --sql --from 52 +--to 53` bundle. + ## The one gate `awa migrate` on a 0.7 binary refuses to apply migrations unless one of these holds: From 1ba226f21c75482c1f4d0b019a8509b563bdfbe6 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 11:51:42 +1300 Subject: [PATCH 05/18] fix(worker): restrict statistics refresh to ring parents, cover standby leaders, report lock skips - Analyse only the known ring parents, not every partitioned table in the schema, so a substrate sharing a schema with application tables never analyses them. - A leader that resolved canonical at start but holds the queue-storage runtime as standby refreshes the queue-storage schema once routing flips to it, checked on every pass. - Take the parent's SHARE UPDATE EXCLUSIVE lock before ANALYZE (SKIP_LOCKED), so a conflicting lock on the parent fails the call with lock_not_available (recorded as skipped_lock) instead of skipping every partition and reporting analyzed. Correct the documentation of what SKIP_LOCKED covers. - Link the pull request from the migration header. - Tests: refresh on election before the first interval, both rejection guards of the helper from the runtime role, and a locked parent. --- CHANGELOG.md | 5 +- .../v053_analyze_queue_storage_statistics.sql | 17 +++- awa-worker/src/maintenance.rs | 6 +- awa-worker/src/statistics.rs | 52 ++++++++++- awa/tests/statistics_refresh_test.rs | 88 +++++++++++++++++++ docs/configuration.md | 2 +- 6 files changed, 162 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d1f27174..1287eca44 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,8 +17,9 @@ Notable changes between releases. Detailed migration notes for storage transitio minutes. `ClientBuilder::statistics_refresh_interval` (Python `statistics_refresh_interval_ms`) changes the cadence; zero disables it. One pass takes under a second on a small ring and a few seconds on a large one. - It analyses one relation per transaction with a 2 s `lock_timeout`, skips - children that prune holds, and never makes prune wait: a prune attempt that + It analyses one relation per transaction with a 2 s `lock_timeout` (a call that + times out waiting for a lock is retried next interval and recorded as + `skipped_lock`) and never makes prune wait: a prune attempt that lands while ANALYZE samples the oldest slot's tables returns `blocked` and retries on its normal backoff. Results are in `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration` diff --git a/awa-model/migrations/v053_analyze_queue_storage_statistics.sql b/awa-model/migrations/v053_analyze_queue_storage_statistics.sql index a82a52f8a..93d556d12 100644 --- a/awa-model/migrations/v053_analyze_queue_storage_statistics.sql +++ b/awa-model/migrations/v053_analyze_queue_storage_statistics.sql @@ -1,4 +1,4 @@ --- v053: Planner statistics for the queue-storage tables autovacuum never analyses +-- v053: Planner statistics for the queue-storage tables autovacuum never analyses (#631) -- -- Autovacuum does not analyse partitioned tables, so the ring parents -- (ready_entries, done_entries, leases, ...) have no inherited statistics. It @@ -16,8 +16,16 @@ -- The function analyses one relation per call and the caller runs each call as -- its own transaction, so the locks ANALYZE takes on a partitioned table's -- children are held for that table rather than for the whole refresh. --- SKIP_LOCKED skips a child that ring prune holds ACCESS EXCLUSIVE instead of --- waiting for it and aborting the statistics gathered so far. +-- +-- Locking on a partitioned table, per the ANALYZE documentation: SKIP_LOCKED +-- applies when ANALYZE starts work on a relation, so a child that ring prune +-- holds ACCESS EXCLUSIVE is skipped during the per-child analysis, but +-- sampling the children for the parent's statistics can still wait for it. +-- That wait is bounded by the caller's lock_timeout and fails the call with +-- SQLSTATE 55P03, which the caller records as skipped_lock. SKIP_LOCKED also +-- skips every partition silently when the parent itself is locked, so the +-- function takes the parent's SHARE UPDATE EXCLUSIVE lock first and lets that +-- wait fail instead. -- -- Definer boundary: the function only analyses a head table or a partitioned -- table of a schema that awa.queue_storage_schemas() reports, which requires @@ -94,6 +102,9 @@ BEGIN END IF; IF v_relkind = 'p' THEN + EXECUTE format( + 'LOCK TABLE ONLY %I.%I IN SHARE UPDATE EXCLUSIVE MODE', p_schema, p_relation + ); EXECUTE format('ANALYZE (SKIP_LOCKED) %I.%I', p_schema, p_relation); ELSE EXECUTE format('ANALYZE %I.%I', p_schema, p_relation); diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index 042936b9a..782ca622b 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -1336,12 +1336,16 @@ impl MaintenanceService { if self.statistics_refresh_interval.is_zero() { return None; } - let runtime = self.storage.queue_storage()?; + let (runtime, standby) = match self.storage.queue_storage() { + Some(runtime) => (runtime, false), + None => (self.standby_queue_storage.as_ref()?, true), + }; Some(RefreshGuard::spawn( StatisticsRefresher::new( self.pool.clone(), runtime.store.schema().to_string(), self.metrics.clone(), + standby, ), self.statistics_refresh_interval, &self.cancel, diff --git a/awa-worker/src/statistics.rs b/awa-worker/src/statistics.rs index 4fe0aab59..68994c6da 100644 --- a/awa-worker/src/statistics.rs +++ b/awa-worker/src/statistics.rs @@ -26,6 +26,25 @@ const STATEMENT_TIMEOUT: Duration = Duration::from_secs(120); const HEAD_TABLES: [&str; 2] = ["queue_claim_heads", "queue_enqueue_heads"]; +/// The partitioned ring parents a queue-storage substrate installs. A schema +/// can hold other partitioned tables (a substrate in an application schema), +/// which this refresh must never analyse. +const RING_PARENTS: [&str; 13] = [ + "done_entries", + "lease_claim_batches", + "lease_claim_closure_batches", + "lease_claim_closures", + "lease_claims", + "leases", + "queue_terminal_count_deltas", + "ready_claim_attempt_batches", + "ready_entries", + "ready_segments", + "ready_tombstones", + "receipt_completion_batches", + "receipt_completion_tombstones", +]; + const SQLSTATE_INSUFFICIENT_PRIVILEGE: &str = "42501"; const SQLSTATE_UNDEFINED_FUNCTION: &str = "42883"; const SQLSTATE_LOCK_NOT_AVAILABLE: &str = "55P03"; @@ -99,16 +118,26 @@ pub(crate) struct StatisticsRefresher { pool: PgPool, schema: String, metrics: AwaMetrics, + /// Set for a leader whose own storage is canonical and whose queue-storage + /// runtime is only a standby: a pass runs only once routing has flipped to + /// `schema`. + require_active_schema: bool, warned_privilege: AtomicBool, warned_unavailable: AtomicBool, } impl StatisticsRefresher { - pub(crate) fn new(pool: PgPool, schema: String, metrics: AwaMetrics) -> Self { + pub(crate) fn new( + pool: PgPool, + schema: String, + metrics: AwaMetrics, + require_active_schema: bool, + ) -> Self { Self { pool, schema, metrics, + require_active_schema, warned_privilege: AtomicBool::new(false), warned_unavailable: AtomicBool::new(false), } @@ -136,6 +165,25 @@ impl StatisticsRefresher { let started = Instant::now(); let mut summary = PassSummary::default(); + if self.require_active_schema { + match sqlx::query_scalar::<_, Option>( + "SELECT awa.active_queue_storage_schema()", + ) + .fetch_one(&self.pool) + .await + { + Ok(active) if active.as_deref() == Some(self.schema.as_str()) => {} + Ok(_) => { + debug!(schema = %self.schema, "Routing has not flipped to this schema yet; skipping statistics refresh"); + return summary; + } + Err(err) => { + warn!(error = %err, "Could not resolve the active queue-storage schema for statistics refresh"); + return summary; + } + } + } + match awa_model::storage::queue_storage_schema_ready(&self.pool, &self.schema).await { Ok(true) => {} Ok(false) => { @@ -209,10 +257,12 @@ impl StatisticsRefresher { FROM pg_class AS class JOIN pg_namespace AS namespace ON namespace.oid = class.relnamespace WHERE namespace.nspname = $1 AND class.relkind = 'p' + AND class.relname = ANY($2) ORDER BY class.relname "#, ) .bind(&self.schema) + .bind(&RING_PARENTS[..]) .fetch_all(&self.pool) .await } diff --git a/awa/tests/statistics_refresh_test.rs b/awa/tests/statistics_refresh_test.rs index 7215fcf53..aac55c100 100644 --- a/awa/tests/statistics_refresh_test.rs +++ b/awa/tests/statistics_refresh_test.rs @@ -307,6 +307,26 @@ async fn leader_refreshes_statistics_on_election_and_every_interval() { database.drop().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn leader_refreshes_statistics_on_election_before_the_first_interval() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 200).await; + + let client = probe_client(pool.clone(), Duration::from_secs(300), None); + client.start().await.expect("start client"); + wait_for( + "a statistics pass on election", + Duration::from_secs(30), + || async { stats_rows(pool, "awa", "ready_entries", true).await > 0 }, + ) + .await; + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn zero_interval_disables_the_refresh() { let database = TestDatabase::create(None, &[]).await; @@ -354,6 +374,30 @@ async fn non_owner_runtime_role_refreshes_statistics() { "the runtime role must not be able to analyse the tables itself" ); + for (schema, relation, message) in [ + ( + "public", + "queue_claim_heads", + "is not a queue-storage schema", + ), + ("awa", "jobs", "is not a partitioned or head table"), + ] { + let error = + sqlx::query_scalar::<_, String>("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(schema) + .bind(relation) + .fetch_one(&runtime_pool) + .await + .expect_err("the helper rejects what it may not analyse"); + let database_error = error.as_database_error().expect("database error"); + assert_eq!(database_error.code().as_deref(), Some("22023")); + assert!( + database_error.message().contains(message), + "{schema}.{relation}: {}", + database_error.message() + ); + } + let client = probe_client(runtime_pool.clone(), Duration::from_secs(1), None); client .start() @@ -374,6 +418,50 @@ async fn non_owner_runtime_role_refreshes_statistics() { database.drop().await; } +/// A conflicting lock on a ring parent fails the call with `lock_not_available` +/// rather than letting `ANALYZE (SKIP_LOCKED)` skip every partition and report +/// success. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn locked_ring_parent_fails_the_call_instead_of_reporting_analyzed() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + + let mut blocker = pool.begin().await.expect("begin blocker"); + sqlx::query("LOCK TABLE ONLY awa.ready_entries IN SHARE UPDATE EXCLUSIVE MODE") + .execute(blocker.as_mut()) + .await + .expect("lock the parent"); + + let mut call = pool.begin().await.expect("begin call"); + sqlx::query("SELECT set_config('lock_timeout', '200ms', true)") + .execute(call.as_mut()) + .await + .expect("set lock_timeout"); + let error = sqlx::query_scalar::<_, String>( + "SELECT awa.analyze_queue_storage_relation('awa', 'ready_entries')", + ) + .fetch_one(call.as_mut()) + .await + .expect_err("the locked parent must not report analyzed"); + assert_eq!( + error + .as_database_error() + .and_then(|error| error.code()) + .as_deref(), + Some("55P03") + ); + call.rollback().await.expect("rollback call"); + + blocker.rollback().await.expect("release the parent"); + let status: String = + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation('awa', 'ready_entries')") + .fetch_one(pool) + .await + .expect("analyse once the lock is gone"); + assert_eq!(status, "analyzed"); + database.drop().await; +} + fn sum_by_outcome( resource_metrics: &[opentelemetry_sdk::metrics::data::ResourceMetrics], outcome: &str, diff --git a/docs/configuration.md b/docs/configuration.md index 6db6cec59..6d61b5755 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -538,7 +538,7 @@ Autovacuum does not analyse partitioned tables, so the ring parents (`ready_entr The maintenance leader runs `ANALYZE` on both when it is elected and every `statistics_refresh_interval` (default `5m`). One pass takes well under a second on a small ring and a few seconds on a ring with millions of retained rows. It has no measurable effect on claim or completion throughput. -- **Locks.** Each relation is analysed in its own transaction with a 2 s `lock_timeout` and a 120 s `statement_timeout`. Children that ring prune holds are skipped (`SKIP_LOCKED`). Prune never waits for the refresh: its own 10 ms lock timeout applies, so an attempt that lands while `ANALYZE` samples the oldest slot's tables returns `blocked` and retries on its usual backoff. `ANALYZE` of a partitioned table samples every child, so that window cannot be avoided. +- **Locks.** Each relation is analysed in its own transaction with a 2 s `lock_timeout` and a 120 s `statement_timeout`. `SKIP_LOCKED` skips a child that ring prune holds during the per-child analysis, but sampling for the parent's statistics waits for it up to the `lock_timeout`; a call that times out is recorded as `skipped_lock` and retried on the next interval. Prune never waits for the refresh: its own 10 ms lock timeout applies, so an attempt that lands while `ANALYZE` samples the oldest slot's tables returns `blocked` and retries on its usual backoff. `ANALYZE` of a partitioned table samples every child, so that window cannot be avoided. - **Privileges.** The refresh calls `awa.analyze_queue_storage_relation(schema, relation)` (schema v053), a `SECURITY DEFINER` function owned by the role that ran the migration, so the runtime role needs no `MAINTAIN` grant and does not have to own the tables. If the function is missing (run `awa migrate`) or its owner may not analyse the tables, the leader logs one warning that names what is missing and skips the pass; job processing is unaffected. - **Metrics.** `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration`, by `awa.analyze.scope` (`heads`, `ring_parent`) and `awa.analyze.outcome` (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`). A sustained `skipped_lock` rate means something else holds the tables; any `skipped_privilege` or `unavailable` means the refresh is not running. From c824b056e231fe84cd5efc0781dfc637396c2a6b Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 12:06:36 +1300 Subject: [PATCH 06/18] perf(insert): fall back to the per-row insert path on a schema without the batch function insert_many and insert_many_copy used awa.insert_jobs_compat_batch unconditionally, so a binary running ahead of migration v053 failed every batch insert. They now look for the function in the catalog and use the per-row path until it exists. The per-row statement asks the catalog as it inserts, the process keeps the answer, and a call that finds a cached function gone fails once and resets it. insert_many_skip_unique_conflicts needs the function and reports a schema without it. Tests run a binary against a schema without the function, through the migration appearing and the function disappearing again, and mix producers on the per-row and batch functions on one lane. --- CHANGELOG.md | 16 +- .../v053_insert_jobs_batch_compat.sql | 7 +- awa-model/src/insert.rs | 556 ++++++++++++++++-- awa/tests/insert_many_batch_test.rs | 97 +++ awa/tests/insert_many_schema_fallback_test.rs | 206 +++++++ docs/architecture.md | 2 +- docs/configuration.md | 2 + docs/troubleshooting.md | 2 +- docs/upgrade-0.6-to-0.7.md | 13 +- 9 files changed, 837 insertions(+), 64 deletions(-) create mode 100644 awa/tests/insert_many_schema_fallback_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a09350c9e..9cd219c09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,9 @@ Notable changes between releases. Detailed migration notes for storage transitio with a unique violation, so a producer that replays work had to insert the batch again one job at a time, taking the lane lock once per job. The new function returns the inserted rows in input order, with skipped jobs absent. + It needs migration v053 and returns a validation error on an older schema; + `insert_many_copy` (and Python's `insert_many_copy`) skips unique conflicts on + every schema. ### Changed @@ -26,11 +29,14 @@ Notable changes between releases. Detailed migration notes for storage transitio one lane holds the lock for a fraction of the time it did, which raises the per-lane ceiling several times over and cuts the time producers, promotion and rescue spend waiting for the lane. `insert_many_copy` skips unique - conflicts in one pass instead of one savepoint per row. Run `awa migrate` - before upgrading producers: a binary that calls the new function against a - database without v053 fails with `function awa.insert_jobs_compat_batch ... - does not exist`. Binaries before this release keep calling - `awa.insert_job_compat` and share the lane lock with the new path. + conflicts in one pass instead of one savepoint per row. A binary that runs + ahead of the migration keeps working: `insert_many` and `insert_many_copy` + look for the function in the catalog and use the per-row path until it + exists, then switch to the batch function without a restart (the process + keeps the answer; a call that finds a cached function gone fails once with + `undefined_function` and the next call probes again). Binaries before this + release keep calling `awa.insert_job_compat` and share the lane lock with the + new path, so a mixed fleet keeps each batch contiguous and in order. ### Fixed diff --git a/awa-model/migrations/v053_insert_jobs_batch_compat.sql b/awa-model/migrations/v053_insert_jobs_batch_compat.sql index a6f67912a..8463a9ba6 100644 --- a/awa-model/migrations/v053_insert_jobs_batch_compat.sql +++ b/awa-model/migrations/v053_insert_jobs_batch_compat.sql @@ -22,7 +22,12 @@ -- N-1 binaries: unchanged. They keep calling awa.insert_job_compat, whose -- definition this migration does not touch, and the two paths share the lane -- lock, so producers on either version serialise exactly as before. --- Binaries that call the new function need this migration first. +-- Binaries that know the function: insert_many and insert_many_copy look for it +-- in the catalog and use the per-row function until it exists, so they run +-- unchanged before this migration and pick the function up after it, without a +-- restart. insert_many_skip_unique_conflicts needs it and reports a schema +-- without it. +-- External runners: apply this file; nothing else reads it. -- Live load: one CREATE OR REPLACE FUNCTION, no table DDL. CREATE OR REPLACE FUNCTION awa.insert_jobs_compat_batch( diff --git a/awa-model/src/insert.rs b/awa-model/src/insert.rs index 15b5bf18b..79fdde6cb 100644 --- a/awa-model/src/insert.rs +++ b/awa-model/src/insert.rs @@ -4,7 +4,8 @@ use crate::queue_storage::ordering_key_hash64; use crate::unique::compute_unique_key; use crate::JobArgs; use sqlx::postgres::PgConnection; -use sqlx::{PgExecutor, PgPool}; +use sqlx::{FromRow, PgExecutor, PgPool, Row}; +use std::sync::atomic::{AtomicU8, Ordering}; // ── Lane lock order ───────────────────────────────────────────────────── // @@ -12,13 +13,19 @@ use sqlx::{PgExecutor, PgPool}; // advisory lock on its `(queue, priority, enqueue_shard)` lane until the // enclosing transaction ends, so two transactions that enqueue into the same // lanes in opposite orders can deadlock. The batch inserts here go through -// `awa.insert_jobs_compat_batch`, which does the lane-independent work first -// and then reserves its lanes in one canonical order: `(queue, priority, -// shard)` ascending, queues compared bytewise (`COLLATE "C"`, the order of -// Rust's `str`). `QueueStorage` batch enqueues group rows in a `BTreeMap` -// keyed the same way. Holding the lanes only across the reservation and the -// ready-row writes keeps a lane's lock window independent of the batch's -// unique-claim and payload work. +// `awa.insert_jobs_compat_batch` (migration v053), which does the +// lane-independent work first and then reserves its lanes in one canonical +// order: `(queue, priority, shard)` ascending, queues compared bytewise +// (`COLLATE "C"`, the order of Rust's `str`). `QueueStorage` batch enqueues +// group rows in a `BTreeMap` keyed the same way. Holding the lanes only across +// the reservation and the ready-row writes keeps a lane's lock window +// independent of the batch's unique-claim and payload work. +// +// A schema without the function (a binary running ahead of the migration) is +// served by the per-row path: rows are sorted by the same key, stably, before +// they reach `awa.insert_job_compat`, and returned in input order. Both paths +// take the same lane lock, so producers on either, and older binaries, share a +// lane without reordering it. // // The shard is derived as `insert_job_compat` derives it: rows with an // ordering key use `ordering_key_hash64 % enqueue_shards`, other rows @@ -235,6 +242,68 @@ const INSERT_MANY_BATCH_SQL: &str = r#" ) "#; +/// The per-row form of [`INSERT_MANY_BATCH_SQL`], for schemas without +/// `awa.insert_jobs_compat_batch`: one `awa.insert_job_compat` call per row, +/// rows inserted in canonical lane order (see "Lane lock order" above) and +/// returned in input order. `batch_insert_available` reports whether the batch +/// function exists, so the caller can switch to it. +const INSERT_MANY_ROW_FUNCTION_SQL: &str = r#" + WITH input AS MATERIALIZED ( + SELECT + raw.*, + MOD( + CASE WHEN raw.ordering_hash IS NULL + THEN pg_backend_pid()::numeric + ELSE raw.ordering_hash::numeric + END, + COALESCE( + (SELECT meta.enqueue_shards FROM awa.queue_meta AS meta WHERE meta.queue = raw.queue), + 1 + )::numeric + ) AS lane_shard + FROM unnest( + $1::text[], + $2::text[], + $3::jsonb[], + $4::text[], + $5::smallint[], + $6::smallint[], + $7::timestamptz[], + $8::jsonb[], + $9::jsonb[], + $10::bytea[], + $11::text[], + $12::bytea[], + $13::text[] + ) WITH ORDINALITY AS raw( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ordering_key, ordering_hash, ord + ) + ORDER BY raw.queue COLLATE "C", raw.priority, lane_shard, raw.ord + ) + SELECT inserted.*, (SELECT to_regprocedure('awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],smallint[],smallint[],timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') IS NOT NULL ) AS batch_insert_available + FROM input + CROSS JOIN LATERAL awa.insert_job_compat( + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, ord) + ORDER BY tag.ord + ), + input.unique_key, + input.unique_states::bit(8), + input.ordering_key + ) AS inserted + ORDER BY input.ord +"#; + /// Compute unique_key and unique_states from opts. fn compute_unique_fields( kind: &str, @@ -440,8 +509,14 @@ fn ordering_hash_text(row: &PreparedRow) -> Option { /// /// Supports uniqueness constraints — jobs with `unique` opts will have their /// `unique_key` and `unique_states` computed and included. A job whose unique -/// key is already claimed fails the whole batch with a unique violation; use +/// key is already claimed fails the whole batch with a unique violation +/// (SQLSTATE `23505`, as [`AwaError::Database`]); use /// [`insert_many_skip_unique_conflicts`] to insert the rest instead. +/// +/// Uses `awa.insert_jobs_compat_batch` when the schema has migration v053 and +/// the per-row `awa.insert_job_compat` path when it does not, deciding from the +/// schema itself. The first calls of a process use the per-row path until one +/// of them has seen the function. #[tracing::instrument( skip(executor, jobs), fields( @@ -464,7 +539,16 @@ where /// is already claimed. /// /// Returns the rows that were inserted, in input order; a skipped job has no -/// entry in the result. Jobs without `unique` opts are always inserted. +/// entry in the result. Jobs without `unique` opts are always inserted. A key +/// repeated inside the batch inserts its first job. Unlike [`insert_many`], +/// which fails the batch with a unique violation, this is what a caller that +/// replays work wants, and it takes each lane lock once for the batch rather +/// than once per job as a row-by-row retry does. +/// +/// Requires migration v053 (`awa migrate`); on an older schema it returns +/// [`AwaError::Validation`]. [`insert_many_copy`] skips unique conflicts on +/// every schema, through a connection, and is the Python +/// `insert_many_copy`. #[tracing::instrument( skip(executor, jobs), fields( @@ -499,7 +583,128 @@ where } let rows = precompute_rows(jobs)?; - let results = sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL) + + // Skipping conflicts has no per-row equivalent in one statement, so it + // always calls the batch function and reports a schema without it. + if skip_unique_conflicts { + return match bind_batch_columns(sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL), &rows) + .bind(true) + .fetch_all(executor) + .await + { + Ok(results) => { + BATCH_INSERT_FUNCTION.record(BatchFunction::Present); + Ok(results) + } + Err(err) if is_missing_batch_function(&err) => { + BATCH_INSERT_FUNCTION.record(BatchFunction::Absent); + Err(AwaError::Validation( + "insert_many_skip_unique_conflicts needs migration v053; run `awa migrate`" + .to_string(), + )) + } + Err(err) => Err(err.into()), + }; + } + + if BATCH_INSERT_FUNCTION.load() == BatchFunction::Present { + return match bind_batch_columns(sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL), &rows) + .bind(false) + .fetch_all(executor) + .await + { + Ok(results) => Ok(results), + Err(err) => { + // The schema no longer has the function (restored from an + // older backup, or a different database): probe again on the + // next call. This call's executor is spent, and inside a + // transaction it would be aborted anyway. + if is_missing_batch_function(&err) { + BATCH_INSERT_FUNCTION.record(BatchFunction::Unknown); + } + Err(err.into()) + } + }; + } + + // Not known to be present: insert through the per-row function, which + // every schema has, and learn from the same statement whether the batch + // function exists so the next call can use it. + let rows_out = bind_row_function_columns(sqlx::query(INSERT_MANY_ROW_FUNCTION_SQL), &rows) + .fetch_all(executor) + .await?; + let mut results = Vec::with_capacity(rows_out.len()); + let mut available = false; + for row in &rows_out { + results.push(JobRow::from_row(row)?); + available = row.try_get::("batch_insert_available")?; + } + BATCH_INSERT_FUNCTION.record(if available { + BatchFunction::Present + } else { + BatchFunction::Absent + }); + Ok(results) +} + +/// Whether `awa.insert_jobs_compat_batch` exists in the schema this process +/// writes to, as far as it has seen. +/// +/// The executors the insert functions take are consumed by one statement, so +/// the answer comes from the per-row statement, which asks the catalog as it +/// inserts, and is cached for the process: present stays present until a +/// call finds the function missing, absent is asked again by every per-row +/// call and so is picked up as soon as the migration has run. A process that +/// writes to databases at different migration levels may see one failed call +/// when it switches to one that lacks the function. +struct BatchFunctionState(AtomicU8); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BatchFunction { + Unknown, + Present, + Absent, +} + +static BATCH_INSERT_FUNCTION: BatchFunctionState = BatchFunctionState(AtomicU8::new(0)); + +impl BatchFunctionState { + fn load(&self) -> BatchFunction { + match self.0.load(Ordering::Relaxed) { + 1 => BatchFunction::Present, + 2 => BatchFunction::Absent, + _ => BatchFunction::Unknown, + } + } + + fn record(&self, state: BatchFunction) { + self.0.store( + match state { + BatchFunction::Unknown => 0, + BatchFunction::Present => 1, + BatchFunction::Absent => 2, + }, + Ordering::Relaxed, + ); + } +} + +/// `undefined_function` naming the batch function: the schema predates +/// migration v053 or the function was dropped. +fn is_missing_batch_function(err: &sqlx::Error) -> bool { + matches!( + err, + sqlx::Error::Database(db_err) + if db_err.code().as_deref() == Some("42883") + && db_err.message().contains("insert_jobs_compat_batch") + ) +} + +fn bind_batch_columns<'q, O>( + query: sqlx::query::QueryAs<'q, sqlx::Postgres, O, sqlx::postgres::PgArguments>, + rows: &'q [PreparedRow], +) -> sqlx::query::QueryAs<'q, sqlx::Postgres, O, sqlx::postgres::PgArguments> { + query .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) .bind( rows.iter() @@ -532,17 +737,58 @@ where .collect::>(), ) .bind(rows.iter().map(ordering_hash_text).collect::>()) - .bind(skip_unique_conflicts) - .fetch_all(executor) - .await?; +} - Ok(results) +fn bind_row_function_columns<'q>( + query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>, + rows: &'q [PreparedRow], +) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> { + query + .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) + .bind( + rows.iter() + .map(|row| row.queue.as_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| &row.args).collect::>()) + .bind( + rows.iter() + .map(|row| row.state_db_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.priority).collect::>()) + .bind(rows.iter().map(|row| row.max_attempts).collect::>()) + .bind(rows.iter().map(|row| row.run_at).collect::>()) + .bind(rows.iter().map(|row| &row.metadata).collect::>()) + .bind( + rows.iter() + .map(|row| serde_json::json!(row.tags)) + .collect::>(), + ) + .bind( + rows.iter() + .map(|row| row.unique_key.as_deref()) + .collect::>(), + ) + .bind( + rows.iter() + .map(|row| row.unique_states.as_deref()) + .collect::>(), + ) + .bind( + rows.iter() + .map(|row| row.ordering_key.as_deref()) + .collect::>(), + ) + .bind(rows.iter().map(ordering_hash_text).collect::>()) } /// Insert many jobs through the compatibility COPY surface. /// /// Uses a temp staging table with no constraints for fast COPY ingestion, -/// then routes staged rows through `awa.insert_job_compat`. Accepts +/// then inserts them through `awa.insert_jobs_compat_batch`, or, on a schema +/// without migration v053, through `awa.insert_job_compat` row by row. Jobs +/// whose unique key is already claimed are skipped. Accepts /// `&mut PgConnection` so callers can use pool connections or transactions /// (Transaction derefs to PgConnection). For high-volume queue-storage /// producers, prefer [`QueueStorage::enqueue_params_copy`](crate::QueueStorage::enqueue_params_copy) @@ -614,10 +860,77 @@ pub async fn insert_many_copy( copy_in.finish().await?; } - // 3. Insert the staged rows through the batch function, which skips rows - // whose unique key is already claimed when the batch has any. + // 3. Insert the staged rows. With the batch function, unique conflicts are + // skipped in the one statement; without it, each unique row goes through + // the per-row function in its own savepoint. let has_unique = rows.iter().any(|r| r.unique_key.is_some()); - let results = sqlx::query_as::<_, JobRow>( + let results = if batch_function_available(conn).await? { + insert_staged_rows_with_batch_function(conn, has_unique).await? + } else { + insert_staged_rows_per_row(conn, has_unique).await? + }; + + // Keep the session-local staging table reusable across multiple COPY calls + // within the same outer transaction. + sqlx::query("DELETE FROM pg_temp.awa_copy_staging") + .execute(&mut *conn) + .await?; + + Ok(results) +} + +/// Convenience wrapper that acquires a connection from the pool. +/// +/// Wraps the operation in a transaction so the staging rows are cleaned up at +/// commit time even if the caller does not reuse the connection afterward. +#[tracing::instrument( + skip(pool, jobs), + fields( + job.count = jobs.len(), + otel.name = "send", + otel.kind = "producer", + messaging.system = "awa", + messaging.operation.r#type = "send", + messaging.batch.message_count = jobs.len(), + ) +)] +pub async fn insert_many_copy_from_pool( + pool: &PgPool, + jobs: &[InsertParams], +) -> Result, AwaError> { + if jobs.is_empty() { + return Ok(Vec::new()); + } + + let mut tx = pool.begin().await?; + let results = insert_many_copy(&mut tx, jobs).await?; + tx.commit().await?; + + Ok(results) +} + +async fn batch_function_available(conn: &mut PgConnection) -> Result { + if BATCH_INSERT_FUNCTION.load() == BatchFunction::Present { + return Ok(true); + } + let available: bool = sqlx::query_scalar( + "SELECT to_regprocedure('awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],smallint[],smallint[],timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') IS NOT NULL", + ) + .fetch_one(&mut *conn) + .await?; + BATCH_INSERT_FUNCTION.record(if available { + BatchFunction::Present + } else { + BatchFunction::Absent + }); + Ok(available) +} + +async fn insert_staged_rows_with_batch_function( + conn: &mut PgConnection, + skip_unique_conflicts: bool, +) -> Result, AwaError> { + let result = sqlx::query_as::<_, JobRow>( r#" WITH staged AS ( SELECT *, row_number() OVER () AS ord @@ -658,46 +971,187 @@ pub async fn insert_many_copy( ) AS inserted "#, ) - .bind(has_unique) + .bind(skip_unique_conflicts) .fetch_all(&mut *conn) .await?; + Ok(result) +} - // Keep the session-local staging table reusable across multiple COPY calls - // within the same outer transaction. - sqlx::query("DELETE FROM pg_temp.awa_copy_staging") - .execute(&mut *conn) - .await?; - - Ok(results) +/// CTEs that number the staged rows in staging order and re-sort them into +/// canonical lane order (see "Lane lock order" above); `staged_in_lane_order` +/// carries the staging-order position as `ord`. +macro_rules! staged_in_lane_order_ctes { + () => { + r#" + staged AS ( + SELECT *, row_number() OVER () AS ord + FROM pg_temp.awa_copy_staging + ), + staged_in_lane_order AS MATERIALIZED ( + SELECT + staged.*, + MOD( + CASE WHEN staged.ordering_hash IS NULL + THEN pg_backend_pid()::numeric + ELSE staged.ordering_hash::numeric + END, + COALESCE( + (SELECT meta.enqueue_shards FROM awa.queue_meta AS meta WHERE meta.queue = staged.queue), + 1 + )::numeric + ) AS lane_shard + FROM staged + ORDER BY staged.queue COLLATE "C", staged.priority, lane_shard, staged.ord + ) + "# + }; } -/// Convenience wrapper that acquires a connection from the pool. -/// -/// Wraps the operation in a transaction so the staging rows are cleaned up at -/// commit time even if the caller does not reuse the connection afterward. -#[tracing::instrument( - skip(pool, jobs), - fields( - job.count = jobs.len(), - otel.name = "send", - otel.kind = "producer", - messaging.system = "awa", - messaging.operation.r#type = "send", - messaging.batch.message_count = jobs.len(), - ) -)] -pub async fn insert_many_copy_from_pool( - pool: &PgPool, - jobs: &[InsertParams], +async fn insert_staged_rows_per_row( + conn: &mut PgConnection, + has_unique: bool, ) -> Result, AwaError> { - if jobs.is_empty() { - return Ok(Vec::new()); - } + let results: Vec = if has_unique { + // The compatibility `awa.jobs` surface is now a view backed by hot and + // deferred tables, so the old `ON CONFLICT` path is no longer available + // here. Keep COPY for staging/parsing, then insert unique rows one at a + // time and skip duplicates explicitly. + let staged_rows = sqlx::query_as::< + _, + ( + i64, + String, + String, + serde_json::Value, + String, + i16, + i16, + Option>, + serde_json::Value, + Vec, + Option>, + Option, + Option>, + ), + >(concat!( + "WITH", + staged_in_lane_order_ctes!(), + r#" + SELECT + ord, + kind, + queue, + args, + state::text, + priority, + max_attempts, + run_at, + metadata, + tags, + unique_key, + unique_states::text, + ordering_key + FROM staged_in_lane_order + ORDER BY queue COLLATE "C", priority, lane_shard, ord + "# + )) + .fetch_all(&mut *conn) + .await?; - let mut tx = pool.begin().await?; - let results = insert_many_copy(&mut tx, jobs).await?; - tx.commit().await?; + let mut inserted = Vec::with_capacity(staged_rows.len()); + for ( + staged_position, + kind, + queue, + args, + state, + priority, + max_attempts, + run_at, + metadata, + tags, + unique_key, + unique_states, + ordering_key, + ) in staged_rows + { + sqlx::query("SAVEPOINT awa_copy_unique_row") + .execute(&mut *conn) + .await?; + + let result = sqlx::query_as::<_, JobRow>(POSTGRES_INSERT_JOB_SQL) + .bind(&kind) + .bind(&queue) + .bind(&args) + .bind(&state) + .bind(priority) + .bind(max_attempts) + .bind(run_at) + .bind(&metadata) + .bind(&tags) + .bind(&unique_key) + .bind(&unique_states) + .bind(&ordering_key) + .fetch_one(&mut *conn) + .await; + + match result { + Ok(row) => { + inserted.push((staged_position, row)); + sqlx::query("RELEASE SAVEPOINT awa_copy_unique_row") + .execute(&mut *conn) + .await?; + } + Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("23505") => { + sqlx::query("ROLLBACK TO SAVEPOINT awa_copy_unique_row") + .execute(&mut *conn) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_copy_unique_row") + .execute(&mut *conn) + .await?; + continue; + } + Err(err) => { + sqlx::query("ROLLBACK TO SAVEPOINT awa_copy_unique_row") + .execute(&mut *conn) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_copy_unique_row") + .execute(&mut *conn) + .await?; + return Err(AwaError::Database(err)); + } + } + } + inserted.sort_unstable_by_key(|(staged_position, _)| *staged_position); + inserted.into_iter().map(|(_, row)| row).collect() + } else { + sqlx::query_as::<_, JobRow>(concat!( + "WITH", + staged_in_lane_order_ctes!(), + r#" + SELECT inserted.* + FROM staged_in_lane_order AS staged + CROSS JOIN LATERAL awa.insert_job_compat( + staged.kind, + staged.queue, + staged.args, + staged.state, + staged.priority, + staged.max_attempts, + staged.run_at, + staged.metadata, + staged.tags, + staged.unique_key, + staged.unique_states, + staged.ordering_key + ) AS inserted + ORDER BY staged.ord + "# + )) + .fetch_all(&mut *conn) + .await? + }; Ok(results) } diff --git a/awa/tests/insert_many_batch_test.rs b/awa/tests/insert_many_batch_test.rs index aafdf9202..d90c081ce 100644 --- a/awa/tests/insert_many_batch_test.rs +++ b/awa/tests/insert_many_batch_test.rs @@ -424,3 +424,100 @@ async fn batch_waits_for_the_lane_after_its_lane_independent_work() { assert_eq!(inserted.len(), 50); assert_eq!(ready_lane_seqs(&pool, &queue).await.len(), 51); } + +/// Producers on the per-row function, as binaries before v053 are, and +/// producers on the batch function share one lane at the same time. Each batch +/// still occupies a contiguous run of `lane_seq`, in input order, and the lane +/// has no duplicate or missing positions. +#[tokio::test] +async fn old_and_new_producers_share_a_lane_without_reordering() { + const TASKS: usize = 8; + const BATCHES: usize = 25; + const BATCH: i64 = 40; + + let _guard = ENGINE_GUARD.lock().await; + let pool = setup::setup(TASKS as u32 + 4).await; + setup::activate_queue_storage(&pool).await; + let queue = unique_queue("mixed_fleet"); + + // Make sure the batch function is the one `insert_many` uses. + insert_many(&pool, &[job(&queue, -1, false)]) + .await + .expect("warm up"); + + let mut tasks = Vec::new(); + for task in 0..TASKS { + let pool = pool.clone(); + let queue = queue.clone(); + tasks.push(tokio::spawn(async move { + for batch in 0..BATCHES { + let jobs: Vec = (0..BATCH) + .map(|seq| InsertParams { + kind: "insert_many_batch_job".to_string(), + args: serde_json::json!({ "seq": seq, "batch": format!("{task}-{batch}") }), + opts: InsertOpts { + queue: queue.clone(), + ..Default::default() + }, + }) + .collect(); + if task % 2 == 0 { + insert_through_row_function(&pool, &jobs).await; + } else { + insert_many(&pool, &jobs) + .await + .expect("batch function insert"); + } + } + })); + } + for task in tasks { + task.await.expect("producer task"); + } + + let rows: Vec<(String, i64, i64)> = sqlx::query_as( + "SELECT args->>'batch', (args->>'seq')::bigint, lane_seq FROM awa.ready_entries + WHERE queue = $1 AND args ? 'batch' ORDER BY lane_seq", + ) + .bind(&queue) + .fetch_all(&pool) + .await + .expect("read lane"); + assert_eq!(rows.len(), TASKS * BATCHES * BATCH as usize); + for pair in rows.windows(2) { + assert!(pair[1].2 > pair[0].2, "lane_seq is unique and ascending"); + } + let mut by_batch: std::collections::HashMap<&str, Vec<(i64, i64)>> = + std::collections::HashMap::new(); + for (batch, seq, lane_seq) in &rows { + by_batch.entry(batch).or_default().push((*seq, *lane_seq)); + } + for (batch, members) in by_batch { + for (offset, (seq, lane_seq)) in members.iter().enumerate() { + assert_eq!(*seq, offset as i64, "batch {batch} keeps input order"); + assert_eq!( + *lane_seq, + members[0].1 + offset as i64, + "batch {batch} occupies a contiguous run of the lane" + ); + } + } +} + +/// What a binary from before v053 sends: one `insert_job_compat` call per row. +async fn insert_through_row_function(pool: &PgPool, jobs: &[InsertParams]) { + sqlx::query( + r#" + SELECT inserted.* + FROM unnest($1::text[], $2::text[], $3::jsonb[]) WITH ORDINALITY AS input(kind, queue, args, ord) + CROSS JOIN LATERAL awa.insert_job_compat(input.kind, input.queue, input.args) AS inserted + ORDER BY input.ord + "#, + ) + .bind(jobs.iter().map(|job| job.kind.as_str()).collect::>()) + .bind(jobs.iter().map(|job| job.opts.queue.as_str()).collect::>()) + .bind(jobs.iter().map(|job| &job.args).collect::>()) + .fetch_all(pool) + .await + .expect("per-row function insert"); +} diff --git a/awa/tests/insert_many_schema_fallback_test.rs b/awa/tests/insert_many_schema_fallback_test.rs new file mode 100644 index 000000000..7c24c954a --- /dev/null +++ b/awa/tests/insert_many_schema_fallback_test.rs @@ -0,0 +1,206 @@ +//! `insert_many` and `insert_many_copy` against a schema without +//! `awa.insert_jobs_compat_batch`. +//! +//! A binary may run ahead of migration v053. It must keep inserting through +//! the per-row function, start using the batch function once it exists, and +//! recover when the function goes away again. The path a call took is read +//! from `ready_segments`: the per-row function writes one segment per row, the +//! batch function one per lane run. +//! +//! The process-wide answer about the function is part of what is tested, so +//! this binary has one test, against a database of its own. +//! +//! Requires a running Postgres instance. +//! Run with: `cargo test --package awa --test insert_many_schema_fallback_test` + +use awa::audited_sql; +use awa::model::{ + insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, migrations, + AwaError, +}; +use awa::{InsertOpts, InsertParams, UniqueOpts}; +use awa_testing::setup; +use sqlx::PgPool; +use uuid::Uuid; + +const DROP_BATCH_FUNCTION: &str = "DROP FUNCTION awa.insert_jobs_compat_batch(text[], text[], \ + jsonb[], text[], smallint[], smallint[], timestamptz[], jsonb[], jsonb[], bytea[], text[], \ + text[], boolean)"; +const MIGRATION_V053: &str = + include_str!("../../awa-model/migrations/v053_insert_jobs_batch_compat.sql"); + +fn database_url_for(database: &str) -> String { + let url = setup::database_url(); + let (without_query, query) = match url.split_once('?') { + Some((prefix, query)) => (prefix, Some(query)), + None => (url.as_str(), None), + }; + let (base, _) = without_query + .rsplit_once('/') + .expect("database name in URL"); + match query { + Some(query) => format!("{base}/{database}?{query}"), + None => format!("{base}/{database}"), + } +} + +fn jobs(queue: &str, range: std::ops::Range, unique: bool) -> Vec { + range + .map(|seq| InsertParams { + kind: "schema_fallback_job".to_string(), + args: serde_json::json!({ "seq": seq }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + }) + .collect() +} + +async fn segments(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar("SELECT count(*)::bigint FROM awa.ready_segments WHERE queue = $1") + .bind(queue) + .fetch_one(pool) + .await + .expect("count segments") +} + +async fn ready_seqs(pool: &PgPool, queue: &str) -> Vec { + sqlx::query_scalar( + "SELECT (args->>'seq')::bigint FROM awa.ready_entries WHERE queue = $1 ORDER BY lane_seq", + ) + .bind(queue) + .fetch_all(pool) + .await + .expect("read ready rows") +} + +fn seqs(rows: &[awa::JobRow]) -> Vec { + rows.iter() + .map(|row| row.args["seq"].as_i64().expect("seq")) + .collect() +} + +#[tokio::test] +async fn binary_ahead_of_the_migration_falls_back_and_recovers() { + let admin = setup::pool(2).await; + let database = format!("awa_batchfn_{}", &Uuid::new_v4().simple().to_string()[..12]); + sqlx::query(audited_sql(format!("CREATE DATABASE {database}"))) + .execute(&admin) + .await + .expect("create scratch database"); + + let pool = setup::pool_with_url(&database_url_for(&database), 4).await; + migrations::run(&pool).await.expect("migrate"); + setup::activate_queue_storage(&pool).await; + + // A schema from before v053. + sqlx::query(DROP_BATCH_FUNCTION) + .execute(&pool) + .await + .expect("drop batch function"); + + let first = insert_many(&pool, &jobs("fallback_a", 0..3, false)) + .await + .expect("insert_many without the batch function"); + assert_eq!(seqs(&first), vec![0, 1, 2]); + assert_eq!(segments(&pool, "fallback_a").await, 3, "per-row path"); + insert_many(&pool, &jobs("fallback_a", 3..6, false)) + .await + .expect("a second call without the function"); + assert_eq!(segments(&pool, "fallback_a").await, 6, "still per-row"); + assert_eq!( + ready_seqs(&pool, "fallback_a").await, + vec![0, 1, 2, 3, 4, 5] + ); + + // Unique handling on the per-row path: a duplicate fails `insert_many` + // and is skipped by `insert_many_copy`. + insert_many(&pool, &jobs("fallback_u", 0..2, true)) + .await + .expect("seed unique jobs"); + let duplicate = insert_many(&pool, &jobs("fallback_u", 1..3, true)) + .await + .expect_err("a unique conflict fails the batch"); + assert!( + matches!(&duplicate, AwaError::Database(sqlx::Error::Database(e)) if e.code().as_deref() == Some("23505")), + "{duplicate:?}" + ); + let copied = insert_many_copy_from_pool(&pool, &jobs("fallback_u", 1..4, true)) + .await + .expect("insert_many_copy without the batch function"); + assert_eq!(seqs(&copied), vec![2, 3], "duplicate skipped"); + + // The skip API is new with v053 and says so. + let skipped = insert_many_skip_unique_conflicts(&pool, &jobs("fallback_u", 0..1, true)) + .await + .expect_err("needs the migration"); + assert!( + matches!(&skipped, AwaError::Validation(message) if message.contains("v053")), + "{skipped:?}" + ); + + // The migration runs while the process is up: the next call still takes + // the per-row path, sees the function, and the one after uses it. + sqlx::raw_sql(MIGRATION_V053) + .execute(&pool) + .await + .expect("apply v053"); + insert_many(&pool, &jobs("fallback_b", 0..3, false)) + .await + .expect("call that finds the function"); + assert_eq!(segments(&pool, "fallback_b").await, 3, "per-row path"); + insert_many(&pool, &jobs("fallback_b", 3..6, false)) + .await + .expect("call through the batch function"); + assert_eq!( + segments(&pool, "fallback_b").await, + 4, + "the batch function wrote one segment for its run" + ); + assert_eq!( + ready_seqs(&pool, "fallback_b").await, + vec![0, 1, 2, 3, 4, 5] + ); + + let copied = insert_many_copy_from_pool(&pool, &jobs("fallback_u", 2..6, true)) + .await + .expect("insert_many_copy with the batch function"); + assert_eq!(seqs(&copied), vec![4, 5], "duplicates skipped in one pass"); + let skipped = insert_many_skip_unique_conflicts(&pool, &jobs("fallback_u", 0..8, true)) + .await + .expect("skip with the batch function"); + assert_eq!(seqs(&skipped), vec![6, 7]); + + // The function goes away under a process that has seen it: one call fails + // with `undefined_function`, and the next probes again and succeeds. + sqlx::query(DROP_BATCH_FUNCTION) + .execute(&pool) + .await + .expect("drop batch function again"); + let failed = insert_many(&pool, &jobs("fallback_c", 0..3, false)) + .await + .expect_err("the cached function is gone"); + assert!( + matches!(&failed, AwaError::Database(sqlx::Error::Database(e)) if e.code().as_deref() == Some("42883")), + "{failed:?}" + ); + assert!(ready_seqs(&pool, "fallback_c").await.is_empty()); + insert_many(&pool, &jobs("fallback_c", 0..3, false)) + .await + .expect("the call after the failure falls back"); + assert_eq!(segments(&pool, "fallback_c").await, 3, "per-row path"); + + pool.close().await; + sqlx::query(audited_sql(format!( + "DROP DATABASE {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop scratch database"); +} diff --git a/docs/architecture.md b/docs/architecture.md index d3c527f61..e68e9b14a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -152,7 +152,7 @@ Enqueue is transactional: if the producer's outer transaction rolls back, the jo There are two COPY-shaped producer paths: - Direct queue-storage COPY is the high-throughput path: Rust producers use a configured `QueueStorage::enqueue_params_copy()`, and Python producers use `enqueue_many_copy()`. Each ready or deferred write group uses array-backed INSERT below 2,048 rows and CSV COPY for larger groups. Direct-copy producers must use the same queue-storage configuration as the worker fleet, especially `queue_stripe_count` / `queue_storage_queue_stripe_count`. -- `insert_many_copy()` is the compatibility path. It stages rows with an array-backed INSERT below 2,048 rows and CSV COPY at or above that cutoff, then takes one of two branches. If any staged row has a `unique_key`, every row is replayed through `POSTGRES_INSERT_JOB_SQL` inside a per-row savepoint so SQLSTATE `23505` can skip that duplicate without aborting the batch; that statement itself calls `awa.insert_job_compat()`. An all-non-unique batch uses one lateral `awa.insert_job_compat()` call per staged row without savepoints. Both branches therefore preserve cross-storage behavior. The planned `awa.insert_job()` public capability becomes the clean stable SQL contract after #342 ships its migration and conformance artifacts. +- `insert_many_copy()` is the compatibility path. It stages rows with an array-backed INSERT below 2,048 rows and CSV COPY at or above that cutoff, then inserts them through `awa.insert_jobs_compat_batch` (migration v053; each lane's lock is held only for its reservation and the ready-row write, and unique conflicts are skipped in the same statement). On a schema without the function it takes one of two per-row branches. If any staged row has a `unique_key`, every row is replayed through `POSTGRES_INSERT_JOB_SQL` inside a per-row savepoint so SQLSTATE `23505` can skip that duplicate without aborting the batch; that statement itself calls `awa.insert_job_compat()`. An all-non-unique batch uses one lateral `awa.insert_job_compat()` call per staged row without savepoints. Both branches therefore preserve cross-storage behavior. The planned `awa.insert_job()` public capability becomes the clean stable SQL contract after #342 ships its migration and conformance artifacts. Claim is cursor-based rather than heap-scan based: diff --git a/docs/configuration.md b/docs/configuration.md index 930bc6bde..477eb58fb 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -521,6 +521,8 @@ Direct-copy producers must use the same queue-storage routing config as the work `insert_many_copy()` is the compatibility insert API. It is still useful when a caller needs the canonical insert surface, but in queue-storage mode it routes through the compatibility function rather than being the primary producer fast path. +The Rust compatibility inserts differ in how they treat a unique-key conflict. `insert_many` fails the whole batch with a unique violation (SQLSTATE `23505`). `insert_many_skip_unique_conflicts` inserts the rest of the batch and returns only the rows it inserted, in input order, so a producer that replays work does not have to insert the batch again one job at a time (each of those inserts takes the lane lock separately, which on a busy lane is slower than the batch itself). `insert_many_copy` also skips conflicting jobs, and is what Python's `insert_many_copy()` calls; Python has no separate skip function because that is already its behaviour. `insert_many_skip_unique_conflicts` needs migration v053; the other two use `awa.insert_jobs_compat_batch` when the schema has it and the per-row function when it does not. + ### Terminal count rollup Queue storage keeps exact terminal counts without updating mutable counter rows on the completion hot path. `done_entries` terminal inserts append positive rows to `queue_terminal_count_deltas_*`; retry, discard, DLQ move, and compatibility delete of `done_entries` rows append negative rows. Compact receipt completions do not write count deltas: exact count reads add retained `receipt_completion_batches_*` counts and subtract retained `receipt_completion_tombstones_*`. diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index ff6085d91..67bcc18df 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -450,7 +450,7 @@ Cause: Fix: -- run `awa migrate` so producers use the batch function +- run `awa migrate`: producers on the per-row path (a schema before v053) hold the lane for the whole batch, and switch to the batch function as soon as it exists - when `insert_many` fails on a unique conflict, retry with `insert_many_skip_unique_conflicts` rather than inserting the batch again row by row: each single-row insert queues for the lane separately - raise `awa.queue_meta.enqueue_shards` for a queue whose producers still saturate one lane, accepting partitioned rather than strict FIFO ([ADR-025](adr/025-sharded-enqueue-heads.md)) diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 0a114d970..de44cbfab 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -266,11 +266,14 @@ The runtime role needs `EXECUTE` on the new function; the [default privileges](security/database-roles.md#runtime-grants) cover functions the migrator creates. -A 0.7 binary that calls the new function before the migration has run fails -its `insert_many` and `insert_many_copy` calls with `function -awa.insert_jobs_compat_batch(...) does not exist`. Migrate before rolling -producers. External migration runners apply the `awa migrate --sql --from 52 ---to 53` bundle. +A 0.7 binary running before the migration keeps inserting: `insert_many` and +`insert_many_copy` look for the function in the catalog, use the per-row +function until it exists, and switch to the batch function without a restart +once it does. `insert_many_skip_unique_conflicts` needs the function and returns +a validation error before the migration. The per-row and batch paths take the +same lane lock, so a fleet that mixes them keeps each batch contiguous and in +order. External migration runners apply the `awa migrate --sql --from 52 --to +53` bundle. ## The one gate From 629274cd6f90e321c034e3f0df4a3ff45fab825f Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 13:27:19 +1300 Subject: [PATCH 07/18] test(discovery): refuse through the newest migration that runs the discovery preflight The fail-closed test rewound the schema by the newest migration and expected the rerun to hit the substrate-ownership preflight. A migration that touches no custom substrate, such as the batch insert function, carries no preflight, so the test now picks the newest migration that does. --- awa-model/tests/schema_discovery_owner_test.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/awa-model/tests/schema_discovery_owner_test.rs b/awa-model/tests/schema_discovery_owner_test.rs index 3dfc52bbd..914414343 100644 --- a/awa-model/tests/schema_discovery_owner_test.rs +++ b/awa-model/tests/schema_discovery_owner_test.rs @@ -556,8 +556,15 @@ async fn assert_configured_substrates_fail_closed( ) .fetch_one(migrator_pool) .await?; - sqlx::query("DELETE FROM awa.schema_version WHERE version = $1") - .bind(migrations::CURRENT_VERSION) + // The newest migration that runs the discovery preflight. Later ones + // that touch no custom substrate carry no preflight to refuse. + let (preflight_version, _, pending_sql) = migrations::migration_sql() + .into_iter() + .rev() + .find(|(_, _, sql)| sql.contains("unsafe configured queue-storage schema")) + .expect("a migration that runs the discovery preflight"); + sqlx::query("DELETE FROM awa.schema_version WHERE version >= $1") + .bind(preflight_version) .execute(migrator_pool) .await?; let previous_version: i32 = sqlx::query_scalar("SELECT max(version) FROM awa.schema_version") @@ -602,7 +609,6 @@ async fn assert_configured_substrates_fail_closed( "CREATE SCHEMA awa_exp AUTHORIZATION {outsider}" )))) .await?; - let (_, _, pending_sql) = migrations::migration_sql().pop().unwrap(); for (details, expected_schema) in [ ( serde_json::json!({"schema": "injected_custom"}), From 0cb61a66798e58492dd60b626af26bd13b09282a Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 13:27:37 +1300 Subject: [PATCH 08/18] fix(queue-storage): create queue lanes as the schema owner (v053) A lane's enqueue and claim sequences were created by whichever role first touched the lane, through SECURITY INVOKER helpers, so only that role and the roles its default privileges named could use them. Another runtime role then failed with "permission denied for sequence" on enqueue, claim or promotion. Each queue-storage schema now has two SECURITY DEFINER functions owned by the schema owner, provision_lane and provision_lane_sequences. The lane helpers create a missing sequence through them, and each new sequence gets the privileges of the head table it backs. ensure_lane_inserts calls provision_lane when it exists and is executable, inside a savepoint, and otherwise writes the lane rows as before. v053 refreshes the helpers and repairs existing lane sequences one at a time, reassigning and granting only what the migrating role may and reporting the rest. awa.lane_sequence_problems and the new `awa storage lane-sequences` command list what remains with the statements to run; awa.repair_lane_sequences and `--repair` rerun the repair. --- CHANGELOG.md | 24 + awa-cli/src/main.rs | 59 +- .../v023_install_queue_storage_substrate.sql | 416 +++++++++- .../v053_definer_lane_provisioning.sql | 63 ++ awa-model/src/migrations.rs | 8 +- awa-model/src/queue_storage.rs | 106 ++- awa-model/src/storage.rs | 76 ++ .../tests/least_privilege_runtime_test.rs | 41 +- .../fixtures/v052_lane_sequence_helpers.sql | 190 +++++ awa/tests/lane_ownership_roles_test.rs | 735 ++++++++++++++++++ docs/queue-storage-substrate.md | 2 +- docs/security/database-roles.md | 50 +- docs/upgrade-0.6-to-0.7.md | 39 +- skills/awa-operations/SKILL.md | 12 +- 14 files changed, 1734 insertions(+), 87 deletions(-) create mode 100644 awa-model/migrations/v053_definer_lane_provisioning.sql create mode 100644 awa/tests/fixtures/v052_lane_sequence_helpers.sql create mode 100644 awa/tests/lane_ownership_roles_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 065d50db1..6780d6113 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,30 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **A queue lane is usable by every runtime role, whichever role creates it + first (migration v053).** A lane's enqueue and claim sequences belonged to + the role that first enqueued into, claimed from or promoted into it, so with + several runtime roles sharing a group role, any role other than the creator + whose default privileges did not cover it failed with `permission denied for + sequence`. That included the maintenance leader promoting into a shard + created after `enqueue_shards` was raised. Lanes are now created through two + `SECURITY DEFINER` functions per queue-storage schema, `provision_lane` and + `provision_lane_sequences`, owned by the schema owner. Each new sequence gets + the privileges of the head table it backs, explicitly rather than through + default privileges. A runtime without schema `CREATE` can therefore create + lanes itself, and `awa storage prepare-queue` is optional and may run as any + runtime role. v053 gives existing lane sequences to the schema owner and adds + their missing grants, one sequence at a time and only where the migrating + role is allowed to. It never fails on a sequence it cannot reassign. The new + read-only `awa storage lane-sequences` (or + `awa.lane_sequence_problems(schema)`) lists what remains, with the statements + to run and the roles to run them as, and exits 2 while any remain; + `--repair` (or `awa.repair_lane_sequences(schema)`) repairs what the + connected role may. `awa migrate` prints a summary of the leftovers. A binary + running ahead of v053 detects that `provision_lane` is missing and creates + lanes as before, and older binaries create correctly owned lanes once v053 + has refreshed the helpers they call. See "Lane ownership and v053" in the + 0.6 to 0.7 upgrade guide. - **Batch enqueues take their lane locks in one order, so they no longer deadlock with each other.** On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index 405fc3023..5e53264c5 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -328,11 +328,13 @@ fn access_for(command: &Commands) -> context::Access { }, Commands::Storage { command } => match command { StorageCommands::Status => ReadOnly, + StorageCommands::LaneSequences { repair: false, .. } => ReadOnly, StorageCommands::Finalize { check: true, .. } => ReadOnly, StorageCommands::FlipRingAuthority { check: true, .. } => ReadOnly, StorageCommands::Prepare { .. } | StorageCommands::PrepareQueueStorageSchema { .. } | StorageCommands::PrepareQueue { .. } + | StorageCommands::LaneSequences { .. } | StorageCommands::Abort | StorageCommands::EnterMixedTransition { .. } | StorageCommands::Finalize { .. } @@ -811,7 +813,7 @@ enum StorageCommands { #[arg(long)] reset: bool, }, - /// Provision lane sequences as the migrator, before starting a runtime without DDL + /// Provision every priority and shard of a queue ahead of its first enqueue PrepareQueue { /// Queue-storage schema to provision #[arg(long, default_value = "awa")] @@ -826,6 +828,16 @@ enum StorageCommands { #[arg(long, default_value_t = 1)] queue_stripe_count: usize, }, + /// List lane sequences not owned by the schema owner or missing runtime + /// grants, with the statements that fix them. Exits 2 while any remain. + LaneSequences { + /// Queue-storage schema to check + #[arg(long, default_value = "awa")] + schema: String, + /// Reassign and grant what the connected role is allowed to first + #[arg(long)] + repair: bool, + }, /// Abort a prepared or mixed-transition storage rollout before final activation Abort, /// Enter mixed transition and begin routing new writes to the prepared engine @@ -1182,6 +1194,16 @@ async fn run() -> Result<(), Box> { }; awa_model::migrations::run_with_options(&pool, options).await?; println!("Migrations applied successfully."); + for (schema, count) in + awa_model::storage::lane_sequence_problem_counts(&pool).await? + { + eprintln!( + "warning: {count} lane sequence(s) in schema {schema} are not owned by \ + the schema owner or lack runtime grants, and the migrating role could \ + not repair them. Run `awa storage lane-sequences --schema {schema}` for \ + the statements to run and the roles to run them as." + ); + } } } @@ -1933,6 +1955,41 @@ async fn run() -> Result<(), Box> { }))? ); } + StorageCommands::LaneSequences { schema, repair } => { + let remaining = if repair { + let results = + awa_model::storage::repair_lane_sequences(&pool, &schema).await?; + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema": schema, "results": results, + }))? + ); + results + .iter() + .filter(|result| result.outcome != "repaired") + .count() + } else { + let problems = + awa_model::storage::lane_sequence_problems(&pool, &schema).await?; + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema": schema, "problems": problems, + }))? + ); + problems.len() + }; + if remaining > 0 { + eprintln!( + "{remaining} lane sequence(s) in schema {schema} need repair. Run \ + each missing_grants statement as the sequence owner (or a role \ + that can SET ROLE to it), and each reassign_sql statement as a \ + role that can SET ROLE to both the owner and expected_owner." + ); + std::process::exit(2); + } + } StorageCommands::Abort => { awa_model::storage::abort(&pool).await?; let report = awa_model::storage::status_report(&pool).await?; diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index d852b2ac5..b5526cccc 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -42,6 +42,12 @@ LANGUAGE plpgsql SECURITY INVOKER SET search_path = pg_catalog AS $refresh$ +DECLARE + v_schema_owner OID; + v_function REGPROCEDURE; + v_function_owner OID; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; BEGIN EXECUTE format( $ddl$ @@ -69,28 +75,40 @@ BEGIN ); BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_enqueue_seq)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_enqueue_seq USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_enqueue_seq + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_enqueue_seq - ); END IF; IF to_regclass(format('%%I.%%I', %1$L, v_claim_seq)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_claim_seq USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_claim_seq + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_claim_seq - ); END IF; UPDATE %1$I.queue_enqueue_heads @@ -131,16 +149,22 @@ BEGIN v_start BIGINT; BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(NEW.queue, NEW.priority, NEW.enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_seq_name - ); END IF; NEW.seq_name := v_seq_name; @@ -197,16 +221,22 @@ BEGIN ); BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(NEW.queue, NEW.priority, NEW.enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_seq_name - ); END IF; NEW.seq_name := v_seq_name; @@ -221,11 +251,341 @@ BEGIN p_schema ); + -- Lane sequences are created by the schema owner through these two + -- SECURITY DEFINER functions, whichever role first touches the lane, so + -- every runtime role can use them. Their grants mirror the head table each + -- sequence backs (see awa.lane_sequence_problems); they never depend on + -- the creating role's default privileges. + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.provision_lane_sequences( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path = pg_catalog, pg_temp + AS $func$ + DECLARE + v_seq_name TEXT; + v_grant TEXT; + BEGIN + IF p_queue IS NULL OR p_queue = '' + OR p_priority IS NULL OR p_priority NOT BETWEEN 1 AND 4 + OR p_enqueue_shard IS NULL OR p_enqueue_shard NOT BETWEEN 0 AND 63 THEN + RAISE EXCEPTION 'invalid queue lane: queue %%, priority %%, enqueue shard %%', + quote_nullable(p_queue), p_priority, p_enqueue_shard + USING ERRCODE = '22023'; + END IF; + + -- The lane lock QueueStorage takes before it writes lane rows, so + -- concurrent first touches of one lane create its sequences once. + PERFORM pg_advisory_xact_lock(hashtextextended( + format('%%s:%%s:%%s:%%s', %1$L, p_queue, p_priority, p_enqueue_shard), + 0 + )); + + FOREACH v_seq_name IN ARRAY ARRAY[ + %1$I.queue_lane_sequence_name('queue_enqueue_seq', p_queue, p_priority, p_enqueue_shard), + %1$I.queue_lane_sequence_name('queue_claim_seq', p_queue, p_priority, p_enqueue_shard) + ] LOOP + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + FOR v_grant IN + SELECT unnest(problem.missing_grants) + FROM awa.lane_sequence_problems(%1$L, v_seq_name) AS problem + WHERE problem.owner = current_user + LOOP + EXECUTE v_grant; + END LOOP; + END IF; + END LOOP; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.provision_lane( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path = pg_catalog, pg_temp + AS $func$ + BEGIN + -- Validates the lane and takes its lock before any row is written. + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + + INSERT INTO %1$I.queue_lanes (queue, priority) + VALUES (p_queue, p_priority) + ON CONFLICT (queue, priority) DO NOTHING; + + -- The head triggers fire before ON CONFLICT and set the lane's + -- sequence, so rows that already exist are skipped outright. + INSERT INTO %1$I.queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT p_queue, p_priority, p_enqueue_shard + WHERE NOT EXISTS ( + SELECT 1 FROM %1$I.queue_enqueue_heads AS heads + WHERE heads.queue = p_queue + AND heads.priority = p_priority + AND heads.enqueue_shard = p_enqueue_shard + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + INSERT INTO %1$I.queue_claim_heads (queue, priority, enqueue_shard) + SELECT p_queue, p_priority, p_enqueue_shard + WHERE NOT EXISTS ( + SELECT 1 FROM %1$I.queue_claim_heads AS heads + WHERE heads.queue = p_queue + AND heads.priority = p_priority + AND heads.enqueue_shard = p_enqueue_shard + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + PERFORM %1$I.ensure_lane_sequences(p_queue, p_priority, p_enqueue_shard); + END; + $func$ + $ddl$, + p_schema + ); + + SELECT namespace.nspowner INTO v_schema_owner + FROM pg_namespace AS namespace + WHERE namespace.nspname = p_schema; + + FOREACH v_function IN ARRAY ARRAY[ + format('%I.provision_lane_sequences(text,smallint,smallint)', p_schema)::regprocedure, + format('%I.provision_lane(text,smallint,smallint)', p_schema)::regprocedure + ] LOOP + EXECUTE format('GRANT EXECUTE ON FUNCTION %s TO PUBLIC', v_function); + + SELECT routine.proowner INTO v_function_owner + FROM pg_proc AS routine + WHERE routine.oid = v_function; + + IF v_function_owner <> v_schema_owner THEN + IF pg_has_role(current_user, v_function_owner, 'USAGE') + AND pg_has_role(current_user, v_schema_owner, v_set_mode) THEN + EXECUTE format( + 'ALTER FUNCTION %s OWNER TO %I', + v_function, + pg_get_userbyid(v_schema_owner) + ); + ELSE + RAISE WARNING '% is owned by %, not the owner % of schema %; lane sequences it creates will be owned by % too', + v_function, pg_get_userbyid(v_function_owner), pg_get_userbyid(v_schema_owner), + p_schema, pg_get_userbyid(v_function_owner) + USING HINT = format( + 'Run awa migrate as a role that can SET ROLE to %s, or as %s itself.', + quote_ident(pg_get_userbyid(v_schema_owner)), + quote_ident(pg_get_userbyid(v_schema_owner)) + ); + END IF; + END IF; + END LOOP; + END; $refresh$; REVOKE EXECUTE ON FUNCTION awa.refresh_lane_sequence_helpers(TEXT) FROM PUBLIC; +-- Lane sequences in p_schema (or the one named p_sequence) that are not owned +-- by the schema owner or lack a runtime grant. A lane sequence carries a head +-- table's cursor, so whoever may write that head table (INSERT or UPDATE) gets +-- USAGE, SELECT and UPDATE on it, and whoever may read it gets SELECT: +-- queue_enqueue_seq_* mirror queue_enqueue_heads, queue_claim_seq_* mirror +-- queue_claim_heads. missing_grants holds the statements to run as the +-- sequence owner; reassign_sql, when the owner is wrong, needs a role that can +-- SET ROLE to both the current owner and the schema owner. Read-only. +CREATE OR REPLACE FUNCTION awa.lane_sequence_problems(p_schema TEXT, p_sequence TEXT DEFAULT NULL) +RETURNS TABLE ( + sequence_name TEXT, + owner TEXT, + expected_owner TEXT, + missing_grants TEXT[], + reassign_sql TEXT +) +LANGUAGE sql +STABLE +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +AS $problems$ + WITH target AS ( + SELECT namespace.oid, namespace.nspowner + FROM pg_namespace AS namespace + WHERE namespace.nspname = p_schema + ), + lane_sequence AS ( + SELECT sequence.oid, + sequence.relname::text AS relname, + sequence.relowner, + sequence.relacl, + target.nspowner, + CASE WHEN starts_with(sequence.relname::text, 'queue_enqueue_seq_') + THEN 'queue_enqueue_heads' ELSE 'queue_claim_heads' END AS head_table + FROM target + JOIN pg_class AS sequence ON sequence.relnamespace = target.oid + WHERE sequence.relkind = 'S' + AND sequence.relname ~ '^queue_(enqueue|claim)_seq_[0-9a-f]{32}$' + AND (p_sequence IS NULL OR sequence.relname = p_sequence) + ), + wanted AS ( + SELECT DISTINCT lane_sequence.oid, head_privilege.grantee, sequence_privilege.privilege_type + FROM lane_sequence + JOIN pg_class AS head + ON head.relnamespace = (SELECT target.oid FROM target) + AND head.relname = lane_sequence.head_table + CROSS JOIN LATERAL aclexplode(COALESCE(head.relacl, acldefault('r', head.relowner))) AS head_privilege + CROSS JOIN LATERAL unnest( + CASE head_privilege.privilege_type + WHEN 'SELECT' THEN ARRAY['SELECT'] + WHEN 'INSERT' THEN ARRAY['USAGE', 'SELECT', 'UPDATE'] + WHEN 'UPDATE' THEN ARRAY['USAGE', 'SELECT', 'UPDATE'] + ELSE ARRAY[]::TEXT[] + END + ) AS sequence_privilege(privilege_type) + WHERE head_privilege.grantee <> lane_sequence.nspowner + ), + missing AS ( + SELECT wanted.oid, wanted.grantee, wanted.privilege_type + FROM wanted + JOIN lane_sequence ON lane_sequence.oid = wanted.oid + WHERE NOT EXISTS ( + SELECT 1 + FROM aclexplode(COALESCE(lane_sequence.relacl, acldefault('S', lane_sequence.relowner))) AS held + WHERE held.grantee = wanted.grantee + AND held.privilege_type = wanted.privilege_type + ) + ) + SELECT lane_sequence.relname, + pg_get_userbyid(lane_sequence.relowner)::text, + pg_get_userbyid(lane_sequence.nspowner)::text, + ARRAY( + SELECT format( + 'GRANT %s ON SEQUENCE %I.%I TO %s', + string_agg(missing.privilege_type, ', ' ORDER BY missing.privilege_type DESC), + p_schema, + lane_sequence.relname, + CASE missing.grantee WHEN 0 THEN 'PUBLIC' + ELSE quote_ident(pg_get_userbyid(missing.grantee)) END + ) + FROM missing + WHERE missing.oid = lane_sequence.oid + GROUP BY missing.grantee + ORDER BY 1 + ), + CASE WHEN lane_sequence.relowner <> lane_sequence.nspowner THEN format( + 'ALTER SEQUENCE %I.%I OWNER TO %I', + p_schema, lane_sequence.relname, pg_get_userbyid(lane_sequence.nspowner) + ) END + FROM lane_sequence + WHERE lane_sequence.relowner <> lane_sequence.nspowner + OR EXISTS (SELECT 1 FROM missing WHERE missing.oid = lane_sequence.oid) + ORDER BY lane_sequence.relname +$problems$; + +-- Give every lane sequence in p_schema to the schema owner and add its missing +-- runtime grants, one sequence at a time, doing only what the calling role is +-- allowed to: a sequence it cannot reassign keeps its owner and still gets +-- the grants when the caller can act as that owner. Never raises for a +-- sequence it cannot repair. Returns one row per sequence that had a problem: +-- outcome is 'repaired', 'owner_unchanged' (grants complete, owner still +-- wrong) or 'unrepaired', with the statements still to run. +CREATE OR REPLACE FUNCTION awa.repair_lane_sequences(p_schema TEXT) +RETURNS TABLE ( + sequence_name TEXT, + owner TEXT, + expected_owner TEXT, + outcome TEXT, + missing_grants TEXT[], + reassign_sql TEXT +) +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +AS $repair$ +DECLARE + v_problem RECORD; + v_remaining RECORD; + v_has_remaining BOOLEAN; + v_owner OID; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; + v_statement TEXT; +BEGIN + FOR v_problem IN + SELECT * FROM awa.lane_sequence_problems(p_schema) AS problem + LOOP + SELECT sequence.relowner INTO v_owner + FROM pg_class AS sequence + WHERE sequence.oid = to_regclass(format('%I.%I', p_schema, v_problem.sequence_name)); + + -- ALTER ... OWNER TO needs the current owner's privileges, SET ROLE to + -- the new owner, and the new owner's CREATE on the schema (which a + -- schema owner has). Check first so an unrelated owner is skipped + -- rather than raising. + IF v_problem.reassign_sql IS NOT NULL + AND pg_has_role(current_user, v_owner, 'USAGE') + AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.expected_owner)), v_set_mode) THEN + BEGIN + EXECUTE v_problem.reassign_sql; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END IF; + + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, v_problem.sequence_name) AS problem; + v_has_remaining := FOUND; + + IF v_has_remaining + AND cardinality(v_remaining.missing_grants) > 0 + AND pg_has_role(current_user, to_regrole(quote_ident(v_remaining.owner)), 'USAGE') THEN + FOREACH v_statement IN ARRAY v_remaining.missing_grants LOOP + BEGIN + EXECUTE v_statement; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END LOOP; + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, v_problem.sequence_name) AS problem; + v_has_remaining := FOUND; + END IF; + + sequence_name := v_problem.sequence_name; + expected_owner := v_problem.expected_owner; + IF v_has_remaining THEN + owner := v_remaining.owner; + missing_grants := v_remaining.missing_grants; + reassign_sql := v_remaining.reassign_sql; + outcome := CASE WHEN cardinality(v_remaining.missing_grants) = 0 + THEN 'owner_unchanged' ELSE 'unrepaired' END; + ELSE + owner := v_problem.expected_owner; + missing_grants := ARRAY[]::TEXT[]; + reassign_sql := NULL; + outcome := 'repaired'; + END IF; + RETURN NEXT; + END LOOP; +END +$repair$; + +REVOKE EXECUTE ON FUNCTION awa.repair_lane_sequences(TEXT) FROM PUBLIC; + -- Retained job output (ADR-048): add the output columns to an existing -- substrate and expose them through terminal_jobs. Called by the installer -- and by migration v052, which does not replay the installer. Each column is diff --git a/awa-model/migrations/v053_definer_lane_provisioning.sql b/awa-model/migrations/v053_definer_lane_provisioning.sql new file mode 100644 index 000000000..a56d96978 --- /dev/null +++ b/awa-model/migrations/v053_definer_lane_provisioning.sql @@ -0,0 +1,63 @@ +-- v053: Create queue lanes as the schema owner, whichever role touches them first. +-- Lane sequences used to belong to the role that first enqueued into, claimed +-- from or promoted into a lane, and only that role (plus whoever its default +-- privileges named) could use them. refresh_lane_sequence_helpers now installs +-- two SECURITY DEFINER functions, provision_lane and provision_lane_sequences, +-- owned by the schema owner, and the lane helpers create a missing sequence +-- through them. Each new sequence gets the grants of the head table it backs. +-- This migration then repairs existing lane sequences: it gives each one to +-- the schema owner and adds its missing grants, one sequence at a time and +-- only where the migrating role may. Sequences it cannot reassign are reported +-- in a warning and by awa.lane_sequence_problems(); the migration still commits. +-- +-- N-1 binaries call the same helpers with the same arguments and cursor +-- semantics, so lanes they create are owned by the schema owner too. A binary +-- that predates provision_lane keeps using the helpers; this binary probes for +-- provision_lane and uses the helpers when it is missing. +-- +-- Live load: function replacement, plus ALTER SEQUENCE OWNER on the lane +-- sequences that need it. That takes each such sequence's ACCESS EXCLUSIVE +-- lock until commit, so enqueues and claims on those lanes wait for the +-- migration to commit; correctly owned lanes are not touched. Reads the +-- catalog only; no table DDL or row scans. + +DO $$ +DECLARE + v_schema TEXT; + v_leftovers TEXT; + v_count BIGINT; +BEGIN + FOR v_schema IN + SELECT discovered.schema_name FROM awa.queue_storage_schemas() AS discovered + LOOP + IF to_regclass(format('%I.queue_enqueue_heads', v_schema)) IS NULL + OR to_regclass(format('%I.queue_claim_heads', v_schema)) IS NULL + OR to_regclass(format('%I.queue_ring_state', v_schema)) IS NULL + OR to_regclass(format('%I.lease_ring_state', v_schema)) IS NULL + OR to_regclass(format('%I.claim_ring_state', v_schema)) IS NULL THEN + CONTINUE; + END IF; + + PERFORM awa.refresh_lane_sequence_helpers(v_schema); + + SELECT count(*), + string_agg(format('%s (owner %s, %s)', repair.sequence_name, repair.owner, repair.outcome), + ', ' ORDER BY repair.sequence_name) + INTO v_count, v_leftovers + FROM awa.repair_lane_sequences(v_schema) AS repair + WHERE repair.outcome <> 'repaired'; + + IF v_count > 0 THEN + RAISE WARNING '% lane sequence(s) in schema % still need repair: %', v_count, v_schema, v_leftovers + USING HINT = format( + 'Run awa storage lane-sequences --schema %s for the statements to run and the role to run them as.', + v_schema + ); + END IF; + END LOOP; +END; +$$; + +INSERT INTO awa.schema_version (version, description) +VALUES (53, 'Create queue lanes as the schema owner') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 385035b65..51c423222 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -5,7 +5,7 @@ use sqlx::{Connection, PgPool}; use tracing::{info, warn}; /// Current schema version. -pub const CURRENT_VERSION: i32 = 52; +pub const CURRENT_VERSION: i32 = 53; /// Migrations that require an exclusive (no-live-runtime) upgrade window. /// @@ -292,6 +292,11 @@ const MIGRATIONS: &[(i32, &str, &[&str])] = &[ ), (51, "Per-slot claim functions", &[V23_UP, V51_UP]), (52, "Retained job output", &[V23_UP, V52_UP]), + ( + 53, + "Create queue lanes as the schema owner", + &[V23_UP, V53_UP], + ), ]; const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql"); @@ -352,6 +357,7 @@ const V49_UP: &str = include_str!("../migrations/v049_prepared_lane_sequences.sq const V50_UP: &str = include_str!("../migrations/v050_drop_ready_lane_shard_index.sql"); const V51_UP: &str = include_str!("../migrations/v051_claim_slot_functions.sql"); const V52_UP: &str = include_str!("../migrations/v052_retained_job_output.sql"); +const V53_UP: &str = include_str!("../migrations/v053_definer_lane_provisioning.sql"); /// Old version numbers from pre-0.4 releases that used V3/V4/V5 numbering. /// Also tolerates the unreleased inline-V6 branch numbering used during review. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 339001705..dafbbaa14 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -13,7 +13,7 @@ use sqlx::{PgPool, Postgres}; use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::hash::{Hash, Hasher}; -use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU16, AtomicUsize, Ordering}; use std::sync::Mutex; use std::time::{Duration, Instant}; use uuid::Uuid; @@ -3461,6 +3461,12 @@ pub struct QueueStorage { /// TRUNCATEs queue_lanes; `advance_enqueue_head` repairs a stale entry /// (head row gone after a rolled-back ensure_lane). ensured_lanes: Mutex>, + /// Whether `{schema}.provision_lane` (migration v053) was callable the + /// last time a lane was created. Only a positive answer is kept: lane + /// creation is rare, and the probe rides on the lane-lock round trip. + lane_provisioner_present: AtomicBool, + /// Set once a provisioner failure has been logged. + lane_provisioner_warned: AtomicBool, slot_fill_observations: Mutex>, prune_lock_timeout: Duration, cancel_phase_hook: CancelPhaseHookSlot, @@ -3516,6 +3522,8 @@ impl QueueStorage { enqueue_shards_cache: Mutex::new(HashMap::new()), claimer_gates: Mutex::new(HashMap::new()), ensured_lanes: Mutex::new(HashSet::new()), + lane_provisioner_present: AtomicBool::new(false), + lane_provisioner_warned: AtomicBool::new(false), slot_fill_observations: Mutex::new(HashMap::new()), prune_lock_timeout: DEFAULT_PRUNE_LOCK_TIMEOUT, cancel_phase_hook: CancelPhaseHookSlot(None), @@ -4594,9 +4602,10 @@ impl QueueStorage { .await } - /// Provision every priority and shard of a queue before deploying a runtime - /// role without schema CREATE. Run as the schema owner/migrator. This does - /// not change enqueue routing or the configured shard count. + /// Provision every priority and shard of a queue ahead of its first + /// enqueue. From migration v053 any runtime role may run it: the lanes are + /// created as the schema owner. This does not change enqueue routing or + /// the configured shard count. pub async fn prepare_queue( &self, pool: &PgPool, @@ -4637,11 +4646,35 @@ impl QueueStorage { ) -> Result<(), AwaError> { let schema = self.schema(); let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}"); - sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") - .bind(lane_lock_key) - .execute(tx.as_mut()) + let provisioner_callable = if self.lane_provisioner_present.load(Ordering::Relaxed) { + sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") + .bind(&lane_lock_key) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + true + } else { + sqlx::query_scalar( + r#" + SELECT COALESCE(has_function_privilege(to_regprocedure($2), 'EXECUTE'), false) + FROM (SELECT pg_advisory_xact_lock(hashtextextended($1, 0))) AS lane_lock + "#, + ) + .bind(&lane_lock_key) + .bind(format!("{schema}.provision_lane(text,smallint,smallint)")) + .fetch_one(tx.as_mut()) .await - .map_err(map_sqlx_error)?; + .map_err(map_sqlx_error)? + }; + + if provisioner_callable + && self + .provision_lane(tx, queue, priority, enqueue_shard) + .await? + { + self.mark_lane_ensured(queue, priority, enqueue_shard); + return Ok(()); + } sqlx::query(audited_sql(format!( r#" @@ -4712,6 +4745,63 @@ impl QueueStorage { Ok(()) } + /// Create the lane through `{schema}.provision_lane`, which runs as the + /// schema owner, inside a savepoint. Returns false, with the transaction + /// as it was, when the function has gone (`42883`) or the schema owner + /// may not write the lane (`42501`); the caller then writes the rows + /// itself, which is what a binary that predates the function does. + async fn provision_lane<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + queue: &str, + priority: i16, + enqueue_shard: i16, + ) -> Result { + let schema = self.schema(); + sqlx::query("SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let provisioned = sqlx::query(audited_sql(format!( + "SELECT {schema}.provision_lane($1, $2, $3)" + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(tx.as_mut()) + .await; + match provisioned { + Ok(_) => { + sqlx::query("RELEASE SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.lane_provisioner_present.store(true, Ordering::Relaxed); + Ok(true) + } + Err(sqlx::Error::Database(db_err)) + if matches!(db_err.code().as_deref(), Some("42883" | "42501")) => + { + sqlx::query("ROLLBACK TO SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.lane_provisioner_present + .store(false, Ordering::Relaxed); + if !self.lane_provisioner_warned.swap(true, Ordering::Relaxed) { + tracing::warn!( + schema, + code = db_err.code().as_deref().unwrap_or_default(), + error = %db_err.message(), + "Lane provisioning function failed; creating lane rows as the connected role" + ); + } + Ok(false) + } + Err(err) => Err(map_sqlx_error(err)), + } + } + /// Pick the enqueue shard for a row on `queue`. /// /// Reads `awa.queue_meta.enqueue_shards` (default 1, cached diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index b7a0ff1ad..d67dc2048 100644 --- a/awa-model/src/storage.rs +++ b/awa-model/src/storage.rs @@ -469,3 +469,79 @@ pub async fn status_report(pool: &PgPool) -> Result, + pub reassign_sql: Option, +} + +/// What [`repair_lane_sequences`] did to one lane sequence that had a problem. +#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq, Eq)] +pub struct LaneSequenceRepair { + pub sequence_name: String, + /// Owner after the repair. + pub owner: String, + pub expected_owner: String, + /// `repaired`; `owner_unchanged` when the grants are complete but the + /// connected role could not reassign the sequence; or `unrepaired` when + /// grants are still missing too. + pub outcome: String, + pub missing_grants: Vec, + pub reassign_sql: Option, +} + +/// Lane sequences in `schema` that need repair. Read-only; any role may run it. +pub async fn lane_sequence_problems<'e, E>( + executor: E, + schema: &str, +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + Ok(sqlx::query_as( + "SELECT sequence_name, owner, expected_owner, missing_grants, reassign_sql \ + FROM awa.lane_sequence_problems($1)", + ) + .bind(schema) + .fetch_all(executor) + .await?) +} + +/// Give each lane sequence in `schema` to the schema owner and add its missing +/// grants, as far as the connected role is allowed to. A sequence the role +/// cannot reassign is returned with the statements still to run, never an +/// error. +pub async fn repair_lane_sequences( + pool: &PgPool, + schema: &str, +) -> Result, AwaError> { + Ok(sqlx::query_as( + "SELECT sequence_name, owner, expected_owner, outcome, missing_grants, reassign_sql \ + FROM awa.repair_lane_sequences($1)", + ) + .bind(schema) + .fetch_all(pool) + .await?) +} + +/// How many lane sequences need repair in each queue-storage schema that has +/// any. Read-only. +pub async fn lane_sequence_problem_counts(pool: &PgPool) -> Result, AwaError> { + Ok(sqlx::query_as( + "SELECT discovered.schema_name, count(*) \ + FROM awa.queue_storage_schemas() AS discovered \ + CROSS JOIN LATERAL awa.lane_sequence_problems(discovered.schema_name) AS problem \ + GROUP BY discovered.schema_name \ + ORDER BY discovered.schema_name", + ) + .fetch_all(pool) + .await?) +} diff --git a/awa-model/tests/least_privilege_runtime_test.rs b/awa-model/tests/least_privilege_runtime_test.rs index 99e69d536..db66046a4 100644 --- a/awa-model/tests/least_privilege_runtime_test.rs +++ b/awa-model/tests/least_privilege_runtime_test.rs @@ -250,21 +250,26 @@ async fn check_runtime(pool: &PgPool, schema: &str) -> TestResult { INSERT INTO {schema}.queue_claim_heads(queue, priority, enqueue_shard) VALUES ('prepared_queue', 2, 0) ON CONFLICT DO NOTHING; SELECT {schema}.reserve_enqueue_seq('prepared_queue', 2::smallint, 0::smallint, 3); "#)))).await?; - let error = pool - .execute(sqlx::raw_sql(awa_model::audited_sql(format!( - "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint)" - )))) - .await - .unwrap_err(); - assert_eq!( - error.as_database_error().unwrap().code().as_deref(), - Some("42501") - ); - assert!(error - .as_database_error() - .unwrap() - .message() - .contains("has not been provisioned")); + // An unprepared lane is created through the schema owner's provisioning + // function, so the runtime still issues no DDL of its own. + pool.execute(sqlx::raw_sql(awa_model::audited_sql(format!( + "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint); \ + SELECT {schema}.provision_lane('unprepared_queue', 2::smallint, 0::smallint);" + )))) + .await?; + let foreign_owned: i64 = sqlx::query_scalar(awa_model::audited_sql(format!( + "SELECT count(*) FROM pg_class \ + WHERE relnamespace = '{schema}'::regnamespace AND relkind = 'S' \ + AND relowner <> (SELECT nspowner FROM pg_namespace WHERE nspname = '{schema}')" + ))) + .fetch_one(pool) + .await?; + assert_eq!(foreign_owned, 0); + sqlx::query_scalar::<_, i64>(awa_model::audited_sql(format!( + "SELECT {schema}.reserve_enqueue_seq('unprepared_queue', 2::smallint, 0::smallint, 1)" + ))) + .fetch_one(pool) + .await?; let error = pool .execute(sqlx::raw_sql(awa_model::audited_sql(format!( "CREATE TABLE {schema}.must_not_create(id integer)" @@ -331,10 +336,14 @@ async fn exercise_lazy_lane(options: PgConnectOptions, owner: &str, runtime: &st "GRANT USAGE, CREATE ON SCHEMA awa TO {runtime}; \ GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {runtime}; \ GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {runtime}; \ - GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {runtime};" + GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {runtime}; \ + REVOKE EXECUTE ON FUNCTION awa.provision_lane_sequences(text, smallint, smallint) \ + FROM PUBLIC, {runtime};" ), ) .await?; + // Without the provisioning function the runtime creates the lane's + // sequences itself and owns them, as a 0.6 runtime with CREATE did. let runtime_pool = PgPoolOptions::new() .max_connections(2) .connect_with(options.username(runtime)) diff --git a/awa/tests/fixtures/v052_lane_sequence_helpers.sql b/awa/tests/fixtures/v052_lane_sequence_helpers.sql new file mode 100644 index 000000000..b7d396dad --- /dev/null +++ b/awa/tests/fixtures/v052_lane_sequence_helpers.sql @@ -0,0 +1,190 @@ +-- The lane-sequence helpers as migration v052 left them, before lane +-- provisioning functions existed. Installed under a separate name so a test +-- can put a schema back in its v052 shape. +CREATE OR REPLACE FUNCTION awa.v052_refresh_lane_sequence_helpers(p_schema TEXT) +RETURNS VOID +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog +AS $refresh$ +BEGIN + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.ensure_lane_sequences( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_enqueue_seq TEXT := %1$I.queue_lane_sequence_name( + 'queue_enqueue_seq', + p_queue, + p_priority, + p_enqueue_shard + ); + v_claim_seq TEXT := %1$I.queue_lane_sequence_name( + 'queue_claim_seq', + p_queue, + p_priority, + p_enqueue_shard + ); + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_enqueue_seq)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_enqueue_seq + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_enqueue_seq + ); + END IF; + IF to_regclass(format('%%I.%%I', %1$L, v_claim_seq)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_claim_seq + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_claim_seq + ); + END IF; + + UPDATE %1$I.queue_enqueue_heads + SET seq_name = v_enqueue_seq + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard + AND seq_name IS DISTINCT FROM v_enqueue_seq; + + UPDATE %1$I.queue_claim_heads + SET seq_name = v_claim_seq + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard + AND seq_name IS DISTINCT FROM v_claim_seq; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.queue_enqueue_head_sequence_sync() + RETURNS TRIGGER + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_seq_name TEXT := %1$I.queue_lane_sequence_name( + 'queue_enqueue_seq', + NEW.queue, + NEW.priority, + NEW.enqueue_shard + ); + v_qualified_seq TEXT := format('%%I.%%I', %1$L, v_seq_name); + v_count BIGINT; + v_start BIGINT; + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + END IF; + NEW.seq_name := v_seq_name; + + IF TG_OP = 'UPDATE' + AND NEW.next_seq IS DISTINCT FROM OLD.next_seq + AND NEW.next_seq > OLD.next_seq + THEN + v_count := NEW.next_seq - OLD.next_seq; + PERFORM pg_advisory_lock(hashtextextended(v_qualified_seq, 0)); + BEGIN + EXECUTE format( + 'SELECT nextval(%%L::regclass)::bigint', + v_qualified_seq + ) + INTO v_start; + + IF v_count > 1 THEN + EXECUTE format( + 'SELECT setval(%%L::regclass, %%s, true)', + v_qualified_seq, + v_start + v_count - 1 + ); + END IF; + EXCEPTION WHEN OTHERS THEN + PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); + RAISE; + END; + PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); + NEW.next_seq := v_start + v_count; + ELSE + PERFORM %1$I.set_sequence_next(v_seq_name, NEW.next_seq); + END IF; + + RETURN NEW; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.queue_claim_head_sequence_sync() + RETURNS TRIGGER + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_seq_name TEXT := %1$I.queue_lane_sequence_name( + 'queue_claim_seq', + NEW.queue, + NEW.priority, + NEW.enqueue_shard + ); + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + END IF; + NEW.seq_name := v_seq_name; + + IF TG_OP = 'INSERT' THEN + PERFORM %1$I.set_sequence_next(v_seq_name, NEW.claim_seq); + END IF; + + RETURN NEW; + END; + $func$ + $ddl$, + p_schema + ); + +END; +$refresh$; diff --git a/awa/tests/lane_ownership_roles_test.rs b/awa/tests/lane_ownership_roles_test.rs new file mode 100644 index 000000000..e5aadbd39 --- /dev/null +++ b/awa/tests/lane_ownership_roles_test.rs @@ -0,0 +1,735 @@ +//! Queue lanes under a split-role install: a schema owner, a migration role +//! that is a member of it, and several application roles that share a runtime +//! group role. Only the owner and the migration role have default privileges, +//! so a lane sequence is usable by every runtime role only if it is created as +//! the owner or granted explicitly, whichever role touches the lane first. + +use awa::model::{migrations, storage, AwaError, InsertOpts, InsertParams, QueueStorage}; +use awa::model::{JobState, QueueStorageConfig}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::{Executor, PgPool}; +use std::str::FromStr; +use std::time::Duration; + +type TestResult = Result<(), Box>; + +fn sql(text: impl Into) -> sqlx::AssertSqlSafe { + awa::audited_sql(text) +} + +/// Roles, a database owned by the schema owner, and a pool per role. +struct Cluster { + admin: PgPool, + database: String, + roles: Vec, + owner_name: String, + migrator_name: String, + runtime_name: String, + owner: PgPool, + migrator: PgPool, + app_a: PgPool, + app_b: PgPool, + app_c: PgPool, +} + +#[derive(Clone, Copy, PartialEq)] +enum Schema { + /// Every migration, including the lane provisioning functions. + Current, + /// The release before lane provisioning (v052). + BeforeProvisioning, +} + +impl Cluster { + /// `runtime_create` grants the runtime group CREATE on the schema, as a + /// 0.6 install does so runtimes can create lane sequences themselves. + async fn new( + schema: Schema, + runtime_create: bool, + ) -> Result> { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let admin = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await?; + let suffix = uuid::Uuid::new_v4().simple().to_string()[..12].to_string(); + let owner_name = format!("lane_owner_{suffix}"); + let migrator_name = format!("lane_migrator_{suffix}"); + let runtime_name = format!("lane_runtime_{suffix}"); + let app_names = ["a", "b", "c"].map(|app| format!("lane_app_{app}_{suffix}")); + let database = format!("lane_roles_{suffix}"); + admin + .execute(sqlx::raw_sql(sql(format!( + r#" + CREATE ROLE {owner_name} LOGIN PASSWORD 'lane_test'; + CREATE ROLE {migrator_name} LOGIN PASSWORD 'lane_test' IN ROLE {owner_name}; + CREATE ROLE {runtime_name} NOLOGIN; + CREATE ROLE {a} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + CREATE ROLE {b} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + CREATE ROLE {c} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + "#, + a = app_names[0], + b = app_names[1], + c = app_names[2], + )))) + .await?; + admin + .execute(sqlx::raw_sql(sql(format!( + "CREATE DATABASE {database} OWNER {owner_name}" + )))) + .await?; + let mut roles = vec![ + owner_name.clone(), + migrator_name.clone(), + runtime_name.clone(), + ]; + roles.extend(app_names.iter().cloned()); + let options = PgConnectOptions::from_str(&url)? + .database(&database) + .password("lane_test"); + let pool = |user: &str| { + PgPoolOptions::new() + .max_connections(3) + .connect_with(options.clone().username(user)) + }; + let owner = pool(&owner_name).await?; + let migrator = pool(&migrator_name).await?; + + owner + .execute(sqlx::raw_sql(sql(format!( + r#" + CREATE SCHEMA awa; + GRANT USAGE ON SCHEMA awa TO {runtime_name}; + "# + )))) + .await?; + for (creator, creator_pool) in [(&owner_name, &owner), (&migrator_name, &migrator)] { + creator_pool + .execute(sqlx::raw_sql(sql(format!( + r#" + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON TABLES TO {runtime_name}; + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO {runtime_name}; + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT EXECUTE ON FUNCTIONS TO {runtime_name}; + "# + )))) + .await?; + } + // The migration role runs every migration with the owner as its + // effective role, as the database-roles guide describes. + let migrate_as_owner = PgPoolOptions::new() + .max_connections(1) + .connect_with( + options + .clone() + .username(&migrator_name) + .options([("role", owner_name.as_str())]), + ) + .await?; + let target = match schema { + Schema::Current => migrations::CURRENT_VERSION, + Schema::BeforeProvisioning => 52, + }; + for (_, _, statement) in migrations::migration_sql_range(0, target) { + migrate_as_owner + .execute(sqlx::raw_sql(sql(statement))) + .await?; + } + if schema == Schema::BeforeProvisioning { + // Every migration bundle carries the current helper definitions, + // so put the v052 helpers back and drop what v052 did not have. + migrate_as_owner + .execute(sqlx::raw_sql(sql(include_str!( + "fixtures/v052_lane_sequence_helpers.sql" + )))) + .await?; + migrate_as_owner + .execute(sqlx::raw_sql(sql(r#" + SELECT awa.v052_refresh_lane_sequence_helpers('awa'); + DROP FUNCTION awa.v052_refresh_lane_sequence_helpers(text); + DROP FUNCTION awa.provision_lane(text, smallint, smallint); + DROP FUNCTION awa.provision_lane_sequences(text, smallint, smallint); + DROP FUNCTION awa.repair_lane_sequences(text); + DROP FUNCTION awa.lane_sequence_problems(text, text); + "#))) + .await?; + } + migrate_as_owner.close().await; + if runtime_create { + owner + .execute(sqlx::raw_sql(sql(format!( + "GRANT CREATE ON SCHEMA awa TO {runtime_name}" + )))) + .await?; + } + + let app_a = pool(&app_names[0]).await?; + let app_b = pool(&app_names[1]).await?; + let app_c = pool(&app_names[2]).await?; + Ok(Self { + admin, + database, + roles, + owner_name, + migrator_name, + runtime_name, + owner, + migrator, + app_a, + app_b, + app_c, + }) + } + + async fn teardown(self) -> TestResult { + for pool in [ + &self.owner, + &self.migrator, + &self.app_a, + &self.app_b, + &self.app_c, + ] { + pool.close().await; + } + self.admin + .execute(sqlx::raw_sql(sql(format!( + "DROP DATABASE {} WITH (FORCE)", + self.database + )))) + .await?; + for role in self.roles.iter().rev() { + self.admin + .execute(sqlx::raw_sql(sql(format!("DROP ROLE {role}")))) + .await?; + } + self.admin.close().await; + Ok(()) + } + + /// `(sequence, owner)` for every lane sequence of `queue`. + async fn lane_sequence_owners( + &self, + queue: &str, + ) -> Result, sqlx::Error> { + sqlx::query_as( + r#" + SELECT sequence.relname::text, pg_get_userbyid(sequence.relowner)::text + FROM awa.queue_enqueue_heads AS heads + CROSS JOIN LATERAL (VALUES (heads.seq_name), (( + SELECT claims.seq_name FROM awa.queue_claim_heads AS claims + WHERE claims.queue = heads.queue + AND claims.priority = heads.priority + AND claims.enqueue_shard = heads.enqueue_shard + ))) AS lane(seq_name) + JOIN pg_class AS sequence + ON sequence.relname = lane.seq_name + AND sequence.relnamespace = 'awa'::regnamespace + WHERE heads.queue = $1 + ORDER BY 1 + "#, + ) + .bind(queue) + .fetch_all(&self.owner) + .await + } + + async fn assert_lanes_owned_by_owner(&self, queue: &str, lanes: usize) -> TestResult { + let owners = self.lane_sequence_owners(queue).await?; + assert_eq!( + owners.len(), + lanes * 2, + "lane sequences of {queue}: {owners:?}" + ); + for (sequence, owner) in owners { + assert_eq!(owner, self.owner_name, "{sequence} of {queue}"); + } + Ok(()) + } + + async fn set_enqueue_shards(&self, queue: &str, shards: i16) -> Result<(), sqlx::Error> { + sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) \ + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(&self.owner) + .await?; + Ok(()) + } +} + +fn store() -> QueueStorage { + QueueStorage::new(QueueStorageConfig::default()).expect("default queue storage config") +} + +/// Enqueue `count` jobs through a fresh store, as a producer process would. +async fn enqueue(pool: &PgPool, queue: &str, priority: i16, count: i64) -> Result { + store().enqueue_batch(pool, queue, priority, count).await +} + +/// Claim and complete everything on `queue`, as a worker does; returns how +/// many jobs completed. +async fn drain(pool: &PgPool, queue: &str) -> Result { + let worker = store(); + let mut completed = 0; + loop { + let claimed = worker + .claim_runtime_batch(pool, queue, 100, Duration::from_secs(60)) + .await?; + if claimed.is_empty() { + return Ok(completed); + } + completed += worker.complete_runtime_batch(pool, &claimed).await?.len(); + } +} + +/// Schedule jobs that become due shortly, one per priority. +async fn schedule(pool: &PgPool, queue: &str, priorities: &[i16]) -> Result<(), AwaError> { + let run_at = chrono::Utc::now() + chrono::Duration::milliseconds(200); + let jobs: Vec = priorities + .iter() + .map(|&priority| InsertParams { + kind: "lane_ownership_job".into(), + args: serde_json::json!({}), + opts: InsertOpts { + queue: queue.into(), + priority, + run_at: Some(run_at), + ..Default::default() + }, + }) + .collect(); + store().enqueue_params_batch(pool, &jobs).await?; + tokio::time::sleep(Duration::from_millis(300)).await; + Ok(()) +} + +async fn first_touch_by_one_app_role(runtime_create: bool) -> TestResult { + let cluster = Cluster::new(Schema::Current, runtime_create).await?; + + // Role A creates the lane; role B enqueues, claims and completes on it. + assert_eq!(enqueue(&cluster.app_a, "fresh_queue", 3, 2).await?, 2); + assert_eq!(enqueue(&cluster.app_b, "fresh_queue", 3, 2).await?, 2); + assert_eq!(drain(&cluster.app_b, "fresh_queue").await?, 4); + + // The maintenance leader, running as role C, promotes scheduled jobs into + // lanes nobody has created yet: two priorities over two shards. + cluster.set_enqueue_shards("fresh_queue", 2).await?; + schedule(&cluster.app_a, "fresh_queue", &[1, 4]).await?; + assert_eq!( + store() + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 2 + ); + let promoted_shards: Vec = sqlx::query_scalar( + "SELECT enqueue_shard FROM awa.queue_enqueue_heads \ + WHERE queue = 'fresh_queue' AND priority IN (1, 4) ORDER BY 1", + ) + .fetch_all(&cluster.owner) + .await?; + assert_eq!(promoted_shards, vec![0, 1]); + + // A and B both use the lanes the leader created; each producer's rotor + // visits both shards. + let (producer_a, producer_b) = (store(), store()); + for _ in 0..2 { + producer_a + .enqueue_batch(&cluster.app_a, "fresh_queue", 4, 1) + .await?; + producer_b + .enqueue_batch(&cluster.app_b, "fresh_queue", 1, 1) + .await?; + } + assert_eq!(drain(&cluster.app_a, "fresh_queue").await?, 6); + + cluster + .assert_lanes_owned_by_owner("fresh_queue", 5) + .await?; + assert!(storage::lane_sequence_problems(&cluster.owner, "awa") + .await? + .is_empty()); + cluster.teardown().await +} + +#[tokio::test] +async fn lane_first_touched_by_an_app_role_with_schema_create_serves_every_runtime_role( +) -> TestResult { + first_touch_by_one_app_role(true).await +} + +#[tokio::test] +async fn lane_first_touched_by_an_app_role_without_schema_create_serves_every_runtime_role( +) -> TestResult { + first_touch_by_one_app_role(false).await +} + +#[tokio::test] +async fn raising_enqueue_shards_live_lets_the_leader_create_a_shard_first() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + enqueue(&cluster.app_a, "sharded_queue", 2, 3).await?; + + cluster.set_enqueue_shards("sharded_queue", 2).await?; + // The leader's store has already routed a promotion to shard 0, as a + // long-running leader has, so its next pick is the new shard. + let leader = store(); + schedule(&cluster.app_a, "sharded_queue", &[2]).await?; + assert_eq!( + leader + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 1 + ); + schedule(&cluster.app_b, "sharded_queue", &[2]).await?; + assert_eq!( + leader + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 1 + ); + let shards: Vec = sqlx::query_scalar( + "SELECT enqueue_shard FROM awa.queue_enqueue_heads WHERE queue = 'sharded_queue' ORDER BY 1", + ) + .fetch_all(&cluster.owner) + .await?; + assert_eq!(shards, vec![0, 1]); + + // Producers on both roles spread over both shards; B drains them. + let producer_a = store(); + let producer_b = store(); + for _ in 0..4 { + producer_a + .enqueue_batch(&cluster.app_a, "sharded_queue", 2, 1) + .await?; + producer_b + .enqueue_batch(&cluster.app_b, "sharded_queue", 2, 1) + .await?; + } + assert_eq!(drain(&cluster.app_b, "sharded_queue").await?, 13); + cluster + .assert_lanes_owned_by_owner("sharded_queue", 2) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn prepare_queue_as_an_app_role_provisions_lanes_every_runtime_role_can_use() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + store() + .prepare_queue(&cluster.app_a, "prepared_queue", 2) + .await?; + cluster + .assert_lanes_owned_by_owner("prepared_queue", 8) + .await?; + + cluster.set_enqueue_shards("prepared_queue", 2).await?; + let producer = store(); + for priority in 1..=4 { + producer + .enqueue_batch(&cluster.app_b, "prepared_queue", priority, 1) + .await?; + producer + .enqueue_batch(&cluster.app_c, "prepared_queue", priority, 1) + .await?; + } + assert_eq!(drain(&cluster.app_b, "prepared_queue").await?, 8); + cluster + .assert_lanes_owned_by_owner("prepared_queue", 8) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn mixed_fleet_shares_lanes_whichever_path_creates_them() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + // Role A cannot call provision_lane, so its stores take the path a binary + // that predates the function takes: they write the lane rows themselves. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE EXECUTE ON FUNCTION awa.provision_lane(text, smallint, smallint) FROM PUBLIC, {}", + cluster.runtime_name + )))) + .await?; + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "GRANT EXECUTE ON FUNCTION awa.provision_lane(text, smallint, smallint) TO {}, {}", + cluster.roles[4], cluster.roles[5] + )))) + .await?; + let can_provision: bool = sqlx::query_scalar( + "SELECT has_function_privilege('awa.provision_lane(text,smallint,smallint)', 'EXECUTE')", + ) + .fetch_one(&cluster.app_a) + .await?; + assert!(!can_provision); + + enqueue(&cluster.app_a, "fleet_queue", 1, 2).await?; + enqueue(&cluster.app_b, "fleet_queue", 2, 2).await?; + enqueue(&cluster.app_b, "fleet_queue", 1, 2).await?; + enqueue(&cluster.app_a, "fleet_queue", 2, 2).await?; + assert_eq!(drain(&cluster.app_a, "fleet_queue").await?, 8); + cluster + .assert_lanes_owned_by_owner("fleet_queue", 2) + .await?; + + // A provisioner that disappears after a store has used it: the next lane + // falls back inside the same transaction instead of failing the enqueue. + let producer = store(); + producer + .enqueue_batch(&cluster.app_b, "fleet_queue", 3, 1) + .await?; + cluster + .owner + .execute("DROP FUNCTION awa.provision_lane(text, smallint, smallint)") + .await?; + producer + .enqueue_batch(&cluster.app_b, "fleet_queue", 4, 1) + .await?; + enqueue(&cluster.app_c, "fleet_queue", 4, 1).await?; + assert_eq!(drain(&cluster.app_c, "fleet_queue").await?, 3); + cluster + .assert_lanes_owned_by_owner("fleet_queue", 4) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn binary_ahead_of_the_schema_creates_lanes_as_before() -> TestResult { + let cluster = Cluster::new(Schema::BeforeProvisioning, false).await?; + let has_provisioner: bool = sqlx::query_scalar( + "SELECT to_regprocedure('awa.provision_lane(text,smallint,smallint)') IS NOT NULL", + ) + .fetch_one(&cluster.owner) + .await?; + assert!(!has_provisioner); + + // The owner creates lanes lazily, as on v052. + enqueue(&cluster.owner, "ahead_queue", 2, 2).await?; + enqueue(&cluster.app_b, "ahead_queue", 2, 1).await?; + assert_eq!(drain(&cluster.app_b, "ahead_queue").await?, 3); + // A runtime without schema CREATE still gets v052's provisioning hint. + let error = enqueue(&cluster.app_a, "ahead_queue", 3, 1) + .await + .unwrap_err(); + let AwaError::Database(sqlx::Error::Database(db_error)) = &error else { + panic!("expected a database error, got {error:?}"); + }; + assert_eq!(db_error.code().as_deref(), Some("42501")); + assert!(db_error.message().contains("has not been provisioned")); + cluster.teardown().await +} + +#[tokio::test] +async fn concurrent_first_touch_of_one_lane_by_two_roles() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + + // A holds a brand-new lane's creation open; B's first touch waits for it + // and then uses A's sequences. + let mut holding = cluster.app_a.begin().await?; + sqlx::query("SELECT awa.provision_lane('raced_queue', 2::smallint, 0::smallint)") + .execute(holding.as_mut()) + .await?; + let app_b = cluster.app_b.clone(); + let waiting = tokio::spawn(async move { enqueue(&app_b, "raced_queue", 2, 1).await }); + let mut blocked = false; + for _ in 0..100 { + blocked = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND NOT granted \ + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()))", + ) + .fetch_one(&cluster.owner) + .await?; + if blocked { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!(blocked, "B's first touch should wait on A's lane lock"); + holding.commit().await?; + assert_eq!(waiting.await??, 1); + + // Unordered races on many fresh lanes. + for round in 0..20 { + let queue = format!("raced_queue_{round}"); + let (from_a, from_b) = tokio::join!( + enqueue(&cluster.app_a, &queue, 1, 1), + enqueue(&cluster.app_b, &queue, 1, 1) + ); + assert_eq!(from_a? + from_b?, 2); + cluster.assert_lanes_owned_by_owner(&queue, 1).await?; + assert_eq!(drain(&cluster.app_c, &queue).await?, 2); + } + assert_eq!(drain(&cluster.app_c, "raced_queue").await?, 1); + cluster + .assert_lanes_owned_by_owner("raced_queue", 1) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { + let mut cluster = Cluster::new(Schema::BeforeProvisioning, true).await?; + // Lanes as a 0.6 fleet left them: one each created by an app role, by the + // migration role acting as itself, and by the owner without the runtime's + // grant. + enqueue(&cluster.app_a, "app_lane", 1, 1).await?; + enqueue(&cluster.migrator, "migrator_lane", 1, 1).await?; + enqueue(&cluster.owner, "owner_lane", 1, 1).await?; + for (sequence, _) in cluster.lane_sequence_owners("owner_lane").await? { + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE ALL ON SEQUENCE awa.{sequence} FROM {}", + cluster.runtime_name + )))) + .await?; + } + let failed = drain(&cluster.app_b, "owner_lane").await.unwrap_err(); + assert!( + failed + .to_string() + .contains("permission denied for sequence"), + "{failed}" + ); + let failed = drain(&cluster.app_b, "app_lane").await.unwrap_err(); + assert!( + failed + .to_string() + .contains("permission denied for sequence"), + "{failed}" + ); + // 0.7 refuses to migrate while app roles hold CREATE on the schema. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE CREATE ON SCHEMA awa FROM {}", + cluster.runtime_name + )))) + .await?; + + // The migration runs as the migration role itself. It is a member of the + // owner, so it may reassign its own and the owner's sequences, but not + // the app role's. + for (_, _, statement) in migrations::migration_sql_range(52, migrations::CURRENT_VERSION) { + cluster + .migrator + .execute(sqlx::raw_sql(sql(statement))) + .await?; + } + assert_eq!( + migrations::current_version(&cluster.owner).await?, + migrations::CURRENT_VERSION + ); + cluster + .assert_lanes_owned_by_owner("migrator_lane", 1) + .await?; + cluster.assert_lanes_owned_by_owner("owner_lane", 1).await?; + assert_eq!(drain(&cluster.app_b, "migrator_lane").await?, 1); + assert_eq!(drain(&cluster.app_b, "owner_lane").await?, 1); + + let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; + assert_eq!(problems.len(), 2, "{problems:?}"); + let app_a = cluster.roles[3].clone(); + let app_a = &app_a; + for problem in &problems { + assert_eq!(&problem.owner, app_a); + assert_eq!(problem.expected_owner, cluster.owner_name); + assert_eq!( + problem.missing_grants, + vec![format!( + "GRANT USAGE, UPDATE, SELECT ON SEQUENCE awa.{} TO {}", + problem.sequence_name, cluster.runtime_name + )] + ); + assert_eq!( + problem.reassign_sql.as_deref(), + Some( + format!( + "ALTER SEQUENCE awa.{} OWNER TO {}", + problem.sequence_name, cluster.owner_name + ) + .as_str() + ) + ); + } + + // Re-running the repair as the migration role changes nothing and fails + // nothing. The migration role created the repair function, so it grants + // its use. + cluster + .migrator + .execute(sqlx::raw_sql(sql(format!( + "GRANT EXECUTE ON FUNCTION awa.repair_lane_sequences(text) TO {}", + cluster.runtime_name + )))) + .await?; + let rerun = storage::repair_lane_sequences(&cluster.migrator, "awa").await?; + assert_eq!(rerun.len(), 2); + assert!(rerun.iter().all(|result| result.outcome == "unrepaired")); + + // A role that can act as the app role but cannot become the owner adds + // the grants and leaves the owner for someone who can. + let delegate = format!("{}_delegate", cluster.runtime_name); + cluster + .admin + .execute(sqlx::raw_sql(sql(format!( + "CREATE ROLE {delegate} LOGIN PASSWORD 'lane_test' IN ROLE {app_a}; \ + GRANT CONNECT ON DATABASE {} TO {delegate}", + cluster.database + )))) + .await?; + let delegate_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + PgConnectOptions::from_str(&std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }))? + .database(&cluster.database) + .username(&delegate) + .password("lane_test"), + ) + .await?; + let delegated = storage::repair_lane_sequences(&delegate_pool, "awa").await?; + delegate_pool.close().await; + assert_eq!(delegated.len(), 2); + for result in &delegated { + assert_eq!(result.outcome, "owner_unchanged"); + assert_eq!(&result.owner, app_a); + assert!(result.missing_grants.is_empty()); + assert!(result.reassign_sql.is_some()); + } + assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); + enqueue(&cluster.app_c, "app_lane", 1, 1).await?; + + // A role that is a member of both the app role and the owner reassigns. + cluster + .admin + .execute(sqlx::raw_sql(sql(format!( + "GRANT {app_a} TO {migrator}", + migrator = cluster.migrator_name + )))) + .await?; + let repaired = storage::repair_lane_sequences(&cluster.migrator, "awa").await?; + assert_eq!(repaired.len(), 2); + assert!( + repaired.iter().all(|result| result.outcome == "repaired"), + "{repaired:?}" + ); + assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") + .await? + .is_empty()); + assert!(storage::repair_lane_sequences(&cluster.migrator, "awa") + .await? + .is_empty()); + cluster.assert_lanes_owned_by_owner("app_lane", 1).await?; + assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); + + cluster.roles.push(delegate); + cluster.teardown().await +} diff --git a/docs/queue-storage-substrate.md b/docs/queue-storage-substrate.md index 4476712e4..9542e3914 100644 --- a/docs/queue-storage-substrate.md +++ b/docs/queue-storage-substrate.md @@ -40,7 +40,7 @@ awa storage prepare-queue-storage-schema \ The command invokes the idempotent `awa.install_queue_storage_substrate(...)` helper under a per-schema advisory transaction lock. The helper is activation-neutral: preparing a schema does not route work to it. Activation is a separate staged transition described in [Upgrading from 0.5 to 0.6](upgrade-0.5-to-0.6.md). -The installer is `SECURITY INVOKER`; its caller needs DDL privileges on the target schema. After the migrator provisions each queue with `awa storage prepare-queue`, workers need runtime DML, sequence `USAGE`/`SELECT`/`UPDATE`, and `TRUNCATE` for guarded ring reclamation. They do not need schema CREATE. See [Database roles](security/database-roles.md). +The installer is `SECURITY INVOKER`; its caller needs DDL privileges on the target schema. Workers need runtime DML, sequence `USAGE`/`SELECT`/`UPDATE`, and `TRUNCATE` for guarded ring reclamation. They do not need schema CREATE: lanes are created through the schema owner's provisioning functions. See [Database roles](security/database-roles.md). ## Operator rules diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index b6c1557e9..1b5e0091f 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -63,36 +63,44 @@ ALTER DEFAULT PRIVILEGES FOR ROLE awa_owner IN SCHEMA awa If an existing installation created objects as `awa_migrator` without first setting the owner role, transfer every existing schema object to `awa_owner` before relying on owner-scoped defaults. New migrations should use the effective-role command above. -## Provision queues before starting restricted runtimes +## Lane sequences -From schema v049, runtime calls use existing lane sequences without issuing -`CREATE SEQUENCE`. Run provisioning as the migrator for every queue before -starting its producers or workers: +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. From schema v053 they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after validating the queue, priority (1–4) and shard (0–63). A runtime therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. + +A new lane sequence gets the privileges of the head table it backs. A role that may `INSERT` or `UPDATE` `queue_enqueue_heads` gets `USAGE, SELECT, UPDATE` on the enqueue sequences, and one that may only `SELECT` it gets `SELECT`; `queue_claim_heads` does the same for the claim sequences. Granting the runtime the table privileges above is enough, and no sequence default privileges are needed for lanes. + +Run migrations with the schema owner as the effective role, as above, so the provisioning functions belong to it. A migrator that cannot `SET ROLE` to the owner leaves them owned by itself and the migration warns. + +`awa storage prepare-queue` still creates every priority and shard of a queue in one transaction, and any runtime role may run it. Use it to create lanes before a deploy rather than on first use. + +### Check and repair lane ownership + +Lanes created before v053 belong to whichever role touched them first, and only that role, or the roles its default privileges named, can use them. A different runtime role then fails with `permission denied for sequence` on that lane. List lane sequences that are not owned by the schema owner or lack a runtime grant (read-only, any role): + +```bash +awa --database-url "$AWA_DATABASE_URL" storage lane-sequences --schema awa +``` + +```sql +SELECT * FROM awa.lane_sequence_problems('awa'); +``` + +The command exits 2 while any remain. Migration v053 runs the repair as the migrating role: for each sequence it changes the owner to the schema owner and adds the missing grants, wherever that role is allowed to. A sequence it is not allowed to change does not fail the migration. It is reported in a warning, and `awa migrate` prints a summary. Run the repair again at any time; it is idempotent: ```bash PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-queue \ - --queue email --enqueue-shards 1 + awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair ``` -The command provisions all four priorities in one transaction and is safe to -repeat on a live queue: existing sequence cursors keep their positions. It does -not change routing or the queue's configured shard count. Prepare additional -shards before increasing `enqueue_shards`; match `--queue-stripe-count` to the -runtime's queue-storage configuration. With `PartitionedQueue`, provision each -physical queue. Use `--schema` for a custom storage schema. - -Apply the sequence grants above after provisioning, or set the owner's default -privileges first. Lane cursors need `UPDATE` for `setval`, in addition to `USAGE` -and `SELECT`. `CREATE SEQUENCE IF NOT EXISTS` requires schema CREATE even when -the sequence already exists, so provisioning alone does not repair pre-v049 -helpers: apply the migration first. Schema owners retain lazy creation for -single-role installations; a restricted runtime encountering an unprepared -lane fails with a provisioning hint. +Changing a sequence's owner requires a role that can `SET ROLE` to both the current owner and the schema owner. A migrator usually can for sequences it or the owner created, but not for one an unrelated application role created, and managed PostgreSQL offers no superuser. For each sequence that remains, the output gives the statements to run: + +- `missing_grants`: run them as the current owner, or as a role that can `SET ROLE` to it. Every runtime role can then use the lane, although the owner stays unchanged. +- `reassign_sql`: run it as a role that is a member of both roles, for example after `GRANT app_role TO awa_migrator` from a role with `ADMIN` on `app_role`. Revoke that membership afterwards. +## Custom queue-storage schemas ## Custom queue-storage schemas -Repeat the schema, table, sequence, function, and default-privilege grants for a custom queue-storage schema. Prepare custom schemas as the migrator; the runtime needs DML, sequence access, function execution, and `TRUNCATE`, but never DDL. See [Queue storage](../queue-storage-substrate.md). +Repeat the schema, table, sequence, function, and default-privilege grants for a custom queue-storage schema. Prepare custom schemas as the migrator; the runtime needs DML, sequence access, function execution, and `TRUNCATE`, but never DDL. Check lanes in a custom schema with `awa storage lane-sequences --schema `. See [Queue storage](../queue-storage-substrate.md). ## Why the grants are broad diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 0afc48afb..7fff6173f 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -14,14 +14,41 @@ A [migrate-first rehearsal with released `awa-pg==0.6.7`](https://github.com/har verified a restricted worker completing jobs on the upgraded schema. That rehearsal does not cover concurrent mixed-version rollouts. -For a runtime without DDL privileges, apply v049 as the migrator, run -`awa storage prepare-queue --queue ` for every physical queue, and grant -sequence `USAGE, SELECT, UPDATE` before starting producers or workers. Match -shard and stripe counts and provision extra lanes before increasing them. External migration runners must apply the complete generated `awa migrate --sql --from 48 --to 49` bundle: it refreshes the installer -definition before replacing lane helpers in the installed schemas. -See [Database roles](security/database-roles.md#provision-queues-before-starting-restricted-runtimes). +definition before replacing lane helpers in the installed schemas. From v053, +runtimes without DDL privileges create lanes themselves through the schema +owner's provisioning functions, so `awa storage prepare-queue` is optional. + +## Lane ownership and v053 + +Before v053 a lane's sequences belonged to the role that first enqueued into, +claimed from or promoted into it. With several runtime roles, only roles whose +default privileges covered that role could use the lane; any other role failed +with `permission denied for sequence`. v053 installs `SECURITY DEFINER` +`provision_lane` and `provision_lane_sequences` functions, owned by the schema +owner, in each queue-storage schema. The lane helpers create sequences through +them, with grants that mirror the lane's head table. v053 then gives every +existing lane sequence to the schema owner and adds its missing grants, as far +as the migrating role is allowed to. + +0.6 and earlier 0.7 binaries call the same helpers, so lanes they create after +v053 are owned by the schema owner too. A 0.7 binary running ahead of v053 +detects that `provision_lane` is missing and creates lanes as before. + +To adopt it: + +1. Check the current state: `awa storage lane-sequences` (read-only, exits 2 + while lanes need repair). +2. Run `awa migrate` with the schema owner as the effective role, as described + in [Database roles](security/database-roles.md). +3. Read the migrate output. Any lane sequence it could not reassign is listed; + `awa storage lane-sequences` gives the statements to run and the role to run + them as. See + [Check and repair lane ownership](security/database-roles.md#check-and-repair-lane-ownership). + +External migration runners apply the generated `awa migrate --sql --from 52 +--to 53` bundle, then run `awa storage lane-sequences` to check for leftovers. ## Custom-schema discovery during migration diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index 4cfb039b0..5124fbc3f 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -91,11 +91,13 @@ Do not skip a major version. 0.5 → 0.7 is unsupported; step through 0.6. ## Restricted Database Roles -When deploying a runtime without schema CREATE, apply v049 and provision every -queue with `awa storage prepare-queue` as the migrator before starting producers -or workers. Match shard/stripe counts; provision additional lanes before changing -routing. Use the sequence grants and owner defaults in -[Database roles](../../docs/security/database-roles.md#provision-queues-before-starting-restricted-runtimes). +From v053 any runtime role creates lanes through the schema owner's +`SECURITY DEFINER` provisioning functions, so runtimes need neither schema +CREATE nor sequence default privileges, and `awa storage prepare-queue` is +optional. Run migrations with the schema owner as the effective role. Check +lane ownership with `awa storage lane-sequences` (exits 2 while lanes need +repair) and fix it with `--repair` or the statements it prints; see +[Database roles](../../docs/security/database-roles.md#check-and-repair-lane-ownership). ## Rolling Upgrades From 70fe5ff66ae289189b3b3340737280f2d316372e Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 13:35:46 +1300 Subject: [PATCH 09/18] docs(migrations): link v053 to its pull request --- awa-model/migrations/v053_definer_lane_provisioning.sql | 1 + 1 file changed, 1 insertion(+) diff --git a/awa-model/migrations/v053_definer_lane_provisioning.sql b/awa-model/migrations/v053_definer_lane_provisioning.sql index a56d96978..758627703 100644 --- a/awa-model/migrations/v053_definer_lane_provisioning.sql +++ b/awa-model/migrations/v053_definer_lane_provisioning.sql @@ -1,4 +1,5 @@ -- v053: Create queue lanes as the schema owner, whichever role touches them first. +-- Tracking: https://github.com/hardbyte/awa/pull/637 -- Lane sequences used to belong to the role that first enqueued into, claimed -- from or promoted into a lane, and only that role (plus whoever its default -- privileges named) could use them. refresh_lane_sequence_helpers now installs From 80178b39b6b2c40af8430809667a0a8a2b701069 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 14:04:15 +1300 Subject: [PATCH 10/18] fix(queue-storage): accept every lane the invoker path accepts in provision_lane Queue storage takes any queue name, including an empty one, and any smallint priority, so the provisioning functions only reject a null argument or a shard outside 0-63. --- awa-model/migrations/v023_install_queue_storage_substrate.sql | 4 ++-- docs/security/database-roles.md | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index b5526cccc..39df93e86 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -272,8 +272,8 @@ BEGIN v_seq_name TEXT; v_grant TEXT; BEGIN - IF p_queue IS NULL OR p_queue = '' - OR p_priority IS NULL OR p_priority NOT BETWEEN 1 AND 4 + IF p_queue IS NULL + OR p_priority IS NULL OR p_enqueue_shard IS NULL OR p_enqueue_shard NOT BETWEEN 0 AND 63 THEN RAISE EXCEPTION 'invalid queue lane: queue %%, priority %%, enqueue shard %%', quote_nullable(p_queue), p_priority, p_enqueue_shard diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index 1b5e0091f..5f15793e7 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -65,7 +65,7 @@ If an existing installation created objects as `awa_migrator` without first sett ## Lane sequences -Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. From schema v053 they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after validating the queue, priority (1–4) and shard (0–63). A runtime therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. From schema v053 they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after checking that the queue, priority and shard are set and the shard is 0–63. A runtime therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. A new lane sequence gets the privileges of the head table it backs. A role that may `INSERT` or `UPDATE` `queue_enqueue_heads` gets `USAGE, SELECT, UPDATE` on the enqueue sequences, and one that may only `SELECT` it gets `SELECT`; `queue_claim_heads` does the same for the claim sequences. Granting the runtime the table privileges above is enough, and no sequence default privileges are needed for lanes. From 4e51915e07c2ee23726396ed0139ac797d80621c Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 17:49:27 +1300 Subject: [PATCH 11/18] fix(worker): restrict who may call the statistics helper and what it may analyse - Revoke EXECUTE on awa.analyze_queue_storage_relation from PUBLIC. The SHARE UPDATE EXCLUSIVE lock it takes lasts until the caller's transaction ends, so any role could hold it open and block ring prune. The migration grants EXECUTE to the roles that already hold INSERT, UPDATE and DELETE on the awa head tables, and the function requires the calling session user to hold those privileges on the target schema. - Give the function its own lock_timeout. - Accept only the two head tables and the thirteen ring parents inside the function, not any partitioned table of the schema. - Probe for the function once per pass, so a binary ahead of the schema makes one catalog query instead of fifteen failing calls and server ERROR lines. Pass-level outcomes use a new "pass" metric scope. - Correct the header, which claimed the function takes no lock. - Tests: a role without EXECUTE and a role without head-table privileges are denied, the migration grants EXECUTE to existing grantees, non-ring tables are rejected, and a missing helper costs one probe and no calls. --- CHANGELOG.md | 23 ++- awa-metrics/src/lib.rs | 6 +- .../v053_analyze_queue_storage_statistics.sql | 107 +++++++++--- awa-worker/src/statistics.rs | 135 +++++++++++++-- awa/tests/statistics_refresh_test.rs | 159 +++++++++++++++++- docs/configuration.md | 4 +- docs/grafana/README.md | 2 +- docs/security/database-roles.md | 4 +- docs/upgrade-0.6-to-0.7.md | 6 +- 9 files changed, 391 insertions(+), 55 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1287eca44..699c9bb16 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,21 +23,26 @@ Notable changes between releases. Detailed migration notes for storage transitio lands while ANALYZE samples the oldest slot's tables returns `blocked` and retries on its normal backoff. Results are in `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration` - (scope `heads` or `ring_parent`; outcome `ok`, `skipped_privilege`, + (scope `heads`, `ring_parent` or `pass`; outcome `ok`, `skipped_privilege`, `skipped_lock`, `unavailable` or `failed`). - **Migration v053 adds `awa.analyze_queue_storage_relation(schema, relation)`, so a runtime role that does not own the tables can refresh statistics.** `ANALYZE` skips a table its caller may not analyse, and the documented runtime grants do not include `MAINTAIN`, so the refresh runs through this `SECURITY DEFINER` function owned by the migration role. It analyses only a - head table or partitioned table of a schema that - `awa.queue_storage_schemas()` reports as trusted, and `EXECUTE` is granted - to `PUBLIC` like the other runtime helpers. Run `awa migrate` as the schema - owner (see [database roles](docs/security/database-roles.md)); no grant - changes are needed. A binary running ahead of the migration, or against a - function whose owner may not analyse the tables, logs one warning naming what - is missing and skips the refresh without affecting job processing. Earlier - binaries never call the function. + head table or ring parent of a schema that `awa.queue_storage_schemas()` + reports as trusted. It is not executable by `PUBLIC`, because the lock it + takes lasts until the caller's transaction ends and would block ring prune: + the migration grants `EXECUTE` to the roles that already hold INSERT, UPDATE + and DELETE on the head tables (the runtime grantees), and the function + requires the calling session user to hold them. A role created later needs + the documented `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa`. Run + `awa migrate` as the schema owner (see + [database roles](docs/security/database-roles.md)). A binary running ahead of + the migration probes for the function once per pass and, like one whose + function's owner may not analyse the tables, logs one warning naming what is + missing and skips the refresh without making failing calls or affecting job + processing. Earlier binaries never call the function. ### Fixed diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 14e3cb540..3f1eee208 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -279,7 +279,7 @@ pub struct AwaMetrics { /// rather than one per tick. Fleets alert on this rather than /// scraping the matching `tracing::warn!` line. Issue #242. pub maintenance_branch_overrun_total: Counter, - /// Planner-statistics `ANALYZE` calls by scope (`heads`, `ring_parent`) + /// Planner-statistics `ANALYZE` calls by scope (`heads`, `ring_parent`, `pass`) /// and outcome (`ok`, `skipped_privilege`, `skipped_lock`, /// `unavailable`, `failed`). One increment per relation per refresh pass. pub maintenance_analyze_attempts: Counter, @@ -1304,8 +1304,8 @@ impl AwaMetrics { self.maintenance_branch_overrun_total.add(1, &attrs); } - /// Record one planner-statistics `ANALYZE` call. `scope` is `"heads"` or - /// `"ring_parent"`; `outcome` is `"ok"`, `"skipped_privilege"`, + /// Record one planner-statistics `ANALYZE` call. `scope` is `"heads"`, + /// `"ring_parent"` or `"pass"` (a pass-level probe or failure); `outcome` is `"ok"`, `"skipped_privilege"`, /// `"skipped_lock"`, `"unavailable"` or `"failed"`. pub fn record_statistics_analyze( &self, diff --git a/awa-model/migrations/v053_analyze_queue_storage_statistics.sql b/awa-model/migrations/v053_analyze_queue_storage_statistics.sql index 93d556d12..d9ec60d0f 100644 --- a/awa-model/migrations/v053_analyze_queue_storage_statistics.sql +++ b/awa-model/migrations/v053_analyze_queue_storage_statistics.sql @@ -21,24 +21,35 @@ -- applies when ANALYZE starts work on a relation, so a child that ring prune -- holds ACCESS EXCLUSIVE is skipped during the per-child analysis, but -- sampling the children for the parent's statistics can still wait for it. --- That wait is bounded by the caller's lock_timeout and fails the call with +-- That wait is bounded by the function's lock_timeout and fails the call with -- SQLSTATE 55P03, which the caller records as skipped_lock. SKIP_LOCKED also -- skips every partition silently when the parent itself is locked, so the -- function takes the parent's SHARE UPDATE EXCLUSIVE lock first and lets that -- wait fail instead. -- --- Definer boundary: the function only analyses a head table or a partitioned --- table of a schema that awa.queue_storage_schemas() reports, which requires --- the schema's ownership, grants and triggers to be trusted, and builds SQL --- from nothing but quoted identifiers. ANALYZE writes statistics and nothing --- else. +-- Definer boundary: only a head table or ring parent of a schema that +-- awa.queue_storage_schemas() reports is analysed, which requires the schema's +-- ownership, grants and triggers to be trusted, and SQL is built from nothing +-- but quoted identifiers. ANALYZE writes statistics and nothing else. +-- Only the fixed list of ring parents and the two head tables is accepted; +-- a migration that adds a partitioned table to the substrate must extend it. +-- +-- Who may call it: the lock above lasts until the caller's transaction ends +-- and blocks ring prune of that table, so EXECUTE is not granted to PUBLIC. It +-- goes to the roles that already hold INSERT, UPDATE and DELETE on the head +-- tables of the awa schema (the runtime grantees), and the function itself +-- requires the calling session user to hold those privileges on the target +-- schema's queue_claim_heads. A role created later needs the documented +-- `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa` (or default privileges). -- -- Returns 'analyzed'; 'skipped_privilege' when the function owner may not -- analyse the relation; or 'missing' when the relation does not exist. -- --- N-1 binaries never call it. A binary newer than this schema logs once that --- the function is missing and skips the refresh. Completes in milliseconds and --- takes no lock on any queue-storage table. +-- N-1 binaries never call it. A binary newer than this schema probes for the +-- function, logs once that it is missing and skips the refresh. Creating or +-- replacing it takes no table lock and completes in milliseconds; calling it +-- takes SHARE UPDATE EXCLUSIVE on the relation until the caller's transaction +-- ends. CREATE OR REPLACE FUNCTION awa.analyze_queue_storage_relation( p_schema TEXT, @@ -48,10 +59,13 @@ RETURNS TEXT LANGUAGE plpgsql SECURITY DEFINER SET search_path = pg_catalog, pg_temp +SET lock_timeout = '2s' AS $$ DECLARE + v_heads REGCLASS; v_relid OID; v_relkind "char"; + v_relowner OID; v_may_analyze BOOLEAN; BEGIN IF NOT EXISTS ( @@ -62,26 +76,51 @@ BEGIN USING ERRCODE = '22023'; END IF; + IF p_relation <> ALL (ARRAY[ + 'queue_claim_heads', + 'queue_enqueue_heads', + 'done_entries', + 'lease_claim_batches', + 'lease_claim_closure_batches', + 'lease_claim_closures', + 'lease_claims', + 'leases', + 'queue_terminal_count_deltas', + 'ready_claim_attempt_batches', + 'ready_entries', + 'ready_segments', + 'ready_tombstones', + 'receipt_completion_batches', + 'receipt_completion_tombstones' + ]) THEN + RAISE EXCEPTION '%.% is not a ring parent or head table of a queue-storage schema', + p_schema, p_relation + USING ERRCODE = '22023'; + END IF; + + v_heads := to_regclass(format('%I.queue_claim_heads', p_schema)); + IF v_heads IS NULL + OR NOT has_table_privilege(session_user, v_heads, 'INSERT') + OR NOT has_table_privilege(session_user, v_heads, 'UPDATE') + OR NOT has_table_privilege(session_user, v_heads, 'DELETE') THEN + RAISE EXCEPTION 'permission denied: % may not maintain queue storage in schema %', + session_user, p_schema + USING ERRCODE = '42501'; + END IF; + v_relid := to_regclass(format('%I.%I', p_schema, p_relation)); IF v_relid IS NULL THEN RETURN 'missing'; END IF; - SELECT class.relkind INTO v_relkind FROM pg_class AS class WHERE class.oid = v_relid; - IF NOT ( - v_relkind = 'p' - OR p_relation IN ('queue_claim_heads', 'queue_enqueue_heads') - ) THEN - RAISE EXCEPTION '%.% is not a partitioned or head table of a queue-storage schema', - p_schema, p_relation - USING ERRCODE = '22023'; - END IF; + SELECT class.relkind, class.relowner INTO v_relkind, v_relowner + FROM pg_class AS class WHERE class.oid = v_relid; -- ANALYZE skips, with a warning rather than an error, a table the caller -- may not analyse: the owner, a MAINTAIN holder (PostgreSQL 17+), a -- member of pg_maintain, or the database owner may analyse it. SELECT - pg_has_role(current_user, class.relowner, 'USAGE') + pg_has_role(current_user, v_relowner, 'USAGE') OR pg_has_role( current_user, (SELECT database.datdba FROM pg_database AS database @@ -90,12 +129,10 @@ BEGIN ) OR CASE WHEN current_setting('server_version_num')::int >= 170000 - THEN has_table_privilege(class.oid, 'MAINTAIN') + THEN has_table_privilege(v_relid, 'MAINTAIN') ELSE FALSE END - INTO v_may_analyze - FROM pg_class AS class - WHERE class.oid = v_relid; + INTO v_may_analyze; IF NOT v_may_analyze THEN RETURN 'skipped_privilege'; @@ -113,7 +150,29 @@ BEGIN END $$; -GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(TEXT, TEXT) TO PUBLIC; +REVOKE EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(TEXT, TEXT) FROM PUBLIC; + +DO $grant$ +DECLARE + v_grantee TEXT; +BEGIN + FOR v_grantee IN + SELECT pg_get_userbyid(privilege.grantee) + FROM pg_class AS heads + CROSS JOIN LATERAL aclexplode(heads.relacl) AS privilege + WHERE heads.oid = to_regclass('awa.queue_claim_heads') + AND privilege.grantee <> 0 + AND privilege.privilege_type IN ('INSERT', 'UPDATE', 'DELETE') + GROUP BY privilege.grantee + HAVING count(DISTINCT privilege.privilege_type) = 3 + LOOP + EXECUTE format( + 'GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO %I', + v_grantee + ); + END LOOP; +END +$grant$; INSERT INTO awa.schema_version (version, description) VALUES (53, 'Planner statistics refresh for queue-storage tables autovacuum never analyses') diff --git a/awa-worker/src/statistics.rs b/awa-worker/src/statistics.rs index 68994c6da..352b97b64 100644 --- a/awa-worker/src/statistics.rs +++ b/awa-worker/src/statistics.rs @@ -54,6 +54,8 @@ const SQLSTATE_LOCK_NOT_AVAILABLE: &str = "55P03"; pub(crate) enum RelationScope { Heads, RingParent, + /// The pass as a whole: a probe or listing failed, or the helper is absent. + Pass, } impl RelationScope { @@ -61,6 +63,7 @@ impl RelationScope { match self { RelationScope::Heads => "heads", RelationScope::RingParent => "ring_parent", + RelationScope::Pass => "pass", } } } @@ -193,7 +196,31 @@ impl StatisticsRefresher { Err(err) => { warn!(error = %err, schema = %self.schema, "Could not check the queue-storage schema for statistics refresh"); self.record( - RelationScope::RingParent, + RelationScope::Pass, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + return summary; + } + } + + match self.helper_installed().await { + Ok(true) => {} + Ok(false) => { + self.warn_unavailable_once(); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Unavailable, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Unavailable); + return summary; + } + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not probe for the statistics helper"); + self.record( + RelationScope::Pass, AnalyzeOutcome::Failed, started.elapsed(), ); @@ -207,7 +234,7 @@ impl StatisticsRefresher { Err(err) => { warn!(error = %err, schema = %self.schema, "Could not list ring tables for statistics refresh"); self.record( - RelationScope::RingParent, + RelationScope::Pass, AnalyzeOutcome::Failed, started.elapsed(), ); @@ -250,6 +277,18 @@ impl StatisticsRefresher { .record_statistics_analyze(scope.as_str(), outcome.as_str(), elapsed); } + /// One catalog lookup per pass. A missing helper (a binary running ahead + /// of the schema) then costs this query instead of a failing call, and a + /// server-side ERROR line, per relation; the pass interval is the + /// re-probe interval. + async fn helper_installed(&self) -> Result { + sqlx::query_scalar( + "SELECT to_regprocedure('awa.analyze_queue_storage_relation(text,text)') IS NOT NULL", + ) + .fetch_one(&self.pool) + .await + } + async fn ring_parents(&self) -> Result, sqlx::Error> { sqlx::query_scalar( r#" @@ -327,19 +366,14 @@ impl StatisticsRefresher { AnalyzeOutcome::SkippedLock } Some(SQLSTATE_UNDEFINED_FUNCTION) => { - if !self.warned_unavailable.swap(true, Ordering::Relaxed) { - warn!( - schema = %self.schema, - "Planner statistics refresh is disabled: awa.analyze_queue_storage_relation() \ - is missing. Run `awa migrate` (schema v053 or later) to install it" - ); - } + self.warn_unavailable_once(); AnalyzeOutcome::Unavailable } Some(SQLSTATE_INSUFFICIENT_PRIVILEGE) => { self.warn_privilege_once( relation, - "the runtime role may not execute awa.analyze_queue_storage_relation(); \ + "the runtime role may not execute awa.analyze_queue_storage_relation(), or lacks \ + INSERT, UPDATE and DELETE on the queue head tables; \ GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) to it", ); AnalyzeOutcome::SkippedPrivilege @@ -351,6 +385,17 @@ impl StatisticsRefresher { } } + fn warn_unavailable_once(&self) { + if self.warned_unavailable.swap(true, Ordering::Relaxed) { + return; + } + warn!( + schema = %self.schema, + "Planner statistics refresh is disabled: awa.analyze_queue_storage_relation() \ + is missing. Run `awa migrate` (schema v053 or later) to install it" + ); + } + fn warn_privilege_once(&self, relation: &str, missing: &str) { if self.warned_privilege.swap(true, Ordering::Relaxed) { return; @@ -395,3 +440,73 @@ impl Drop for RefreshGuard { self.cancel.cancel(); } } + +#[cfg(test)] +mod tests { + use super::*; + use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; + use std::str::FromStr; + + /// A binary ahead of the schema must probe once and attempt no calls: a + /// failing call per relation would count as `unavailable` fifteen times + /// and write fifteen server-side ERROR lines per pass. + #[tokio::test] + async fn missing_helper_costs_one_probe_and_no_calls() { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let options = PgConnectOptions::from_str(&url).expect("parse DATABASE_URL"); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect_with(options.clone()) + .await + .expect("connect admin"); + let database = format!("awa_stats_unit_{}", std::process::id()); + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "CREATE DATABASE {database}" + ))) + .execute(&admin) + .await + .expect("create database"); + let pool = PgPoolOptions::new() + .max_connections(2) + .connect_with(options.database(&database)) + .await + .expect("connect test database"); + awa_model::migrations::run(&pool).await.expect("migrate"); + + let refresher = StatisticsRefresher::new( + pool.clone(), + "awa".to_string(), + AwaMetrics::from_global(), + false, + ); + let present = refresher.run_pass().await; + assert_eq!(present.ok, 15, "{present:?}"); + + sqlx::raw_sql("DROP FUNCTION awa.analyze_queue_storage_relation(text, text)") + .execute(&pool) + .await + .expect("drop the helper"); + let absent = refresher.run_pass().await; + assert_eq!( + absent, + PassSummary { + unavailable: 1, + ..PassSummary::default() + } + ); + pool.close().await; + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop database"); + } +} diff --git a/awa/tests/statistics_refresh_test.rs b/awa/tests/statistics_refresh_test.rs index aac55c100..bc9ac8525 100644 --- a/awa/tests/statistics_refresh_test.rs +++ b/awa/tests/statistics_refresh_test.rs @@ -32,6 +32,13 @@ const RING_PARENTS: [&str; 13] = [ static DATABASE_COUNTER: AtomicU32 = AtomicU32::new(0); +async fn run_sql(pool: &PgPool, sql: String) -> Result<(), sqlx::Error> { + sqlx::raw_sql(audited_sql(sql)) + .execute(pool) + .await + .map(|_| ()) +} + fn base_database_url() -> String { std::env::var("DATABASE_URL") .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) @@ -380,7 +387,7 @@ async fn non_owner_runtime_role_refreshes_statistics() { "queue_claim_heads", "is not a queue-storage schema", ), - ("awa", "jobs", "is not a partitioned or head table"), + ("awa", "jobs", "is not a ring parent or head table"), ] { let error = sqlx::query_scalar::<_, String>("SELECT awa.analyze_queue_storage_relation($1, $2)") @@ -418,6 +425,153 @@ async fn non_owner_runtime_role_refreshes_statistics() { database.drop().await; } +async fn call_helper( + pool: &PgPool, + schema: &str, + relation: &str, +) -> Result { + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(schema) + .bind(relation) + .fetch_one(pool) + .await +} + +fn denied(error: &sqlx::Error, message: &str) { + let database_error = error.as_database_error().expect("database error"); + assert_eq!(database_error.code().as_deref(), Some("42501")); + assert!( + database_error.message().contains(message), + "unexpected message: {}", + database_error.message() + ); +} + +/// EXECUTE is not public: holding the lock the helper takes until the end of +/// a transaction would block ring prune of that table. A role needs EXECUTE and +/// write privileges on the head tables, and the migration grants EXECUTE to +/// the roles that already have them. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn execute_is_limited_to_roles_that_may_maintain_queue_storage() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let bystander = format!("stats_bystander_{suffix}"); + let reader = format!("stats_reader_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime, &bystander, &reader]).await; + let owner_pool = &database.owner_pool; + run_sql( + owner_pool, + format!( + "GRANT USAGE ON SCHEMA awa TO {bystander}, {reader}, {runtime}; + GRANT SELECT ON ALL TABLES IN SCHEMA awa TO {reader}; + GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {reader}; + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {runtime}" + ), + ) + .await + .expect("grant privileges"); + let connect = |role: String| { + let options = database.options.clone(); + async move { + PgPoolOptions::new() + .max_connections(2) + .connect_with(options.username(&role).password(ROLE_PASSWORD)) + .await + .expect("connect role pool") + } + }; + let bystander_pool = connect(bystander.clone()).await; + let reader_pool = connect(reader.clone()).await; + let runtime_pool = connect(runtime.clone()).await; + + denied( + &call_helper(&bystander_pool, "awa", "queue_claim_heads") + .await + .expect_err("a role without EXECUTE is denied"), + "permission denied for function", + ); + denied( + &call_helper(&reader_pool, "awa", "queue_claim_heads") + .await + .expect_err("EXECUTE without write privileges on the head tables is denied"), + "may not maintain queue storage", + ); + + // The runtime role holds write privileges but was granted no function + // EXECUTE: re-running the migration grants it, as an upgrade does. + denied( + &call_helper(&runtime_pool, "awa", "queue_claim_heads") + .await + .expect_err("write privileges alone do not grant EXECUTE"), + "permission denied for function", + ); + for (_, _, sql) in migrations::migration_sql_range( + migrations::CURRENT_VERSION - 1, + migrations::CURRENT_VERSION, + ) { + run_sql(owner_pool, sql) + .await + .expect("re-run the migration"); + } + assert_eq!( + call_helper(&runtime_pool, "awa", "queue_claim_heads") + .await + .expect("the migration granted EXECUTE to the runtime grantee"), + "analyzed" + ); + + for pool in [bystander_pool, reader_pool, runtime_pool] { + pool.close().await; + } + database.drop().await; +} + +/// The helper analyses the ring parents and head tables only, not any other +/// partitioned table of the schema. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn helper_rejects_a_partitioned_table_that_is_not_a_ring_parent() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime]).await; + run_sql( + &database.owner_pool, + "CREATE TABLE awa.application_events (id int) PARTITION BY LIST (id)".to_string(), + ) + .await + .expect("create a partitioned application table"); + let runtime_pool = database.runtime_pool(&runtime).await; + + for relation in ["application_events", "ready_entries_0", "jobs_hot"] { + let error = call_helper(&runtime_pool, "awa", relation) + .await + .expect_err("only ring parents and head tables are accepted"); + let database_error = error.as_database_error().expect("database error"); + assert_eq!( + database_error.code().as_deref(), + Some("22023"), + "{relation}" + ); + assert!( + database_error + .message() + .contains("is not a ring parent or head table"), + "{relation}: {}", + database_error.message() + ); + } + assert_eq!( + call_helper(&runtime_pool, "awa", "ready_entries") + .await + .expect("a ring parent is accepted"), + "analyzed" + ); + + runtime_pool.close().await; + database.drop().await; +} + /// A conflicting lock on a ring parent fails the call with `lock_not_available` /// rather than letting `ANALYZE (SKIP_LOCKED)` skip every partition and report /// success. @@ -515,7 +669,8 @@ async fn unrunnable_refresh_is_recorded_and_does_not_fail_the_runtime() { .expect("let the schema owner hand over the function"); sqlx::raw_sql(audited_sql(format!( "GRANT USAGE, CREATE ON SCHEMA awa TO {helper}; - ALTER FUNCTION awa.analyze_queue_storage_relation(text, text) OWNER TO {helper}" + ALTER FUNCTION awa.analyze_queue_storage_relation(text, text) OWNER TO {helper}; + GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {owner}" ))) .execute(&database.owner_pool) .await diff --git a/docs/configuration.md b/docs/configuration.md index 6d61b5755..154c54801 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -539,8 +539,8 @@ Autovacuum does not analyse partitioned tables, so the ring parents (`ready_entr The maintenance leader runs `ANALYZE` on both when it is elected and every `statistics_refresh_interval` (default `5m`). One pass takes well under a second on a small ring and a few seconds on a ring with millions of retained rows. It has no measurable effect on claim or completion throughput. - **Locks.** Each relation is analysed in its own transaction with a 2 s `lock_timeout` and a 120 s `statement_timeout`. `SKIP_LOCKED` skips a child that ring prune holds during the per-child analysis, but sampling for the parent's statistics waits for it up to the `lock_timeout`; a call that times out is recorded as `skipped_lock` and retried on the next interval. Prune never waits for the refresh: its own 10 ms lock timeout applies, so an attempt that lands while `ANALYZE` samples the oldest slot's tables returns `blocked` and retries on its usual backoff. `ANALYZE` of a partitioned table samples every child, so that window cannot be avoided. -- **Privileges.** The refresh calls `awa.analyze_queue_storage_relation(schema, relation)` (schema v053), a `SECURITY DEFINER` function owned by the role that ran the migration, so the runtime role needs no `MAINTAIN` grant and does not have to own the tables. If the function is missing (run `awa migrate`) or its owner may not analyse the tables, the leader logs one warning that names what is missing and skips the pass; job processing is unaffected. -- **Metrics.** `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration`, by `awa.analyze.scope` (`heads`, `ring_parent`) and `awa.analyze.outcome` (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`). A sustained `skipped_lock` rate means something else holds the tables; any `skipped_privilege` or `unavailable` means the refresh is not running. +- **Privileges.** The refresh calls `awa.analyze_queue_storage_relation(schema, relation)` (schema v053), a `SECURITY DEFINER` function owned by the role that ran the migration, so the runtime role needs no `MAINTAIN` grant and does not have to own the tables. It is executable only by roles that hold INSERT, UPDATE and DELETE on the queue head tables (the migration grants it to the existing ones). Each pass first probes for the function; if it is missing (run `awa migrate`) or its owner may not analyse the tables, the leader logs one warning that names what is missing and skips the pass without making failing calls; job processing is unaffected. +- **Metrics.** `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration`, by `awa.analyze.scope` (`heads`, `ring_parent`, `pass`) and `awa.analyze.outcome` (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`). A sustained `skipped_lock` rate means something else holds the tables; any `skipped_privilege` or `unavailable` means the refresh is not running. ### Rust diff --git a/docs/grafana/README.md b/docs/grafana/README.md index 4b8fa278a..500b13047 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -207,7 +207,7 @@ In Prometheus they appear under the OpenTelemetry [OTLP-to-Prometheus translatio | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | | `awa.maintenance.prune.attempts` | Counter | ring, outcome, reason | Ring prune outcomes, including destructive `pruned` and no-DDL `already_pruned` | | `awa.maintenance.prune.duration` | Histogram (s) | ring, phase | Successful destructive prune time split into `lock`, `truncate`, and `commit` | -| `awa.maintenance.analyze.attempts` | Counter | scope, outcome | Planner-statistics `ANALYZE` calls on the head tables and ring parents (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`) | +| `awa.maintenance.analyze.attempts` | Counter | scope (`heads`, `ring_parent`, `pass`), outcome | Planner-statistics `ANALYZE` calls on the head tables and ring parents (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`) | | `awa.maintenance.analyze.duration` | Histogram (s) | scope, outcome | Duration of one planner-statistics `ANALYZE` call | ### Heartbeat diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index 2e2d19a60..150c2fe81 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -65,9 +65,9 @@ If an existing installation created objects as `awa_migrator` without first sett ## Planner statistics -The maintenance leader runs `ANALYZE` on the queue-storage head tables and partitioned ring parents (see [Planner statistics](../configuration.md#planner-statistics)). `ANALYZE` skips a table its caller may not analyse, so a runtime that does not own the tables cannot do this itself; the leader calls `awa.analyze_queue_storage_relation(schema, relation)` instead, a `SECURITY DEFINER` function from migration v053 owned by whichever role ran the migration. Nothing is needed beyond the grants above: the function is executable by `PUBLIC`, and the runtime does not need `MAINTAIN`. +The maintenance leader runs `ANALYZE` on the queue-storage head tables and partitioned ring parents (see [Planner statistics](../configuration.md#planner-statistics)). `ANALYZE` skips a table its caller may not analyse, so a runtime that does not own the tables cannot do this itself; the leader calls `awa.analyze_queue_storage_relation(schema, relation)` instead, a `SECURITY DEFINER` function from migration v053 owned by whichever role ran the migration. The runtime does not need `MAINTAIN`. The function is not executable by `PUBLIC`, because the lock it takes lasts until the caller's transaction ends and would block ring prune: the migration grants `EXECUTE` to the roles that already hold `INSERT`, `UPDATE` and `DELETE` on the head tables, and the function requires the calling session user to hold them. A runtime role created later gets it from the `GRANT EXECUTE ON ALL FUNCTIONS` above or the default privileges. For a custom queue-storage schema, the role needs write privileges on that schema's `queue_claim_heads`. -Run migrations with the schema owner as the effective role, as shown above, so the function's owner owns the tables. A function owned by a role that may not analyse them (it is not the table owner, the database owner, or a `MAINTAIN` holder) makes every pass report `skipped_privilege` and logs one warning. The function only analyses a head table or partitioned table of a schema that `awa.queue_storage_schemas()` reports as trusted, and `ANALYZE` changes statistics only. +Run migrations with the schema owner as the effective role, as shown above, so the function's owner owns the tables. A function owned by a role that may not analyse them (it is not the table owner, the database owner, or a `MAINTAIN` holder) makes every pass report `skipped_privilege` and logs one warning. The function only analyses the two head tables and the thirteen ring parents of a schema that `awa.queue_storage_schemas()` reports as trusted, and `ANALYZE` changes statistics only. ## Provision queues before starting restricted runtimes diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 50f6cc587..5c0f47993 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -260,8 +260,10 @@ v053 adds `awa.analyze_queue_storage_relation(schema, relation)`, a maintenance leader refresh planner statistics for the head tables and ring parents when the runtime role does not own them (see [Planner statistics](configuration.md#planner-statistics)). It creates one -function and grants `EXECUTE` to `PUBLIC`; it takes no lock on any table and -completes in milliseconds. Run the migration as the schema owner so the +function, revokes `EXECUTE` from `PUBLIC` and grants it to the roles that +already hold INSERT, UPDATE and DELETE on the `awa` head tables (the runtime +grantees); it takes no table lock and completes in milliseconds. Roles created +later need `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa`. Run the migration as the schema owner so the function's owner owns the tables. Binaries before v053 never call it. A 0.7 binary running before the migration From 9a7eed99ae2fa439b9488e728b9129dac9c4917f Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 18:21:00 +1300 Subject: [PATCH 12/18] fix(queue-storage): repair lane sequences per transaction with a bounded lock wait v053 reassigned lane sequences inside the migration transaction. ALTER SEQUENCE ... OWNER takes the sequence's ACCESS EXCLUSIVE lock until commit, with no lock_timeout, so an open enqueue on one lane could stall the whole migration and every repaired lane waited for the upgrade to commit. v053 now only refreshes the helpers. After it commits, awa migrate repairs each lane sequence with its own statement through awa.repair_lane_sequence(schema, sequence), which has lock_timeout = 300ms; on 55P03 the sequence is left unchanged and reported as busy. `awa storage lane-sequences --repair` uses the same path, and the set-returning awa.repair_lane_sequences(schema) is gone. Reassigning needs the current owner's privileges, which the schema owner lacks for sequences the migration role or an app role created. A migration role connected with role set to the owner therefore could not reassign its own sequences. The repair now retries such a sequence as the login role, lane_sequence_problems reports repair_roles (the login roles that can repair each sequence), and leftovers are reported with the next step to take. The docs state which role to run the repair as. --- CHANGELOG.md | 13 +- awa-cli/src/main.rs | 61 +++-- .../v023_install_queue_storage_substrate.sql | 122 +++++---- .../v053_definer_lane_provisioning.sql | 30 +-- awa-model/src/migrations.rs | 62 ++++- awa-model/src/storage.rs | 237 ++++++++++++++++-- awa/tests/lane_ownership_roles_test.rs | 194 ++++++++++---- docs/security/database-roles.md | 27 +- docs/upgrade-0.6-to-0.7.md | 17 +- skills/awa-operations/SKILL.md | 3 +- 10 files changed, 577 insertions(+), 189 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6780d6113..7f8ce530a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,14 +18,17 @@ Notable changes between releases. Detailed migration notes for storage transitio the privileges of the head table it backs, explicitly rather than through default privileges. A runtime without schema `CREATE` can therefore create lanes itself, and `awa storage prepare-queue` is optional and may run as any - runtime role. v053 gives existing lane sequences to the schema owner and adds - their missing grants, one sequence at a time and only where the migrating - role is allowed to. It never fails on a sequence it cannot reassign. The new + runtime role. After v053 commits, `awa migrate` gives existing lane + sequences to the schema owner and adds their missing grants, one sequence per + transaction with a 300 ms lock wait, and only where it is allowed to, trying + the effective role and then the login role. A sequence an open transaction + holds is skipped and reported as busy; one it may not change is reported + with the roles that can repair it. Neither fails the migration. The new read-only `awa storage lane-sequences` (or `awa.lane_sequence_problems(schema)`) lists what remains, with the statements to run and the roles to run them as, and exits 2 while any remain; - `--repair` (or `awa.repair_lane_sequences(schema)`) repairs what the - connected role may. `awa migrate` prints a summary of the leftovers. A binary + `--repair` (or `awa.repair_lane_sequence(schema, sequence)` per sequence) + repairs what the connected role may. `awa migrate` prints a summary of the leftovers. A binary running ahead of v053 detects that `provision_lane` is missing and creates lanes as before, and older binaries create correctly owned lanes once v053 has refreshed the helpers they call. See "Lane ownership and v053" in the diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index 5e53264c5..e111c3574 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -273,6 +273,23 @@ fn sanitize_filename(description: &str) -> String { /// mutating commands require an explicit `--context` or `--database-url` /// once more than one context is defined, and are confirmation-gated on /// contexts marked `production = true`. +/// After a migration, say which lane sequences still need repair and how to +/// finish it. Read-only. +async fn print_lane_sequence_leftovers( + pool: &sqlx::PgPool, +) -> Result<(), Box> { + let mut problems = Vec::new(); + for (schema, _) in awa_model::storage::lane_sequence_problem_counts(pool).await? { + problems.extend(awa_model::storage::lane_sequence_problems(pool, &schema).await?); + } + for next_step in awa_model::storage::lane_sequence_next_steps( + problems.iter().map(|problem| (problem, false)), + ) { + eprintln!("warning: {next_step}"); + } + Ok(()) +} + fn access_for(command: &Commands) -> context::Access { use context::Access::{Mutating, ReadOnly}; match command { @@ -834,7 +851,8 @@ enum StorageCommands { /// Queue-storage schema to check #[arg(long, default_value = "awa")] schema: String, - /// Reassign and grant what the connected role is allowed to first + /// Reassign and grant first, as far as the connected role (or, failing + /// that, its login role) is allowed to #[arg(long)] repair: bool, }, @@ -1194,16 +1212,7 @@ async fn run() -> Result<(), Box> { }; awa_model::migrations::run_with_options(&pool, options).await?; println!("Migrations applied successfully."); - for (schema, count) in - awa_model::storage::lane_sequence_problem_counts(&pool).await? - { - eprintln!( - "warning: {count} lane sequence(s) in schema {schema} are not owned by \ - the schema owner or lack runtime grants, and the migrating role could \ - not repair them. Run `awa storage lane-sequences --schema {schema}` for \ - the statements to run and the roles to run them as." - ); - } + print_lane_sequence_leftovers(&pool).await?; } } @@ -1956,19 +1965,18 @@ async fn run() -> Result<(), Box> { ); } StorageCommands::LaneSequences { schema, repair } => { - let remaining = if repair { - let results = - awa_model::storage::repair_lane_sequences(&pool, &schema).await?; + let next_steps = if repair { + let mut conn = pool.acquire().await?; + let repairs = + awa_model::storage::repair_lane_sequences(&mut conn, &schema) + .await?; println!( "{}", serde_json::to_string_pretty(&serde_json::json!({ - "schema": schema, "results": results, + "schema": schema, "results": repairs, }))? ); - results - .iter() - .filter(|result| result.outcome != "repaired") - .count() + awa_model::storage::lane_sequence_repair_next_steps(&repairs) } else { let problems = awa_model::storage::lane_sequence_problems(&pool, &schema).await?; @@ -1978,15 +1986,14 @@ async fn run() -> Result<(), Box> { "schema": schema, "problems": problems, }))? ); - problems.len() + awa_model::storage::lane_sequence_next_steps( + problems.iter().map(|problem| (problem, false)), + ) }; - if remaining > 0 { - eprintln!( - "{remaining} lane sequence(s) in schema {schema} need repair. Run \ - each missing_grants statement as the sequence owner (or a role \ - that can SET ROLE to it), and each reassign_sql statement as a \ - role that can SET ROLE to both the owner and expected_owner." - ); + if !next_steps.is_empty() { + for next_step in next_steps { + eprintln!("{next_step}"); + } std::process::exit(2); } } diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index 39df93e86..f3736cef1 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -407,15 +407,18 @@ REVOKE EXECUTE ON FUNCTION awa.refresh_lane_sequence_helpers(TEXT) FROM PUBLIC; -- USAGE, SELECT and UPDATE on it, and whoever may read it gets SELECT: -- queue_enqueue_seq_* mirror queue_enqueue_heads, queue_claim_seq_* mirror -- queue_claim_heads. missing_grants holds the statements to run as the --- sequence owner; reassign_sql, when the owner is wrong, needs a role that can --- SET ROLE to both the current owner and the schema owner. Read-only. +-- sequence owner; reassign_sql, when the owner is wrong, needs a role that +-- holds the current owner's privileges and can SET ROLE to the schema owner. +-- repair_roles lists the login roles, superusers aside, that can run both. +-- Read-only. CREATE OR REPLACE FUNCTION awa.lane_sequence_problems(p_schema TEXT, p_sequence TEXT DEFAULT NULL) RETURNS TABLE ( sequence_name TEXT, owner TEXT, expected_owner TEXT, missing_grants TEXT[], - reassign_sql TEXT + reassign_sql TEXT, + repair_roles TEXT[] ) LANGUAGE sql STABLE @@ -489,55 +492,76 @@ AS $problems$ CASE WHEN lane_sequence.relowner <> lane_sequence.nspowner THEN format( 'ALTER SEQUENCE %I.%I OWNER TO %I', p_schema, lane_sequence.relname, pg_get_userbyid(lane_sequence.nspowner) - ) END + ) END, + ARRAY( + SELECT candidate.rolname::text + FROM pg_roles AS candidate + WHERE candidate.rolcanlogin + AND NOT candidate.rolsuper + AND pg_has_role(candidate.oid, lane_sequence.relowner, 'USAGE') + AND (lane_sequence.relowner = lane_sequence.nspowner + OR (pg_has_role(candidate.oid, lane_sequence.nspowner, 'USAGE') + AND pg_has_role( + candidate.oid, + lane_sequence.nspowner, + CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END + ))) + ORDER BY candidate.rolname + ) FROM lane_sequence WHERE lane_sequence.relowner <> lane_sequence.nspowner OR EXISTS (SELECT 1 FROM missing WHERE missing.oid = lane_sequence.oid) ORDER BY lane_sequence.relname $problems$; --- Give every lane sequence in p_schema to the schema owner and add its missing --- runtime grants, one sequence at a time, doing only what the calling role is --- allowed to: a sequence it cannot reassign keeps its owner and still gets --- the grants when the caller can act as that owner. Never raises for a --- sequence it cannot repair. Returns one row per sequence that had a problem: --- outcome is 'repaired', 'owner_unchanged' (grants complete, owner still --- wrong) or 'unrepaired', with the statements still to run. -CREATE OR REPLACE FUNCTION awa.repair_lane_sequences(p_schema TEXT) +-- Give one lane sequence to the schema owner and add its missing runtime +-- grants, doing only what the calling role is allowed to: a sequence it +-- cannot reassign keeps its owner and still gets the grants when the caller +-- holds that owner's privileges. Never raises for a sequence it cannot +-- repair. Reassigning takes the sequence's ACCESS EXCLUSIVE lock, which an +-- open enqueue or claim on the lane conflicts with, so the lock wait is +-- bounded: on timeout nothing changes and the outcome is 'busy'. The lock is +-- held until the caller's transaction ends, so callers repair one sequence +-- per transaction. Returns no row when the sequence needs nothing, otherwise +-- one row: outcome is 'repaired', 'owner_unchanged' (grants complete, owner +-- still wrong), 'unrepaired' or 'busy', with the statements still to run. +CREATE OR REPLACE FUNCTION awa.repair_lane_sequence(p_schema TEXT, p_sequence TEXT) RETURNS TABLE ( sequence_name TEXT, owner TEXT, expected_owner TEXT, outcome TEXT, missing_grants TEXT[], - reassign_sql TEXT + reassign_sql TEXT, + repair_roles TEXT[] ) LANGUAGE plpgsql SECURITY INVOKER SET search_path = pg_catalog, pg_temp +SET lock_timeout = '300ms' AS $repair$ DECLARE v_problem RECORD; v_remaining RECORD; - v_has_remaining BOOLEAN; - v_owner OID; + v_busy BOOLEAN := false; v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 THEN 'SET' ELSE 'MEMBER' END; v_statement TEXT; BEGIN - FOR v_problem IN - SELECT * FROM awa.lane_sequence_problems(p_schema) AS problem - LOOP - SELECT sequence.relowner INTO v_owner - FROM pg_class AS sequence - WHERE sequence.oid = to_regclass(format('%I.%I', p_schema, v_problem.sequence_name)); + SELECT * INTO v_problem + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + IF NOT FOUND THEN + RETURN; + END IF; + BEGIN -- ALTER ... OWNER TO needs the current owner's privileges, SET ROLE to -- the new owner, and the new owner's CREATE on the schema (which a -- schema owner has). Check first so an unrelated owner is skipped - -- rather than raising. + -- without waiting for the lock. IF v_problem.reassign_sql IS NOT NULL - AND pg_has_role(current_user, v_owner, 'USAGE') + AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.owner)), 'USAGE') AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.expected_owner)), v_set_mode) THEN BEGIN EXECUTE v_problem.reassign_sql; @@ -547,10 +571,8 @@ BEGIN END IF; SELECT * INTO v_remaining - FROM awa.lane_sequence_problems(p_schema, v_problem.sequence_name) AS problem; - v_has_remaining := FOUND; - - IF v_has_remaining + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + IF FOUND AND cardinality(v_remaining.missing_grants) > 0 AND pg_has_role(current_user, to_regrole(quote_ident(v_remaining.owner)), 'USAGE') THEN FOREACH v_statement IN ARRAY v_remaining.missing_grants LOOP @@ -560,31 +582,35 @@ BEGIN NULL; END; END LOOP; - SELECT * INTO v_remaining - FROM awa.lane_sequence_problems(p_schema, v_problem.sequence_name) AS problem; - v_has_remaining := FOUND; END IF; - - sequence_name := v_problem.sequence_name; - expected_owner := v_problem.expected_owner; - IF v_has_remaining THEN - owner := v_remaining.owner; - missing_grants := v_remaining.missing_grants; - reassign_sql := v_remaining.reassign_sql; - outcome := CASE WHEN cardinality(v_remaining.missing_grants) = 0 - THEN 'owner_unchanged' ELSE 'unrepaired' END; - ELSE - owner := v_problem.expected_owner; - missing_grants := ARRAY[]::TEXT[]; - reassign_sql := NULL; - outcome := 'repaired'; - END IF; - RETURN NEXT; - END LOOP; + EXCEPTION WHEN lock_not_available THEN + v_busy := true; + END; + + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + sequence_name := v_problem.sequence_name; + expected_owner := v_problem.expected_owner; + IF FOUND THEN + owner := v_remaining.owner; + missing_grants := v_remaining.missing_grants; + reassign_sql := v_remaining.reassign_sql; + repair_roles := v_remaining.repair_roles; + outcome := CASE WHEN v_busy THEN 'busy' + WHEN cardinality(v_remaining.missing_grants) = 0 THEN 'owner_unchanged' + ELSE 'unrepaired' END; + ELSE + owner := v_problem.expected_owner; + missing_grants := ARRAY[]::TEXT[]; + reassign_sql := NULL; + repair_roles := ARRAY[]::TEXT[]; + outcome := 'repaired'; + END IF; + RETURN NEXT; END $repair$; -REVOKE EXECUTE ON FUNCTION awa.repair_lane_sequences(TEXT) FROM PUBLIC; +REVOKE EXECUTE ON FUNCTION awa.repair_lane_sequence(TEXT, TEXT) FROM PUBLIC; -- Retained job output (ADR-048): add the output columns to an existing -- substrate and expose them through terminal_jobs. Called by the installer diff --git a/awa-model/migrations/v053_definer_lane_provisioning.sql b/awa-model/migrations/v053_definer_lane_provisioning.sql index 758627703..2eba66ff6 100644 --- a/awa-model/migrations/v053_definer_lane_provisioning.sql +++ b/awa-model/migrations/v053_definer_lane_provisioning.sql @@ -6,26 +6,24 @@ -- two SECURITY DEFINER functions, provision_lane and provision_lane_sequences, -- owned by the schema owner, and the lane helpers create a missing sequence -- through them. Each new sequence gets the grants of the head table it backs. --- This migration then repairs existing lane sequences: it gives each one to --- the schema owner and adds its missing grants, one sequence at a time and --- only where the migrating role may. Sequences it cannot reassign are reported --- in a warning and by awa.lane_sequence_problems(); the migration still commits. +-- This migration does not reassign existing lane sequences: ALTER SEQUENCE +-- OWNER takes the sequence's ACCESS EXCLUSIVE lock until commit, and the +-- migration runs in one transaction. After it commits, `awa migrate` repairs +-- each lane sequence in its own short transaction with a bounded lock wait +-- (awa.repair_lane_sequence). Another runner can do the same with +-- `awa storage lane-sequences --repair`. -- -- N-1 binaries call the same helpers with the same arguments and cursor -- semantics, so lanes they create are owned by the schema owner too. A binary -- that predates provision_lane keeps using the helpers; this binary probes for -- provision_lane and uses the helpers when it is missing. -- --- Live load: function replacement, plus ALTER SEQUENCE OWNER on the lane --- sequences that need it. That takes each such sequence's ACCESS EXCLUSIVE --- lock until commit, so enqueues and claims on those lanes wait for the --- migration to commit; correctly owned lanes are not touched. Reads the --- catalog only; no table DDL or row scans. +-- Live load: function replacement only. Reads the catalog; no table DDL, row +-- scans or sequence locks. DO $$ DECLARE v_schema TEXT; - v_leftovers TEXT; v_count BIGINT; BEGIN FOR v_schema IN @@ -41,17 +39,13 @@ BEGIN PERFORM awa.refresh_lane_sequence_helpers(v_schema); - SELECT count(*), - string_agg(format('%s (owner %s, %s)', repair.sequence_name, repair.owner, repair.outcome), - ', ' ORDER BY repair.sequence_name) - INTO v_count, v_leftovers - FROM awa.repair_lane_sequences(v_schema) AS repair - WHERE repair.outcome <> 'repaired'; + SELECT count(*) INTO v_count + FROM awa.lane_sequence_problems(v_schema) AS problem; IF v_count > 0 THEN - RAISE WARNING '% lane sequence(s) in schema % still need repair: %', v_count, v_schema, v_leftovers + RAISE NOTICE '% lane sequence(s) in schema % are not owned by the schema owner or lack runtime grants', v_count, v_schema USING HINT = format( - 'Run awa storage lane-sequences --schema %s for the statements to run and the role to run them as.', + 'awa migrate repairs them after this commits. Otherwise run awa storage lane-sequences --schema %s --repair as the migration role itself, without SET ROLE.', v_schema ); END IF; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 51c423222..f30033198 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -423,9 +423,9 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< // only sound because every migration is transaction-safe (no // `CREATE INDEX CONCURRENTLY`, `VACUUM`, or explicit transaction control). let mut retry = 0; - loop { + let version_before = loop { match run_migration_transaction(pool, &options, lock_key).await { - Ok(()) => break, + Ok(version_before) => break version_before, Err(error) if is_migration_deadlock(&error) && retry < MIGRATION_DEADLOCK_RETRIES => { retry += 1; let delay = std::time::Duration::from_millis(50 * (1 << (retry - 1))); @@ -439,6 +439,10 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< } Err(error) => return Err(error), } + }; + + if version_before < LANE_PROVISIONING_VERSION { + repair_lane_sequences_after_migration(pool).await; } // Best-effort admin-metadata cache warmup, deliberately outside the @@ -449,29 +453,73 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< Ok(()) } +/// Returns the schema version the transaction started from. async fn run_migration_transaction( pool: &PgPool, options: &MigrateOptions, lock_key: i64, -) -> Result<(), AwaError> { +) -> Result { let mut tx = pool.begin().await?; sqlx::query("SELECT pg_advisory_xact_lock($1)") .bind(lock_key) .execute(&mut *tx) .await?; - apply_migrations(&mut tx, options).await?; + let version_before = apply_migrations(&mut tx, options).await?; tx.commit().await?; - Ok(()) + Ok(version_before) +} + +/// The migration that makes lane sequences belong to the schema owner. +const LANE_PROVISIONING_VERSION: i32 = 53; + +/// Repair existing lane sequences once v053 has committed. +/// +/// `ALTER SEQUENCE ... OWNER` takes the sequence's ACCESS EXCLUSIVE lock until +/// its transaction ends, so it cannot run inside the migration transaction +/// without stalling every enqueue and claim on those lanes until the whole +/// upgrade commits. Each sequence is repaired in its own short transaction +/// with a bounded lock wait instead. Best-effort: the schema is committed, and +/// whatever is left is reported with the command that finishes it. +/// +/// Uses a detached connection for the reason given on +/// [`warm_admin_metadata_cache`]. +async fn repair_lane_sequences_after_migration(pool: &PgPool) { + let mut conn = match pool.acquire().await { + Ok(conn) => conn.detach(), + Err(error) => { + warn!(%error, "Could not connect to repair lane sequences; run `awa storage lane-sequences --repair`"); + return; + } + }; + match crate::storage::repair_lane_sequences_in_every_schema(&mut conn).await { + Ok(repairs) => { + let repaired = repairs + .iter() + .filter(|repair| repair.outcome == crate::storage::LaneSequenceOutcome::Repaired) + .count(); + if repaired > 0 { + info!(repaired, "Gave lane sequences to the schema owner"); + } + for next_step in crate::storage::lane_sequence_repair_next_steps(&repairs) { + warn!("{next_step}"); + } + } + Err(error) => { + warn!(%error, "Could not repair lane sequences; run `awa storage lane-sequences --repair`"); + } + } + let _ = conn.close().await; } fn is_migration_deadlock(error: &AwaError) -> bool { error.is_deadlock() } +/// Returns the schema version before this run. async fn apply_migrations( conn: &mut PgConnection, options: &MigrateOptions, -) -> Result<(), AwaError> { +) -> Result { let has_schema: bool = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM pg_namespace WHERE nspname = 'awa')") .fetch_one(&mut *conn) @@ -564,7 +612,7 @@ async fn apply_migrations( info!(version = current, "Schema is up to date"); } - Ok(()) + Ok(current) } /// Warm the admin-metadata cache after migrations commit. diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index d67dc2048..14770c40a 100644 --- a/awa-model/src/storage.rs +++ b/awa-model/src/storage.rs @@ -2,7 +2,7 @@ use crate::admin; use crate::error::AwaError; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use sqlx::{FromRow, PgExecutor, PgPool}; +use sqlx::{Connection, FromRow, PgConnection, PgExecutor, PgPool}; use std::collections::BTreeMap; #[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq)] @@ -472,30 +472,65 @@ pub async fn status_report(pool: &PgPool) -> Result, pub reassign_sql: Option, + /// Login roles, superusers aside, that can run every statement above. + /// Empty when no such role exists yet. + pub repair_roles: Vec, } -/// What [`repair_lane_sequences`] did to one lane sequence that had a problem. +/// What a repair did to one lane sequence. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum LaneSequenceOutcome { + /// Owned by the schema owner with every runtime grant. + Repaired, + /// Grants complete, but the role could not reassign the sequence. + OwnerUnchanged, + /// The role could neither reassign the sequence nor add its grants. + Unrepaired, + /// An open transaction held the sequence's lock past the repair's lock + /// timeout, so nothing changed. Rerun the repair. + Busy, +} + +impl TryFrom for LaneSequenceOutcome { + type Error = String; + + fn try_from(outcome: String) -> Result { + match outcome.as_str() { + "repaired" => Ok(Self::Repaired), + "owner_unchanged" => Ok(Self::OwnerUnchanged), + "unrepaired" => Ok(Self::Unrepaired), + "busy" => Ok(Self::Busy), + _ => Err(format!("unknown lane sequence repair outcome {outcome:?}")), + } + } +} + +/// What [`repair_lane_sequences`] did to one lane sequence that had a problem, +/// with what is still wrong afterwards. #[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq, Eq)] pub struct LaneSequenceRepair { + pub schema: String, pub sequence_name: String, /// Owner after the repair. pub owner: String, pub expected_owner: String, - /// `repaired`; `owner_unchanged` when the grants are complete but the - /// connected role could not reassign the sequence; or `unrepaired` when - /// grants are still missing too. - pub outcome: String, + #[sqlx(try_from = "String")] + pub outcome: LaneSequenceOutcome, pub missing_grants: Vec, pub reassign_sql: Option, + /// Login roles, superusers aside, that can finish the repair. + pub repair_roles: Vec, } /// Lane sequences in `schema` that need repair. Read-only; any role may run it. @@ -507,7 +542,8 @@ where E: PgExecutor<'e>, { Ok(sqlx::query_as( - "SELECT sequence_name, owner, expected_owner, missing_grants, reassign_sql \ + "SELECT $1::text AS schema, sequence_name, owner, expected_owner, missing_grants, \ + reassign_sql, repair_roles \ FROM awa.lane_sequence_problems($1)", ) .bind(schema) @@ -516,25 +552,182 @@ where } /// Give each lane sequence in `schema` to the schema owner and add its missing -/// grants, as far as the connected role is allowed to. A sequence the role -/// cannot reassign is returned with the statements still to run, never an -/// error. +/// grants, as far as the connection's roles are allowed to. A sequence that +/// cannot be repaired is returned with the statements still to run and the +/// roles that can run them, never an error. +/// +/// Each sequence is repaired by its own statement, outside any transaction +/// the caller holds, so a reassignment's ACCESS EXCLUSIVE lock is released +/// before the next sequence and waits at most the repair function's lock +/// timeout. A sequence locked for longer by an open enqueue or claim comes +/// back [`LaneSequenceOutcome::Busy`]. +/// +/// Reassigning needs the current owner's privileges. A migration role that +/// connects as itself and sets its role to the schema owner (`role` in the +/// connection options) loses its own privileges while acting as the owner, +/// so a sequence the current role cannot repair is retried once as the login +/// role. pub async fn repair_lane_sequences( - pool: &PgPool, + conn: &mut PgConnection, schema: &str, ) -> Result, AwaError> { - Ok(sqlx::query_as( - "SELECT sequence_name, owner, expected_owner, outcome, missing_grants, reassign_sql \ - FROM awa.repair_lane_sequences($1)", - ) - .bind(schema) - .fetch_all(pool) - .await?) + let (current_role, login_role): (String, String) = + sqlx::query_as("SELECT current_user::text, session_user::text") + .fetch_one(&mut *conn) + .await?; + let problems = lane_sequence_problems(&mut *conn, schema).await?; + let mut repairs = Vec::with_capacity(problems.len()); + for problem in problems { + let mut repair: Option = sqlx::query_as(REPAIR_LANE_SEQUENCE_SQL) + .bind(schema) + .bind(&problem.sequence_name) + .fetch_optional(&mut *conn) + .await?; + let role_was_the_obstacle = matches!( + repair.as_ref().map(|repair| repair.outcome), + Some(LaneSequenceOutcome::OwnerUnchanged | LaneSequenceOutcome::Unrepaired) + ); + if role_was_the_obstacle && login_role != current_role { + if let Some(retried) = + repair_lane_sequence_as_login_role(conn, schema, &problem.sequence_name).await? + { + repair = retried; + } + } + repairs.extend(repair); + } + Ok(repairs) +} + +const REPAIR_LANE_SEQUENCE_SQL: &str = + "SELECT $1::text AS schema, sequence_name, owner, expected_owner, outcome, missing_grants, \ + reassign_sql, repair_roles \ + FROM awa.repair_lane_sequence($1, $2)"; + +/// Repair one sequence with the login role as the current role, for this +/// transaction only. `None` when the login role may not run the repair. +async fn repair_lane_sequence_as_login_role( + conn: &mut PgConnection, + schema: &str, + sequence_name: &str, +) -> Result>, AwaError> { + let mut transaction = conn.begin().await?; + sqlx::query("SET LOCAL ROLE NONE") + .execute(&mut *transaction) + .await?; + let repair = sqlx::query_as(REPAIR_LANE_SEQUENCE_SQL) + .bind(schema) + .bind(sequence_name) + .fetch_optional(&mut *transaction) + .await; + match repair { + Ok(repair) => { + transaction.commit().await?; + Ok(Some(repair)) + } + Err(sqlx::Error::Database(error)) if error.code().as_deref() == Some("42501") => { + transaction.rollback().await?; + Ok(None) + } + Err(error) => Err(error.into()), + } +} + +/// [`repair_lane_sequences`] for every queue-storage schema with a lane +/// sequence that needs repair. +pub async fn repair_lane_sequences_in_every_schema( + conn: &mut PgConnection, +) -> Result, AwaError> { + let mut repairs = Vec::new(); + for (schema, _) in lane_sequence_problem_counts(&mut *conn).await? { + repairs.extend(repair_lane_sequences(conn, &schema).await?); + } + Ok(repairs) +} + +/// One line per group of lane sequences that still need repair, saying what +/// to do next. `busy` marks sequences a repair found locked. +pub fn lane_sequence_next_steps<'a>( + leftovers: impl IntoIterator, +) -> Vec { + let mut groups: BTreeMap<(String, bool, String, String, Vec), usize> = BTreeMap::new(); + for (problem, busy) in leftovers { + *groups + .entry(( + problem.schema.clone(), + busy, + problem.owner.clone(), + problem.expected_owner.clone(), + problem.repair_roles.clone(), + )) + .or_default() += 1; + } + groups + .into_iter() + .map( + |((schema, busy, owner, expected_owner, repair_roles), count)| { + let rerun = format!("awa storage lane-sequences --schema {schema} --repair"); + if busy { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} are busy: an \ + open transaction held the lock, so nothing changed. Rerun `{rerun}`." + ) + } else if repair_roles.is_empty() { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} still need \ + repair, and no login role can both act as {owner} and SET ROLE to \ + {expected_owner}. Grant {owner} to the migration role, then rerun `{rerun}` \ + connected as the migration role, or run the statements from \ + `awa storage lane-sequences --schema {schema}` as {owner}." + ) + } else { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} still need \ + repair. Rerun `{rerun}` connected as one of: {}.", + repair_roles.join(", ") + ) + } + }, + ) + .collect() +} + +impl LaneSequenceRepair { + /// The problem left after this repair, or `None` when it was repaired. + pub fn remaining_problem(&self) -> Option { + (self.outcome != LaneSequenceOutcome::Repaired).then(|| LaneSequenceProblem { + schema: self.schema.clone(), + sequence_name: self.sequence_name.clone(), + owner: self.owner.clone(), + expected_owner: self.expected_owner.clone(), + missing_grants: self.missing_grants.clone(), + reassign_sql: self.reassign_sql.clone(), + repair_roles: self.repair_roles.clone(), + }) + } +} + +/// [`lane_sequence_next_steps`] for the sequences a repair left behind. +pub fn lane_sequence_repair_next_steps(repairs: &[LaneSequenceRepair]) -> Vec { + let leftovers: Vec<(LaneSequenceProblem, bool)> = repairs + .iter() + .filter_map(|repair| { + repair + .remaining_problem() + .map(|problem| (problem, repair.outcome == LaneSequenceOutcome::Busy)) + }) + .collect(); + lane_sequence_next_steps(leftovers.iter().map(|(problem, busy)| (problem, *busy))) } /// How many lane sequences need repair in each queue-storage schema that has /// any. Read-only. -pub async fn lane_sequence_problem_counts(pool: &PgPool) -> Result, AwaError> { +pub async fn lane_sequence_problem_counts<'e, E>( + executor: E, +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ Ok(sqlx::query_as( "SELECT discovered.schema_name, count(*) \ FROM awa.queue_storage_schemas() AS discovered \ @@ -542,6 +735,6 @@ pub async fn lane_sequence_problem_counts(pool: &PgPool) -> Result Result { + PgPoolOptions::new() + .max_connections(1) + .connect_with( + self.connect_options + .clone() + .username(&self.migrator_name) + .options([("role", self.owner_name.as_str())]), + ) + .await + } + + /// The `queue_enqueue_seq_*` sequence of `queue`'s first lane. + async fn enqueue_sequence(&self, queue: &str) -> Result { + Ok(self + .lane_sequence_owners(queue) + .await? + .into_iter() + .map(|(sequence, _)| sequence) + .find(|sequence| sequence.starts_with("queue_enqueue_seq_")) + .expect("lane has an enqueue sequence")) + } + /// `(sequence, owner)` for every lane sequence of `queue`. async fn lane_sequence_owners( &self, @@ -590,20 +618,15 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { )))) .await?; } - let failed = drain(&cluster.app_b, "owner_lane").await.unwrap_err(); - assert!( - failed - .to_string() - .contains("permission denied for sequence"), - "{failed}" - ); - let failed = drain(&cluster.app_b, "app_lane").await.unwrap_err(); - assert!( - failed - .to_string() - .contains("permission denied for sequence"), - "{failed}" - ); + for lane in ["owner_lane", "app_lane"] { + let failed = drain(&cluster.app_b, lane).await.unwrap_err(); + assert!( + failed + .to_string() + .contains("permission denied for sequence"), + "{failed}" + ); + } // 0.7 refuses to migrate while app roles hold CREATE on the schema. cluster .owner @@ -613,19 +636,17 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { )))) .await?; - // The migration runs as the migration role itself. It is a member of the - // owner, so it may reassign its own and the owner's sequences, but not - // the app role's. - for (_, _, statement) in migrations::migration_sql_range(52, migrations::CURRENT_VERSION) { - cluster - .migrator - .execute(sqlx::raw_sql(sql(statement))) - .await?; - } + // `awa migrate` connects as the migration role and acts as the owner. As + // the owner it adds the grants the owner's sequences lack; the migration + // role's own sequences need the migration role's privileges, so they are + // reassigned as the login role. Neither can act as the app role. + let migrate_pool = cluster.migrator_as_owner().await?; + tokio::time::timeout(Duration::from_secs(120), migrations::run(&migrate_pool)).await??; assert_eq!( migrations::current_version(&cluster.owner).await?, migrations::CURRENT_VERSION ); + migrate_pool.close().await; cluster .assert_lanes_owned_by_owner("migrator_lane", 1) .await?; @@ -633,10 +654,10 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { assert_eq!(drain(&cluster.app_b, "migrator_lane").await?, 1); assert_eq!(drain(&cluster.app_b, "owner_lane").await?, 1); - let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; - assert_eq!(problems.len(), 2, "{problems:?}"); let app_a = cluster.roles[3].clone(); let app_a = &app_a; + let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; + assert_eq!(problems.len(), 2, "{problems:?}"); for problem in &problems { assert_eq!(&problem.owner, app_a); assert_eq!(problem.expected_owner, cluster.owner_name); @@ -657,21 +678,31 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { .as_str() ) ); + assert!(problem.repair_roles.is_empty(), "{problem:?}"); } + let next_steps = + storage::lane_sequence_next_steps(problems.iter().map(|problem| (problem, false))); + assert_eq!(next_steps.len(), 1, "{next_steps:?}"); + assert!( + next_steps[0].contains(&format!("no login role can both act as {app_a}")), + "{next_steps:?}" + ); // Re-running the repair as the migration role changes nothing and fails - // nothing. The migration role created the repair function, so it grants - // its use. + // nothing. cluster - .migrator + .owner .execute(sqlx::raw_sql(sql(format!( - "GRANT EXECUTE ON FUNCTION awa.repair_lane_sequences(text) TO {}", + "GRANT EXECUTE ON FUNCTION awa.repair_lane_sequence(text, text) TO {}", cluster.runtime_name )))) .await?; - let rerun = storage::repair_lane_sequences(&cluster.migrator, "awa").await?; + let rerun = + storage::repair_lane_sequences(&mut *cluster.migrator.acquire().await?, "awa").await?; assert_eq!(rerun.len(), 2); - assert!(rerun.iter().all(|result| result.outcome == "unrepaired")); + assert!(rerun + .iter() + .all(|result| result.outcome == LaneSequenceOutcome::Unrepaired)); // A role that can act as the app role but cannot become the owner adds // the grants and leaves the owner for someone who can. @@ -686,20 +717,14 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { .await?; let delegate_pool = PgPoolOptions::new() .max_connections(1) - .connect_with( - PgConnectOptions::from_str(&std::env::var("DATABASE_URL").unwrap_or_else(|_| { - "postgres://postgres:test@localhost:15432/awa_test".to_string() - }))? - .database(&cluster.database) - .username(&delegate) - .password("lane_test"), - ) + .connect_with(cluster.connect_options.clone().username(&delegate)) .await?; - let delegated = storage::repair_lane_sequences(&delegate_pool, "awa").await?; + let delegated = + storage::repair_lane_sequences(&mut *delegate_pool.acquire().await?, "awa").await?; delegate_pool.close().await; assert_eq!(delegated.len(), 2); for result in &delegated { - assert_eq!(result.outcome, "owner_unchanged"); + assert_eq!(result.outcome, LaneSequenceOutcome::OwnerUnchanged); assert_eq!(&result.owner, app_a); assert!(result.missing_grants.is_empty()); assert!(result.reassign_sql.is_some()); @@ -707,7 +732,8 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); enqueue(&cluster.app_c, "app_lane", 1, 1).await?; - // A role that is a member of both the app role and the owner reassigns. + // Once the migration role is a member of the app role too, the check names + // it, and a repair connected as it reassigns. cluster .admin .execute(sqlx::raw_sql(sql(format!( @@ -715,21 +741,93 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { migrator = cluster.migrator_name )))) .await?; - let repaired = storage::repair_lane_sequences(&cluster.migrator, "awa").await?; + let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; + assert!(problems + .iter() + .all(|problem| problem.repair_roles == vec![cluster.migrator_name.clone()])); + let repaired = + storage::repair_lane_sequences(&mut *cluster.migrator.acquire().await?, "awa").await?; assert_eq!(repaired.len(), 2); assert!( - repaired.iter().all(|result| result.outcome == "repaired"), + repaired + .iter() + .all(|result| result.outcome == LaneSequenceOutcome::Repaired), "{repaired:?}" ); assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") .await? .is_empty()); - assert!(storage::repair_lane_sequences(&cluster.migrator, "awa") - .await? - .is_empty()); cluster.assert_lanes_owned_by_owner("app_lane", 1).await?; assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); cluster.roles.push(delegate); cluster.teardown().await } + +/// Reassigning a lane sequence waits on any open enqueue or claim on the lane, +/// so the repair bounds that wait and holds each sequence's lock only for its +/// own short transaction. A held sequence is skipped and reported; the rest +/// are repaired, and a later repair finishes it. +#[tokio::test] +async fn migration_skips_a_lane_sequence_an_open_transaction_holds() -> TestResult { + let cluster = Cluster::new(Schema::BeforeProvisioning, true).await?; + enqueue(&cluster.migrator, "held_lane", 1, 1).await?; + enqueue(&cluster.migrator, "free_lane", 1, 1).await?; + let held_sequence = cluster.enqueue_sequence("held_lane").await?; + // 0.7 refuses to migrate while app roles hold CREATE on the schema. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE CREATE ON SCHEMA awa FROM {}", + cluster.runtime_name + )))) + .await?; + + // An enqueue in flight on held_lane holds its sequence's lock. + let mut holder = cluster.app_b.begin().await?; + sqlx::query("SELECT nextval($1::regclass)") + .bind(format!("awa.{held_sequence}")) + .execute(&mut *holder) + .await?; + + let migrate_pool = cluster.migrator_as_owner().await?; + // The migration completes while the lock is held. + tokio::time::timeout(Duration::from_secs(120), migrations::run(&migrate_pool)).await??; + for (sequence, owner) in cluster.lane_sequence_owners("held_lane").await? { + let expected = if sequence == held_sequence { + &cluster.migrator_name + } else { + &cluster.owner_name + }; + assert_eq!(&owner, expected, "{sequence}"); + } + cluster.assert_lanes_owned_by_owner("free_lane", 1).await?; + + // A repair while the transaction is still open reports the sequence busy. + let mut conn = migrate_pool.acquire().await?; + let busy = storage::repair_lane_sequences(&mut conn, "awa").await?; + assert_eq!(busy.len(), 1, "{busy:?}"); + assert_eq!(busy[0].sequence_name, held_sequence); + assert_eq!(busy[0].outcome, LaneSequenceOutcome::Busy); + assert_eq!(busy[0].owner, cluster.migrator_name); + let next_steps = storage::lane_sequence_repair_next_steps(&busy); + assert_eq!(next_steps.len(), 1); + assert!(next_steps[0].contains("busy"), "{next_steps:?}"); + assert!( + next_steps[0].contains("awa storage lane-sequences --schema awa --repair"), + "{next_steps:?}" + ); + + holder.rollback().await?; + let repaired = storage::repair_lane_sequences(&mut conn, "awa").await?; + assert_eq!(repaired.len(), 1, "{repaired:?}"); + assert_eq!(repaired[0].outcome, LaneSequenceOutcome::Repaired); + drop(conn); + migrate_pool.close().await; + cluster.assert_lanes_owned_by_owner("held_lane", 1).await?; + assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") + .await? + .is_empty()); + assert_eq!(drain(&cluster.app_b, "held_lane").await?, 1); + cluster.teardown().await +} diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index 5f15793e7..d8d9f5497 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -85,17 +85,32 @@ awa --database-url "$AWA_DATABASE_URL" storage lane-sequences --schema awa SELECT * FROM awa.lane_sequence_problems('awa'); ``` -The command exits 2 while any remain. Migration v053 runs the repair as the migrating role: for each sequence it changes the owner to the schema owner and adds the missing grants, wherever that role is allowed to. A sequence it is not allowed to change does not fail the migration. It is reported in a warning, and `awa migrate` prints a summary. Run the repair again at any time; it is idempotent: +The command exits 2 while any remain. + +`awa migrate` repairs them once the migration commits, one sequence per short transaction: it gives each sequence to the schema owner and adds its missing grants, wherever it is allowed to. Reassigning takes the sequence's `ACCESS EXCLUSIVE` lock, so the repair waits at most 300 ms for it. A sequence that an open enqueue or claim holds for longer is skipped, left unchanged and reported as busy. A sequence the repair is not allowed to change is reported with the roles that can finish it. Neither fails the migration. Run the repair again at any time; it is idempotent: ```bash -PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair +awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair ``` -Changing a sequence's owner requires a role that can `SET ROLE` to both the current owner and the schema owner. A migrator usually can for sequences it or the owner created, but not for one an unrelated application role created, and managed PostgreSQL offers no superuser. For each sequence that remains, the output gives the statements to run: +Which role can repair a sequence: + +- Reassigning it needs a role that holds the current owner's privileges (a member of the owner, with inheritance) and can `SET ROLE` to the schema owner. +- Adding its grants needs the current owner's privileges. + +The schema owner holds neither the migration role's nor an application role's privileges, so a repair running as `awa_owner` cannot reassign sequences those roles created. That includes the migration role connected with `PGOPTIONS='-c role=awa_owner'`. Run the repair as the migration role itself, without `PGOPTIONS`, as above. `awa migrate` and `--repair` retry each sequence the effective role could not repair as the login role, so `awa migrate` run with `PGOPTIONS='-c role=awa_owner'`, as documented, still repairs the migration role's sequences. + +For each sequence that remains, the output lists `repair_roles`, the login roles other than superusers that can repair it, and the statements to run. A busy sequence only needs another `--repair`. When `repair_roles` is empty, which is usual for a sequence an application role created on managed PostgreSQL, either: -- `missing_grants`: run them as the current owner, or as a role that can `SET ROLE` to it. Every runtime role can then use the lane, although the owner stays unchanged. -- `reassign_sql`: run it as a role that is a member of both roles, for example after `GRANT app_role TO awa_migrator` from a role with `ADMIN` on `app_role`. Revoke that membership afterwards. +- run `missing_grants` as the current owner. Every runtime role can then use the lane, although the owner stays unchanged; or +- give the migration role the owner's privileges, for example `GRANT app_role TO awa_migrator` from a role with `ADMIN` on `app_role`, rerun `--repair` as the migration role, then revoke the membership. + +Without the CLI, repair one sequence per transaction from `psql` in autocommit mode, connected as the migration role: + +```sql +SELECT format('SELECT * FROM awa.repair_lane_sequence(%L, %L)', 'awa', sequence_name) +FROM awa.lane_sequence_problems('awa') \gexec +``` ## Custom queue-storage schemas ## Custom queue-storage schemas diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 7fff6173f..59ef9ba89 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -28,9 +28,10 @@ default privileges covered that role could use the lane; any other role failed with `permission denied for sequence`. v053 installs `SECURITY DEFINER` `provision_lane` and `provision_lane_sequences` functions, owned by the schema owner, in each queue-storage schema. The lane helpers create sequences through -them, with grants that mirror the lane's head table. v053 then gives every -existing lane sequence to the schema owner and adds its missing grants, as far -as the migrating role is allowed to. +them, with grants that mirror the lane's head table. Once v053 commits, +`awa migrate` gives every existing lane sequence to the schema owner and adds +its missing grants, as far as it is allowed to, one sequence per short +transaction with a bounded lock wait. 0.6 and earlier 0.7 binaries call the same helpers, so lanes they create after v053 are owned by the schema owner too. A 0.7 binary running ahead of v053 @@ -42,13 +43,15 @@ To adopt it: while lanes need repair). 2. Run `awa migrate` with the schema owner as the effective role, as described in [Database roles](security/database-roles.md). -3. Read the migrate output. Any lane sequence it could not reassign is listed; - `awa storage lane-sequences` gives the statements to run and the role to run - them as. See +3. Read the migrate output. Any lane sequence it could not repair is listed in + a warning. A busy one needs another run of + `awa storage lane-sequences --repair` as the migration role (without + `PGOPTIONS`); the others name the roles that can repair them. See [Check and repair lane ownership](security/database-roles.md#check-and-repair-lane-ownership). External migration runners apply the generated `awa migrate --sql --from 52 ---to 53` bundle, then run `awa storage lane-sequences` to check for leftovers. +--to 53` bundle, which does not repair existing lane sequences, then run +`awa storage lane-sequences --repair` as the migration role itself. ## Custom-schema discovery during migration diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index 5124fbc3f..64a23de16 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -96,7 +96,8 @@ From v053 any runtime role creates lanes through the schema owner's CREATE nor sequence default privileges, and `awa storage prepare-queue` is optional. Run migrations with the schema owner as the effective role. Check lane ownership with `awa storage lane-sequences` (exits 2 while lanes need -repair) and fix it with `--repair` or the statements it prints; see +repair) and fix it with `--repair`, connected as the migration role itself +rather than as the schema owner, or with the statements it prints; see [Database roles](../../docs/security/database-roles.md#check-and-repair-lane-ownership). ## Rolling Upgrades From bfccec55abf279ce6ac9c1344b3f77cd1ed29490 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 20:58:37 +1300 Subject: [PATCH 13/18] docs(security): clarify lane sequence repair outcomes --- awa-model/migrations/v053_definer_lane_provisioning.sql | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/awa-model/migrations/v053_definer_lane_provisioning.sql b/awa-model/migrations/v053_definer_lane_provisioning.sql index 2eba66ff6..ced1a7642 100644 --- a/awa-model/migrations/v053_definer_lane_provisioning.sql +++ b/awa-model/migrations/v053_definer_lane_provisioning.sql @@ -45,7 +45,7 @@ BEGIN IF v_count > 0 THEN RAISE NOTICE '% lane sequence(s) in schema % are not owned by the schema owner or lack runtime grants', v_count, v_schema USING HINT = format( - 'awa migrate repairs them after this commits. Otherwise run awa storage lane-sequences --schema %s --repair as the migration role itself, without SET ROLE.', + 'awa migrate attempts repair after this commits; review its warnings for any sequences left unrepaired. To retry, run awa storage lane-sequences --schema %s --repair as the migration role itself, without SET ROLE.', v_schema ); END IF; From 0f84fa30aa5b1c73bad598d6c32b17c0b008b27b Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 21:11:39 +1300 Subject: [PATCH 14/18] Check statistics helper privileges before each refresh pass --- CHANGELOG.md | 2 + .../v053_analyze_queue_storage_statistics.sql | 1 + awa-worker/src/statistics.rs | 162 ++++++++++++++++-- docs/configuration.md | 2 +- docs/security/database-roles.md | 2 + skills/awa-operations/SKILL.md | 7 +- 6 files changed, 159 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 699c9bb16..f915fefb5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,8 @@ Notable changes between releases. Detailed migration notes for storage transitio function's owner may not analyse the tables, logs one warning naming what is missing and skips the refresh without making failing calls or affecting job processing. Earlier binaries never call the function. + Missing EXECUTE or head-table privileges are checked before each pass, + avoiding a failing helper call for every relation. ### Fixed diff --git a/awa-model/migrations/v053_analyze_queue_storage_statistics.sql b/awa-model/migrations/v053_analyze_queue_storage_statistics.sql index d9ec60d0f..0591b8cc8 100644 --- a/awa-model/migrations/v053_analyze_queue_storage_statistics.sql +++ b/awa-model/migrations/v053_analyze_queue_storage_statistics.sql @@ -68,6 +68,7 @@ DECLARE v_relowner OID; v_may_analyze BOOLEAN; BEGIN + -- awa_statistics_helper_v1: bump when the definition or grants change. IF NOT EXISTS ( SELECT 1 FROM awa.queue_storage_schemas() AS substrate WHERE substrate.schema_name = p_schema diff --git a/awa-worker/src/statistics.rs b/awa-worker/src/statistics.rs index 352b97b64..f09940bb9 100644 --- a/awa-worker/src/statistics.rs +++ b/awa-worker/src/statistics.rs @@ -116,6 +116,13 @@ impl PassSummary { } } +#[derive(Debug, PartialEq, Eq)] +enum HelperAccess { + Available, + Missing, + Denied, +} + /// Refreshes planner statistics for one queue-storage schema. pub(crate) struct StatisticsRefresher { pool: PgPool, @@ -205,9 +212,9 @@ impl StatisticsRefresher { } } - match self.helper_installed().await { - Ok(true) => {} - Ok(false) => { + match self.helper_access().await { + Ok(HelperAccess::Available) => {} + Ok(HelperAccess::Missing) => { self.warn_unavailable_once(); self.record( RelationScope::Pass, @@ -217,8 +224,23 @@ impl StatisticsRefresher { summary.record(AnalyzeOutcome::Unavailable); return summary; } + Ok(HelperAccess::Denied) => { + self.warn_privilege_once( + "queue_claim_heads", + "the runtime role needs EXECUTE on awa.analyze_queue_storage_relation(text, text), \ + and the login role needs INSERT, UPDATE and DELETE on this schema's queue_claim_heads; \ + restore those grants as the schema owner", + ); + self.record( + RelationScope::Pass, + AnalyzeOutcome::SkippedPrivilege, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::SkippedPrivilege); + return summary; + } Err(err) => { - warn!(error = %err, schema = %self.schema, "Could not probe for the statistics helper"); + warn!(error = %err, schema = %self.schema, "Could not probe access to the statistics helper"); self.record( RelationScope::Pass, AnalyzeOutcome::Failed, @@ -277,16 +299,33 @@ impl StatisticsRefresher { .record_statistics_analyze(scope.as_str(), outcome.as_str(), elapsed); } - /// One catalog lookup per pass. A missing helper (a binary running ahead - /// of the schema) then costs this query instead of a failing call, and a - /// server-side ERROR line, per relation; the pass interval is the - /// re-probe interval. - async fn helper_installed(&self) -> Result { - sqlx::query_scalar( - "SELECT to_regprocedure('awa.analyze_queue_storage_relation(text,text)') IS NOT NULL", + /// Check the helper and its caller privileges once per pass, so a missing + /// migration or grant does not produce a failing call for every relation. + async fn helper_access(&self) -> Result { + let (installed, allowed): (bool, bool) = sqlx::query_as( + r#" + SELECT helper IS NOT NULL, + COALESCE( + has_function_privilege(helper, 'EXECUTE') + AND has_table_privilege(session_user, heads, 'INSERT') + AND has_table_privilege(session_user, heads, 'UPDATE') + AND has_table_privilege(session_user, heads, 'DELETE'), + false + ) + FROM ( + SELECT to_regprocedure('awa.analyze_queue_storage_relation(text,text)') AS helper, + to_regclass(format('%I.queue_claim_heads', $1::text)) AS heads + ) AS objects + "#, ) + .bind(&self.schema) .fetch_one(&self.pool) - .await + .await?; + Ok(match (installed, allowed) { + (false, _) => HelperAccess::Missing, + (true, false) => HelperAccess::Denied, + (true, true) => HelperAccess::Available, + }) } async fn ring_parents(&self) -> Result, sqlx::Error> { @@ -314,8 +353,8 @@ impl StatisticsRefresher { self.warn_privilege_once( relation, "the owner of awa.analyze_queue_storage_relation() may not ANALYZE this table; \ - run migrations as the role that owns the queue-storage tables, or grant the \ - function owner MAINTAIN on them (PostgreSQL 17+)", + change the function's owner to the table owner with ALTER FUNCTION, \ + or grant the function owner MAINTAIN on them (PostgreSQL 17+)", ); AnalyzeOutcome::SkippedPrivilege } @@ -509,4 +548,99 @@ mod tests { .await .expect("drop database"); } + + /// Missing EXECUTE or login-role DML is a single skipped pass, rather + /// than fifteen server-side permission errors. + #[tokio::test] + async fn missing_helper_privileges_skip_the_pass_without_calling_it() { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let options = PgConnectOptions::from_str(&url).expect("parse DATABASE_URL"); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect_with(options.clone()) + .await + .expect("connect admin"); + let database = format!("awa_stats_access_unit_{}", std::process::id()); + let role = format!("awa_stats_access_{}", std::process::id()); + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(awa_model::audited_sql(format!("DROP ROLE IF EXISTS {role}; CREATE ROLE {role} LOGIN PASSWORD 'statistics_access_test'"))) + .execute(&admin) + .await + .expect("create test role"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "CREATE DATABASE {database}" + ))) + .execute(&admin) + .await + .expect("create database"); + let owner = PgPoolOptions::new() + .max_connections(2) + .connect_with(options.clone().database(&database)) + .await + .expect("connect owner"); + awa_model::migrations::run(&owner).await.expect("migrate"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "GRANT USAGE ON SCHEMA awa TO {role}; \ + GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA awa TO {role}; \ + REVOKE EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) FROM {role}" + ))) + .execute(&owner) + .await + .expect("grant head-table privileges only"); + let runtime = PgPoolOptions::new() + .max_connections(2) + .connect_with( + options + .database(&database) + .username(&role) + .password("statistics_access_test"), + ) + .await + .expect("connect runtime"); + let refresher = StatisticsRefresher::new( + runtime.clone(), + "awa".to_string(), + AwaMetrics::from_global(), + false, + ); + let denied = PassSummary { + skipped_privilege: 1, + ..PassSummary::default() + }; + assert_eq!(refresher.run_pass().await, denied); + sqlx::raw_sql(awa_model::audited_sql(format!( + "GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {role}" + ))) + .execute(&owner) + .await + .expect("grant helper EXECUTE"); + assert_eq!(refresher.run_pass().await.ok, 15); + sqlx::raw_sql(awa_model::audited_sql(format!( + "REVOKE INSERT ON awa.queue_claim_heads FROM {role}" + ))) + .execute(&owner) + .await + .expect("revoke login-role head privilege"); + assert_eq!(refresher.run_pass().await, denied); + + runtime.close().await; + owner.close().await; + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop database"); + sqlx::raw_sql(awa_model::audited_sql(format!("DROP ROLE {role}"))) + .execute(&admin) + .await + .expect("drop test role"); + admin.close().await; + } } diff --git a/docs/configuration.md b/docs/configuration.md index 154c54801..4640dda6d 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -539,7 +539,7 @@ Autovacuum does not analyse partitioned tables, so the ring parents (`ready_entr The maintenance leader runs `ANALYZE` on both when it is elected and every `statistics_refresh_interval` (default `5m`). One pass takes well under a second on a small ring and a few seconds on a ring with millions of retained rows. It has no measurable effect on claim or completion throughput. - **Locks.** Each relation is analysed in its own transaction with a 2 s `lock_timeout` and a 120 s `statement_timeout`. `SKIP_LOCKED` skips a child that ring prune holds during the per-child analysis, but sampling for the parent's statistics waits for it up to the `lock_timeout`; a call that times out is recorded as `skipped_lock` and retried on the next interval. Prune never waits for the refresh: its own 10 ms lock timeout applies, so an attempt that lands while `ANALYZE` samples the oldest slot's tables returns `blocked` and retries on its usual backoff. `ANALYZE` of a partitioned table samples every child, so that window cannot be avoided. -- **Privileges.** The refresh calls `awa.analyze_queue_storage_relation(schema, relation)` (schema v053), a `SECURITY DEFINER` function owned by the role that ran the migration, so the runtime role needs no `MAINTAIN` grant and does not have to own the tables. It is executable only by roles that hold INSERT, UPDATE and DELETE on the queue head tables (the migration grants it to the existing ones). Each pass first probes for the function; if it is missing (run `awa migrate`) or its owner may not analyse the tables, the leader logs one warning that names what is missing and skips the pass without making failing calls; job processing is unaffected. +- **Privileges.** The refresh calls `awa.analyze_queue_storage_relation(schema, relation)`, a `SECURITY DEFINER` function installed by migration v053. The runtime needs EXECUTE on the helper, and the login role needs INSERT, UPDATE and DELETE on the target schema's `queue_claim_heads`; the migration grants EXECUTE to existing runtime grantees. A missing helper or caller privilege skips the pass before making calls and logs one warning. Run `awa migrate` for a missing helper; restore the grants as the schema owner for `skipped_privilege`. If the helper owner may not analyse the tables, transfer the function to the table owner or grant its owner `MAINTAIN` on PostgreSQL 17+. See the [database roles guide](security/database-roles.md) for the ownership command. - **Metrics.** `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration`, by `awa.analyze.scope` (`heads`, `ring_parent`, `pass`) and `awa.analyze.outcome` (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`). A sustained `skipped_lock` rate means something else holds the tables; any `skipped_privilege` or `unavailable` means the refresh is not running. ### Rust diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index 150c2fe81..2a7ed1c14 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -69,6 +69,8 @@ The maintenance leader runs `ANALYZE` on the queue-storage head tables and parti Run migrations with the schema owner as the effective role, as shown above, so the function's owner owns the tables. A function owned by a role that may not analyse them (it is not the table owner, the database owner, or a `MAINTAIN` holder) makes every pass report `skipped_privilege` and logs one warning. The function only analyses the two head tables and the thirteen ring parents of a schema that `awa.queue_storage_schemas()` reports as trusted, and `ANALYZE` changes statistics only. +If the helper's owner may not analyse the tables, run `ALTER FUNCTION awa.analyze_queue_storage_relation(text, text) OWNER TO awa_owner` as a role permitted to transfer ownership, or grant its owner `MAINTAIN` on the target tables on PostgreSQL 17+. Replacing the function body leaves its owner unchanged. + ## Provision queues before starting restricted runtimes From schema v049, runtime calls use existing lane sequences without issuing diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index 90ad16015..7f81928e5 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -321,8 +321,11 @@ On queue storage the leader runs `ANALYZE` on `queue_claim_heads`, analyses them. It goes through the `SECURITY DEFINER` function `awa.analyze_queue_storage_relation` from migration v053, so a non-owner runtime role needs no `MAINTAIN` grant. Watch `awa.maintenance.analyze.attempts`: any -`skipped_privilege` or `unavailable` means it is not running (run `awa migrate` -as the schema owner). See `docs/configuration.md#planner-statistics`. +`unavailable` means the helper is missing: run `awa migrate` as the schema +owner. For `skipped_privilege`, restore helper EXECUTE and the login role's +head-table DML, or correct the helper owner; migration reruns preserve an +existing function's owner. See `docs/configuration.md#planner-statistics` and +`docs/security/database-roles.md`. ## Managed Postgres From 7d459b50dc2837d4e9629c249a49f13a639adb31 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 21:13:24 +1300 Subject: [PATCH 15/18] fix(insert): keep schema availability local to each connection --- CHANGELOG.md | 21 +- awa-model/src/insert.rs | 182 ++++-------------- awa/tests/enqueue_lane_contention_test.rs | 2 +- awa/tests/insert_many_batch_test.rs | 8 +- awa/tests/insert_many_schema_fallback_test.rs | 127 +++++++++--- docs/architecture.md | 4 +- docs/configuration.md | 2 +- docs/upgrade-0.6-to-0.7.md | 13 +- skills/awa-jobs/SKILL.md | 5 +- 9 files changed, 154 insertions(+), 210 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9cd219c09..c1db94b96 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,26 +17,7 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Changed -- **Batch enqueues hold a lane lock for the lane write only.** `insert_many` - and `insert_many_copy` called `awa.insert_job_compat` per row, which takes - the `(queue, priority, shard)` lane lock on the batch's first row and holds - it, because the lock lasts to commit, while every remaining row is - processed: concurrent producers on one lane queued behind whole batches. - Migration v053 adds `awa.insert_jobs_compat_batch`, which allocates job ids, - claims unique keys (in key order), registers new lanes and builds the rows - first, then reserves each lane's sequence range once, in lane order, and - writes the ready rows and segments in two statements. A 185-row batch into - one lane holds the lock for a fraction of the time it did, which raises the - per-lane ceiling several times over and cuts the time producers, promotion - and rescue spend waiting for the lane. `insert_many_copy` skips unique - conflicts in one pass instead of one savepoint per row. A binary that runs - ahead of the migration keeps working: `insert_many` and `insert_many_copy` - look for the function in the catalog and use the per-row path until it - exists, then switch to the batch function without a restart (the process - keeps the answer; a call that finds a cached function gone fails once with - `undefined_function` and the next call probes again). Binaries before this - release keep calling `awa.insert_job_compat` and share the lane lock with the - new path, so a mixed fleet keeps each batch contiguous and in order. +- **COPY and opt-in skip-conflict batch enqueues hold lane locks only for reservation, ready-row writes and commit.** Migration v053 adds `awa.insert_jobs_compat_batch`, which allocates job ids, claims unique keys in key order, registers lanes and prepares rows before reserving each lane's sequence range in lane order. `insert_many_copy` probes its own connection on every call and uses the function when installed, falling back to per-row inserts on older schemas. `insert_many_skip_unique_conflicts` requires the function. Generic `insert_many` keeps its single-statement per-row path on every supported schema, preserving caller transactions when one process writes to databases at different migration levels. All paths return rows in input order and share the same per-lane lock with older binaries. ### Fixed diff --git a/awa-model/src/insert.rs b/awa-model/src/insert.rs index 79fdde6cb..61a163843 100644 --- a/awa-model/src/insert.rs +++ b/awa-model/src/insert.rs @@ -4,16 +4,16 @@ use crate::queue_storage::ordering_key_hash64; use crate::unique::compute_unique_key; use crate::JobArgs; use sqlx::postgres::PgConnection; -use sqlx::{FromRow, PgExecutor, PgPool, Row}; -use std::sync::atomic::{AtomicU8, Ordering}; +use sqlx::{PgExecutor, PgPool}; // ── Lane lock order ───────────────────────────────────────────────────── // // On queue storage every available-state insert holds a transaction-scoped // advisory lock on its `(queue, priority, enqueue_shard)` lane until the // enclosing transaction ends, so two transactions that enqueue into the same -// lanes in opposite orders can deadlock. The batch inserts here go through -// `awa.insert_jobs_compat_batch` (migration v053), which does the +// lanes in opposite orders can deadlock. `insert_many_copy` and +// `insert_many_skip_unique_conflicts` use `awa.insert_jobs_compat_batch` +// (migration v053), which does the // lane-independent work first and then reserves its lanes in one canonical // order: `(queue, priority, shard)` ascending, queues compared bytewise // (`COLLATE "C"`, the order of Rust's `str`). `QueueStorage` batch enqueues @@ -21,8 +21,8 @@ use std::sync::atomic::{AtomicU8, Ordering}; // the reservation and the ready-row writes keeps a lane's lock window // independent of the batch's unique-claim and payload work. // -// A schema without the function (a binary running ahead of the migration) is -// served by the per-row path: rows are sorted by the same key, stably, before +// Generic `insert_many` always uses the per-row path, as does +// `insert_many_copy` before the migration: rows are sorted by the same key before // they reach `awa.insert_job_compat`, and returned in input order. Both paths // take the same lane lock, so producers on either, and older binaries, share a // lane without reordering it. @@ -245,8 +245,8 @@ const INSERT_MANY_BATCH_SQL: &str = r#" /// The per-row form of [`INSERT_MANY_BATCH_SQL`], for schemas without /// `awa.insert_jobs_compat_batch`: one `awa.insert_job_compat` call per row, /// rows inserted in canonical lane order (see "Lane lock order" above) and -/// returned in input order. `batch_insert_available` reports whether the batch -/// function exists, so the caller can switch to it. +/// returned in input order. It works before and after the batch migration, +/// including when one process writes to databases at different schema versions. const INSERT_MANY_ROW_FUNCTION_SQL: &str = r#" WITH input AS MATERIALIZED ( SELECT @@ -281,7 +281,7 @@ const INSERT_MANY_ROW_FUNCTION_SQL: &str = r#" ) ORDER BY raw.queue COLLATE "C", raw.priority, lane_shard, raw.ord ) - SELECT inserted.*, (SELECT to_regprocedure('awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],smallint[],smallint[],timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') IS NOT NULL ) AS batch_insert_available + SELECT inserted.* FROM input CROSS JOIN LATERAL awa.insert_job_compat( input.kind, @@ -513,10 +513,12 @@ fn ordering_hash_text(row: &PreparedRow) -> Option { /// (SQLSTATE `23505`, as [`AwaError::Database`]); use /// [`insert_many_skip_unique_conflicts`] to insert the rest instead. /// -/// Uses `awa.insert_jobs_compat_batch` when the schema has migration v053 and -/// the per-row `awa.insert_job_compat` path when it does not, deciding from the -/// schema itself. The first calls of a process use the per-row path until one -/// of them has seen the function. +/// Calls `awa.insert_job_compat` per row, in canonical lane order, on every +/// supported schema. A generic executor is consumed by one statement and +/// cannot be probed separately, so this path never assumes another database +/// has the same schema. [`insert_many_copy`] probes its concrete connection +/// and uses the batch function when available; +/// [`insert_many_skip_unique_conflicts`] requires that function. #[tracing::instrument( skip(executor, jobs), fields( @@ -532,7 +534,16 @@ pub async fn insert_many<'e, E>(executor: E, jobs: &[InsertParams]) -> Result, { - insert_many_batch(executor, jobs, false).await + if jobs.is_empty() { + return Ok(Vec::new()); + } + let rows = precompute_rows(jobs)?; + Ok(bind_row_function_columns( + sqlx::query_as::<_, JobRow>(INSERT_MANY_ROW_FUNCTION_SQL), + &rows, + ) + .fetch_all(executor) + .await?) } /// Insert multiple jobs in a single statement, skipping those whose unique key @@ -564,128 +575,23 @@ pub async fn insert_many_skip_unique_conflicts<'e, E>( executor: E, jobs: &[InsertParams], ) -> Result, AwaError> -where - E: PgExecutor<'e>, -{ - insert_many_batch(executor, jobs, true).await -} - -async fn insert_many_batch<'e, E>( - executor: E, - jobs: &[InsertParams], - skip_unique_conflicts: bool, -) -> Result, AwaError> where E: PgExecutor<'e>, { if jobs.is_empty() { return Ok(Vec::new()); } - let rows = precompute_rows(jobs)?; - - // Skipping conflicts has no per-row equivalent in one statement, so it - // always calls the batch function and reports a schema without it. - if skip_unique_conflicts { - return match bind_batch_columns(sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL), &rows) - .bind(true) - .fetch_all(executor) - .await - { - Ok(results) => { - BATCH_INSERT_FUNCTION.record(BatchFunction::Present); - Ok(results) - } - Err(err) if is_missing_batch_function(&err) => { - BATCH_INSERT_FUNCTION.record(BatchFunction::Absent); - Err(AwaError::Validation( - "insert_many_skip_unique_conflicts needs migration v053; run `awa migrate`" - .to_string(), - )) - } - Err(err) => Err(err.into()), - }; - } - - if BATCH_INSERT_FUNCTION.load() == BatchFunction::Present { - return match bind_batch_columns(sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL), &rows) - .bind(false) - .fetch_all(executor) - .await - { - Ok(results) => Ok(results), - Err(err) => { - // The schema no longer has the function (restored from an - // older backup, or a different database): probe again on the - // next call. This call's executor is spent, and inside a - // transaction it would be aborted anyway. - if is_missing_batch_function(&err) { - BATCH_INSERT_FUNCTION.record(BatchFunction::Unknown); - } - Err(err.into()) - } - }; - } - - // Not known to be present: insert through the per-row function, which - // every schema has, and learn from the same statement whether the batch - // function exists so the next call can use it. - let rows_out = bind_row_function_columns(sqlx::query(INSERT_MANY_ROW_FUNCTION_SQL), &rows) + match bind_batch_columns(sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL), &rows) + .bind(true) .fetch_all(executor) - .await?; - let mut results = Vec::with_capacity(rows_out.len()); - let mut available = false; - for row in &rows_out { - results.push(JobRow::from_row(row)?); - available = row.try_get::("batch_insert_available")?; - } - BATCH_INSERT_FUNCTION.record(if available { - BatchFunction::Present - } else { - BatchFunction::Absent - }); - Ok(results) -} - -/// Whether `awa.insert_jobs_compat_batch` exists in the schema this process -/// writes to, as far as it has seen. -/// -/// The executors the insert functions take are consumed by one statement, so -/// the answer comes from the per-row statement, which asks the catalog as it -/// inserts, and is cached for the process: present stays present until a -/// call finds the function missing, absent is asked again by every per-row -/// call and so is picked up as soon as the migration has run. A process that -/// writes to databases at different migration levels may see one failed call -/// when it switches to one that lacks the function. -struct BatchFunctionState(AtomicU8); - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum BatchFunction { - Unknown, - Present, - Absent, -} - -static BATCH_INSERT_FUNCTION: BatchFunctionState = BatchFunctionState(AtomicU8::new(0)); - -impl BatchFunctionState { - fn load(&self) -> BatchFunction { - match self.0.load(Ordering::Relaxed) { - 1 => BatchFunction::Present, - 2 => BatchFunction::Absent, - _ => BatchFunction::Unknown, - } - } - - fn record(&self, state: BatchFunction) { - self.0.store( - match state { - BatchFunction::Unknown => 0, - BatchFunction::Present => 1, - BatchFunction::Absent => 2, - }, - Ordering::Relaxed, - ); + .await + { + Ok(results) => Ok(results), + Err(err) if is_missing_batch_function(&err) => Err(AwaError::Validation( + "insert_many_skip_unique_conflicts needs migration v053; run `awa migrate`".to_string(), + )), + Err(err) => Err(err.into()), } } @@ -740,9 +646,9 @@ fn bind_batch_columns<'q, O>( } fn bind_row_function_columns<'q>( - query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>, + query: sqlx::query::QueryAs<'q, sqlx::Postgres, JobRow, sqlx::postgres::PgArguments>, rows: &'q [PreparedRow], -) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> { +) -> sqlx::query::QueryAs<'q, sqlx::Postgres, JobRow, sqlx::postgres::PgArguments> { query .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) .bind( @@ -909,21 +815,15 @@ pub async fn insert_many_copy_from_pool( Ok(results) } +/// Probe this connection on every COPY call. A process can write to databases +/// at different migration levels, and the generic executor path cannot safely +/// seed a shared answer for them. async fn batch_function_available(conn: &mut PgConnection) -> Result { - if BATCH_INSERT_FUNCTION.load() == BatchFunction::Present { - return Ok(true); - } - let available: bool = sqlx::query_scalar( + Ok(sqlx::query_scalar( "SELECT to_regprocedure('awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],smallint[],smallint[],timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') IS NOT NULL", ) .fetch_one(&mut *conn) - .await?; - BATCH_INSERT_FUNCTION.record(if available { - BatchFunction::Present - } else { - BatchFunction::Absent - }); - Ok(available) + .await?) } async fn insert_staged_rows_with_batch_function( diff --git a/awa/tests/enqueue_lane_contention_test.rs b/awa/tests/enqueue_lane_contention_test.rs index 6495213a8..cd6e05b84 100644 --- a/awa/tests/enqueue_lane_contention_test.rs +++ b/awa/tests/enqueue_lane_contention_test.rs @@ -667,7 +667,7 @@ fn scenario( ) -> Scenario { Scenario { name: name.to_string(), - mode: Mode::InsertMany, + mode: Mode::InsertManySkip, producers, batch_size: 185, unique, diff --git a/awa/tests/insert_many_batch_test.rs b/awa/tests/insert_many_batch_test.rs index d90c081ce..e3e5884e1 100644 --- a/awa/tests/insert_many_batch_test.rs +++ b/awa/tests/insert_many_batch_test.rs @@ -1,4 +1,4 @@ -//! `insert_many` and `insert_many_copy` on the batch insert function. +//! Compatibility inserts and the opt-in batch insert function. //! //! A batch does its lane-independent work (job ids, unique claims, lane //! registration) before it reserves its lanes, and holds a lane lock only for @@ -341,7 +341,7 @@ async fn batch_waits_for_the_lane_after_its_lane_independent_work() { let producer = { let pool = pool.clone(); let jobs = jobs_for_probe.clone(); - tokio::spawn(async move { insert_many(&pool, &jobs).await }) + tokio::spawn(async move { insert_many_skip_unique_conflicts(&pool, &jobs).await }) }; let started = Instant::now(); @@ -440,7 +440,7 @@ async fn old_and_new_producers_share_a_lane_without_reordering() { setup::activate_queue_storage(&pool).await; let queue = unique_queue("mixed_fleet"); - // Make sure the batch function is the one `insert_many` uses. + // Provision the lane before mixing generic and batch producers. insert_many(&pool, &[job(&queue, -1, false)]) .await .expect("warm up"); @@ -464,7 +464,7 @@ async fn old_and_new_producers_share_a_lane_without_reordering() { if task % 2 == 0 { insert_through_row_function(&pool, &jobs).await; } else { - insert_many(&pool, &jobs) + insert_many_skip_unique_conflicts(&pool, &jobs) .await .expect("batch function insert"); } diff --git a/awa/tests/insert_many_schema_fallback_test.rs b/awa/tests/insert_many_schema_fallback_test.rs index 7c24c954a..7063c3ed7 100644 --- a/awa/tests/insert_many_schema_fallback_test.rs +++ b/awa/tests/insert_many_schema_fallback_test.rs @@ -1,22 +1,18 @@ //! `insert_many` and `insert_many_copy` against a schema without //! `awa.insert_jobs_compat_batch`. //! -//! A binary may run ahead of migration v053. It must keep inserting through -//! the per-row function, start using the batch function once it exists, and -//! recover when the function goes away again. The path a call took is read -//! from `ready_segments`: the per-row function writes one segment per row, the -//! batch function one per lane run. -//! -//! The process-wide answer about the function is part of what is tested, so -//! this binary has one test, against a database of its own. +//! A binary may run ahead of migration v053. Generic `insert_many` keeps using +//! the per-row function; COPY probes its own connection and uses the batch +//! function when installed. Different databases in one process can have +//! different schema versions without aborting their caller's transactions. //! //! Requires a running Postgres instance. //! Run with: `cargo test --package awa --test insert_many_schema_fallback_test` use awa::audited_sql; use awa::model::{ - insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, migrations, - AwaError, + insert_many, insert_many_copy, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, + migrations, AwaError, }; use awa::{InsertOpts, InsertParams, UniqueOpts}; use awa_testing::setup; @@ -145,23 +141,23 @@ async fn binary_ahead_of_the_migration_falls_back_and_recovers() { "{skipped:?}" ); - // The migration runs while the process is up: the next call still takes - // the per-row path, sees the function, and the one after uses it. + // Generic insert_many stays on the per-row path when the migration runs; + // COPY chooses the batch function on its next call. sqlx::raw_sql(MIGRATION_V053) .execute(&pool) .await .expect("apply v053"); insert_many(&pool, &jobs("fallback_b", 0..3, false)) .await - .expect("call that finds the function"); + .expect("generic insert after migration"); assert_eq!(segments(&pool, "fallback_b").await, 3, "per-row path"); insert_many(&pool, &jobs("fallback_b", 3..6, false)) .await - .expect("call through the batch function"); + .expect("second generic insert after migration"); assert_eq!( segments(&pool, "fallback_b").await, - 4, - "the batch function wrote one segment for its run" + 6, + "generic insert_many keeps its schema-independent per-row path" ); assert_eq!( ready_seqs(&pool, "fallback_b").await, @@ -177,24 +173,23 @@ async fn binary_ahead_of_the_migration_falls_back_and_recovers() { .expect("skip with the batch function"); assert_eq!(seqs(&skipped), vec![6, 7]); - // The function goes away under a process that has seen it: one call fails - // with `undefined_function`, and the next probes again and succeeds. + // Removing the function after a successful batch does not poison the next + // generic or COPY call with a stale process-wide answer. sqlx::query(DROP_BATCH_FUNCTION) .execute(&pool) .await .expect("drop batch function again"); - let failed = insert_many(&pool, &jobs("fallback_c", 0..3, false)) + let mut tx = pool.begin().await.expect("caller transaction"); + insert_many(&mut *tx, &jobs("fallback_c", 0..3, false)) .await - .expect_err("the cached function is gone"); - assert!( - matches!(&failed, AwaError::Database(sqlx::Error::Database(e)) if e.code().as_deref() == Some("42883")), - "{failed:?}" - ); - assert!(ready_seqs(&pool, "fallback_c").await.is_empty()); - insert_many(&pool, &jobs("fallback_c", 0..3, false)) + .expect("generic insert still works"); + insert_many_copy(&mut tx, &jobs("fallback_c", 3..6, false)) + .await + .expect("COPY re-probes its own connection"); + tx.commit() .await - .expect("the call after the failure falls back"); - assert_eq!(segments(&pool, "fallback_c").await, 3, "per-row path"); + .expect("caller transaction remains usable"); + assert_eq!(segments(&pool, "fallback_c").await, 6, "per-row paths"); pool.close().await; sqlx::query(audited_sql(format!( @@ -204,3 +199,79 @@ async fn binary_ahead_of_the_migration_falls_back_and_recovers() { .await .expect("drop scratch database"); } + +/// Availability learned from one database must never select SQL for another +/// database. Both paths join the caller's transaction, including its rollback. +#[tokio::test] +async fn mixed_schema_databases_do_not_abort_caller_transactions() { + let current = setup::TestDatabase::canonical().await; + let older = setup::TestDatabase::canonical().await; + for db in [¤t, &older] { + setup::activate_queue_storage(db.pool()).await; + sqlx::query("CREATE TABLE application_writes (id bigint PRIMARY KEY)") + .execute(db.pool()) + .await + .expect("create application table"); + } + sqlx::query(DROP_BATCH_FUNCTION) + .execute(older.pool()) + .await + .expect("older database has no batch function"); + + // This call previously seeded the shared cache with Present. + insert_many_skip_unique_conflicts(current.pool(), &jobs("mixed_warm", 0..2, true)) + .await + .expect("batch API on current schema"); + + for (index, db) in [¤t, &older, ¤t].into_iter().enumerate() { + let pool = db.pool(); + let queue = format!("mixed_database_{index}"); + let mut tx = pool.begin().await.expect("caller transaction"); + sqlx::query("INSERT INTO application_writes (id) VALUES ($1)") + .bind(index as i64) + .execute(&mut *tx) + .await + .expect("application write"); + let inserted = insert_many(&mut *tx, &jobs(&queue, 0..3, true)) + .await + .expect("generic insert must use this database's supported SQL"); + assert_eq!(seqs(&inserted), vec![0, 1, 2]); + let copied = insert_many_copy(&mut tx, &jobs(&queue, 2..5, true)) + .await + .expect("COPY must probe this database"); + assert_eq!(seqs(&copied), vec![3, 4], "duplicate skipped"); + let written: i64 = + sqlx::query_scalar("SELECT count(*) FROM application_writes WHERE id = $1") + .bind(index as i64) + .fetch_one(&mut *tx) + .await + .expect("caller transaction was not aborted"); + assert_eq!(written, 1); + tx.commit() + .await + .expect("commit application and jobs together"); + assert_eq!(ready_seqs(pool, &queue).await, vec![0, 1, 2, 3, 4]); + + let mut tx = pool.begin().await.expect("rollback transaction"); + sqlx::query("INSERT INTO application_writes (id) VALUES ($1)") + .bind(index as i64 + 100) + .execute(&mut *tx) + .await + .expect("application write to roll back"); + insert_many(&mut *tx, &jobs(&queue, 5..7, true)) + .await + .expect("generic insert to roll back"); + insert_many_copy(&mut tx, &jobs(&queue, 7..9, true)) + .await + .expect("COPY to roll back"); + tx.rollback().await.expect("roll back caller transaction"); + assert_eq!(ready_seqs(pool, &queue).await, vec![0, 1, 2, 3, 4]); + let rolled_back: i64 = + sqlx::query_scalar("SELECT count(*) FROM application_writes WHERE id = $1") + .bind(index as i64 + 100) + .fetch_one(pool) + .await + .expect("check application rollback"); + assert_eq!(rolled_back, 0); + } +} diff --git a/docs/architecture.md b/docs/architecture.md index e68e9b14a..79d9940df 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -152,7 +152,9 @@ Enqueue is transactional: if the producer's outer transaction rolls back, the jo There are two COPY-shaped producer paths: - Direct queue-storage COPY is the high-throughput path: Rust producers use a configured `QueueStorage::enqueue_params_copy()`, and Python producers use `enqueue_many_copy()`. Each ready or deferred write group uses array-backed INSERT below 2,048 rows and CSV COPY for larger groups. Direct-copy producers must use the same queue-storage configuration as the worker fleet, especially `queue_stripe_count` / `queue_storage_queue_stripe_count`. -- `insert_many_copy()` is the compatibility path. It stages rows with an array-backed INSERT below 2,048 rows and CSV COPY at or above that cutoff, then inserts them through `awa.insert_jobs_compat_batch` (migration v053; each lane's lock is held only for its reservation and the ready-row write, and unique conflicts are skipped in the same statement). On a schema without the function it takes one of two per-row branches. If any staged row has a `unique_key`, every row is replayed through `POSTGRES_INSERT_JOB_SQL` inside a per-row savepoint so SQLSTATE `23505` can skip that duplicate without aborting the batch; that statement itself calls `awa.insert_job_compat()`. An all-non-unique batch uses one lateral `awa.insert_job_compat()` call per staged row without savepoints. Both branches therefore preserve cross-storage behavior. The planned `awa.insert_job()` public capability becomes the clean stable SQL contract after #342 ships its migration and conformance artifacts. +- `insert_many_copy()` is the compatibility path. It stages rows with an array-backed INSERT below 2,048 rows and CSV COPY at or above that cutoff, then inserts them through `awa.insert_jobs_compat_batch` (migration v053; each lane's lock is held only for its reservation and the ready-row write, and unique conflicts are skipped in the same statement). It probes its own connection on every call; on a schema without the function it takes one of two per-row branches. If any staged row has a `unique_key`, every row is replayed through `POSTGRES_INSERT_JOB_SQL` inside a per-row savepoint so SQLSTATE `23505` can skip that duplicate without aborting the batch; that statement itself calls `awa.insert_job_compat()`. An all-non-unique batch uses one lateral `awa.insert_job_compat()` call per staged row without savepoints. Both branches therefore preserve cross-storage behavior. The planned `awa.insert_job()` public capability becomes the clean stable SQL contract after #342 ships its migration and conformance artifacts. + +Generic `insert_many` always uses one statement that calls `awa.insert_job_compat` per row in canonical lane order. Its `PgExecutor` is consumed by that statement, so it cannot probe the connection first; it does not cache another database's schema availability. `insert_many_skip_unique_conflicts` explicitly requires the batch function. Claim is cursor-based rather than heap-scan based: diff --git a/docs/configuration.md b/docs/configuration.md index 477eb58fb..231537923 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -521,7 +521,7 @@ Direct-copy producers must use the same queue-storage routing config as the work `insert_many_copy()` is the compatibility insert API. It is still useful when a caller needs the canonical insert surface, but in queue-storage mode it routes through the compatibility function rather than being the primary producer fast path. -The Rust compatibility inserts differ in how they treat a unique-key conflict. `insert_many` fails the whole batch with a unique violation (SQLSTATE `23505`). `insert_many_skip_unique_conflicts` inserts the rest of the batch and returns only the rows it inserted, in input order, so a producer that replays work does not have to insert the batch again one job at a time (each of those inserts takes the lane lock separately, which on a busy lane is slower than the batch itself). `insert_many_copy` also skips conflicting jobs, and is what Python's `insert_many_copy()` calls; Python has no separate skip function because that is already its behaviour. `insert_many_skip_unique_conflicts` needs migration v053; the other two use `awa.insert_jobs_compat_batch` when the schema has it and the per-row function when it does not. +The Rust compatibility inserts differ in how they treat a unique-key conflict. `insert_many` fails the whole batch with a unique violation (SQLSTATE `23505`). `insert_many_skip_unique_conflicts` inserts the rest of the batch and returns only the rows it inserted, in input order, so a producer that replays work does not have to insert the batch again one job at a time (each of those inserts takes the lane lock separately, which on a busy lane is slower than the batch itself). `insert_many_copy` also skips conflicting jobs, and is what Python's `insert_many_copy()` calls; Python has no separate skip function because that is already its behaviour. `insert_many_skip_unique_conflicts` needs migration v053. `insert_many_copy` probes its own connection on every call and uses `awa.insert_jobs_compat_batch` when installed, otherwise the per-row function. Generic `insert_many` always uses the per-row function in one statement, so a process can write to databases at different migration levels without assuming they share a schema. ### Terminal count rollup diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index de44cbfab..0c6315572 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -255,8 +255,8 @@ External migration runners apply the complete `awa migrate --sql --from 51 ## v053: batch producer insert -v053 adds `awa.insert_jobs_compat_batch`, the function `insert_many` and -`insert_many_copy` call. It is one `CREATE OR REPLACE FUNCTION`: no table DDL, +v053 adds `awa.insert_jobs_compat_batch`, the function `insert_many_copy` and +`insert_many_skip_unique_conflicts` call. It is one `CREATE OR REPLACE FUNCTION`: no table DDL, no scan, no lock beyond the catalog, so it applies instantly under live load. `awa.insert_job_compat` is unchanged, and binaries before v053 keep calling it; both paths take the same per-lane reservation lock, so producers on either @@ -266,14 +266,7 @@ The runtime role needs `EXECUTE` on the new function; the [default privileges](security/database-roles.md#runtime-grants) cover functions the migrator creates. -A 0.7 binary running before the migration keeps inserting: `insert_many` and -`insert_many_copy` look for the function in the catalog, use the per-row -function until it exists, and switch to the batch function without a restart -once it does. `insert_many_skip_unique_conflicts` needs the function and returns -a validation error before the migration. The per-row and batch paths take the -same lane lock, so a fleet that mixes them keeps each batch contiguous and in -order. External migration runners apply the `awa migrate --sql --from 52 --to -53` bundle. +A 0.7 binary running before the migration keeps inserting. `insert_many_copy` probes its own connection on every call, uses the per-row function until the batch function exists, and switches without a restart. Generic `insert_many` always uses the per-row function in one statement, so one process can write to databases at different migration levels without aborting a caller transaction on an unavailable function. `insert_many_skip_unique_conflicts` needs the function and returns a validation error before the migration. The per-row and batch paths take the same lane lock, so a fleet that mixes them keeps each batch contiguous and in order. External migration runners apply the `awa migrate --sql --from 52 --to 53` bundle. ## The one gate diff --git a/skills/awa-jobs/SKILL.md b/skills/awa-jobs/SKILL.md index 3895bf63a..e28f20704 100644 --- a/skills/awa-jobs/SKILL.md +++ b/skills/awa-jobs/SKILL.md @@ -124,10 +124,7 @@ low-priority work cannot starve indefinitely. **Unique jobs** deduplicate on a BLAKE3 key over kind + optionally queue + args + a time-period bucket. Defaults: dedup by args, not by queue. Conflict handling -depends on the API: a single-row `insert` / `insert_with` and the row-by-row -`insert_many` **raise `AwaError::UniqueConflict`** (Python raises the -corresponding error) on a duplicate — you must handle it, do not assume a silent -no-op. Only the COPY bulk path (`insert_many_copy`) silently skips duplicates. +depends on the API: a single-row `insert` / `insert_with` raises `AwaError::UniqueConflict` (Python raises the corresponding error), and `insert_many` fails the whole batch with SQLSTATE `23505`. `insert_many_copy` skips duplicates on every supported schema; Rust's `insert_many_skip_unique_conflicts` also skips them but requires migration v053. Both return only inserted rows, in input order. Generic `insert_many` always uses one per-row SQL statement; COPY probes its own connection and uses the batch function when installed. Cancel without storing the id via `cancel_by_unique_key(kind, queue=, args=, period_bucket=)` — the components must match those used at insert. It cancels the oldest live job holding the key. On queue storage, a `by_queue` key is From 459f0bcd7d8d8351b96715ec5a2219e189eba58e Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 21:18:39 +1300 Subject: [PATCH 16/18] docs(insert): clarify which producer paths use batch inserts --- awa-model/migrations/v053_insert_jobs_batch_compat.sql | 10 +++++----- docs/troubleshooting.md | 8 ++++---- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/awa-model/migrations/v053_insert_jobs_batch_compat.sql b/awa-model/migrations/v053_insert_jobs_batch_compat.sql index 8463a9ba6..034477373 100644 --- a/awa-model/migrations/v053_insert_jobs_batch_compat.sql +++ b/awa-model/migrations/v053_insert_jobs_batch_compat.sql @@ -1,5 +1,6 @@ -- v053: Batch producer insert that holds each enqueue lane lock only for the -- lane write. +-- Tracking: https://github.com/hardbyte/awa/pull/635 -- -- awa.insert_job_compat takes the (queue, priority, shard) lane lock in -- reserve_enqueue_seq on a batch's first row and, because the lock is @@ -22,11 +23,10 @@ -- N-1 binaries: unchanged. They keep calling awa.insert_job_compat, whose -- definition this migration does not touch, and the two paths share the lane -- lock, so producers on either version serialise exactly as before. --- Binaries that know the function: insert_many and insert_many_copy look for it --- in the catalog and use the per-row function until it exists, so they run --- unchanged before this migration and pick the function up after it, without a --- restart. insert_many_skip_unique_conflicts needs it and reports a schema --- without it. +-- Binaries that know the function: insert_many_copy probes its own connection +-- and uses the per-row function until it exists, picking up the migration +-- without a restart. Generic insert_many always uses the per-row function; +-- insert_many_skip_unique_conflicts requires the batch function. -- External runners: apply this file; nothing else reads it. -- Live load: one CREATE OR REPLACE FUNCTION, no table DDL. diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 67bcc18df..d9e7fc7f2 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -186,7 +186,7 @@ Reading them together separates the two usual failure modes: **1. Producer is going through the compatibility insert path.** -`insert_many_copy_from_pool` / `client.insert_many_copy` routes every row through `awa.insert_job_compat()` once per row. On any database with non-trivial per-statement latency (auth-proxy hop, managed Postgres, network) the per-row function call dominates — measured at ~100–150 ms per row through a Cloud SQL Auth Proxy in staging, which caps a single producer at ~7 rows/sec regardless of batch or chunk size. +Before migration v053, `insert_many_copy_from_pool` / `client.insert_many_copy` calls `awa.insert_job_compat()` per staged row. With unique jobs, per-row savepoints add database round trips that can dominate behind a proxy or on managed PostgreSQL. After migration, COPY probes its own connection and uses the batch function. Generic `insert_many` retains its per-row statement on every supported schema. Switch the producer to the direct queue-storage COPY entry point: @@ -445,12 +445,12 @@ Symptoms: Cause: - on queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until it commits, so concurrent producers for one lane run one at a time; a lane's ceiling is a batch's rows divided by the time the lock is held -- from migration v053, `insert_many` and `insert_many_copy` hold the lock across the lane write only, a few milliseconds for a 185-row batch; before it, the lock was held while every row of the batch was processed +- from migration v053, `insert_many_copy` and `insert_many_skip_unique_conflicts` prepare the batch before taking lane locks; generic `insert_many` continues to hold the lane while subsequent rows are processed - maintenance promotion moves due scheduled and retryable jobs into the lane under the same lock; a promotion batch is one transaction across every queue it touches, and the leader runs its branches one at a time, so a batch parked behind a hot lane delays the due jobs of other queues in that batch and the leader's other branches. Rescue itself does not take the lane lock; it writes retryable rows that promotion moves later Fix: -- run `awa migrate`: producers on the per-row path (a schema before v053) hold the lane for the whole batch, and switch to the batch function as soon as it exists +- run `awa migrate` and use `insert_many_copy` or `insert_many_skip_unique_conflicts` for the batch function; generic `insert_many` stays on its compatible per-row path - when `insert_many` fails on a unique conflict, retry with `insert_many_skip_unique_conflicts` rather than inserting the batch again row by row: each single-row insert queues for the lane separately - raise `awa.queue_meta.enqueue_shards` for a queue whose producers still saturate one lane, accepting partitioned rather than strict FIFO ([ADR-025](adr/025-sharded-enqueue-heads.md)) @@ -459,7 +459,7 @@ Fix: Cause: - on queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until the transaction commits; the `DETAIL` names two 64-bit advisory keys, and two transactions that took lanes in opposite orders hold one each and wait on the other's -- `insert_many`, `insert_many_copy` and the `QueueStorage` batch enqueues take their lane locks in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, so batches never deadlock with each other +- `insert_many`, `insert_many_copy` and the `QueueStorage` batch enqueues take their lane locks in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, avoiding deadlocks caused by opposing lane-lock order; uniqueness and other application locks can still form cycles - a transaction that enqueues with several separate statements (`insert` per job, handler-staged follow-ups, `Transaction.insert` in a loop) takes locks in statement order and can deadlock against a batch or another such transaction that touches the same lanes in a different order Fix: From 934d0f6539123975a56408150c79ba1e948ee1f7 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 00:25:11 +1300 Subject: [PATCH 17/18] Replay statistics migration explicitly in grant regression --- awa/tests/statistics_refresh_test.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/awa/tests/statistics_refresh_test.rs b/awa/tests/statistics_refresh_test.rs index bc9ac8525..44fcca251 100644 --- a/awa/tests/statistics_refresh_test.rs +++ b/awa/tests/statistics_refresh_test.rs @@ -506,10 +506,9 @@ async fn execute_is_limited_to_roles_that_may_maintain_queue_storage() { .expect_err("write privileges alone do not grant EXECUTE"), "permission denied for function", ); - for (_, _, sql) in migrations::migration_sql_range( - migrations::CURRENT_VERSION - 1, - migrations::CURRENT_VERSION, - ) { + // Replay the statistics migration itself: later migrations do not repeat + // its EXECUTE grants for a runtime role created after initial migration. + for (_, _, sql) in migrations::migration_sql_range(53, 54) { run_sql(owner_pool, sql) .await .expect("re-run the migration"); From d504269cbc087ae93cc29d6b8383f04677b3ef82 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 02:53:36 +1300 Subject: [PATCH 18/18] Map all prune transaction lock timeouts to Blocked --- CHANGELOG.md | 20 + awa-model/src/queue_storage.rs | 1387 +++++++++++------------ awa/tests/queue_storage_runtime_test.rs | 175 +++ 3 files changed, 876 insertions(+), 706 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5aa39f220..7c176b4bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -88,6 +88,26 @@ Notable changes between releases. Detailed migration notes for storage transitio lanes as before, and older binaries create correctly owned lanes once v053 has refreshed the helpers they call. See "Lane ownership and v053" in the 0.6 to 0.7 upgrade guide. + +- **Ring prune and terminal-count rollup report `Blocked` when any of their + statements times out on a lock.** Prune sets a short `lock_timeout` (10 ms + by default) before it takes `ACCESS EXCLUSIVE` on a slot's partitions, but + only those `LOCK TABLE` statements and the final `TRUNCATE` turned a + timeout into `PruneOutcome::Blocked`. The timeout covers every later + statement in the transaction too: the post-lock liveness checks read the + lease and claim partitions, the failed-row carry inserts into the current + slot's `done_entries` and terminal-count deltas, and the fold appends + `queue_terminal_rollup_deltas` and deletes `queue_terminal_live_counts` + rows. When one of those waited out the timeout, for example behind + autovacuum truncating a relation or a hosted Postgres's own background + workers, prune failed with a `lock_not_available` (`55P03`) database error, + was logged as an error and skipped the prune backoff. A lock timeout + anywhere in a queue, lease or claim prune, or in a terminal-count rollup + slot, now rolls the whole transaction back and reports `Blocked` + (`blocked_slots` for rollup), so maintenance backs off and retries; nothing + the attempt wrote survives the rollback. A rollup pass stops at its first + blocked slot rather than waiting out the timeout again on each remaining + one. - **Batch enqueues take their lane locks in one order, so they no longer deadlock with each other.** On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index dafbbaa14..3d784fc27 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -713,6 +713,35 @@ async fn set_prune_lock_timeout_tx( Ok(()) } +/// Commit or roll back a maintenance transaction that has set the prune +/// `lock_timeout`. +/// +/// The timeout applies to every later statement in the transaction, not only +/// the explicit `LOCK TABLE`: an insert, update or read of another relation can +/// also queue behind a conflicting lock (autovacuum's heap truncation takes +/// `ACCESS EXCLUSIVE`, for example). A lock timeout anywhere therefore rolls the +/// whole transaction back and reports `blocked`, so the caller backs off and +/// retries instead of surfacing an error. +async fn settle_prune_tx( + tx: sqlx::Transaction<'_, Postgres>, + result: Result, + blocked: T, +) -> Result { + match result { + Ok(outcome) => { + tx.commit().await.map_err(map_sqlx_error)?; + Ok(outcome) + } + Err(err) => { + let _ = tx.rollback().await; + match err { + AwaError::Database(ref db_err) if is_lock_contention_error(db_err) => Ok(blocked), + err => Err(err), + } + } + } +} + fn validate_ident(ident: &str) -> Result<(), AwaError> { let mut chars = ident.chars(); match chars.next() { @@ -18253,7 +18282,10 @@ impl QueueStorage { break; } TerminalDeltaSlotRollup::Blocked => { + // The next slot would most likely wait out the same lock + // timeout, and the leader loop is blocked while it does. outcome.blocked_slots += 1; + break; } } } @@ -18440,237 +18472,215 @@ impl QueueStorage { return Ok(TerminalDeltaSlotRollup::SkippedActive); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - - let lock_delta = sqlx::query(audited_sql(format!( - "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - - match lock_delta { - Ok(_) => {} - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - return Ok(TerminalDeltaSlotRollup::Blocked); - } - Err(err) => { - let _ = tx.rollback().await; - return Err(map_sqlx_error(err)); - } - } + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let active_leases = - queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; + sqlx::query(audited_sql(format!( + "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::SkippedActive); - } + let active_leases = + queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; - let unclosed_claim_refs = - queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; - - if unclosed_claim_refs { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::SkippedActive); - } + if active_leases { + return Ok(TerminalDeltaSlotRollup::SkippedActive); + } - let delta_rows: i64 = sqlx::query_scalar(audited_sql(format!( - r#" - SELECT count(*)::bigint - FROM {delta_child} - WHERE ready_generation = $1 - "# - ))) - .bind(generation) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; + let unclosed_claim_refs = + queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; - if delta_rows == 0 { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::Empty); - } + if unclosed_claim_refs { + return Ok(TerminalDeltaSlotRollup::SkippedActive); + } - // In ledger authority each touched lane collapses to one bucket-0 row - // holding its old total (summed over every bucket) plus the signed - // delta, which keeps a negative delta exact when the matching positive - // count sits in another bucket. Before the flip, N-1 maintenance can - // still win the election and folds per job-id bucket, so the fold - // keeps that shape until the flip fences it (ADR-041). - let grouped_keys: i64 = match authority { - RingAuthority::Ledger => sqlx::query_scalar(audited_sql(format!( + let delta_rows: i64 = sqlx::query_scalar(audited_sql(format!( r#" - WITH grouped AS MATERIALIZED ( - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - SUM(terminal_delta)::bigint AS delta + SELECT count(*)::bigint FROM {delta_child} WHERE ready_generation = $1 - GROUP BY ready_slot, queue, priority, enqueue_shard - HAVING SUM(terminal_delta) <> 0 - ), - collapsed AS ( - DELETE FROM {schema}.queue_terminal_live_counts AS counts - USING grouped - WHERE counts.ready_slot = grouped.ready_slot - AND counts.queue = grouped.queue - AND counts.priority = grouped.priority - AND counts.enqueue_shard = grouped.enqueue_shard - RETURNING - counts.ready_slot, - counts.queue, - counts.priority, - counts.enqueue_shard, - counts.live_terminal_count - ), - folded AS ( - SELECT - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - (grouped.delta + COALESCE(SUM(collapsed.live_terminal_count), 0))::bigint - AS live_terminal_count - FROM grouped - LEFT JOIN collapsed - ON collapsed.ready_slot = grouped.ready_slot - AND collapsed.queue = grouped.queue - AND collapsed.priority = grouped.priority - AND collapsed.enqueue_shard = grouped.enqueue_shard - GROUP BY - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - grouped.delta - ), - inserted AS ( - INSERT INTO {schema}.queue_terminal_live_counts ( - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - live_terminal_count - ) - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - 0::smallint, - live_terminal_count - FROM folded - WHERE live_terminal_count <> 0 - ORDER BY ready_slot, queue, priority, enqueue_shard - RETURNING 1 - ) - SELECT count(*)::bigint FROM grouped - "# - ))) - .bind(generation) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?, - RingAuthority::Columns => sqlx::query_scalar(audited_sql(format!( - r#" - WITH grouped AS MATERIALIZED ( - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - SUM(terminal_delta)::bigint AS delta - FROM {delta_child} - WHERE ready_generation = $1 - GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket - HAVING SUM(terminal_delta) <> 0 - ), - updated AS ( - UPDATE {schema}.queue_terminal_live_counts AS counts - SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta) - FROM grouped - WHERE counts.ready_slot = grouped.ready_slot - AND counts.queue = grouped.queue - AND counts.priority = grouped.priority - AND counts.enqueue_shard = grouped.enqueue_shard - AND counts.counter_bucket = grouped.counter_bucket - RETURNING - counts.ready_slot, - counts.queue, - counts.priority, - counts.enqueue_shard, - counts.counter_bucket - ), - inserted AS ( - INSERT INTO {schema}.queue_terminal_live_counts AS counts ( - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - live_terminal_count - ) - SELECT - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - grouped.counter_bucket, - grouped.delta - FROM grouped - WHERE grouped.delta > 0 - AND NOT EXISTS ( - SELECT 1 - FROM updated - WHERE updated.ready_slot = grouped.ready_slot - AND updated.queue = grouped.queue - AND updated.priority = grouped.priority - AND updated.enqueue_shard = grouped.enqueue_shard - AND updated.counter_bucket = grouped.counter_bucket - ) - ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket - ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE - SET live_terminal_count = - GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count) - RETURNING 1 - ) - SELECT count(*)::bigint FROM grouped - "# + "# ))) .bind(generation) .fetch_one(tx.as_mut()) .await - .map_err(map_sqlx_error)?, - }; - - let truncate_delta = sqlx::query(audited_sql(format!("TRUNCATE TABLE {delta_child}"))) - .execute(tx.as_mut()) - .await; + .map_err(map_sqlx_error)?; - match truncate_delta { - Ok(_) => { - tx.commit().await.map_err(map_sqlx_error)?; - Ok(TerminalDeltaSlotRollup::Rolled { - delta_rows, - grouped_keys, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(TerminalDeltaSlotRollup::Blocked) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) + if delta_rows == 0 { + return Ok(TerminalDeltaSlotRollup::Empty); } + + // In ledger authority each touched lane collapses to one bucket-0 row + // holding its old total (summed over every bucket) plus the signed + // delta, which keeps a negative delta exact when the matching positive + // count sits in another bucket. Before the flip, N-1 maintenance can + // still win the election and folds per job-id bucket, so the fold + // keeps that shape until the flip fences it (ADR-041). + let grouped_keys: i64 = match authority { + RingAuthority::Ledger => sqlx::query_scalar(audited_sql(format!( + r#" + WITH grouped AS MATERIALIZED ( + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + SUM(terminal_delta)::bigint AS delta + FROM {delta_child} + WHERE ready_generation = $1 + GROUP BY ready_slot, queue, priority, enqueue_shard + HAVING SUM(terminal_delta) <> 0 + ), + collapsed AS ( + DELETE FROM {schema}.queue_terminal_live_counts AS counts + USING grouped + WHERE counts.ready_slot = grouped.ready_slot + AND counts.queue = grouped.queue + AND counts.priority = grouped.priority + AND counts.enqueue_shard = grouped.enqueue_shard + RETURNING + counts.ready_slot, + counts.queue, + counts.priority, + counts.enqueue_shard, + counts.live_terminal_count + ), + folded AS ( + SELECT + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + (grouped.delta + COALESCE(SUM(collapsed.live_terminal_count), 0))::bigint + AS live_terminal_count + FROM grouped + LEFT JOIN collapsed + ON collapsed.ready_slot = grouped.ready_slot + AND collapsed.queue = grouped.queue + AND collapsed.priority = grouped.priority + AND collapsed.enqueue_shard = grouped.enqueue_shard + GROUP BY + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + grouped.delta + ), + inserted AS ( + INSERT INTO {schema}.queue_terminal_live_counts ( + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + live_terminal_count + ) + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + 0::smallint, + live_terminal_count + FROM folded + WHERE live_terminal_count <> 0 + ORDER BY ready_slot, queue, priority, enqueue_shard + RETURNING 1 + ) + SELECT count(*)::bigint FROM grouped + "# + ))) + .bind(generation) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?, + RingAuthority::Columns => sqlx::query_scalar(audited_sql(format!( + r#" + WITH grouped AS MATERIALIZED ( + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + SUM(terminal_delta)::bigint AS delta + FROM {delta_child} + WHERE ready_generation = $1 + GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket + HAVING SUM(terminal_delta) <> 0 + ), + updated AS ( + UPDATE {schema}.queue_terminal_live_counts AS counts + SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta) + FROM grouped + WHERE counts.ready_slot = grouped.ready_slot + AND counts.queue = grouped.queue + AND counts.priority = grouped.priority + AND counts.enqueue_shard = grouped.enqueue_shard + AND counts.counter_bucket = grouped.counter_bucket + RETURNING + counts.ready_slot, + counts.queue, + counts.priority, + counts.enqueue_shard, + counts.counter_bucket + ), + inserted AS ( + INSERT INTO {schema}.queue_terminal_live_counts AS counts ( + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + live_terminal_count + ) + SELECT + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + grouped.counter_bucket, + grouped.delta + FROM grouped + WHERE grouped.delta > 0 + AND NOT EXISTS ( + SELECT 1 + FROM updated + WHERE updated.ready_slot = grouped.ready_slot + AND updated.queue = grouped.queue + AND updated.priority = grouped.priority + AND updated.enqueue_shard = grouped.enqueue_shard + AND updated.counter_bucket = grouped.counter_bucket + ) + ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket + ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE + SET live_terminal_count = + GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count) + RETURNING 1 + ) + SELECT count(*)::bigint FROM grouped + "# + ))) + .bind(generation) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?, + }; + + sqlx::query(audited_sql(format!("TRUNCATE TABLE {delta_child}"))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(TerminalDeltaSlotRollup::Rolled { + delta_rows, + grouped_keys, + }) } + .await; + + settle_prune_tx(tx, locked, TerminalDeltaSlotRollup::Blocked).await } /// Fold pending `queue_terminal_rollup_deltas` rows into the @@ -18920,172 +18930,119 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let lock_started = Instant::now(); - // Claims read segments before ready rows, while producers write in - // the opposite order. Queue behind segment readers, then abandon - // the transaction rather than wait with an exclusive child held. - let lock_tables = async { - sqlx::query(audited_sql(format!( - "LOCK TABLE {segment_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await?; - sqlx::query(audited_sql(format!( - "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE NOWAIT" - ))) - .execute(tx.as_mut()) + let lock_started = Instant::now(); + // Claims read segments before ready rows, while producers write in + // the opposite order. Queue behind segment readers, then abandon + // the transaction rather than wait with an exclusive child held. + async { + sqlx::query(audited_sql(format!( + "LOCK TABLE {segment_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await?; + sqlx::query(audited_sql(format!( + "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE NOWAIT" + ))) + .execute(tx.as_mut()) + .await + } .await - } - .await; - let lock_duration = lock_started.elapsed(); - - if let Err(err) = lock_tables { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); + + // The pre-lock gates are only an MVCC snapshot. If an in-flight + // producer or claimer committed while the exclusive partition + // lock waited, these post-lock gates see it before TRUNCATE. + let active_leases = + queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; + if active_leases { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueueActiveLeases, + count: busy_indicator(active_leases), + }); } - return Err(map_sqlx_error(err)); - } - - // The pre-lock gates are only an MVCC snapshot. If an in-flight - // producer or claimer committed while the exclusive partition - // lock waited, these post-lock gates see it before TRUNCATE. - let active_leases = - queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueueActiveLeases, - count: busy_indicator(active_leases), - }); - } - let pending = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; - if pending { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueuePendingReady, - count: busy_indicator(pending), - }); - } + let pending = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; + if pending { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueuePendingReady, + count: busy_indicator(pending), + }); + } - let unclosed_claim_refs = - queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; - if unclosed_claim_refs { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueueUnclosedClaimRefs, - count: busy_indicator(unclosed_claim_refs), - }); - } + let unclosed_claim_refs = + queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; + if unclosed_claim_refs { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueueUnclosedClaimRefs, + count: busy_indicator(unclosed_claim_refs), + }); + } - let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX); - let retention_floor = failed_retention_secs > 0; - - // #337: failed terminal rows still inside the retention floor - // are re-homed into the live done segment before the TRUNCATE - // so they stay retryable. The ring rotation lock is held above, - // which serializes against rotate — the current slot (and its - // generation, read from the ledger cursor) cannot move while - // carried rows are re-inserted into it. - let carried_failed_rows = if retention_floor { - // Carried rows are widened to the self-contained synthetic - // done-row shape (the COALESCE chain mirrors - // `done_row_projection`) because the ready rows they would - // otherwise hydrate from are truncated with this slot. - // Deliberately no ON CONFLICT: lane_seq comes from - // never-resetting sequences, so a PK collision means - // corrupted terminal state and must abort the prune loudly - // rather than silently drop a terminal fact. - let carried = sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.done_entries ( - ready_slot, ready_generation, job_id, kind, queue, - args, state, priority, attempt, run_lease, - max_attempts, lane_seq, enqueue_shard, run_at, - attempted_at, finalized_at, created_at, unique_key, - unique_states, payload - ) - SELECT - $2, - $3, - done.job_id, - done.kind, - done.queue, - COALESCE(done.args, ready.args, '{{}}'::jsonb), - done.state, - done.priority, - done.attempt, - done.run_lease, - COALESCE(done.max_attempts, ready.max_attempts, 25::smallint), - done.lane_seq, - done.enqueue_shard, - COALESCE(done.run_at, ready.run_at, done.finalized_at), - COALESCE(done.attempted_at, ready.attempted_at), - done.finalized_at, - COALESCE(done.created_at, ready.created_at, done.finalized_at), - COALESCE(done.unique_key, ready.unique_key), - COALESCE(done.unique_states, ready.unique_states), - COALESCE(done.payload, ready.payload, '{{}}'::jsonb) - FROM {done_child} AS done - LEFT JOIN {ready_child} AS ready - ON ready.ready_slot = done.ready_slot - AND ready.ready_generation = done.ready_generation - AND ready.queue = done.queue - AND ready.priority = done.priority - AND ready.enqueue_shard = done.enqueue_shard - AND ready.lane_seq = done.lane_seq - WHERE done.state = 'failed' - AND done.finalized_at >= now() - make_interval(secs => $1::bigint) - "# - ))) - .bind(failed_retention_secs) - .bind(state.0) - .bind(current_generation) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)? - .rows_affected(); - - if carried > 0 { - // This slot's live counter rows and delta segment are - // wiped below, so the carried rows' exact-count - // evidence moves with them: re-append positive deltas - // under the current slot, grouped exactly like the - // completion path's delta append. - sqlx::query(audited_sql(format!( + let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX); + let retention_floor = failed_retention_secs > 0; + + // #337: failed terminal rows still inside the retention floor + // are re-homed into the live done segment before the TRUNCATE + // so they stay retryable. The ring rotation lock is held above, + // which serializes against rotate — the current slot (and its + // generation, read from the ledger cursor) cannot move while + // carried rows are re-inserted into it. + let carried_failed_rows = if retention_floor { + // Carried rows are widened to the self-contained synthetic + // done-row shape (the COALESCE chain mirrors + // `done_row_projection`) because the ready rows they would + // otherwise hydrate from are truncated with this slot. + // Deliberately no ON CONFLICT: lane_seq comes from + // never-resetting sequences, so a PK collision means + // corrupted terminal state and must abort the prune loudly + // rather than silently drop a terminal fact. + let carried = sqlx::query(audited_sql(format!( r#" - INSERT INTO {schema}.queue_terminal_count_deltas ( - ready_slot, - ready_generation, - queue, - priority, - enqueue_shard, - counter_bucket, - terminal_delta + INSERT INTO {schema}.done_entries ( + ready_slot, ready_generation, job_id, kind, queue, + args, state, priority, attempt, run_lease, + max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, finalized_at, created_at, unique_key, + unique_states, payload ) SELECT $2, $3, - queue, - priority, - enqueue_shard, - mod( - mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint) - + {TERMINAL_COUNTER_BUCKETS}::bigint, - {TERMINAL_COUNTER_BUCKETS}::bigint - )::smallint AS counter_bucket, - count(*)::bigint - FROM {done_child} - WHERE state = 'failed' - AND finalized_at >= now() - make_interval(secs => $1::bigint) - GROUP BY queue, priority, enqueue_shard, counter_bucket - ORDER BY queue, priority, enqueue_shard, counter_bucket + done.job_id, + done.kind, + done.queue, + COALESCE(done.args, ready.args, '{{}}'::jsonb), + done.state, + done.priority, + done.attempt, + done.run_lease, + COALESCE(done.max_attempts, ready.max_attempts, 25::smallint), + done.lane_seq, + done.enqueue_shard, + COALESCE(done.run_at, ready.run_at, done.finalized_at), + COALESCE(done.attempted_at, ready.attempted_at), + done.finalized_at, + COALESCE(done.created_at, ready.created_at, done.finalized_at), + COALESCE(done.unique_key, ready.unique_key), + COALESCE(done.unique_states, ready.unique_states), + COALESCE(done.payload, ready.payload, '{{}}'::jsonb) + FROM {done_child} AS done + LEFT JOIN {ready_child} AS ready + ON ready.ready_slot = done.ready_slot + AND ready.ready_generation = done.ready_generation + AND ready.queue = done.queue + AND ready.priority = done.priority + AND ready.enqueue_shard = done.enqueue_shard + AND ready.lane_seq = done.lane_seq + WHERE done.state = 'failed' + AND done.finalized_at >= now() - make_interval(secs => $1::bigint) "# ))) .bind(failed_retention_secs) @@ -19093,187 +19050,232 @@ impl QueueStorage { .bind(current_generation) .execute(tx.as_mut()) .await - .map_err(map_sqlx_error)?; - } - carried - } else { - 0 - }; - - // #290: scan the about-to-be-truncated partition for the rollup - // fold. The rollup column is *permanent* state, so we MUST fold - // from ground truth (the `{done_child}` partition itself), not - // from `queue_terminal_live_counts` — the counter can drift - // briefly during a rolling upgrade from a pre-counter binary - // and folding drift into the rollup would bake it in forever. - // The counter rows for this slot are still cleaned up after the - // truncate; reads from the counter (queue_counts_exact in #305) - // may transiently disagree with the rollup until the operator - // runs `awa storage rebuild-terminal-counters`, but the rollup - // itself stays authoritative. See PR #304 reviewer finding - // "High: A1 can persist stale counter state before the - // read-switch PR" for the trade-off. - // - // Carried failed rows are excluded from both columns: they are - // still live in `done_entries`, so folding them into the - // permanent rollup would double-count them. - let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = - sqlx::query_as(audited_sql(format!( - r#" - WITH done_counts AS ( - SELECT - queue, - priority, - (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed, - (count(*) FILTER ( + .map_err(map_sqlx_error)? + .rows_affected(); + + if carried > 0 { + // This slot's live counter rows and delta segment are + // wiped below, so the carried rows' exact-count + // evidence moves with them: re-append positive deltas + // under the current slot, grouped exactly like the + // completion path's delta append. + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.queue_terminal_count_deltas ( + ready_slot, + ready_generation, + queue, + priority, + enqueue_shard, + counter_bucket, + terminal_delta + ) + SELECT + $2, + $3, + queue, + priority, + enqueue_shard, + mod( + mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint) + + {TERMINAL_COUNTER_BUCKETS}::bigint, + {TERMINAL_COUNTER_BUCKETS}::bigint + )::smallint AS counter_bucket, + count(*)::bigint + FROM {done_child} WHERE state = 'failed' - AND ( - $2::boolean IS FALSE - OR finalized_at < now() - make_interval(secs => $1::bigint) - ) - ))::bigint AS failed - FROM {done_child} - GROUP BY queue, priority - ), - -- Scan the entire partition being truncated (all generations), - -- mirroring done_counts' whole-child scan above. The rotate guard - -- keeps a partition to one generation at prune time, so this is the - -- same set today; counting the whole child keeps the rollup fold - -- conservative even if that invariant is ever weakened (a stray - -- generation's batches would otherwise be truncated without being - -- folded into the rollup -> terminal-count undercount). - batch_rows AS ( - SELECT - batch.ready_generation, - batch.queue, - batch.priority, - item.job_id, - item.run_lease - FROM {receipt_batch_child} AS batch - CROSS JOIN LATERAL unnest( - batch.job_ids, - batch.run_leases - ) AS item(job_id, run_lease) - ), - batch_counts AS ( + AND finalized_at >= now() - make_interval(secs => $1::bigint) + GROUP BY queue, priority, enqueue_shard, counter_bucket + ORDER BY queue, priority, enqueue_shard, counter_bucket + "# + ))) + .bind(failed_retention_secs) + .bind(state.0) + .bind(current_generation) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + } + carried + } else { + 0 + }; + + // #290: scan the about-to-be-truncated partition for the rollup + // fold. The rollup column is *permanent* state, so we MUST fold + // from ground truth (the `{done_child}` partition itself), not + // from `queue_terminal_live_counts` — the counter can drift + // briefly during a rolling upgrade from a pre-counter binary + // and folding drift into the rollup would bake it in forever. + // The counter rows for this slot are still cleaned up after the + // truncate; reads from the counter (queue_counts_exact in #305) + // may transiently disagree with the rollup until the operator + // runs `awa storage rebuild-terminal-counters`, but the rollup + // itself stays authoritative. See PR #304 reviewer finding + // "High: A1 can persist stale counter state before the + // read-switch PR" for the trade-off. + // + // Carried failed rows are excluded from both columns: they are + // still live in `done_entries`, so folding them into the + // permanent rollup would double-count them. + let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = + sqlx::query_as(audited_sql(format!( + r#" + WITH done_counts AS ( + SELECT + queue, + priority, + (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed, + (count(*) FILTER ( + WHERE state = 'failed' + AND ( + $2::boolean IS FALSE + OR finalized_at < now() - make_interval(secs => $1::bigint) + ) + ))::bigint AS failed + FROM {done_child} + GROUP BY queue, priority + ), + -- Scan the entire partition being truncated (all generations), + -- mirroring done_counts' whole-child scan above. The rotate guard + -- keeps a partition to one generation at prune time, so this is the + -- same set today; counting the whole child keeps the rollup fold + -- conservative even if that invariant is ever weakened (a stray + -- generation's batches would otherwise be truncated without being + -- folded into the rollup -> terminal-count undercount). + batch_rows AS ( + SELECT + batch.ready_generation, + batch.queue, + batch.priority, + item.job_id, + item.run_lease + FROM {receipt_batch_child} AS batch + CROSS JOIN LATERAL unnest( + batch.job_ids, + batch.run_leases + ) AS item(job_id, run_lease) + ), + batch_counts AS ( + SELECT + batch_rows.queue, + batch_rows.priority, + count(*)::bigint AS completed + FROM batch_rows + WHERE NOT EXISTS ( + SELECT 1 + FROM {receipt_tomb_child} AS tomb + WHERE tomb.ready_generation = batch_rows.ready_generation + AND tomb.job_id = batch_rows.job_id + AND tomb.run_lease = batch_rows.run_lease + ) + GROUP BY batch_rows.queue, batch_rows.priority + ), + keys AS ( + SELECT queue, priority FROM done_counts + UNION + SELECT queue, priority FROM batch_counts + ) SELECT - batch_rows.queue, - batch_rows.priority, - count(*)::bigint AS completed - FROM batch_rows - WHERE NOT EXISTS ( - SELECT 1 - FROM {receipt_tomb_child} AS tomb - WHERE tomb.ready_generation = batch_rows.ready_generation - AND tomb.job_id = batch_rows.job_id - AND tomb.run_lease = batch_rows.run_lease + keys.queue, + keys.priority, + ( + COALESCE(done_counts.completed, 0) + + COALESCE(batch_counts.completed, 0) + )::bigint AS pruned_completed_count, + COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count + FROM keys + LEFT JOIN done_counts USING (queue, priority) + LEFT JOIN batch_counts USING (queue, priority) + "# + ))) + .bind(failed_retention_secs) + .bind(retention_floor) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + + let truncate_started = Instant::now(); + self + .truncate_used_children_tx( + &mut tx, + &[ + &ready_child, + &claim_attempt_child, + &done_child, + &tomb_child, + &segment_child, + &receipt_batch_child, + &receipt_tomb_child, + &delta_child, + ], ) - GROUP BY batch_rows.queue, batch_rows.priority - ), - keys AS ( - SELECT queue, priority FROM done_counts - UNION - SELECT queue, priority FROM batch_counts - ) - SELECT - keys.queue, - keys.priority, - ( - COALESCE(done_counts.completed, 0) - + COALESCE(batch_counts.completed, 0) - )::bigint AS pruned_completed_count, - COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count - FROM keys - LEFT JOIN done_counts USING (queue, priority) - LEFT JOIN batch_counts USING (queue, priority) - "# + .await + .map_err(map_sqlx_error)?; + let truncate_duration = truncate_started.elapsed(); + if !pruned_terminal_counts.is_empty() { + // #371: append the pruned counts as delta rows instead + // of upserting `queue_terminal_rollups` here — the + // upsert under a pinned MVCC horizon left one + // unreclaimable dead rollup version per prune. The + // maintenance leader folds the deltas + // (`fold_terminal_rollup_deltas`) when the horizon is + // clear; exact readers add unfolded delta sums to + // their rollup reads. + self.append_terminal_rollup_deltas_tx( + &mut tx, + pruned_terminal_counts.into_iter(), + ) + .await?; + } + + // #290: the live counter rows for this slot are about to + // be orphans (their underlying done_entries rows just got + // truncated). Delete them in the same transaction. The + // rollup fold above already captured ground-truth from + // the partition scan; this just cleans up the counter + // index entries so a future insert into a re-rotated + // slot starts from zero. + sqlx::query(audited_sql(format!( + "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" ))) - .bind(failed_retention_secs) - .bind(retention_floor) - .fetch_all(tx.as_mut()) + .bind(slot) + .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { + slot, + generation, + carried_failed_rows, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) + } + .await; - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx( - &mut tx, - &[ - &ready_child, - &claim_attempt_child, - &done_child, - &tomb_child, - &segment_child, - &receipt_batch_child, - &receipt_tomb_child, - &delta_child, - ], - ) - .await; - let truncate_duration = truncate_started.elapsed(); - - match truncate { - Ok(_) => { - if !pruned_terminal_counts.is_empty() { - // #371: append the pruned counts as delta rows instead - // of upserting `queue_terminal_rollups` here — the - // upsert under a pinned MVCC horizon left one - // unreclaimable dead rollup version per prune. The - // maintenance leader folds the deltas - // (`fold_terminal_rollup_deltas`) when the horizon is - // clear; exact readers add unfolded delta sums to - // their rollup reads. - self.append_terminal_rollup_deltas_tx( - &mut tx, - pruned_terminal_counts.into_iter(), - ) - .await?; - } - - // #290: the live counter rows for this slot are about to - // be orphans (their underlying done_entries rows just got - // truncated). Delete them in the same transaction. The - // rollup fold above already captured ground-truth from - // the partition scan; this just cleans up the counter - // index entries so a future insert into a re-rotated - // slot starts from zero. - sqlx::query(audited_sql(format!( - "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" - ))) - .bind(slot) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - if carried_failed_rows > 0 { - tracing::info!( - slot, - carried_failed_rows, - "Carried failed terminal rows inside the retention floor to the live done segment" - ); - } - Ok(PruneOutcome::Pruned { + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { + carried_failed_rows, + durations, + .. + } = &mut outcome + { + durations.commit = commit_started.elapsed(); + let carried_failed_rows = *carried_failed_rows; + if carried_failed_rows > 0 { + tracing::info!( slot, - generation, carried_failed_rows, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) + "Carried failed terminal rows inside the retention floor to the live done segment" + ); } } + Ok(outcome) } #[tracing::instrument( @@ -19349,82 +19351,68 @@ impl QueueStorage { let lease_child = lease_child_name(schema, slot as usize); - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let lock_started = Instant::now(); - let lock_table = sqlx::query(audited_sql(format!( - "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - let lock_duration = lock_started.elapsed(); + let lock_started = Instant::now(); + sqlx::query(audited_sql(format!( + "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); - if let Err(err) = lock_table { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); + // The ring rotation serialization held above keeps the cursor + // stable, but keep the explicit recheck to document the truncate + // precondition. + let (current_slot, _) = self + .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) + .await?; + + if current_slot == slot { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::LeaseCurrent, + count: 0, + }); } - return Err(map_sqlx_error(err)); - } - // The ring rotation serialization held above keeps the cursor - // stable, but keep the explicit recheck to document the truncate - // precondition. - let (current_slot, _) = self - .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) - .await?; + let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?; - if current_slot == slot { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::LeaseCurrent, - count: 0, - }); - } - - let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?; + if active_leases { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::LeaseActive, + count: busy_indicator(active_leases), + }); + } - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { + let truncate_started = Instant::now(); + self.truncate_used_children_tx(&mut tx, &[&lease_child]) + .await + .map_err(map_sqlx_error)?; + let truncate_duration = truncate_started.elapsed(); + Ok(PruneOutcome::Pruned { slot, - reason: SkipReason::LeaseActive, - count: busy_indicator(active_leases), - }); + generation, + carried_failed_rows: 0, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) } + .await; - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx(&mut tx, &[&lease_child]) - .await; - let truncate_duration = truncate_started.elapsed(); - - match truncate { - Ok(_) => { - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - Ok(PruneOutcome::Pruned { - slot, - generation, - carried_failed_rows: 0, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) - } + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { durations, .. } = &mut outcome { + durations.commit = commit_started.elapsed(); } + Ok(outcome) } pub async fn vacuum_leases(&self, pool: &PgPool) -> Result<(), AwaError> { @@ -19728,116 +19716,103 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let lock_started = Instant::now(); - let lock_tables = sqlx::query(audited_sql(format!( - "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - let lock_duration = lock_started.elapsed(); - - if let Err(err) = lock_tables { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); - } - return Err(map_sqlx_error(err)); - } - - // After taking ACCESS EXCLUSIVE, recheck that the slot is still - // not the current one. The ring rotation serialization should - // already make this stable, but keeping the explicit gate documents - // the truncate precondition. - let (current_slot, _) = self - .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) - .await?; + let lock_started = Instant::now(); + sqlx::query(audited_sql(format!( + "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); - if current_slot == slot { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::ClaimCurrent, - count: 0, - }); - } + // After taking ACCESS EXCLUSIVE, recheck that the slot is still + // not the current one. The ring rotation serialization should + // already make this stable, but keeping the explicit gate documents + // the truncate precondition. + let (current_slot, _) = self + .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) + .await?; - let open_claims = claim_prune_has_open_claims_tx( - &mut tx, - schema, - &claim_child, - &claim_batch_child, - &closure_child, - &closure_batch_child, - ) - .await?; - if open_claims { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::ClaimOpen, - count: busy_indicator(open_claims), - }); - } + if current_slot == slot { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::ClaimCurrent, + count: 0, + }); + } - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx( + let open_claims = claim_prune_has_open_claims_tx( &mut tx, - &[ - &claim_child, - &claim_batch_child, - &closure_child, - &closure_batch_child, - ], + schema, + &claim_child, + &claim_batch_child, + &closure_child, + &closure_batch_child, ) - .await; - let truncate_duration = truncate_started.elapsed(); + .await?; + if open_claims { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::ClaimOpen, + count: busy_indicator(open_claims), + }); + } - match truncate { - Ok(_) => { - sqlx::query(audited_sql(format!( - r#" - UPDATE {schema}.claim_ring_slots - SET rescue_cursor_claimed_at = '-infinity'::timestamptz, - rescue_cursor_job_id = 0, - rescue_cursor_run_lease = 0, - deadline_cursor_deadline_at = '-infinity'::timestamptz, - deadline_cursor_job_id = 0, - deadline_cursor_run_lease = 0, - batch_deadline_cursor_deadline_at = '-infinity'::timestamptz, - batch_deadline_cursor_batch_id = 0 - WHERE slot = $1 - "# - ))) - .bind(slot) - .execute(tx.as_mut()) + let truncate_started = Instant::now(); + self + .truncate_used_children_tx( + &mut tx, + &[ + &claim_child, + &claim_batch_child, + &closure_child, + &closure_batch_child, + ], + ) .await .map_err(map_sqlx_error)?; - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - Ok(PruneOutcome::Pruned { - slot, - generation, - carried_failed_rows: 0, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) - } + let truncate_duration = truncate_started.elapsed(); + sqlx::query(audited_sql(format!( + r#" + UPDATE {schema}.claim_ring_slots + SET rescue_cursor_claimed_at = '-infinity'::timestamptz, + rescue_cursor_job_id = 0, + rescue_cursor_run_lease = 0, + deadline_cursor_deadline_at = '-infinity'::timestamptz, + deadline_cursor_job_id = 0, + deadline_cursor_run_lease = 0, + batch_deadline_cursor_deadline_at = '-infinity'::timestamptz, + batch_deadline_cursor_batch_id = 0 + WHERE slot = $1 + "# + ))) + .bind(slot) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { + slot, + generation, + carried_failed_rows: 0, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) } + .await; + + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { durations, .. } = &mut outcome { + durations.commit = commit_started.elapsed(); + } + Ok(outcome) } fn job_id_sequence(&self) -> String { diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index bf51546fe..3b57f3370 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -15632,6 +15632,181 @@ async fn test_queue_storage_prune_oldest_blocks_on_reader_lock() { ); } +/// The prune `lock_timeout` covers every statement after it is set, not +/// only the `LOCK TABLE` on the slot's children. A statement that waits on +/// another table's lock — here the failed-row carry, its count deltas, the +/// rollup delta append and the live-counter cleanup — must roll the whole prune back +/// and report `Blocked`, leaving the slot to the next attempt. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_prune_lock_timeout_after_child_locks_reports_blocked() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_prune_late_lock_timeout"; + let schema = "awa_qs_prune_late_lock_timeout"; + let store = create_store(&pool, schema).await; + + store + .enqueue_batch(&pool, queue, 1, 1) + .await + .expect("enqueue completed job"); + let claimed = store + .claim_batch(&pool, queue, 1) + .await + .expect("claim completed job"); + store + .complete_batch(&pool, &claimed) + .await + .expect("complete job"); + + let failed_job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim failing job") + .into_iter() + .next() + .expect("missing claimed job"); + store + .fail_terminal(&pool, failed_job_id, claimed.job.run_lease, "boom", None) + .await + .expect("fail job") + .expect("running job should fail"); + + let rotated = store.rotate(&pool).await.expect("rotate"); + assert!( + matches!(rotated, RotateOutcome::Rotated { slot: 1, .. }), + "unexpected rotate outcome: {rotated:?}" + ); + + let retention = Duration::from_secs(3600); + for held in [ + format!("{schema}.done_entries_1"), + format!("{schema}.queue_terminal_count_deltas_1"), + format!("{schema}.queue_terminal_rollup_deltas"), + format!("{schema}.queue_terminal_live_counts"), + ] { + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(audited_sql(format!("LOCK TABLE {held} IN SHARE MODE"))) + .execute(holder.as_mut()) + .await + .expect("hold share lock"); + + let outcome = store.prune_oldest(&pool, retention).await; + assert!( + matches!(outcome, Ok(PruneOutcome::Blocked { slot: 0 })), + "prune waiting on {held} should report Blocked, got {outcome:?}" + ); + + holder.rollback().await.expect("release share lock"); + + assert_eq!( + done_entries_count(&pool, schema, queue).await, + 2, + "blocked prune on {held} must not truncate or carry rows" + ); + let rollup_rows: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*)::bigint FROM {schema}.queue_terminal_rollup_deltas" + ))) + .fetch_one(&pool) + .await + .expect("count rollups"); + assert_eq!( + rollup_rows, 0, + "blocked prune on {held} must not fold rollups" + ); + assert_invariant_holds( + &pool, + schema, + queue, + &format!("after prune blocked on {held}"), + ) + .await; + } + + let pruned = store + .prune_oldest(&pool, retention) + .await + .expect("prune after locks released"); + assert!( + matches!( + pruned, + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } + ), + "unexpected prune outcome after locks released: {pruned:?}" + ); + assert_invariant_holds(&pool, schema, queue, "after prune").await; + let counts = store.queue_counts(&pool, queue).await.expect("counts"); + assert_eq!(counts.terminal, 2, "rollup (1) + carried live row (1)"); +} + +/// Terminal-delta rollup shares prune's lock timeout. A timeout on the +/// live-counter fold, after the delta child is locked, must report a blocked +/// slot rather than fail the rollup pass, and the pass stops there instead of +/// waiting out the same timeout on every remaining slot. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_terminal_count_delta_rollup_lock_timeout_after_child_lock_reports_blocked() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_terminal_delta_rollup_late_lock"; + let schema = "awa_qs_terminal_delta_rollup_late_lock"; + let store = create_store(&pool, schema).await; + + for _ in 0..2 { + store + .enqueue_batch(&pool, queue, 1, 3) + .await + .expect("enqueue"); + let claimed = store.claim_batch(&pool, queue, 3).await.expect("claim"); + store + .complete_batch(&pool, &claimed) + .await + .expect("complete"); + match store.rotate(&pool).await.expect("rotate") { + RotateOutcome::Rotated { .. } => {} + other => panic!("expected Rotated, got {other:?}"), + } + } + + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.queue_terminal_live_counts IN SHARE MODE" + ))) + .execute(holder.as_mut()) + .await + .expect("hold share lock"); + + let outcome = store.rollup_terminal_count_deltas(&pool, 4).await; + holder.rollback().await.expect("release share lock"); + let outcome = outcome.expect("rollup waiting on live counts should not error"); + assert_eq!(outcome.rolled_slots, 0); + assert_eq!( + outcome.blocked_slots, 1, + "the pass stops at the first blocked slot" + ); + assert_eq!(live_count_sum(&pool, schema, queue).await, 0); + assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 6); + + let outcome = store + .rollup_terminal_count_deltas(&pool, 4) + .await + .expect("rollup after lock released"); + assert_eq!(outcome.rolled_slots, 2); + assert_eq!(outcome.blocked_slots, 0); + assert_eq!(live_count_sum(&pool, schema, queue).await, 6); + assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 0); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_prune_allows_segment_reader_to_finish_claim_lock_order() { let (_db_guard, pool) = setup_pool(8).await;