From c7c8f9e7912c161eb04d39f0bae4e2957af5b0a0 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 08:10:57 +1300 Subject: [PATCH 01/16] feat(worker): make the queue-storage prune lock timeout configurable (0.6) Ring prune and terminal-delta rollup take their exclusive locks under a 10 ms lock_timeout that could only be changed through a hidden QueueStorage method. When short readers of the ring partitions hold a conflicting lock at that instant, prune reports blocked every tick and the rings stop recycling. Add ClientBuilder::prune_lock_timeout and the Python queue_storage_prune_lock_timeout_ms start() argument. The default stays 10 ms and zero is rejected, as with_prune_lock_timeout does. Make QueueStorage::with_prune_lock_timeout public and add set_prune_lock_timeout / prune_lock_timeout. awa.maintenance.prune.attempts already counts pruned, noop, blocked and skipped_active per ring; add backed_off (prune skipped by the post-skipped_active backoff) and error (prune failed with a database error) so a ring that has stopped recycling is visible for every cause. --- CHANGELOG.md | 19 ++ awa-metrics/src/lib.rs | 116 ++++++++- awa-model/src/queue_storage.rs | 26 +- awa-python/python/awa/_awa.pyi | 1 + awa-python/python/awa/client.py | 8 + awa-python/src/client.rs | 12 +- awa-python/tests/test_start_config.py | 7 + awa-worker/src/client.rs | 79 +++++- awa-worker/src/maintenance.rs | 6 + awa-worker/src/storage.rs | 13 + awa/tests/queue_storage_runtime_test.rs | 316 ++++++++++++++++++++++++ docs/configuration.md | 4 + docs/deployment.md | 1 + 13 files changed, 602 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a17deb..220649b20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,25 @@ Notable changes between releases. Detailed migration notes for storage transitio ## [Unreleased] +### Added + +- **The queue-storage prune `lock_timeout` is configurable, and prune + attempts are counted by outcome for every ring.** Prune (queue, lease and + claim rings) and terminal-delta rollup take their exclusive locks under a + 10 ms `lock_timeout` that could only be changed through the hidden + `QueueStorage::with_prune_lock_timeout`. When short readers of the ring + partitions held a conflicting lock at that instant, prune reported + `blocked` every tick and the rings stopped recycling. Set it with + `ClientBuilder::prune_lock_timeout` (Rust) or + `queue_storage_prune_lock_timeout_ms` on `client.start()` (Python); the + default stays 10 ms and zero is rejected. `with_prune_lock_timeout` is now + public, with `set_prune_lock_timeout` and `prune_lock_timeout()`. + `awa_maintenance_prune_attempts_total{awa_ring,awa_ring_outcome,awa_ring_reason}` + gains `backed_off` (prune skipped by the post-`skipped_active` backoff) and + `error` (prune failed with a database error) alongside the existing + `pruned`, `noop`, `blocked` and `skipped_active`. See + [Configuration](docs/configuration.md#what-the-knobs-mean). + ### Fixed - **`admin::job_kind_overviews` on queue storage no longer depends on the diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 3f81abce3..ad45a81ea 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -221,8 +221,11 @@ pub struct AwaMetrics { /// pinned by handfuls of stragglers or by mountains of unfinished work. pub maintenance_rotate_skipped_rows: Histogram, /// Maintenance ring prune attempts, attributed to (ring, outcome, reason). - /// reason="none" for Pruned/Noop/Blocked; otherwise carries the - /// SkipReason discriminator (e.g. "queue.active_leases", + /// Outcomes: `pruned`, `noop`, `blocked` (the exclusive locks were not + /// available within the prune lock timeout), `skipped_active`, + /// `backed_off` (the tick skipped prune after repeated `skipped_active` + /// results) and `error`. reason="none" except for `skipped_active`, which + /// carries the SkipReason discriminator (e.g. "queue.active_leases", /// "queue.pending_ready", "claim.open"). pub maintenance_prune_attempts: Counter, /// Magnitude of the reason count when prune returns SkippedActive. @@ -1092,6 +1095,27 @@ impl AwaMetrics { self.maintenance_branch_overrun_total.add(1, &attrs); } + /// Record a maintenance tick that skipped prune because the ring's + /// backoff after repeated `skipped_active` results had not yet elapsed. + pub fn record_prune_backed_off(&self, ring: &'static str) { + self.record_prune_attempt_without_reason(ring, "backed_off"); + } + + /// Record a prune attempt that failed with a database error rather than + /// returning an outcome. + pub fn record_prune_error(&self, ring: &'static str) { + self.record_prune_attempt_without_reason(ring, "error"); + } + + fn record_prune_attempt_without_reason(&self, ring: &'static str, outcome: &'static str) { + let attrs = [ + opentelemetry::KeyValue::new("awa.ring", ring), + opentelemetry::KeyValue::new("awa.ring.outcome", outcome), + opentelemetry::KeyValue::new("awa.ring.reason", "none"), + ]; + self.maintenance_prune_attempts.add(1, &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 { @@ -1162,6 +1186,94 @@ mod tests { metrics.record_job_kind_info("minimal", None, None, None, None, &[]); } + #[test] + fn prune_attempts_count_every_outcome_per_ring() { + use awa_model::{PruneOutcome, SkipReason}; + + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + let metrics = AwaMetrics::new(&provider.meter("awa-metrics-test")); + + for ring in ["queue", "lease", "claim"] { + metrics.record_prune_outcome( + ring, + &PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 0, + }, + ); + metrics.record_prune_outcome(ring, &PruneOutcome::Blocked { slot: 0 }); + metrics.record_prune_outcome(ring, &PruneOutcome::Noop); + metrics.record_prune_backed_off(ring); + metrics.record_prune_error(ring); + } + metrics.record_prune_outcome("claim", &PruneOutcome::Blocked { slot: 1 }); + metrics.record_prune_outcome( + "claim", + &PruneOutcome::SkippedActive { + slot: 0, + reason: SkipReason::ClaimOpen, + count: 3, + }, + ); + provider.force_flush().expect("flush metrics"); + let resource_metrics = exporter.get_finished_metrics().expect("read metrics"); + + let mut counts: std::collections::BTreeMap<(String, String, String), u64> = + std::collections::BTreeMap::new(); + for resource in &resource_metrics { + for scope in resource.scope_metrics() { + for metric in scope.metrics() { + let ( + names::MAINTENANCE_PRUNE_ATTEMPTS, + AggregatedMetrics::U64(MetricData::Sum(sum)), + ) = (metric.name(), metric.data()) + else { + continue; + }; + for point in sum.data_points() { + let label = |key: &str| { + point + .attributes() + .find(|attribute| attribute.key.as_str() == key) + .map(|attribute| attribute.value.as_str().into_owned()) + .unwrap_or_default() + }; + counts.insert( + ( + label("awa.ring"), + label("awa.ring.outcome"), + label("awa.ring.reason"), + ), + point.value(), + ); + } + } + } + } + + let count = |ring: &str, outcome: &str, reason: &str| { + counts + .get(&(ring.to_string(), outcome.to_string(), reason.to_string())) + .copied() + }; + for ring in ["queue", "lease", "claim"] { + for outcome in ["pruned", "noop", "backed_off", "error"] { + assert_eq!(count(ring, outcome, "none"), Some(1), "{ring} {outcome}"); + } + let expected_blocked = if ring == "claim" { 2 } else { 1 }; + assert_eq!( + count(ring, "blocked", "none"), + Some(expected_blocked), + "{ring} blocked" + ); + } + assert_eq!(count("claim", "skipped_active", "claim.open"), Some(1)); + assert_eq!(count("queue", "skipped_active", "claim.open"), None); + } + /// A second-scale observation must not fall in the SDK's default /// millisecond-scale buckets (upper bounds 5, 10, 25, ...), and every /// histogram must export the intended explicit boundaries. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 930aab715..e6073f3c7 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -2503,15 +2503,37 @@ impl QueueStorage { }) } - #[doc(hidden)] + /// Set the PostgreSQL `lock_timeout` applied to the exclusive-lock steps of + /// ring prune (queue, lease and claim rings) and terminal-delta rollup. + /// + /// A prune that cannot take its locks within the timeout rolls back and + /// reports [`PruneOutcome::Blocked`]; the maintenance loop retries on its + /// next tick. The default is 10 ms, which keeps prune from queueing behind + /// readers and so from stalling claims and enqueues behind its pending + /// `ACCESS EXCLUSIVE` request. Raise it when prune is persistently + /// `Blocked` by short-lived readers of the ring partitions. The timeout is + /// applied at millisecond granularity. + /// + /// Returns a validation error for a zero timeout. pub fn with_prune_lock_timeout(mut self, timeout: Duration) -> Result { + self.set_prune_lock_timeout(timeout)?; + Ok(self) + } + + /// In-place form of [`Self::with_prune_lock_timeout`]. + pub fn set_prune_lock_timeout(&mut self, timeout: Duration) -> Result<(), AwaError> { if timeout.is_zero() { return Err(AwaError::Validation( "queue storage prune lock timeout must be greater than zero".into(), )); } self.prune_lock_timeout = timeout; - Ok(self) + Ok(()) + } + + /// The `lock_timeout` applied to prune's exclusive-lock steps. + pub fn prune_lock_timeout(&self) -> Duration { + self.prune_lock_timeout } pub fn from_existing_schema(schema: impl Into) -> Result { diff --git a/awa-python/python/awa/_awa.pyi b/awa-python/python/awa/_awa.pyi index fcd040853..177df916f 100644 --- a/awa-python/python/awa/_awa.pyi +++ b/awa-python/python/awa/_awa.pyi @@ -583,6 +583,7 @@ class Client: queue_storage_queue_rotate_interval_ms: int = 1000, queue_storage_lease_rotate_interval_ms: int = 250, queue_storage_claim_rotate_interval_ms: int | None = None, + queue_storage_prune_lock_timeout_ms: int | None = None, storage_transition_role: str | None = None, ) -> None: ... async def shutdown(self, timeout_ms: int = 2000) -> None: ... diff --git a/awa-python/python/awa/client.py b/awa-python/python/awa/client.py index a237cc838..57346a6b7 100644 --- a/awa-python/python/awa/client.py +++ b/awa-python/python/awa/client.py @@ -867,6 +867,7 @@ async def start( queue_storage_queue_rotate_interval_ms: int = 1000, queue_storage_lease_rotate_interval_ms: int = 250, queue_storage_claim_rotate_interval_ms: int | None = None, + queue_storage_prune_lock_timeout_ms: int | None = None, storage_transition_role: Literal[ "auto", "canonical_drain", "queue_storage_target" ] = "auto", @@ -900,6 +901,12 @@ async def start( counters. Exact reads include pending deltas, so the default favors fewer hot counter updates over immediate folding. + ``queue_storage_prune_lock_timeout_ms`` sets the PostgreSQL + ``lock_timeout`` the maintenance leader uses when pruning the queue, + lease and claim rings (default 10 ms). A prune that cannot lock its + partitions in time is counted as ``blocked`` and retried on the next + tick; raise it when ``blocked`` outcomes persist. Must be > 0. + 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 @@ -929,6 +936,7 @@ async def start( queue_storage_queue_rotate_interval_ms=queue_storage_queue_rotate_interval_ms, queue_storage_lease_rotate_interval_ms=queue_storage_lease_rotate_interval_ms, queue_storage_claim_rotate_interval_ms=queue_storage_claim_rotate_interval_ms, + queue_storage_prune_lock_timeout_ms=queue_storage_prune_lock_timeout_ms, storage_transition_role=storage_transition_role, ) diff --git a/awa-python/src/client.rs b/awa-python/src/client.rs index 7de483c14..6109868e0 100644 --- a/awa-python/src/client.rs +++ b/awa-python/src/client.rs @@ -1666,7 +1666,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, 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, 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_prune_lock_timeout_ms=None, storage_transition_role=None))] #[allow(clippy::too_many_arguments)] fn start<'py>( &self, @@ -1694,6 +1694,7 @@ impl PyClient { queue_storage_queue_rotate_interval_ms: u64, queue_storage_lease_rotate_interval_ms: u64, queue_storage_claim_rotate_interval_ms: Option, + queue_storage_prune_lock_timeout_ms: Option, storage_transition_role: Option, ) -> PyResult> { let entries: Vec<_> = self @@ -1849,6 +1850,12 @@ impl PyClient { } } + if queue_storage_prune_lock_timeout_ms == Some(0) { + return Err(pyo3::exceptions::PyValueError::new_err( + "queue_storage_prune_lock_timeout_ms must be > 0", + )); + } + builder = builder.queue_storage( QueueStorageConfig { schema: queue_storage_schema @@ -1870,6 +1877,9 @@ impl PyClient { if let Some(ms) = queue_storage_claim_rotate_interval_ms { builder = builder.claim_rotate_interval(Duration::from_millis(ms)); } + if let Some(ms) = queue_storage_prune_lock_timeout_ms { + builder = builder.prune_lock_timeout(Duration::from_millis(ms)); + } builder = builder.transition_role(parse_transition_worker_role( storage_transition_role.as_deref(), )?); diff --git a/awa-python/tests/test_start_config.py b/awa-python/tests/test_start_config.py index 8a68c0123..e5ccb8818 100644 --- a/awa-python/tests/test_start_config.py +++ b/awa-python/tests/test_start_config.py @@ -738,6 +738,13 @@ async def handle(job): queue_storage_claim_rotate_interval_ms=0, ) + with pytest.raises(ValueError, match="queue_storage_prune_lock_timeout_ms must be > 0"): + await client.start( + [(queue, 1)], + poll_interval_ms=25, + queue_storage_prune_lock_timeout_ms=0, + ) + with pytest.raises(ValueError, match="terminal_count_rollup_interval_ms must be > 0"): await client.start( [(queue, 1)], diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index e5a498dbb..17eac0f9f 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -164,6 +164,7 @@ pub struct ClientBuilder { storage: RuntimeStorage, transition_role: TransitionWorkerRole, storage_error: Option, + prune_lock_timeout: Option, } impl ClientBuilder { @@ -222,6 +223,7 @@ impl ClientBuilder { storage, transition_role: TransitionWorkerRole::Auto, storage_error, + prune_lock_timeout: None, } } @@ -924,6 +926,25 @@ impl ClientBuilder { self } + /// Override the PostgreSQL `lock_timeout` the maintenance leader applies to + /// the exclusive-lock steps of queue, lease and claim ring prune and of + /// terminal-delta rollup (default 10 ms). + /// + /// A prune that cannot take its locks within the timeout rolls back and is + /// recorded as `blocked` in `awa.maintenance.prune.attempts`; it retries on + /// the next maintenance tick. The short default keeps the pending + /// `ACCESS EXCLUSIVE` request from stalling claims and enqueues behind it. + /// Raise it when `blocked` outcomes persist and the rings stop recycling + /// because short-lived readers (admin queries, dashboards) repeatedly hold + /// the ring partitions at the moment prune tries. A zero timeout makes + /// [`Self::build`] fail. Only takes effect when queue storage is active. + /// May be set before or after [`Self::queue_storage`]; retained across + /// storage configuration changes. + pub fn prune_lock_timeout(mut self, timeout: Duration) -> Self { + self.prune_lock_timeout = Some(timeout); + self + } + /// Force the worker runtime onto canonical storage. /// /// This is primarily useful for migration/testing and benchmark @@ -951,7 +972,7 @@ impl ClientBuilder { } /// Build the client. - pub fn build(self) -> Result { + pub fn build(mut self) -> Result { if self.queues.is_empty() { return Err(BuildError::NoQueuesConfigured); } @@ -960,6 +981,14 @@ impl ClientBuilder { return Err(err); } + if let (RuntimeStorage::QueueStorage(runtime), Some(timeout)) = + (&mut self.storage, self.prune_lock_timeout) + { + runtime + .set_prune_lock_timeout(timeout) + .map_err(|err| BuildError::InvalidQueueStorage(err.to_string()))?; + } + let mut queue_names = HashSet::with_capacity(self.queues.len()); for (queue, _) in &self.queues { if !queue_names.insert(queue.as_str()) { @@ -2844,6 +2873,54 @@ mod tests { builder } + #[tokio::test] + async fn prune_lock_timeout_reaches_the_store_in_either_builder_order() { + let timeout = Duration::from_millis(750); + for configure_first in [true, false] { + let mut builder = Client::builder(lazy_pool()).queue("q", QueueConfig::default()); + if configure_first { + builder = builder.prune_lock_timeout(timeout); + } + builder = builder.queue_storage( + QueueStorageConfig::default(), + Duration::from_secs(3), + Duration::from_secs(7), + ); + if !configure_first { + builder = builder.prune_lock_timeout(timeout); + } + let client = builder.build().unwrap(); + let store = client.storage.queue_storage_store().unwrap(); + assert_eq!(store.prune_lock_timeout(), timeout); + } + + let client = Client::builder(lazy_pool()) + .queue("q", QueueConfig::default()) + .build() + .unwrap(); + assert_eq!( + client + .storage + .queue_storage_store() + .unwrap() + .prune_lock_timeout(), + Duration::from_millis(10) + ); + } + + #[tokio::test] + async fn zero_prune_lock_timeout_is_rejected_at_build() { + let result = Client::builder(lazy_pool()) + .queue("q", QueueConfig::default()) + .prune_lock_timeout(Duration::ZERO) + .build(); + assert!( + matches!(result, Err(BuildError::InvalidQueueStorage(ref message)) if message.contains("prune lock timeout")), + "{:?}", + result.err() + ); + } + #[tokio::test] async fn shutdown_bounds_blocked_dispatcher_exit() { let _guard = test_mutex().lock().await; diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index b7c134177..698fb5c93 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -2612,6 +2612,7 @@ impl MaintenanceService { if prune_tracker.should_skip(PRUNE_BRANCH_QUEUE) { debug!(branch = PRUNE_BRANCH_QUEUE, "Prune backed off this tick"); + self.metrics.record_prune_backed_off("queue"); return; } @@ -2652,6 +2653,7 @@ impl MaintenanceService { } } Err(err) => { + self.metrics.record_prune_error("queue"); error!(error = %err, "Failed to prune queue storage queue segments"); } } @@ -2686,6 +2688,7 @@ impl MaintenanceService { if prune_tracker.should_skip(PRUNE_BRANCH_LEASE) { debug!(branch = PRUNE_BRANCH_LEASE, "Prune backed off this tick"); + self.metrics.record_prune_backed_off("lease"); return; } @@ -2716,6 +2719,7 @@ impl MaintenanceService { } } Err(err) => { + self.metrics.record_prune_error("lease"); error!(error = %err, "Failed to prune queue storage lease segments"); } } @@ -2755,6 +2759,7 @@ impl MaintenanceService { if prune_tracker.should_skip(PRUNE_BRANCH_CLAIM) { debug!(branch = PRUNE_BRANCH_CLAIM, "Prune backed off this tick"); + self.metrics.record_prune_backed_off("claim"); return; } @@ -2785,6 +2790,7 @@ impl MaintenanceService { } } Err(err) => { + self.metrics.record_prune_error("claim"); error!(error = %err, "Failed to prune queue storage claim segments"); } } diff --git a/awa-worker/src/storage.rs b/awa-worker/src/storage.rs index 66720be0f..e56dc3bfa 100644 --- a/awa-worker/src/storage.rs +++ b/awa-worker/src/storage.rs @@ -33,6 +33,19 @@ impl QueueStorageRuntime { self.claim_rotate_interval = claim_rotate_interval; self } + + /// Set the prune `lock_timeout` on the store. Only valid before the store + /// is shared, which holds for a runtime still owned by the builder. + pub fn set_prune_lock_timeout(&mut self, timeout: Duration) -> Result<(), awa_model::AwaError> { + Arc::get_mut(&mut self.store) + .ok_or_else(|| { + awa_model::AwaError::Validation( + "queue storage is already shared; set the prune lock timeout before build" + .into(), + ) + })? + .set_prune_lock_timeout(timeout) + } } #[derive(Debug, Clone, Default)] diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index a060369f8..e64589321 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -11781,6 +11781,322 @@ async fn test_queue_storage_claimer_heartbeat_skips_fresh_lease() { ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_runtime_counts_blocked_then_pruned_queue_prune_attempts() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_prune_attempt_metric"; + let schema = "awa_qs_prune_attempt_metric"; + let store_config = QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + + let mut reader_tx = pool.begin().await.expect("begin reader tx"); + sqlx::query(&format!( + "LOCK TABLE {schema}.ready_segments_0 IN ACCESS SHARE MODE" + )) + .execute(reader_tx.as_mut()) + .await + .expect("reader lock"); + + let (exporter, meter_provider) = install_in_memory_metrics(); + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: 4, + poll_interval: Duration::from_millis(25), + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_millis(100), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .prune_lock_timeout(Duration::from_millis(20)) + .register_worker(CompleteWorker) + .promote_interval(Duration::from_millis(25)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .deadline_rescue_interval(Duration::from_millis(100)) + .callback_rescue_interval(Duration::from_millis(25)) + .build() + .expect("build prune metric client"); + client.start().await.expect("start prune metric client"); + + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 4001 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + wait_for_job_state( + &store, + &pool, + job_id, + &[JobState::Completed], + Duration::from_secs(15), + ) + .await; + + let queue_attempts = |outcome: &str| { + meter_provider.force_flush().expect("flush metrics"); + let resource_metrics = exporter.get_finished_metrics().expect("read metrics"); + 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.prune.attempts" { + continue; + } + let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else { + continue; + }; + total += sum + .data_points() + .filter(|point| { + let has = |key: &str, value: &str| { + point + .attributes() + .any(|kv| kv.key.as_str() == key && kv.value.as_str() == value) + }; + has("awa.ring", "queue") && has("awa.ring.outcome", outcome) + }) + .map(|point| point.value()) + .sum::(); + } + } + } + total + }; + + let deadline = Instant::now() + Duration::from_secs(20); + while queue_attempts("blocked") == 0 { + assert!( + Instant::now() < deadline, + "queue prune never reported blocked while a reader held the queue segment" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert_eq!(queue_attempts("pruned"), 0); + + reader_tx.rollback().await.expect("release reader lock"); + let deadline = Instant::now() + Duration::from_secs(20); + while queue_attempts("pruned") == 0 { + assert!( + Instant::now() < deadline, + "queue prune never completed after the reader released" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } + + client.shutdown(Duration::from_secs(5)).await; +} + +#[derive(Clone, Copy, Debug)] +enum PruneRing { + Queue, + Lease, + Claim, +} + +async fn prune_ring( + store: &QueueStorage, + pool: &sqlx::PgPool, + ring: PruneRing, +) -> Result { + match ring { + PruneRing::Queue => store.prune_oldest(pool, Duration::ZERO).await, + PruneRing::Lease => store.prune_oldest_leases(pool).await, + PruneRing::Claim => store.prune_oldest_claims(pool).await, + } +} + +/// A reader holds `locked_relation` so prune must wait for its exclusive lock. +/// With a short timeout prune gives up as `Blocked` only after that timeout; +/// with a long timeout it waits out a reader that releases part-way and then +/// completes. +async fn assert_prune_waits_for_configured_lock_timeout( + pool: &sqlx::PgPool, + config: &QueueStorageConfig, + ring: PruneRing, + locked_relation: &str, +) { + let short_timeout = Duration::from_millis(300); + let short_store = QueueStorage::new(config.clone()) + .expect("short-timeout store") + .with_prune_lock_timeout(short_timeout) + .expect("short prune lock timeout"); + assert_eq!(short_store.prune_lock_timeout(), short_timeout); + + let mut reader_tx = pool.begin().await.expect("begin reader tx"); + sqlx::query(&format!( + "LOCK TABLE {locked_relation} IN ACCESS SHARE MODE" + )) + .execute(reader_tx.as_mut()) + .await + .expect("reader lock"); + + let started = Instant::now(); + let blocked = prune_ring(&short_store, pool, ring) + .await + .expect("prune under reader lock"); + let elapsed = started.elapsed(); + assert!( + matches!(blocked, PruneOutcome::Blocked { .. }), + "{ring:?} prune must report Blocked while the reader holds its lock, got {blocked:?}" + ); + assert!( + elapsed >= short_timeout, + "{ring:?} prune gave up after {elapsed:?}, before the {short_timeout:?} lock timeout" + ); + assert!( + elapsed < Duration::from_secs(10), + "{ring:?} prune waited {elapsed:?} for a {short_timeout:?} lock timeout" + ); + + let long_store = QueueStorage::new(config.clone()) + .expect("long-timeout store") + .with_prune_lock_timeout(Duration::from_secs(60)) + .expect("long prune lock timeout"); + let reader_hold = Duration::from_millis(600); + let release = tokio::spawn(async move { + tokio::time::sleep(reader_hold).await; + reader_tx.rollback().await.expect("release reader lock"); + }); + let started = Instant::now(); + let outcome = prune_ring(&long_store, pool, ring) + .await + .expect("prune after reader release"); + let elapsed = started.elapsed(); + release.await.expect("release task"); + assert!( + !matches!(outcome, PruneOutcome::Blocked { .. }), + "{ring:?} prune with a long lock timeout must outwait the reader, got {outcome:?}" + ); + assert!( + elapsed >= short_timeout, + "{ring:?} prune returned after {elapsed:?} instead of waiting for the reader" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_prune_lock_timeout_is_configurable() { + let (_db_guard, pool) = setup_pool(8).await; + let schema = "awa_qs_prune_lt_queue"; + let store = create_store(&pool, schema).await; + store.enqueue_batch(&pool, "q", 1, 1).await.unwrap(); + let claims = store.claim_batch(&pool, "q", 1).await.unwrap(); + assert_eq!(store.complete_batch(&pool, &claims).await.unwrap(), 1); + assert!(matches!( + store.rotate(&pool).await.unwrap(), + RotateOutcome::Rotated { .. } + )); + + let config = QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }; + assert_prune_waits_for_configured_lock_timeout( + &pool, + &config, + PruneRing::Queue, + &format!("{schema}.ready_segments_0"), + ) + .await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_lease_prune_lock_timeout_is_configurable() { + let (_db_guard, pool) = setup_pool(8).await; + let schema = "awa_qs_prune_lt_lease"; + let store = create_store(&pool, schema).await; + store.enqueue_batch(&pool, "q", 1, 2).await.unwrap(); + let claimed = store + .claim_runtime_batch(&pool, "q", 2, Duration::from_secs(60)) + .await + .unwrap(); + assert_eq!(claimed.len(), 2); + assert!(matches!( + store.rotate_leases(&pool).await.unwrap(), + RotateOutcome::Rotated { slot: 1, .. } + )); + + let config = QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }; + assert_prune_waits_for_configured_lock_timeout( + &pool, + &config, + PruneRing::Lease, + &format!("{schema}.leases_0"), + ) + .await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_claim_prune_lock_timeout_is_configurable() { + let (_db_guard, pool) = setup_pool(8).await; + let schema = "awa_qs_prune_lt_claim"; + let config = QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 4, + lease_claim_receipts: true, + ..Default::default() + }; + let store = create_store_with_config(&pool, config.clone()).await; + sqlx::query(&format!( + r#" + INSERT INTO {schema}.lease_claims ( + claim_slot, job_id, run_lease, ready_slot, ready_generation, + queue, priority, attempt, max_attempts, lane_seq + ) VALUES (0, 1001, 1, 0, 0, 'synthetic', 2, 1, 25, 1001) + "# + )) + .execute(&pool) + .await + .expect("seed claim"); + sqlx::query(&format!( + "INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome) \ + VALUES (0, 1001, 1, 'completed')" + )) + .execute(&pool) + .await + .expect("seed closure"); + assert!(matches!( + store.rotate_claims(&pool).await.unwrap(), + RotateOutcome::Rotated { slot: 1, .. } + )); + + assert_prune_waits_for_configured_lock_timeout( + &pool, + &config, + PruneRing::Claim, + &format!("{schema}.lease_claims_0"), + ) + .await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_prune_oldest_blocks_on_reader_lock() { let (_db_guard, pool) = setup_pool(10).await; diff --git a/docs/configuration.md b/docs/configuration.md index 868dfb3cf..677a0ddf3 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -537,6 +537,7 @@ let client = Client::builder(pool.clone()) Duration::from_millis(1_000), ) .claim_rotate_interval(Duration::from_millis(1_000)) + .prune_lock_timeout(Duration::from_millis(10)) .build()?; ``` @@ -553,6 +554,7 @@ await client.start( queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=1000, + queue_storage_prune_lock_timeout_ms=10, ) ``` @@ -568,6 +570,7 @@ await client.start( | `queue_rotate_interval` | `1000ms` | How often ready/terminal segments rotate | | `lease_rotate_interval` | `1000ms` | How often lease segments rotate | | `claim_rotate_interval` | matches `queue_rotate_interval` | How often the ADR-023 claim-ring rotates. Set with `ClientBuilder::claim_rotate_interval` (Rust) or `queue_storage_claim_rotate_interval_ms` (Python). Test harnesses sometimes set this to a long interval to pin claim-ring layout for deterministic count assertions. | +| `prune_lock_timeout` | `10ms` | The PostgreSQL `lock_timeout` the maintenance leader applies when it takes the exclusive locks to reclaim a queue, lease or claim ring partition (and when it folds terminal-count deltas). A prune that cannot lock its partitions within the timeout rolls back, is counted as `blocked`, and is retried on the next maintenance tick. Set with `ClientBuilder::prune_lock_timeout` (Rust) or `queue_storage_prune_lock_timeout_ms` (Python). Must be greater than zero. The short default keeps prune's pending `ACCESS EXCLUSIVE` request from stalling claims and enqueues queued behind it. See below for when to raise it. | The benchmark harness in [postgresql-job-queue-benchmarking](https://github.com/hardbyte/postgresql-job-queue-benchmarking) reads `QUEUE_SLOT_COUNT`, `LEASE_SLOT_COUNT`, `CLAIM_SLOT_COUNT`, `QUEUE_STRIPE_COUNT`, and `LEASE_CLAIM_RECEIPTS` from the environment. Those env vars are benchmark configuration, not general worker-runtime configuration. @@ -577,6 +580,7 @@ Use the defaults unless you have a reason not to: - Increase `lease_slot_count` if lease churn is high enough that dead tuples in the lease ring stop collapsing promptly. - Increase `claim_slot_count` if the rotation cadence (`claim_rotate_interval`) plus the slot count combine to a partition retention window shorter than your longest in-flight zero-deadline short job; running out of empty slots forces `rotate_claims` to return `SkippedBusy` and the receipt-plane churn falls back onto a smaller working set of partitions. - Increase `queue_stripe_count` only for measured workloads where many small replicas contend on the same logical queue. This is queue-storage configuration, not a per-queue `QueueConfig` field: it gives each logical queue the same number of physical stripes. Striping spreads a hot logical queue over `queue#N` physical coordination paths, but it weakens perfect global ordering and can regress calmer shapes if overused. For single hot-queue workloads, benchmark `2` stripes first; higher values should be treated as workload-specific tuning rather than a general default. +- Increase `prune_lock_timeout` when `awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}` stays non-zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Raise it in steps (for example `100ms`, then `1s`) and keep it well below the time you can tolerate claims and enqueues queueing behind prune; while prune waits, new accesses to the same partition queue behind its lock request. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for a few ticks after repeated `skipped_active`, and `error` is a database failure. - Increase rotation intervals to reduce partition churn and metadata activity. - Decrease rotation intervals to tighten dead-tuple bounds at the cost of more frequent rotate/prune work. diff --git a/docs/deployment.md b/docs/deployment.md index 99bf4b87e..2a94d584e 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -82,6 +82,7 @@ Recommended practice: - monitor `pg_stat_activity` for long transactions - monitor `pg_stat_user_tables` on the active queue-storage schema; ready, tombstone, done, and terminal-delta segments should stay near zero dead tuples, lease segments may spike within the rotation window but should collapse after prune, and `attempt_state` should roughly track live long-running attempts - tune autovacuum for the database if the lease tables churn heavily +- alert on `awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}` (and `error`) sustained above zero: the rings are not recycling. If short readers of the ring partitions are the cause, raise [`prune_lock_timeout`](configuration.md#what-the-knobs-mean) Queue-storage depth metrics are designed for high-cadence monitoring. The worker records `available` from lane-head cursor differences and probes queue lag from the next claimable row per lane, instead of running the exact `queue_counts` ready-row scan on every metrics tick. Use the admin API or CLI when you need exact queue counts for an operator action. From aa9c2fbe738940b2a40071a204927d224b849a7a Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 09:06:32 +1300 Subject: [PATCH 02/16] test(worker): make the prune metric runtime test independent of ring rotation timing --- awa/tests/queue_storage_runtime_test.rs | 81 ++++++++++++++++--------- 1 file changed, 51 insertions(+), 30 deletions(-) diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index e64589321..181c12b2d 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -36,6 +36,10 @@ struct RawReceiptClaimRow { claim_slot: i32, } +/// `AwaMetrics::from_global` reads the process-wide meter provider when a +/// client starts, so tests that install their own provider must not overlap. +static GLOBAL_METRICS_TEST_LOCK: Mutex<()> = Mutex::const_new(()); + fn install_in_memory_metrics() -> (InMemoryMetricExporter, SdkMeterProvider) { let exporter = InMemoryMetricExporter::default(); let meter_provider = SdkMeterProvider::builder() @@ -5234,6 +5238,7 @@ async fn test_queue_storage_cancelled_compact_claim_closes_via_batch_and_queue_p #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_capacity_wake_drains_after_partial_drain() { + let _metrics_guard = GLOBAL_METRICS_TEST_LOCK.lock().await; let (exporter, meter_provider) = install_in_memory_metrics(); let (_db_guard, pool) = setup_pool(10).await; let queue = "qs_capacity_wake_partial_drain"; @@ -11783,7 +11788,7 @@ async fn test_queue_storage_claimer_heartbeat_skips_fresh_lease() { #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_runtime_counts_blocked_then_pruned_queue_prune_attempts() { - let (_db_guard, pool) = setup_pool(10).await; + let (_db_guard, pool) = setup_pool(20).await; let queue = "qs_prune_attempt_metric"; let schema = "awa_qs_prune_attempt_metric"; let store_config = QueueStorageConfig { @@ -11795,6 +11800,18 @@ async fn test_runtime_counts_blocked_then_pruned_queue_prune_attempts() { }; let store = create_store_with_config(&pool, store_config.clone()).await; + // Enqueue before the client starts so the job's rows are in slot 0, the + // slot the ring prunes first, whatever the maintenance timing. + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 4001 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; let mut reader_tx = pool.begin().await.expect("begin reader tx"); sqlx::query(&format!( "LOCK TABLE {schema}.ready_segments_0 IN ACCESS SHARE MODE" @@ -11803,6 +11820,7 @@ async fn test_runtime_counts_blocked_then_pruned_queue_prune_attempts() { .await .expect("reader lock"); + let _metrics_guard = GLOBAL_METRICS_TEST_LOCK.lock().await; let (exporter, meter_provider) = install_in_memory_metrics(); let client = Client::builder(pool.clone()) .queue( @@ -11831,16 +11849,6 @@ async fn test_runtime_counts_blocked_then_pruned_queue_prune_attempts() { .expect("build prune metric client"); client.start().await.expect("start prune metric client"); - let job_id = enqueue_job( - &pool, - &store, - &CompleteJob { id: 4001 }, - InsertOpts { - queue: queue.to_string(), - ..Default::default() - }, - ) - .await; wait_for_job_state( &store, &pool, @@ -11850,10 +11858,10 @@ async fn test_runtime_counts_blocked_then_pruned_queue_prune_attempts() { ) .await; - let queue_attempts = |outcome: &str| { + let attempts_by_outcome = || { meter_provider.force_flush().expect("flush metrics"); let resource_metrics = exporter.get_finished_metrics().expect("read metrics"); - let mut total = 0; + let mut totals: std::collections::BTreeMap = Default::default(); for resource in &resource_metrics { for scope in resource.scope_metrics() { for metric in scope.metrics() { @@ -11863,40 +11871,53 @@ async fn test_runtime_counts_blocked_then_pruned_queue_prune_attempts() { let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else { continue; }; - total += sum - .data_points() - .filter(|point| { - let has = |key: &str, value: &str| { - point - .attributes() - .any(|kv| kv.key.as_str() == key && kv.value.as_str() == value) - }; - has("awa.ring", "queue") && has("awa.ring.outcome", outcome) - }) - .map(|point| point.value()) - .sum::(); + for point in sum.data_points() { + let attribute = |key: &str| { + point + .attributes() + .find(|kv| kv.key.as_str() == key) + .map(|kv| kv.value.as_str().into_owned()) + .unwrap_or_default() + }; + *totals + .entry(format!( + "{}/{}", + attribute("awa.ring"), + attribute("awa.ring.outcome") + )) + .or_default() += point.value(); + } } } } - total + totals + }; + let queue_attempts = |outcome: &str| { + attempts_by_outcome() + .get(&format!("queue/{outcome}")) + .copied() + .unwrap_or(0) }; - let deadline = Instant::now() + Duration::from_secs(20); + let deadline = Instant::now() + Duration::from_secs(30); while queue_attempts("blocked") == 0 { assert!( Instant::now() < deadline, - "queue prune never reported blocked while a reader held the queue segment" + "queue prune never reported blocked while a reader held the queue segment; \ + attempts so far: {:?}", + attempts_by_outcome() ); tokio::time::sleep(Duration::from_millis(100)).await; } assert_eq!(queue_attempts("pruned"), 0); reader_tx.rollback().await.expect("release reader lock"); - let deadline = Instant::now() + Duration::from_secs(20); + let deadline = Instant::now() + Duration::from_secs(30); while queue_attempts("pruned") == 0 { assert!( Instant::now() < deadline, - "queue prune never completed after the reader released" + "queue prune never completed after the reader released; attempts so far: {:?}", + attempts_by_outcome() ); tokio::time::sleep(Duration::from_millis(100)).await; } From ed89d9fef8cb05709b1f53a913798fac7b5056be Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 11:51:59 +1300 Subject: [PATCH 03/16] fix(worker): bound the prune lock timeout and tighten its docs and tests --- CHANGELOG.md | 3 ++- awa-metrics/src/lib.rs | 7 ++--- awa-model/src/queue_storage.rs | 9 ++++++- awa-worker/src/client.rs | 36 ++++++++++++++++++------- awa/tests/queue_storage_runtime_test.rs | 16 +++++++++-- docs/configuration.md | 2 +- docs/deployment.md | 2 +- 7 files changed, 57 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 220649b20..8ed8acfa6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,7 +18,8 @@ Notable changes between releases. Detailed migration notes for storage transitio default stays 10 ms and zero is rejected. `with_prune_lock_timeout` is now public, with `set_prune_lock_timeout` and `prune_lock_timeout()`. `awa_maintenance_prune_attempts_total{awa_ring,awa_ring_outcome,awa_ring_reason}` - gains `backed_off` (prune skipped by the post-`skipped_active` backoff) and + gains `backed_off` (prune skipped by the backoff that follows consecutive + `skipped_active` or `blocked` results) and `error` (prune failed with a database error) alongside the existing `pruned`, `noop`, `blocked` and `skipped_active`. See [Configuration](docs/configuration.md#what-the-knobs-mean). diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index ad45a81ea..0185be9cd 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -223,8 +223,8 @@ pub struct AwaMetrics { /// Maintenance ring prune attempts, attributed to (ring, outcome, reason). /// Outcomes: `pruned`, `noop`, `blocked` (the exclusive locks were not /// available within the prune lock timeout), `skipped_active`, - /// `backed_off` (the tick skipped prune after repeated `skipped_active` - /// results) and `error`. reason="none" except for `skipped_active`, which + /// `backed_off` (the tick skipped prune while the ring's backoff after + /// consecutive `skipped_active` or `blocked` results ran down) and `error`. reason="none" except for `skipped_active`, which /// carries the SkipReason discriminator (e.g. "queue.active_leases", /// "queue.pending_ready", "claim.open"). pub maintenance_prune_attempts: Counter, @@ -1096,7 +1096,8 @@ impl AwaMetrics { } /// Record a maintenance tick that skipped prune because the ring's - /// backoff after repeated `skipped_active` results had not yet elapsed. + /// backoff after consecutive `skipped_active` or `blocked` results had not + /// yet elapsed. pub fn record_prune_backed_off(&self, ring: &'static str) { self.record_prune_attempt_without_reason(ring, "backed_off"); } diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index e6073f3c7..c50999925 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -2514,7 +2514,8 @@ impl QueueStorage { /// `Blocked` by short-lived readers of the ring partitions. The timeout is /// applied at millisecond granularity. /// - /// Returns a validation error for a zero timeout. + /// Returns a validation error for a zero timeout or one above PostgreSQL's + /// `lock_timeout` maximum (`i32::MAX` milliseconds, about 24.8 days). pub fn with_prune_lock_timeout(mut self, timeout: Duration) -> Result { self.set_prune_lock_timeout(timeout)?; Ok(self) @@ -2527,6 +2528,12 @@ impl QueueStorage { "queue storage prune lock timeout must be greater than zero".into(), )); } + if timeout.as_millis() > i32::MAX as u128 { + return Err(AwaError::Validation(format!( + "queue storage prune lock timeout must not exceed {} ms", + i32::MAX + ))); + } self.prune_lock_timeout = timeout; Ok(()) } diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 17eac0f9f..1d16bb47c 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -936,7 +936,8 @@ impl ClientBuilder { /// `ACCESS EXCLUSIVE` request from stalling claims and enqueues behind it. /// Raise it when `blocked` outcomes persist and the rings stop recycling /// because short-lived readers (admin queries, dashboards) repeatedly hold - /// the ring partitions at the moment prune tries. A zero timeout makes + /// the ring partitions at the moment prune tries. A zero timeout, or one above + /// PostgreSQL's `lock_timeout` maximum (`i32::MAX` ms), makes /// [`Self::build`] fail. Only takes effect when queue storage is active. /// May be set before or after [`Self::queue_storage`]; retained across /// storage configuration changes. @@ -2909,15 +2910,32 @@ mod tests { } #[tokio::test] - async fn zero_prune_lock_timeout_is_rejected_at_build() { - let result = Client::builder(lazy_pool()) + async fn out_of_range_prune_lock_timeout_is_rejected_at_build() { + for timeout in [Duration::ZERO, Duration::from_millis(i32::MAX as u64 + 1)] { + let result = Client::builder(lazy_pool()) + .queue("q", QueueConfig::default()) + .prune_lock_timeout(timeout) + .build(); + assert!( + matches!(result, Err(BuildError::InvalidQueueStorage(ref message)) if message.contains("prune lock timeout")), + "{timeout:?}: {:?}", + result.err() + ); + } + + let at_maximum = Duration::from_millis(i32::MAX as u64); + let client = Client::builder(lazy_pool()) .queue("q", QueueConfig::default()) - .prune_lock_timeout(Duration::ZERO) - .build(); - assert!( - matches!(result, Err(BuildError::InvalidQueueStorage(ref message)) if message.contains("prune lock timeout")), - "{:?}", - result.err() + .prune_lock_timeout(at_maximum) + .build() + .unwrap(); + assert_eq!( + client + .storage + .queue_storage_store() + .unwrap() + .prune_lock_timeout(), + at_maximum ); } diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 181c12b2d..913890eec 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -11953,6 +11953,7 @@ async fn assert_prune_waits_for_configured_lock_timeout( config: &QueueStorageConfig, ring: PruneRing, locked_relation: &str, + expected_after_release: fn(&PruneOutcome) -> bool, ) { let short_timeout = Duration::from_millis(300); let short_store = QueueStorage::new(config.clone()) @@ -12003,8 +12004,8 @@ async fn assert_prune_waits_for_configured_lock_timeout( let elapsed = started.elapsed(); release.await.expect("release task"); assert!( - !matches!(outcome, PruneOutcome::Blocked { .. }), - "{ring:?} prune with a long lock timeout must outwait the reader, got {outcome:?}" + expected_after_release(&outcome), + "{ring:?} prune with a long lock timeout must outwait the reader and finish, got {outcome:?}" ); assert!( elapsed >= short_timeout, @@ -12037,6 +12038,7 @@ async fn test_queue_prune_lock_timeout_is_configurable() { &config, PruneRing::Queue, &format!("{schema}.ready_segments_0"), + |outcome| matches!(outcome, PruneOutcome::Pruned { .. }), ) .await; } @@ -12069,6 +12071,15 @@ async fn test_lease_prune_lock_timeout_is_configurable() { &config, PruneRing::Lease, &format!("{schema}.leases_0"), + |outcome| { + matches!( + outcome, + PruneOutcome::SkippedActive { + reason: SkipReason::LeaseActive, + .. + } + ) + }, ) .await; } @@ -12114,6 +12125,7 @@ async fn test_claim_prune_lock_timeout_is_configurable() { &config, PruneRing::Claim, &format!("{schema}.lease_claims_0"), + |outcome| matches!(outcome, PruneOutcome::Pruned { .. }), ) .await; } diff --git a/docs/configuration.md b/docs/configuration.md index 677a0ddf3..4c0007681 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -580,7 +580,7 @@ Use the defaults unless you have a reason not to: - Increase `lease_slot_count` if lease churn is high enough that dead tuples in the lease ring stop collapsing promptly. - Increase `claim_slot_count` if the rotation cadence (`claim_rotate_interval`) plus the slot count combine to a partition retention window shorter than your longest in-flight zero-deadline short job; running out of empty slots forces `rotate_claims` to return `SkippedBusy` and the receipt-plane churn falls back onto a smaller working set of partitions. - Increase `queue_stripe_count` only for measured workloads where many small replicas contend on the same logical queue. This is queue-storage configuration, not a per-queue `QueueConfig` field: it gives each logical queue the same number of physical stripes. Striping spreads a hot logical queue over `queue#N` physical coordination paths, but it weakens perfect global ordering and can regress calmer shapes if overused. For single hot-queue workloads, benchmark `2` stripes first; higher values should be treated as workload-specific tuning rather than a general default. -- Increase `prune_lock_timeout` when `awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}` stays non-zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Raise it in steps (for example `100ms`, then `1s`) and keep it well below the time you can tolerate claims and enqueues queueing behind prune; while prune waits, new accesses to the same partition queue behind its lock request. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for a few ticks after repeated `skipped_active`, and `error` is a database failure. +- Increase `prune_lock_timeout` when `increase(awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}[15m])` stays above zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Raise it in steps (for example `100ms`, then `1s`) and keep it well below the time you can tolerate claims and enqueues queueing behind prune; while prune waits, new accesses to the same partition queue behind its lock request. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for up to 32 ticks after consecutive `skipped_active` or `blocked` results, so a persistently blocked ring is retried less often over time, and `error` is a database failure. - Increase rotation intervals to reduce partition churn and metadata activity. - Decrease rotation intervals to tighten dead-tuple bounds at the cost of more frequent rotate/prune work. diff --git a/docs/deployment.md b/docs/deployment.md index 2a94d584e..99492b2cc 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -82,7 +82,7 @@ Recommended practice: - monitor `pg_stat_activity` for long transactions - monitor `pg_stat_user_tables` on the active queue-storage schema; ready, tombstone, done, and terminal-delta segments should stay near zero dead tuples, lease segments may spike within the rotation window but should collapse after prune, and `attempt_state` should roughly track live long-running attempts - tune autovacuum for the database if the lease tables churn heavily -- alert on `awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}` (and `error`) sustained above zero: the rings are not recycling. If short readers of the ring partitions are the cause, raise [`prune_lock_timeout`](configuration.md#what-the-knobs-mean) +- alert when `increase(awa_maintenance_prune_attempts_total{awa_ring_outcome=~"blocked|error"}[15m])` stays above zero for the alert's `for:` duration: the rings are not recycling. The counter is cumulative, so its raw value stays above zero after a single blocked attempt. If short readers of the ring partitions are the cause, raise [`prune_lock_timeout`](configuration.md#what-the-knobs-mean) Queue-storage depth metrics are designed for high-cadence monitoring. The worker records `available` from lane-head cursor differences and probes queue lag from the next claimable row per lane, instead of running the exact `queue_counts` ready-row scan on every metrics tick. Use the admin API or CLI when you need exact queue counts for an operator action. From 8a8d6532dd0f3f188fd06cf429401c9594767495 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 13:27:38 +1300 Subject: [PATCH 04/16] fix(queue-storage): create queue lanes as the schema owner (0.6) 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. The definer_lane_provisioning schema patch refreshes the helpers and repairs existing lane sequences one at a time, reassigning and granting only what the applying 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 | 27 + awa-cli/src/main.rs | 62 +- awa-cli/tests/migrate_patch_cli_test.rs | 11 +- .../patches/definer_lane_provisioning.sql | 80 ++ .../v023_install_queue_storage_substrate.sql | 417 +++++++++- awa-model/src/migrations.rs | 76 ++ awa-model/src/queue_storage.rs | 105 ++- awa-model/src/storage.rs | 83 ++ .../fixtures/v0613_lane_sequence_helpers.sql | 195 +++++ awa/tests/lane_ownership_roles_test.rs | 733 ++++++++++++++++++ docs/security.md | 32 +- docs/upgrade-0.6.13-to-0.6.14.md | 45 ++ 12 files changed, 1814 insertions(+), 52 deletions(-) create mode 100644 awa-model/migrations/patches/definer_lane_provisioning.sql create mode 100644 awa/tests/fixtures/v0613_lane_sequence_helpers.sql create mode 100644 awa/tests/lane_ownership_roles_test.rs create mode 100644 docs/upgrade-0.6.13-to-0.6.14.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a17deb..2ff251781 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,33 @@ 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.** 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. The new schema patch + `definer_lane_provisioning` creates lanes 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. The patch gives existing lane sequences to the schema owner + and adds their missing grants, one sequence at a time and only where the + applying 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 before the patch detects that `provision_lane` is missing and + creates lanes as before, and older binaries create correctly owned lanes once + the patch has refreshed the helpers they call. `awa migrate` applies the + patch on any v040 database; external runners apply the exported + `R__definer_lane_provisioning.sql` after `V40`. See + [`docs/upgrade-0.6.13-to-0.6.14.md`](docs/upgrade-0.6.13-to-0.6.14.md). - **`admin::job_kind_overviews` on queue storage no longer depends on the planner's row estimates.** Behind the admin UI's kinds page, it counted every job through the `current_jobs` view. Two nodes of that view are diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index f20634116..4b994611a 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -417,7 +417,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")] @@ -432,6 +432,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 @@ -669,6 +679,21 @@ async fn main() -> Result<(), Box> { if !report.complete() { std::process::exit(1); } + if report.patches.iter().any(|patch| { + patch.name == "definer_lane_provisioning" + && patch.outcome != awa_model::migrations::PatchOutcome::NotApplicable + }) { + 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." + ); + } + } } } @@ -1245,6 +1270,41 @@ async fn main() -> 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-cli/tests/migrate_patch_cli_test.rs b/awa-cli/tests/migrate_patch_cli_test.rs index 199ebeff4..abf84f477 100644 --- a/awa-cli/tests/migrate_patch_cli_test.rs +++ b/awa-cli/tests/migrate_patch_cli_test.rs @@ -321,12 +321,11 @@ fn psql(database_url: &str, sql: &str, lock_timeout: Option<&str>) -> Output { let mut child = command .spawn() .expect("psql must be on PATH for the exported-SQL test"); - child - .stdin - .take() - .unwrap() - .write_all(sql.as_bytes()) - .unwrap(); + // psql stops reading once ON_ERROR_STOP fires, which can be before the + // whole script is written. + if let Err(err) = child.stdin.take().unwrap().write_all(sql.as_bytes()) { + assert_eq!(err.kind(), std::io::ErrorKind::BrokenPipe, "{err}"); + } child.wait_with_output().unwrap() } diff --git a/awa-model/migrations/patches/definer_lane_provisioning.sql b/awa-model/migrations/patches/definer_lane_provisioning.sql new file mode 100644 index 000000000..439630db9 --- /dev/null +++ b/awa-model/migrations/patches/definer_lane_provisioning.sql @@ -0,0 +1,80 @@ +-- Schema patch definer_lane_provisioning: create queue lanes as the schema +-- owner, whichever role touches them first (awa 0.7 migration v053, minus the +-- schema_version row). +-- +-- Applied after the v023 definitions, whose awa.refresh_lane_sequence_helpers +-- now installs two SECURITY DEFINER functions in each queue-storage schema, +-- provision_lane and provision_lane_sequences, owned by the schema owner. The +-- lane helpers create a missing sequence through them, and each new sequence +-- gets the grants of the head table it backs, so a lane is usable by every +-- runtime role however its creator's default privileges are set. The block +-- below refreshes the helpers on each installed schema, 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 applying role may. It +-- never fails on a sequence it cannot reassign; those are reported in a +-- warning and by awa.lane_sequence_problems(). +-- +-- Live load: function replacement, plus ALTER SEQUENCE OWNER on the lane +-- sequences that need it. That holds each such sequence's ACCESS EXCLUSIVE +-- lock until commit, so enqueues and claims on those lanes wait for the patch +-- to commit; correctly owned lanes are not touched. Reads the catalog only. +-- +-- 0.6 delivery: the 0.6 series cannot add a migration version, so this ships as +-- an idempotent schema patch. `awa migrate` applies it when the schema is +-- exactly at v040 and the probe in migrations.rs reports it missing; +-- `awa migrate --sql` prints it as a repeatable R__ script for external +-- runners. Released 0.6 binaries call the same helpers with the same arguments +-- and cursor semantics, and lanes they create are owned by the schema owner +-- too. Run it as the schema owner, or as a migrator that can SET ROLE to it. + +DO $$ +DECLARE + v_schema TEXT; + v_leftovers TEXT; + v_count BIGINT; +BEGIN + FOR v_schema IN + SELECT n.nspname + FROM pg_namespace AS n + WHERE has_schema_privilege(current_user, n.oid, 'USAGE') + AND EXISTS ( + SELECT 1 FROM pg_proc AS awa_p + WHERE awa_p.pronamespace = n.oid + AND awa_p.proname = 'claim_ready_runtime' + AND oidvectortypes(awa_p.proargtypes) + = 'text, bigint, double precision, double precision' + ) + ORDER BY n.nspname + 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_regprocedure(format('%I.queue_lane_sequence_name(text,text,smallint,smallint)', v_schema)) IS NULL + OR to_regprocedure(format('%I.set_sequence_next(text,bigint)', v_schema)) IS NULL THEN + CONTINUE; + END IF; + + -- Serialize with a concurrent prepare_schema() of the same schema. This + -- is an advisory lock, not a table lock: claims and completions do not + -- take it. + PERFORM pg_advisory_xact_lock( + hashtextextended(format('awa.queue_storage.install:%I', v_schema), 0) + ); + 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; +$$; diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index 5b7bfe127..0d170c34f 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -47,6 +47,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 IF p_schema IS NULL OR p_schema !~ '^[a-z_][a-z0-9_]*$' THEN RAISE EXCEPTION 'refresh_lane_sequence_helpers: schema name must match ^[a-z_][a-z0-9_]*$; got %', @@ -80,28 +86,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 @@ -142,16 +160,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; @@ -208,16 +232,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; @@ -231,11 +261,342 @@ BEGIN $ddl$, 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; + CREATE OR REPLACE FUNCTION awa.install_queue_storage_substrate( p_schema TEXT, p_queue_slot_count INT DEFAULT 16, diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 193c465fc..29a65fc93 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -295,6 +295,61 @@ macro_rules! prepared_lane_sequences_pending_sql { }; } +/// Rows naming what still lacks the `definer_lane_provisioning` patch: the +/// lane-sequence check and repair functions, and each installed queue-storage +/// substrate (found as for `prepared_lane_sequences`) without +/// the SECURITY DEFINER provisioning functions or whose lane helpers do not +/// create sequences through them. Lane sequences the patch could not reassign +/// do not keep it pending: `awa storage lane-sequences` reports those. +macro_rules! definer_lane_provisioning_pending_sql { + () => { + "SELECT 'awa.lane_sequence_problems'::text AS target \ + WHERE to_regprocedure('awa.lane_sequence_problems(text,text)') IS NULL \ + UNION ALL \ + SELECT 'awa.repair_lane_sequences' \ + WHERE to_regprocedure('awa.repair_lane_sequences(text)') IS NULL \ + UNION ALL \ + SELECT 'awa.refresh_lane_sequence_helpers' \ + WHERE EXISTS ( \ + SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = 'awa'::regnamespace \ + AND p.proname = 'refresh_lane_sequence_helpers' \ + AND p.prosrc NOT LIKE '%provision_lane_sequences%') \ + UNION ALL \ + SELECT format('lane provisioning in schema %s', n.nspname) \ + FROM pg_namespace AS n \ + WHERE EXISTS ( \ + SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname = 'claim_ready_runtime' \ + AND oidvectortypes(p.proargtypes) \ + = 'text, bigint, double precision, double precision') \ + AND EXISTS (SELECT 1 FROM pg_class AS c \ + WHERE c.relnamespace = n.oid AND c.relname = 'queue_enqueue_heads') \ + AND EXISTS (SELECT 1 FROM pg_class AS c \ + WHERE c.relnamespace = n.oid AND c.relname = 'queue_claim_heads') \ + AND EXISTS (SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname = 'queue_lane_sequence_name' \ + AND oidvectortypes(p.proargtypes) = 'text, text, smallint, smallint') \ + AND EXISTS (SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname = 'set_sequence_next' \ + AND oidvectortypes(p.proargtypes) = 'text, bigint') \ + AND ((SELECT count(DISTINCT p.proname) FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname IN ('provision_lane', 'provision_lane_sequences') \ + AND p.prosecdef \ + AND oidvectortypes(p.proargtypes) = 'text, smallint, smallint') < 2 \ + OR (SELECT count(DISTINCT p.proname) FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname IN ('ensure_lane_sequences', \ + 'queue_enqueue_head_sequence_sync', \ + 'queue_claim_head_sequence_sync') \ + AND p.prosrc LIKE '%provision_lane_sequences%') < 3)" + }; +} + /// SQL returning one `text` row per object or schema that still needs the /// named patch, for patches that touch more than one place. fn patch_targets_sql(name: &str) -> Option<&'static str> { @@ -303,6 +358,10 @@ fn patch_targets_sql(name: &str) -> Option<&'static str> { prepared_lane_sequences_pending_sql!(), " ORDER BY 1" )), + "definer_lane_provisioning" => Some(concat!( + definer_lane_provisioning_pending_sql!(), + " ORDER BY 1" + )), _ => None, } } @@ -347,6 +406,23 @@ pub const SCHEMA_PATCHES: &[SchemaPatch] = &[ ")" ), }, + SchemaPatch { + name: "definer_lane_provisioning", + description: "Create queue lanes as the schema owner (awa 0.7 migration v053)", + // The v023 definitions first, so the installer and the helper refresh + // carry the provisioning functions; the patch then refreshes each + // installed schema and repairs its existing lane sequences. + sql: concat!( + include_str!("../migrations/v023_install_queue_storage_substrate.sql"), + "\n", + include_str!("../migrations/patches/definer_lane_provisioning.sql"), + ), + applied_probe: concat!( + "SELECT NOT EXISTS (", + definer_lane_provisioning_pending_sql!(), + ")" + ), + }, ]; /// Schema patches a database still needs to reach this binary's schema shape. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 930aab715..e63b3943e 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -10,7 +10,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; use uuid::Uuid; @@ -2455,6 +2455,13 @@ 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` (the `definer_lane_provisioning` + /// schema patch) 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, prune_lock_timeout: Duration, } @@ -2499,6 +2506,8 @@ impl QueueStorage { shard_rotor: AtomicU16::new(0), enqueue_shards_cache: Mutex::new(HashMap::new()), ensured_lanes: Mutex::new(HashSet::new()), + lane_provisioner_present: AtomicBool::new(false), + lane_provisioner_warned: AtomicBool::new(false), prune_lock_timeout: DEFAULT_PRUNE_LOCK_TIMEOUT, }) } @@ -3458,9 +3467,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. Once the `definer_lane_provisioning` schema patch is applied, + /// 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, @@ -3499,11 +3509,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(&format!( r#" @@ -3574,6 +3608,61 @@ 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(&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 64fe8be58..85d949ee2 100644 --- a/awa-model/src/storage.rs +++ b/awa-model/src/storage.rs @@ -379,3 +379,86 @@ 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 n.nspname::text, count(*) \ + FROM pg_namespace AS n \ + CROSS JOIN LATERAL awa.lane_sequence_problems(n.nspname) AS problem \ + WHERE has_schema_privilege(current_user, n.oid, 'USAGE') \ + AND EXISTS ( \ + SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname = 'claim_ready_runtime' \ + AND oidvectortypes(p.proargtypes) \ + = 'text, bigint, double precision, double precision') \ + GROUP BY n.nspname \ + ORDER BY n.nspname", + ) + .fetch_all(pool) + .await?) +} diff --git a/awa/tests/fixtures/v0613_lane_sequence_helpers.sql b/awa/tests/fixtures/v0613_lane_sequence_helpers.sql new file mode 100644 index 000000000..ae3797c45 --- /dev/null +++ b/awa/tests/fixtures/v0613_lane_sequence_helpers.sql @@ -0,0 +1,195 @@ +-- The lane-sequence helpers as awa 0.6.13 installs them, before lane +-- provisioning functions existed. Installed under a separate name so a test +-- can put a schema back in its 0.6.13 shape. +CREATE OR REPLACE FUNCTION awa.v0613_refresh_lane_sequence_helpers(p_schema TEXT) +RETURNS VOID +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog +AS $refresh$ +BEGIN + IF p_schema IS NULL OR p_schema !~ '^[a-z_][a-z0-9_]*$' THEN + RAISE EXCEPTION 'refresh_lane_sequence_helpers: schema name must match ^[a-z_][a-z0-9_]*$; got %', + p_schema + USING ERRCODE = '22023'; + END IF; + + 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..c02d2fa04 --- /dev/null +++ b/awa/tests/lane_ownership_roles_test.rs @@ -0,0 +1,733 @@ +//! 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) -> String { + text.into() +} + +/// 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, + /// awa 0.6.13: every other schema patch, and the lane helpers before + /// lane provisioning. + 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?; + for (_, _, statement) in migrations::migration_sql_range(0, migrations::CURRENT_VERSION) { + migrate_as_owner.execute(sqlx::raw_sql(&statement)).await?; + } + for (name, _, statement) in migrations::schema_patch_sql(migrations::SCHEMA_PATCHES) { + if schema == Schema::BeforeProvisioning && name == "definer_lane_provisioning" { + continue; + } + migrate_as_owner.execute(sqlx::raw_sql(&statement)).await?; + } + if schema == Schema::BeforeProvisioning { + // Every patch carries the current v023 definitions, so put the + // 0.6.13 helpers back and drop what 0.6.13 did not have. + migrate_as_owner + .execute(sqlx::raw_sql(include_str!( + "fixtures/v0613_lane_sequence_helpers.sql" + ))) + .await?; + migrate_as_owner + .execute(sqlx::raw_sql( + r#" + SELECT awa.v0613_refresh_lane_sequence_helpers('awa'); + DROP FUNCTION awa.v0613_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}" + ); + assert_eq!( + migrations::pending_schema_patches(&cluster.owner) + .await? + .iter() + .map(|patch| patch.name) + .collect::>(), + vec!["definer_lane_provisioning"] + ); + + // `awa migrate` 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. The leftovers do not keep the patch pending. + let report = migrations::run_with_report(&cluster.migrator).await?; + assert!(report.complete(), "{:?}", report.patches); + assert!(migrations::pending_schema_patches(&cluster.owner) + .await? + .is_empty()); + 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/security.md b/docs/security.md index a146635bb..a6deef52b 100644 --- a/docs/security.md +++ b/docs/security.md @@ -174,23 +174,37 @@ ALTER DEFAULT PRIVILEGES FOR ROLE awa_migrator IN SCHEMA my_qs_schema The default `awa` queue-storage substrate is migrated by `awa migrate`. Custom queue-storage schemas are migrated by `awa storage prepare-queue-storage-schema`, which should also run as the migrator role and creates objects owned by `awa_migrator` / `awa_owner`. The runtime role needs read/write/execute privileges plus `TRUNCATE` for ring-partition rotation; it never needs DDL. -### Provision queues before starting restricted runtimes +### Lane sequences -Lane sequences are created lazily by the first producer or worker that touches a queue, which needs schema `CREATE`. A runtime role without it needs each queue provisioned in advance, and needs the `prepared_lane_sequences` schema patch: before it, the lane helpers issue `CREATE SEQUENCE IF NOT EXISTS` even for sequences that already exist, which requires schema `CREATE` regardless. Apply the patch with `awa migrate` (it applies on any v040 database), or run the exported `R__prepared_lane_sequences.sql` from `awa migrate --sql` / `--extract-to` / `awa.schema_patches()` after `V40`. +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. Once the `definer_lane_provisioning` schema patch is applied, 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 role therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. Apply the patch with `awa migrate` as the schema owner (`PGOPTIONS='-c role=awa_owner'`), or run the exported `R__definer_lane_provisioning.sql` after `V40`. -Then provision every queue as the migrator before starting its producers or workers: +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. The table grants above are enough. + +`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 the patch 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 -PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-queue \ - --queue email --enqueue-shards 1 +awa --database-url "$AWA_DATABASE_URL" storage lane-sequences --schema awa ``` -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`, and match `--queue-stripe-count` to the runtime's queue-storage configuration. With `PartitionedQueue`, provision each physical queue. Use `--schema` for a custom queue-storage schema. The Rust equivalent is `QueueStorage::prepare_queue`. +```sql +SELECT * FROM awa.lane_sequence_problems('awa'); +``` + +The command exits 2 while any remain. Applying the patch runs the repair as the applying 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 patch. 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 lane-sequences --repair +``` -Lane cursors need `UPDATE` on their sequences for `setval`, in addition to `USAGE` and `SELECT`; use the sequence grants above. A restricted runtime that reaches an unprovisioned lane fails with SQLSTATE `42501` and a provisioning hint. Schema owners keep lazy creation, so single-role installations need no extra step. +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: -On a database where a runtime role with `CREATE` has been creating lanes (any 0.6.9 deployment that ran without schema ownership), the runtime owns those lane sequences. `GRANT ... ON ALL SEQUENCES` run by the owner fails on them, so grant per sequence the owner owns, verify, and only then revoke `CREATE`. [Upgrading from 0.6.9 to 0.6.10](upgrade-0.6.9-to-0.6.10.md#optional-run-the-runtime-role-without-schema-create) has the full procedure and the SQL. +- `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. ### 5. Configure your processes diff --git a/docs/upgrade-0.6.13-to-0.6.14.md b/docs/upgrade-0.6.13-to-0.6.14.md new file mode 100644 index 000000000..1cccbe370 --- /dev/null +++ b/docs/upgrade-0.6.13-to-0.6.14.md @@ -0,0 +1,45 @@ +# Upgrading from 0.6.13 to 0.6.14 + +0.6.14 is a rolling patch upgrade. Traffic keeps running throughout, and you can roll binaries back to 0.6.13 with the schema patch left in place. + +## Upgrade + +1. Install the 0.6.14 CLI. +2. Check lane ownership. This is read-only and exits 2 while lanes need repair: + + ```bash + awa --database-url "$AWA_DATABASE_URL" storage lane-sequences + ``` + +3. Run `awa migrate` once, with the schema owner as the effective role, while workers keep running: + + ```bash + PGOPTIONS='-c role=awa_owner' \ + awa --database-url "$AWA_MIGRATOR_DATABASE_URL" migrate + ``` + +4. Read the output. The `definer_lane_provisioning` patch reports `applied`. If it could not reassign some lane sequences, `awa migrate` prints a warning; run `awa storage lane-sequences` for the statements to run and the role to run them as (see [Check and repair lane ownership](security.md#check-and-repair-lane-ownership)). The migration exit status does not depend on them. +5. Roll workers, producers and admin processes to 0.6.14, in any order. + +## What `awa migrate` changes + +The schema stays at v40. The `definer_lane_provisioning` patch installs `provision_lane` and `provision_lane_sequences`, two `SECURITY DEFINER` functions owned by the schema owner, in `awa` and every other queue-storage schema. It changes the three lane helpers to create a missing sequence through them, and adds the read-only `awa.lane_sequence_problems(schema)` and `awa.repair_lane_sequences(schema)`. Every new lane sequence then belongs to the schema owner and carries the privileges of the head table it backs, whichever role creates the lane: a producer enqueueing to a new queue, priority or shard, the maintenance leader promoting into a shard after `enqueue_shards` was raised, or `awa storage prepare-queue` run as any role. + +The patch then repairs existing lane sequences. Each one not owned by the schema owner is reassigned if the applying role can `SET ROLE` to both its current owner and the schema owner. Missing grants are added if the applying role can act as the owner. A sequence it cannot change is reported and left as it was. The patch never fails because of one. + +The patch: + +- replaces functions and changes only the lane sequences that need it. `ALTER SEQUENCE ... OWNER` holds that sequence's lock until the patch commits, so enqueues and claims on those lanes wait for the commit. Correctly owned lanes are not touched; +- keeps every lane's sequence position; +- leaves storage routing unchanged; +- applies as one transaction, so an interrupted run leaves nothing applied; run it again. + +0.6.13 and earlier binaries call the same helpers, so lanes they create after the patch are owned by the schema owner too. A 0.6.14 binary against a database without the patch detects that `provision_lane` is missing and creates lanes as 0.6.13 did. + +### External migration runners + +`awa migrate --sql --pending` prints the `R__definer_lane_provisioning` patch on a 0.6.13 database; with `--extract-to`, apply `R__definer_lane_provisioning.sql` as a repeatable script after `V40`. Run it as the schema owner, then run `awa storage lane-sequences` to check for leftovers. + +## Roll back + +Roll binaries back to 0.6.13 at any time. Leave the patch in place: 0.6.13 calls the same helpers with the same arguments. From 5c03aed269d10f7f4a7cc438ed416f5d97b6c8ae Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 13:35:43 +1300 Subject: [PATCH 05/16] test: expect restricted runtimes to create lanes through the schema owner --- .../tests/least_privilege_runtime_test.rs | 35 ++++++++++++------- 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/awa-model/tests/least_privilege_runtime_test.rs b/awa-model/tests/least_privilege_runtime_test.rs index 717707c3a..11e37b18d 100644 --- a/awa-model/tests/least_privilege_runtime_test.rs +++ b/awa-model/tests/least_privilege_runtime_test.rs @@ -176,20 +176,25 @@ async fn check_runtime(pool: &PgPool, schema: &str) -> TestResult { ), ) .await?; - let error = run_sql( + // An unprepared lane is created through the schema owner's provisioning + // function, so the runtime still issues no DDL of its own. + run_sql( pool, format!( - "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint)" + "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint); \ + SELECT {schema}.provision_lane('unprepared_queue', 2::smallint, 0::smallint); \ + SELECT {schema}.reserve_enqueue_seq('unprepared_queue', 2::smallint, 0::smallint, 1);" ), ) - .await - .unwrap_err(); - assert_eq!(database_error_code(&error).as_deref(), Some("42501")); - assert!(error - .as_database_error() - .unwrap() - .message() - .contains("has not been provisioned")); + .await?; + let foreign_owned: i64 = sqlx::query_scalar(&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); let error = run_sql( pool, format!("CREATE TABLE {schema}.must_not_create(id integer)"), @@ -354,10 +359,12 @@ async fn exercise_stale_custom_schema(options: PgConnectOptions, migrator: &str) .replace("__schema__", "custom_lanes"), ) .await?; + // The old helper also predates lane provisioning, so both patches that + // refresh the helpers name the schema. let pending = migrations::pending_schema_patches(&pool).await?; assert_eq!( pending.iter().map(|patch| patch.name).collect::>(), - vec!["prepared_lane_sequences"] + vec!["prepared_lane_sequences", "definer_lane_provisioning"] ); let status = lane_patch_status(&pool).await?; assert!(!status.applied); @@ -478,10 +485,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.9 runtime with CREATE did. let runtime_pool = PgPoolOptions::new() .max_connections(2) .connect_with(options.username(runtime)) From 32d671b061e602119793b4e4b30f897a319262de Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 14:04:19 +1300 Subject: [PATCH 06/16] 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. The released-binary rehearsal now expects a 0.6.9 producer without CREATE to create a new lane, owned by the schema owner. --- .../v023_install_queue_storage_substrate.sql | 4 +- awa/tests/patch_upgrade_rehearsal_test.rs | 44 ++++++++++++------- docs/security.md | 2 +- 3 files changed, 31 insertions(+), 19 deletions(-) diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index 0d170c34f..bcc8fe040 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -283,8 +283,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/awa/tests/patch_upgrade_rehearsal_test.rs b/awa/tests/patch_upgrade_rehearsal_test.rs index 4bd2a761a..aef1255af 100644 --- a/awa/tests/patch_upgrade_rehearsal_test.rs +++ b/awa/tests/patch_upgrade_rehearsal_test.rs @@ -1108,12 +1108,23 @@ async fn released_0_6_9_patch_upgrade_and_rollback() { killed_sources.insert(worker_c.kill().await); killed_sources.insert(worker_d.kill().await); - let worker_e = spawn_released_worker(&fixture.runtime_url, &log, &queues).await; - let worker_f = spawn_released_worker(&fixture.runtime_url, &log, &queues).await; + let restricted_queues = [ + MAIN, + BANKED, + NEW_AFTER_PATCH, + NEW_AFTER_ROLLBACK, + PROVISIONED, + UNPROVISIONED, + ]; + let worker_e = spawn_released_worker(&fixture.runtime_url, &log, &restricted_queues).await; + let worker_f = spawn_released_worker(&fixture.runtime_url, &log, &restricted_queues).await; + // A released producer without CREATE creates an unprepared queue's lane + // through the schema owner's provisioning function. for (queue, seq_start) in [ (MAIN, 9_000), (PROVISIONED, 10_000), (NEW_AFTER_ROLLBACK, 11_000), + (UNPROVISIONED, 12_000), ] { released_produce( &fixture.runtime_url, @@ -1130,21 +1141,22 @@ async fn released_0_6_9_patch_upgrade_and_rollback() { .await .unwrap_or_else(|error| panic!("restricted released enqueue on {queue}: {error}")); } - let refused = released_produce( - &fixture.runtime_url, - &log, - ProduceSpec { - queue: UNPROVISIONED, - seq_start: 12_000, - count: 1, - modes: "ok", - delay_s: 0.0, - pause_ms: 0, - }, - ) + let new_lane_owners: Vec = sqlx::query_scalar(&format!( + "SELECT sequence.relowner = namespace.nspowner \ + FROM {SCHEMA}.queue_enqueue_heads AS heads \ + JOIN pg_namespace AS namespace ON namespace.nspname = '{SCHEMA}' \ + JOIN pg_class AS sequence \ + ON sequence.relnamespace = namespace.oid AND sequence.relname = heads.seq_name \ + WHERE heads.queue = $1" + )) + .bind(UNPROVISIONED) + .fetch_all(&fixture.admin) .await - .expect_err("an unprovisioned queue needs provisioning once CREATE is revoked"); - assert!(refused.contains("has not been provisioned"), "{refused}"); + .unwrap(); + assert!( + !new_lane_owners.is_empty() && new_lane_owners.iter().all(|owned| *owned), + "the new lane's sequences belong to the schema owner: {new_lane_owners:?}" + ); wait_for_drain(&log, Duration::from_secs(120)).await; let stats = cli_ok(&released_cli, &fixture.runtime_url, &["queue", "stats"]).await; assert!(stats.contains(PROVISIONED), "{stats}"); diff --git a/docs/security.md b/docs/security.md index a6deef52b..8c56004f1 100644 --- a/docs/security.md +++ b/docs/security.md @@ -176,7 +176,7 @@ The default `awa` queue-storage substrate is migrated by `awa migrate`. Custom q ### Lane sequences -Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. Once the `definer_lane_provisioning` schema patch is applied, 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 role therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. Apply the patch with `awa migrate` as the schema owner (`PGOPTIONS='-c role=awa_owner'`), or run the exported `R__definer_lane_provisioning.sql` after `V40`. +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. Once the `definer_lane_provisioning` schema patch is applied, 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 role therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. Apply the patch with `awa migrate` as the schema owner (`PGOPTIONS='-c role=awa_owner'`), or run the exported `R__definer_lane_provisioning.sql` after `V40`. 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. The table grants above are enough. From 2ce3a31cf6344663cdf19bb833d58c8e6e718c96 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 14:56:16 +1300 Subject: [PATCH 07/16] fix(queue-storage): promote each enqueue lane in its own transaction and hand busy lanes to a waiter (0.6) Backport of #636. Promotion re-enqueues due scheduled and retryable rows into their lane, which takes the lane's reservation lock that producers hold until they commit. One promotion batch was a single transaction across every queue it touched, and the maintenance leader runs its branches one after another, so a lane kept busy by producers delayed due jobs of unrelated queues in the same batch and started rescue, rotation and prune late. QueueStorage::promote_due_lanes selects a batch without locking it and promotes each (queue, priority) lane's share in its own transaction, under an optional lock_timeout; a lane that hits it rolls back and is reported as contended. promote_due_lane promotes the oldest due rows of one lane. promote_due keeps its signature and promotes lane by lane with no timeout. Within a lane, promoted rows take lane sequence numbers in due order. The leader's promoter runs each tick under a 50 ms lane lock timeout, skips contended lanes for the rest of the tick, and hands each to a dedicated waiter task (at most four at once) that promotes the lane under a 30 s lock timeout. PostgreSQL grants a lock to its waiters in arrival order, so the waiter is served after the producers queued ahead of it however busy the lane stays. awa.maintenance.promote.lane_contention counts handed_off, waiter_limit and waiter_timed_out by state and queue. The overrunning-promotion test slowed the queue-storage DELETE, which no longer runs when nothing is due; it now slows the canonical-plane promotion DELETE that every body runs. --- CHANGELOG.md | 21 + Cargo.lock | 1 + awa-metrics/src/lib.rs | 87 ++++ awa-model/src/lib.rs | 4 +- awa-model/src/queue_storage.rs | 263 ++++++++-- awa-worker/src/maintenance.rs | 450 +++++++++++++----- awa/Cargo.toml | 1 + .../enqueue_lane_maintenance_coupling_test.rs | 427 +++++++++++++++++ awa/tests/queue_storage_runtime_test.rs | 323 ++++++++++++- docs/architecture.md | 2 +- docs/configuration.md | 1 + docs/grafana/README.md | 1 + 12 files changed, 1412 insertions(+), 169 deletions(-) create mode 100644 awa/tests/enqueue_lane_maintenance_coupling_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a17deb..81e2d5d61 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,27 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **On queue storage, a queue whose enqueue lane is busy no longer holds + back promotion of other queues, nor the maintenance leader's rescue, + rotation and prune.** Promotion re-enqueues due scheduled and retryable + rows into their lane, which takes the lane's reservation lock that + producers hold until they commit. A promotion batch was one transaction + across every queue it touched, and the leader runs its branches one after + another, so a lane kept busy by producers delayed due jobs of unrelated + queues in the same batch, and rescue started seconds late. Promotion now + gives each `(queue, priority)` lane its own transaction under a 50 ms + `lock_timeout`. A lane that times out keeps its rows deferred and is + handed to a dedicated task that waits up to 30 s for the lock, queued + behind the producers already waiting, so the lane is still promoted in + turn however busy it stays; the leader skips it meanwhile. At most four + lanes wait this way at once, and a further contended lane is retried on + the next tick. `awa.maintenance.promote.lane_contention` + (`awa_maintenance_promote_lane_contention_total{awa_job_state,awa_job_queue,awa_promote_outcome}`) + counts `handed_off`, `waiter_limit` and `waiter_timed_out`. Promoted rows + now take lane sequence numbers in due order within their lane. + `QueueStorage::promote_due` keeps its signature and promotes lane by lane + without a lock timeout; `promote_due_lanes` and `promote_due_lane` expose + the per-lane form. - **`admin::job_kind_overviews` on queue storage no longer depends on the planner's row estimates.** Behind the admin UI's kinds page, it counted every job through the `current_jobs` view. Two nodes of that view are diff --git a/Cargo.lock b/Cargo.lock index 30bd194c9..ce2e43e60 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -400,6 +400,7 @@ dependencies = [ "serde_json", "sqlx", "tokio", + "tokio-util", "tracing", "tracing-opentelemetry", "tracing-subscriber", diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 3f81abce3..a3d015b26 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -64,6 +64,7 @@ pub mod names { pub const MAINTENANCE_PROMOTE_BATCHES: &str = "awa.maintenance.promote_batches"; pub const MAINTENANCE_PROMOTE_BATCH_SIZE: &str = "awa.maintenance.promote_batch_size"; pub const MAINTENANCE_PROMOTE_DURATION: &str = "awa.maintenance.promote_duration"; + pub const MAINTENANCE_PROMOTE_LANE_CONTENTION: &str = "awa.maintenance.promote.lane_contention"; pub const MAINTENANCE_BRANCH_DURATION: &str = "awa.maintenance.branch.duration"; pub const MAINTENANCE_BRANCH_OVERRUN: &str = "awa.maintenance.branch.overrun"; pub const MAINTENANCE_ROTATE_ATTEMPTS: &str = "awa.maintenance.rotate.attempts"; @@ -162,6 +163,11 @@ pub struct AwaMetrics { pub promotion_batch_size: Histogram, /// Promotion query duration. pub promotion_duration_seconds: Histogram, + /// Promotion lanes whose locks were busy past the leader's wait budget, + /// attributed to (state, queue, outcome): `handed_off` to a dedicated + /// waiter, `waiter_limit` when every waiter was busy, `waiter_timed_out` + /// when a waiter's own budget ran out. + pub promotion_lane_contention: Counter, /// Current in-flight jobs (can go up and down). pub jobs_in_flight: UpDownCounter, /// Total heartbeat batches sent. @@ -404,6 +410,11 @@ impl AwaMetrics { .with_unit("s") .with_boundaries(LATENCY_BUCKETS_SECONDS.to_vec()) .build(), + promotion_lane_contention: meter + .u64_counter(names::MAINTENANCE_PROMOTE_LANE_CONTENTION) + .with_description("Promotion lanes whose locks were busy, by state/queue/outcome") + .with_unit("{lane}") + .build(), jobs_in_flight: meter .i64_up_down_counter(names::JOB_IN_FLIGHT) .with_description("Current number of in-flight jobs") @@ -749,6 +760,23 @@ impl AwaMetrics { .record(duration.as_secs_f64(), &attrs); } + /// Record a promotion lane whose locks were busy past the leader's wait + /// budget, with what became of it (`handed_off`, `waiter_limit`, + /// `waiter_timed_out`). + pub fn record_promotion_lane_contention( + &self, + state: &str, + queue: &str, + outcome: &'static str, + ) { + let attrs = [ + opentelemetry::KeyValue::new("awa.job.state", state.to_string()), + opentelemetry::KeyValue::new("awa.job.queue", queue.to_string()), + opentelemetry::KeyValue::new("awa.promote.outcome", outcome), + ]; + self.promotion_lane_contention.add(1, &attrs); + } + /// Record in-flight change. pub fn record_in_flight_change(&self, queue: &str, delta: i64) { let attrs = [opentelemetry::KeyValue::new( @@ -1162,6 +1190,65 @@ mod tests { metrics.record_job_kind_info("minimal", None, None, None, None, &[]); } + #[test] + fn promotion_lane_contention_counts_each_outcome_per_state_and_queue() { + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + let metrics = AwaMetrics::new(&provider.meter("awa-metrics-test")); + + metrics.record_promotion_lane_contention("scheduled", "emails", "handed_off"); + metrics.record_promotion_lane_contention("scheduled", "emails", "handed_off"); + metrics.record_promotion_lane_contention("retryable", "emails", "waiter_limit"); + metrics.record_promotion_lane_contention("scheduled", "billing", "waiter_timed_out"); + provider.force_flush().expect("flush metrics"); + let resource_metrics = exporter.get_finished_metrics().expect("read metrics"); + + let mut counts = BTreeMap::new(); + for resource in &resource_metrics { + for scope in resource.scope_metrics() { + for metric in scope.metrics() { + if let ( + names::MAINTENANCE_PROMOTE_LANE_CONTENTION, + AggregatedMetrics::U64(MetricData::Sum(sum)), + ) = (metric.name(), metric.data()) + { + for point in sum.data_points() { + let label = |key: &str| { + point + .attributes() + .find(|kv| kv.key.as_str() == key) + .map(|kv| kv.value.as_str().into_owned()) + .unwrap_or_default() + }; + counts.insert( + ( + label("awa.job.state"), + label("awa.job.queue"), + label("awa.promote.outcome"), + ), + point.value(), + ); + } + } + } + } + } + + let key = |state: &str, queue: &str, outcome: &str| { + (state.to_string(), queue.to_string(), outcome.to_string()) + }; + assert_eq!( + counts, + BTreeMap::from([ + (key("retryable", "emails", "waiter_limit"), 1), + (key("scheduled", "billing", "waiter_timed_out"), 1), + (key("scheduled", "emails", "handed_off"), 2), + ]) + ); + } + /// A second-scale observation must not fall in the SDK's default /// millisecond-scale buckets (upper bounds 5, 10, 25, ...), and every /// histogram must export the intended explicit boundaries. diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index 267495146..88d658d05 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -45,8 +45,8 @@ pub use partitioned_queue::{ partition_for_ordering_key, partition_hash64, PartitionedQueue, PartitionedQueueError, }; pub use queue_storage::{ - ClaimedEntry, ClaimedRuntimeJob, PruneOutcome, QueueCounts, QueueStorage, QueueStorageConfig, - RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, + ClaimedEntry, ClaimedRuntimeJob, PromoteOutcome, PromotionLane, PruneOutcome, QueueCounts, + QueueStorage, QueueStorageConfig, RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, }; pub use storage::StorageStatus; diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 930aab715..f598d909f 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -362,6 +362,29 @@ enum TerminalDeltaSlotRollup { Blocked, } +/// A `(queue, priority)` enqueue lane as promotion addresses it. Promotion +/// picks the shard when it re-enqueues a row, so the shard is not part of it. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PromotionLane { + pub queue: String, + pub priority: i16, +} + +/// Outcome of [`QueueStorage::promote_due_lanes`] and +/// [`QueueStorage::promote_due_lane`]. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct PromoteOutcome { + /// Rows moved into the ready ring. + pub promoted: usize, + /// Due rows the batch selected, including those left deferred on a + /// contended lane. A batch that selected its full size may have left + /// further due rows behind. + pub selected: usize, + /// Lanes whose promotion transaction hit `lock_timeout` and rolled back. + /// Their rows stay deferred. + pub contended: Vec, +} + /// Discriminator for [`PruneOutcome::SkippedActive`]. /// /// Multiple gates can fire `SkippedActive` for the same ring (e.g. queue @@ -420,7 +443,7 @@ fn is_lock_contention_error(err: &sqlx::Error) -> bool { ) } -async fn set_prune_lock_timeout_tx( +async fn set_lock_timeout_tx( tx: &mut sqlx::Transaction<'_, Postgres>, timeout: Duration, ) -> Result<(), AwaError> { @@ -2124,6 +2147,15 @@ struct ExistingReadyRow { payload: serde_json::Value, } +/// Which of a lane's due rows one promotion transaction takes. +#[derive(Debug, Clone, Copy)] +enum DueLaneRows<'a> { + /// These rows, picked by a batch across lanes, where still due. + Selected(&'a [i64]), + /// The oldest due rows, up to this many. + Oldest(i64), +} + #[derive(Debug, Clone, sqlx::FromRow)] struct DeletedLeaseRow { ready_slot: i32, @@ -14433,32 +14465,205 @@ impl QueueStorage { Ok(rescued) } + /// Promote up to `batch_size` due `state` rows into the ready ring. + /// Equivalent to [`Self::promote_due_lanes`] with no lock timeout and no + /// skipped lanes, returning the number of rows promoted. pub async fn promote_due( &self, pool: &PgPool, state: JobState, batch_size: i64, ) -> Result { + Ok(self + .promote_due_lanes(pool, state, batch_size, None, &[]) + .await? + .promoted) + } + + /// Select up to `batch_size` due `state` rows, oldest `run_at` first, + /// outside the lanes in `skip`, and promote each lane's share in its own + /// transaction. + /// + /// Re-enqueueing a row takes its lane's enqueue reservation lock, which + /// producers hold until they commit, so a lane under heavy enqueue load + /// can keep its promotion waiting. A transaction per lane confines that + /// wait to the lane's own rows. With `lane_lock_timeout`, each lane's + /// transaction runs under that `lock_timeout`; a lane that hits it rolls + /// back, keeps its rows deferred and is reported in + /// [`PromoteOutcome::contended`], and the remaining lanes go ahead. + pub async fn promote_due_lanes( + &self, + pool: &PgPool, + state: JobState, + batch_size: i64, + lane_lock_timeout: Option, + skip: &[PromotionLane], + ) -> Result { if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 { - return Ok(0); + return Ok(PromoteOutcome::default()); } let schema = self.schema(); + let due: Vec<(i64, String, i16)> = sqlx::query_as(&format!( + r#" + SELECT job_id, queue, priority + FROM {schema}.deferred_jobs AS deferred + WHERE state = $1 + AND run_at <= clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM awa.queue_meta + WHERE queue = deferred.queue AND paused = TRUE + ) + AND NOT EXISTS ( + SELECT 1 + FROM unnest($3::text[], $4::smallint[]) AS skip (queue, priority) + WHERE skip.queue = deferred.queue AND skip.priority = deferred.priority + ) + ORDER BY run_at ASC, priority ASC, job_id ASC + LIMIT $2 + "# + )) + .bind(state) + .bind(batch_size) + .bind( + skip.iter() + .map(|lane| lane.queue.as_str()) + .collect::>(), + ) + .bind(skip.iter().map(|lane| lane.priority).collect::>()) + .fetch_all(pool) + .await + .map_err(map_sqlx_error)?; + + let mut outcome = PromoteOutcome { + selected: due.len(), + ..PromoteOutcome::default() + }; + let mut lanes: BTreeMap> = BTreeMap::new(); + for (job_id, queue, priority) in due { + lanes + .entry(PromotionLane { queue, priority }) + .or_default() + .push(job_id); + } + for (lane, job_ids) in lanes { + match self + .promote_lane( + pool, + state, + &lane, + DueLaneRows::Selected(&job_ids), + lane_lock_timeout, + ) + .await? + { + Some(promoted) => outcome.promoted += promoted, + None => outcome.contended.push(lane), + } + } + Ok(outcome) + } + + /// Promote up to `batch_size` due `state` rows of one lane, oldest + /// `run_at` first, in one transaction. With `lane_lock_timeout`, a + /// transaction that hits it rolls back and the lane is returned in + /// [`PromoteOutcome::contended`] with nothing selected. + pub async fn promote_due_lane( + &self, + pool: &PgPool, + state: JobState, + lane: &PromotionLane, + batch_size: i64, + lane_lock_timeout: Option, + ) -> Result { + if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 { + return Ok(PromoteOutcome::default()); + } + Ok( + match self + .promote_lane( + pool, + state, + lane, + DueLaneRows::Oldest(batch_size), + lane_lock_timeout, + ) + .await? + { + Some(promoted) => PromoteOutcome { + promoted, + selected: promoted, + contended: Vec::new(), + }, + None => PromoteOutcome { + contended: vec![lane.clone()], + ..PromoteOutcome::default() + }, + }, + ) + } + + /// One lane's promotion transaction. `None` when it hit `lock_timeout` + /// and rolled back. + async fn promote_lane( + &self, + pool: &PgPool, + state: JobState, + lane: &PromotionLane, + rows: DueLaneRows<'_>, + lock_timeout: Option, + ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let moved: Vec = sqlx::query_as(&format!( + match self + .promote_lane_tx(&mut tx, state, lane, rows, lock_timeout) + .await + { + Ok(promoted) => match tx.commit().await { + Ok(()) => Ok(Some(promoted)), + Err(err) if is_lock_contention_error(&err) => Ok(None), + Err(err) => Err(map_sqlx_error(err)), + }, + Err(AwaError::Database(err)) if is_lock_contention_error(&err) => { + tx.rollback().await.map_err(map_sqlx_error)?; + Ok(None) + } + Err(err) => Err(err), + } + } + + async fn promote_lane_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + state: JobState, + lane: &PromotionLane, + rows: DueLaneRows<'_>, + lock_timeout: Option, + ) -> Result { + if let Some(timeout) = lock_timeout { + set_lock_timeout_tx(tx, timeout).await?; + } + + let schema = self.schema(); + let (row_filter, limit) = match rows { + DueLaneRows::Selected(_) => ("AND job_id = ANY($4)", ""), + DueLaneRows::Oldest(_) => ("", "ORDER BY run_at ASC, job_id ASC LIMIT $4"), + }; + let sql = format!( r#" DELETE FROM {schema}.deferred_jobs WHERE job_id IN ( SELECT job_id - FROM {schema}.deferred_jobs + FROM {schema}.deferred_jobs AS deferred WHERE state = $1 + AND queue = $2 + AND priority = $3 + {row_filter} AND run_at <= clock_timestamp() AND NOT EXISTS ( SELECT 1 FROM awa.queue_meta - WHERE queue = {schema}.deferred_jobs.queue AND paused = TRUE + WHERE queue = deferred.queue AND paused = TRUE ) - ORDER BY run_at ASC, priority ASC, job_id ASC - LIMIT $2 + {limit} FOR UPDATE SKIP LOCKED ) RETURNING @@ -14479,21 +14684,26 @@ impl QueueStorage { unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload "# - )) - .bind(state) - .bind(batch_size) - .fetch_all(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - + ); + let query = sqlx::query_as::<_, DeferredJobRow>(&sql) + .bind(state) + .bind(&lane.queue) + .bind(lane.priority); + let query = match rows { + DueLaneRows::Selected(job_ids) => query.bind(job_ids), + DueLaneRows::Oldest(batch_size) => query.bind(batch_size), + }; + let mut moved: Vec = + query.fetch_all(tx.as_mut()).await.map_err(map_sqlx_error)?; if moved.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; return Ok(0); } + // DELETE ... RETURNING has no order; lane sequence follows due order. + moved.sort_unstable_by_key(|row| (row.run_at, row.job_id)); + let promoted = moved.len(); let ready_rows: Vec = moved - .iter() - .cloned() + .into_iter() .map(|row| ExistingReadyRow { job_id: row.job_id, kind: row.kind, @@ -14511,15 +14721,10 @@ impl QueueStorage { payload: row.payload, }) .collect(); - let queues = ready_rows - .iter() - .map(|row| row.queue.clone()) - .collect::>(); - self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(state)) + self.insert_existing_ready_rows_tx(tx, ready_rows, Some(state)) .await?; - self.notify_queues_tx(&mut tx, queues).await?; - tx.commit().await.map_err(map_sqlx_error)?; - Ok(moved.len()) + self.notify_queues_tx(tx, vec![lane.queue.clone()]).await?; + Ok(promoted) } async fn relation_has_rows_tx( @@ -15113,7 +15318,7 @@ impl QueueStorage { return Ok(TerminalDeltaSlotRollup::SkippedActive); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_delta = sqlx::query(&format!( "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" @@ -15411,7 +15616,7 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_tables = sqlx::query(&format!( "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE" @@ -15796,7 +16001,7 @@ impl QueueStorage { return Ok(PruneOutcome::Noop); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_table = sqlx::query(&format!( "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" @@ -16101,7 +16306,7 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_tables = sqlx::query(&format!( "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index b7c134177..e2a5a5eb3 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -1,4 +1,5 @@ use crate::executor::DlqPolicy; +use crate::metrics::AwaMetrics; use crate::runtime::InFlightMap; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; use awa_model::cron::{ @@ -7,7 +8,8 @@ use awa_model::cron::{ #[cfg(test)] use awa_model::SkipReason; use awa_model::{ - JobRow, JobState, PeriodicJob, PruneOutcome, RotateOutcome, TerminalDeltaRollupOutcome, + AwaError, JobRow, JobState, PeriodicJob, PromotionLane, PruneOutcome, QueueStorage, + RotateOutcome, TerminalDeltaRollupOutcome, }; use chrono::Utc; use croner::Cron; @@ -17,7 +19,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; -use tokio::task::JoinHandle; +use tokio::task::{JoinHandle, JoinSet}; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, warn}; use uuid::Uuid; @@ -719,6 +721,224 @@ const CRON_CATCH_UP_LIMIT: usize = 1_000; const TERMINAL_COUNT_ROLLUP_MAX_SLOTS_PER_TICK: usize = 4; type QueueStorageMetricRow = (String, i64, i64, i64, i64, i64, i64, i64, Option); +/// How long a promotion transaction waits on one lane's locks before the +/// leader hands the lane to a dedicated waiter and goes on without it. +const PROMOTE_LANE_LOCK_TIMEOUT: Duration = Duration::from_millis(50); +/// Lock wait of a dedicated lane waiter. PostgreSQL grants a lock to its +/// waiters in arrival order, so a waiter is served after the producers queued +/// ahead of it, however busy the lane stays. +const PROMOTE_LANE_WAITER_LOCK_TIMEOUT: Duration = Duration::from_secs(30); +/// Most lanes with a dedicated waiter at once. A waiter holds a pool +/// connection while it waits. +const PROMOTE_MAX_LANE_WAITERS: usize = 4; + +/// Queue-storage promotion of due scheduled and retryable rows, as the +/// maintenance leader runs it for one store. +/// +/// Each lane's share of a batch is promoted in its own transaction under a +/// short `lock_timeout`, so a lane whose enqueue lock is busy delays only its +/// own rows. A lane that times out keeps its rows deferred and is handed to a +/// dedicated waiter task, which promotes it under a long lock timeout while +/// the leader carries on with other lanes and other maintenance branches. +/// Dropping the promoter aborts its waiters. +#[doc(hidden)] +pub struct QueueStoragePromoter { + pool: PgPool, + store: Arc, + metrics: AwaMetrics, + waiting: Arc>>, + waiters: std::sync::Mutex>, +} + +impl QueueStoragePromoter { + pub fn new(pool: PgPool, store: Arc, metrics: AwaMetrics) -> Self { + Self { + pool, + store, + metrics, + waiting: Arc::default(), + waiters: std::sync::Mutex::new(JoinSet::new()), + } + } + + /// One promotion tick for `state`: batches until one selects fewer rows + /// than the batch size, each skipping lanes that have a waiter or were + /// contended earlier in the tick. Returns the rows promoted by the tick + /// itself, not by waiters. + pub async fn promote_due( + &self, + state: JobState, + cancel: &CancellationToken, + ) -> Result { + { + let mut waiters = self.waiters.lock().expect("lane waiters poisoned"); + while waiters.try_join_next().is_some() {} + } + + let mut contended: Vec = Vec::new(); + let mut promoted = 0usize; + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + if cancel.is_cancelled() { + break; + } + let mut skip = self.waiting_lanes(state); + skip.extend(contended.iter().cloned()); + let batch_started = Instant::now(); + let outcome = self + .store + .promote_due_lanes( + &self.pool, + state, + PROMOTE_BATCH_SIZE, + Some(PROMOTE_LANE_LOCK_TIMEOUT), + &skip, + ) + .await?; + self.metrics.record_promotion_batch( + state.as_str(), + outcome.promoted as u64, + batch_started.elapsed(), + ); + promoted += outcome.promoted; + for lane in outcome.contended { + self.hand_off(state, lane.clone()); + contended.push(lane); + } + if (outcome.selected as i64) < PROMOTE_BATCH_SIZE { + break; + } + } + Ok(promoted) + } + + fn waiting_lanes(&self, state: JobState) -> Vec { + self.waiting + .lock() + .expect("waiting lanes poisoned") + .iter() + .filter(|(waiting_state, _)| *waiting_state == state) + .map(|(_, lane)| lane.clone()) + .collect() + } + + fn hand_off(&self, state: JobState, lane: PromotionLane) { + let key = (state, lane); + { + let mut waiting = self.waiting.lock().expect("waiting lanes poisoned"); + if waiting.len() >= PROMOTE_MAX_LANE_WAITERS { + self.metrics.record_promotion_lane_contention( + state.as_str(), + &key.1.queue, + "waiter_limit", + ); + debug!( + queue = %key.1.queue, + priority = key.1.priority, + "Promotion lane contended with every lane waiter busy; retrying next tick" + ); + return; + } + if !waiting.insert(key.clone()) { + return; + } + } + self.metrics + .record_promotion_lane_contention(state.as_str(), &key.1.queue, "handed_off"); + let guard = WaitingLane { + waiting: Arc::clone(&self.waiting), + key, + }; + let pool = self.pool.clone(); + let store = Arc::clone(&self.store); + let metrics = self.metrics.clone(); + self.waiters + .lock() + .expect("lane waiters poisoned") + .spawn(async move { + promote_waiting_lane(&pool, &store, &metrics, &guard.key).await; + drop(guard); + }); + } +} + +/// Holds a lane in [`QueueStoragePromoter`]'s waiting set until its waiter +/// finishes, panics or is aborted. +struct WaitingLane { + waiting: Arc>>, + key: (JobState, PromotionLane), +} + +impl Drop for WaitingLane { + fn drop(&mut self) { + if let Ok(mut waiting) = self.waiting.lock() { + waiting.remove(&self.key); + } + } +} + +async fn promote_waiting_lane( + pool: &PgPool, + store: &QueueStorage, + metrics: &AwaMetrics, + (state, lane): &(JobState, PromotionLane), +) { + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + let batch_started = Instant::now(); + match store + .promote_due_lane( + pool, + *state, + lane, + PROMOTE_BATCH_SIZE, + Some(PROMOTE_LANE_WAITER_LOCK_TIMEOUT), + ) + .await + { + Ok(outcome) if !outcome.contended.is_empty() => { + metrics.record_promotion_lane_contention( + state.as_str(), + &lane.queue, + "waiter_timed_out", + ); + warn!( + queue = %lane.queue, + priority = lane.priority, + lock_timeout_ms = PROMOTE_LANE_WAITER_LOCK_TIMEOUT.as_millis() as u64, + "Promotion of a contended lane timed out waiting for its locks; retrying next tick" + ); + return; + } + Ok(outcome) => { + metrics.record_promotion_batch( + state.as_str(), + outcome.promoted as u64, + batch_started.elapsed(), + ); + if (outcome.selected as i64) < PROMOTE_BATCH_SIZE { + return; + } + } + Err(err) => { + error!( + queue = %lane.queue, + priority = lane.priority, + error = %err, + "Failed to promote a contended lane" + ); + return; + } + } + } +} + +/// The leader's queue-storage promoters: its own store's, and the configured +/// store a canonical-resolved leader promotes once routing has flipped +/// (#456). Scoped to one leadership term. +struct LeaderPromoters { + primary: Option, + standby: Option, +} + impl MaintenanceService { #[allow(clippy::too_many_arguments)] pub(crate) fn new( @@ -1019,6 +1239,22 @@ impl MaintenanceService { // tick. Local to the leader loop so a re-election resets // backoff to zero. let prune_tracker = PruneBackoffTracker::new(); + let promoters = LeaderPromoters { + primary: self.storage.queue_storage().map(|runtime| { + QueueStoragePromoter::new( + self.pool.clone(), + Arc::clone(&runtime.store), + self.metrics.clone(), + ) + }), + standby: self.standby_queue_storage.as_ref().map(|runtime| { + QueueStoragePromoter::new( + self.pool.clone(), + Arc::clone(&runtime.store), + self.metrics.clone(), + ) + }), + }; // Skip the first immediate tick heartbeat_rescue_timer.tick().await; @@ -1084,7 +1320,7 @@ impl MaintenanceService { } _ = promote_timer.tick() => { if let Some(timer) = branch_tracker.try_begin_without_cooldown("promote_scheduled", self.promote_interval, &self.metrics) { - self.promote_scheduled().await; + self.promote_scheduled(&promoters).await; timer.finish(); } } @@ -2324,13 +2560,20 @@ impl MaintenanceService { } /// Promote scheduled jobs that are now due. - #[tracing::instrument(skip(self), name = "maintenance.promote")] - async fn promote_scheduled(&self) { - if let Err(err) = self.promote_due_state("scheduled", "scheduled jobs").await { + #[tracing::instrument(skip(self, promoters), name = "maintenance.promote")] + async fn promote_scheduled(&self, promoters: &LeaderPromoters) { + if let Err(err) = self + .promote_due_state(JobState::Scheduled, "scheduled jobs", promoters) + .await + { error!(error = %err, "Failed to promote scheduled jobs"); } if let Err(err) = self - .promote_due_state("retryable", "retryable jobs (backoff elapsed)") + .promote_due_state( + JobState::Retryable, + "retryable jobs (backoff elapsed)", + promoters, + ) .await { error!(error = %err, "Failed to promote retryable jobs"); @@ -2339,138 +2582,55 @@ impl MaintenanceService { async fn promote_due_state( &self, - state: &'static str, + state: JobState, label: &'static str, - ) -> Result<(), awa_model::AwaError> { + promoters: &LeaderPromoters, + ) -> Result<(), AwaError> { + let canonical_state = match state { + JobState::Scheduled => "scheduled", + JobState::Retryable => "retryable", + other => { + return Err(AwaError::Validation(format!( + "unsupported promote state: {other:?}" + ))); + } + }; let mut promoted_total = 0usize; let mut notified_queues = HashSet::new(); - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; - } - - match &self.storage { - RuntimeStorage::Canonical => { - let (promoted, queues) = self - .promote_due_batch(state) - .await - .map_err(awa_model::AwaError::Database)?; - if promoted == 0 { - break; - } - - promoted_total += promoted; - notified_queues.extend(queues); - - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - RuntimeStorage::QueueStorage(runtime) => { - let job_state = match state { - "scheduled" => awa_model::JobState::Scheduled, - "retryable" => awa_model::JobState::Retryable, - other => { - return Err(awa_model::AwaError::Validation(format!( - "unsupported queue storage promote state: {other}" - ))); - } - }; - let promote_start = std::time::Instant::now(); - let promoted = runtime - .store - .promote_due(&self.pool, job_state, PROMOTE_BATCH_SIZE) - .await?; - self.metrics.record_promotion_batch( - state, - promoted as u64, - promote_start.elapsed(), - ); - if promoted == 0 { - break; - } - - promoted_total += promoted; - - if promoted < PROMOTE_BATCH_SIZE as usize { - break; + match &self.storage { + RuntimeStorage::Canonical => { + promoted_total += self + .promote_canonical_due(canonical_state, &mut notified_queues) + .await?; + + // #456 mirror: a canonical-resolved (pre-flip) leader must + // promote the queue-storage deferred backlog once routing has + // flipped — re-scheduled canonical jobs migrate there, and + // post-flip runtimes that would promote them may not hold + // leadership. The active-schema guard keeps this a no-op until + // the flip. + if let Some(promoter) = &promoters.standby { + if self.active_standby_queue_storage().await.is_some() { + promoted_total += promoter.promote_due(state, &self.cancel).await?; } } } - } - - // #456: while a storage transition drains, due jobs can still sit in - // the canonical deferred backlog. Drain-only runtimes execute - // canonical work but stop leading maintenance once a post-flip - // runtime wins the election, so a queue-storage leader must promote - // the canonical plane too. Both queries are cheap no-ops once the - // canonical tables are empty (every finalized or fresh install). - if matches!(&self.storage, RuntimeStorage::QueueStorage(_)) { - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; + RuntimeStorage::QueueStorage(_) => { + if let Some(promoter) = &promoters.primary { + promoted_total += promoter.promote_due(state, &self.cancel).await?; } - let (promoted, queues) = self - .promote_due_batch(state) - .await - .map_err(awa_model::AwaError::Database)?; - if promoted == 0 { - break; - } - promoted_total += promoted; - notified_queues.extend(queues); - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - } - // #456 mirror: a canonical-resolved (pre-flip) leader must promote - // the queue-storage deferred backlog once routing has flipped — - // re-scheduled canonical jobs migrate there, and post-flip runtimes - // that would promote them may not hold leadership. The active-schema - // guard keeps this a no-op until the flip. - if matches!(&self.storage, RuntimeStorage::Canonical) { - if let Some(runtime) = &self.standby_queue_storage { - let active_schema: Option = - sqlx::query_scalar("SELECT awa.active_queue_storage_schema()") - .fetch_one(&self.pool) - .await - .map_err(awa_model::AwaError::Database)?; - if active_schema.as_deref() == Some(runtime.store.schema()) { - let job_state = match state { - "scheduled" => awa_model::JobState::Scheduled, - "retryable" => awa_model::JobState::Retryable, - other => { - return Err(awa_model::AwaError::Validation(format!( - "unsupported queue storage promote state: {other}" - ))); - } - }; - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; - } - let promote_start = std::time::Instant::now(); - let promoted = runtime - .store - .promote_due(&self.pool, job_state, PROMOTE_BATCH_SIZE) - .await?; - self.metrics.record_promotion_batch( - state, - promoted as u64, - promote_start.elapsed(), - ); - if promoted == 0 { - break; - } - promoted_total += promoted; - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - } + // #456: while a storage transition drains, due jobs can still + // sit in the canonical deferred backlog. Drain-only runtimes + // execute canonical work but stop leading maintenance once a + // post-flip runtime wins the election, so a queue-storage + // leader must promote the canonical plane too. This is a + // cheap no-op once the canonical tables are empty (every + // finalized or fresh install). + promoted_total += self + .promote_canonical_due(canonical_state, &mut notified_queues) + .await?; } } @@ -2478,7 +2638,7 @@ impl MaintenanceService { debug!( count = promoted_total, queues = notified_queues.len(), - state, + state = canonical_state, "Promoted {label}" ); } @@ -2486,6 +2646,34 @@ impl MaintenanceService { Ok(()) } + /// Promote due canonical `scheduled_jobs` rows in batches until one comes + /// back short. + async fn promote_canonical_due( + &self, + state: &'static str, + notified_queues: &mut HashSet, + ) -> Result { + let mut promoted_total = 0usize; + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + if self.cancel.is_cancelled() { + break; + } + let (promoted, queues) = self + .promote_due_batch(state) + .await + .map_err(AwaError::Database)?; + if promoted == 0 { + break; + } + promoted_total += promoted; + notified_queues.extend(queues); + if promoted < PROMOTE_BATCH_SIZE as usize { + break; + } + } + Ok(promoted_total) + } + /// SQL template for promotion. The state literal is injected directly /// (not as a parameter) so the planner can match the partial index on /// `(run_at, id) WHERE state = ''`. With a parameter, the planner diff --git a/awa/Cargo.toml b/awa/Cargo.toml index 437cb057b..23ba05b62 100644 --- a/awa/Cargo.toml +++ b/awa/Cargo.toml @@ -25,6 +25,7 @@ sqlx.workspace = true serde.workspace = true serde_json.workspace = true tokio = { workspace = true, features = ["full", "test-util"] } +tokio-util.workspace = true async-trait.workspace = true chrono.workspace = true uuid.workspace = true 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..ce7d53040 --- /dev/null +++ b/awa/tests/enqueue_lane_maintenance_coupling_test.rs @@ -0,0 +1,427 @@ +//! 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. The leader +//! runs its branches (promotion, rescue, rotation, ...) one after another on a +//! single task, so a promotion parked behind a hot lane would hold up the due +//! jobs of other queues and every other maintenance branch. The leader's +//! promoter gives each lane its own short transaction and hands a lane whose +//! lock stays busy to a dedicated waiter, so neither should wait on it. +//! +//! `test_maintenance_under_hot_lane` measures, with and without concurrent +//! producers on one hot lane: +//! * how long each promotion tick holds the leader task, 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::worker::maintenance::QueueStoragePromoter; +use awa::worker::AwaMetrics; +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_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 { + /// How long each promotion tick held the leader task. + tick_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; each tick runs the leader's promoter) 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: &Arc, + 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 promoter = + QueueStoragePromoter::new(pool.clone(), Arc::clone(store), AwaMetrics::from_global()); + let cancel = tokio_util::sync::CancellationToken::new(); + + 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; + let tick_started = Instant::now(); + promoter + .promote_due(JobState::Scheduled, &cancel) + .await + .expect("promote_due"); + report.tick_ms.push(tick_started.elapsed().as_secs_f64() * 1e3); + 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 tick on leader ms p50={:.0} p99={:.0} max={:.0} | 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.tick_ms, 0.5), + percentile(&report.tick_ms, 0.99), + max_of(&report.tick_ms), + 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 = Arc::new(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 = Arc::new(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/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index a060369f8..b0093749c 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -5,8 +5,8 @@ use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, - BatchOperationSpec, PruneOutcome, QueueStorage, QueueStorageConfig, RotateOutcome, SkipReason, - SubmitBatchOperation, + BatchOperationSpec, PromotionLane, PruneOutcome, QueueStorage, QueueStorageConfig, + RotateOutcome, SkipReason, SubmitBatchOperation, }; use awa::{ Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, JobRow, JobState, QueueConfig, @@ -7802,9 +7802,10 @@ fn branch_overruns() -> Arc>> { /// snoozed and retry-after jobs runnable, so overrunning promotion bodies must /// never arm a cooldown on it. /// -/// The first three promotion bodies (two `DELETE`s each) take 1.2s against a -/// 25ms promote interval, so the maintenance branch tracker flags promotion as -/// overrunning. The test asserts on that overrun warning rather than on job +/// Every promotion body runs the canonical-plane promotion `DELETE` once per +/// state, even with nothing due. A statement trigger slows the first six, so +/// the first three bodies take 1.2s against a 25ms promote interval and the +/// maintenance branch tracker flags promotion as overrunning. The test asserts on that overrun warning rather than on job /// latency: the promotion warning must report no cooldown, while another /// branch in the same run (cleanup, whose bodies exceed its 1ms interval) /// does report one, so the check cannot pass vacuously. Jobs that fall due @@ -7831,7 +7832,7 @@ async fn test_queue_storage_runtime_promotion_not_suppressed_after_overrunning_b END $$" ), format!( - "CREATE TRIGGER slow_first_promote_deletes BEFORE DELETE ON {schema}.deferred_jobs \ + "CREATE TRIGGER slow_first_promote_deletes BEFORE DELETE ON awa.scheduled_jobs \ FOR EACH STATEMENT EXECUTE FUNCTION {schema}.slow_first_promote_deletes()" ), ] { @@ -7965,6 +7966,316 @@ async fn test_queue_storage_runtime_promotion_not_suppressed_after_overrunning_b client.shutdown(Duration::from_secs(5)).await; } +/// Advisory lock key of a lane's enqueue reservation, as +/// `reserve_enqueue_seq` takes it. The lane must already exist. +async fn enqueue_lane_lock_key( + pool: &sqlx::PgPool, + store: &QueueStorage, + queue: &str, + priority: i16, +) -> i64 { + sqlx::query_scalar(&format!( + "SELECT hashtextextended(format('%I.%I', $1::text, seq_name), 0) \ + FROM {}.queue_enqueue_heads \ + WHERE queue = $2 AND priority = $3 AND enqueue_shard = 0", + store.schema() + )) + .bind(store.schema()) + .bind(queue) + .bind(priority) + .fetch_one(pool) + .await + .expect("derive enqueue lane lock key") +} + +/// A promotion batch that meets a lane whose enqueue lock is held promotes +/// the other lanes, reports the held lane as contended after its lock timeout +/// and leaves its rows deferred. Promoted later, the lane's rows take lane +/// sequence numbers in due order and are all claimable. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_promotion_leaves_a_contended_lane_deferred_and_in_order() { + const PRIORITY: i16 = 2; + let (_db_guard, pool) = setup_pool(6).await; + let hot = "qs_promote_contended_hot"; + let other = "qs_promote_contended_other"; + let store = create_store(&pool, "awa_qs_promote_contended_lane").await; + let opts = |queue: &str, run_at: Option>| InsertOpts { + queue: queue.to_string(), + priority: PRIORITY, + run_at, + ..Default::default() + }; + + let warm_job_id = enqueue_job(&pool, &store, &CompleteJob { id: 0 }, opts(hot, None)).await; + // Later job ids fall due earlier, so due order differs from id order. + let due_base = Utc::now() - chrono::Duration::seconds(60); + let mut hot_due = Vec::new(); + for n in 0..6 { + let run_at = due_base - chrono::Duration::seconds(n); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: n }, + opts(hot, Some(run_at)), + ) + .await; + hot_due.push((run_at, job_id)); + } + hot_due.sort(); + let hot_in_due_order: Vec = hot_due.iter().map(|(_, job_id)| *job_id).collect(); + let other_job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 100 }, + opts(other, Some(due_base)), + ) + .await; + + let hot_lane = PromotionLane { + queue: hot.to_string(), + priority: PRIORITY, + }; + let lock_key = enqueue_lane_lock_key(&pool, &store, hot, PRIORITY).await; + let mut holder = pool.begin().await.expect("begin lane holder"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(lock_key) + .execute(holder.as_mut()) + .await + .expect("hold the hot lane's enqueue lock"); + + let lock_timeout = Duration::from_millis(100); + let started = Instant::now(); + let outcome = store + .promote_due_lanes(&pool, JobState::Scheduled, 100, Some(lock_timeout), &[]) + .await + .expect("promote with the hot lane held"); + let elapsed = started.elapsed(); + assert_eq!(outcome.promoted, 1, "{outcome:?}"); + assert_eq!(outcome.selected, 7, "{outcome:?}"); + assert_eq!(outcome.contended, vec![hot_lane.clone()]); + assert!( + elapsed >= lock_timeout && elapsed < scaled_timeout(Duration::from_secs(1)), + "the held lane should cost one lock timeout, took {elapsed:?}" + ); + let other_job = store + .load_job(&pool, other_job_id) + .await + .expect("load other job") + .expect("other job exists"); + assert_eq!(other_job.state, JobState::Available); + for job_id in &hot_in_due_order { + let job = store + .load_job(&pool, *job_id) + .await + .expect("load hot job") + .expect("hot job exists"); + assert_eq!(job.state, JobState::Scheduled, "job {job_id}"); + } + + let skipped = store + .promote_due_lanes( + &pool, + JobState::Scheduled, + 100, + Some(lock_timeout), + std::slice::from_ref(&hot_lane), + ) + .await + .expect("promote with the hot lane skipped"); + assert_eq!(skipped, Default::default()); + + holder.commit().await.expect("release the hot lane"); + let mut batches = 0; + loop { + let outcome = store + .promote_due_lane(&pool, JobState::Scheduled, &hot_lane, 4, None) + .await + .expect("promote the hot lane"); + assert!(outcome.contended.is_empty(), "{outcome:?}"); + if outcome.promoted == 0 { + break; + } + batches += 1; + } + assert_eq!(batches, 2); + + let in_lane_order: Vec = sqlx::query_scalar(&format!( + "SELECT job_id FROM {}.ready_entries WHERE job_id = ANY($1) ORDER BY lane_seq", + store.schema() + )) + .bind(&hot_in_due_order) + .fetch_all(&pool) + .await + .expect("read hot lane order"); + assert_eq!(in_lane_order, hot_in_due_order); + + let mut claimed: Vec = store + .claim_runtime_batch(&pool, hot, 100, Duration::from_secs(60)) + .await + .expect("claim the hot lane") + .into_iter() + .map(|claimed| claimed.job.id) + .collect(); + claimed.sort_unstable(); + let mut expected = hot_in_due_order.clone(); + expected.push(warm_job_id); + expected.sort_unstable(); + assert_eq!(claimed, expected); +} + +/// The maintenance leader must not wait on a lane whose enqueue lock never +/// comes free. Two sessions take turns holding the hot lane's lock for +/// `HOLD` each, so a lock request queued behind them waits up to two holds. +/// A due job on another queue must still run, and an abandoned attempt must +/// still be rescued, well within one hold; the hot queue's due jobs must +/// still be promoted while the lane stays busy. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_busy_lane_holds_back_neither_promotion_nor_rescue() { + const HOLD: Duration = Duration::from_secs(5); + const PRIORITY: i16 = 2; + let (_db_guard, pool) = setup_pool(12).await; + let hot = "qs_busy_lane_hot"; + let other = "qs_busy_lane_other"; + let abandoned = "qs_busy_lane_abandoned"; + let schema = "awa_qs_busy_lane_leader"; + let store = create_store(&pool, schema).await; + let opts = |queue: &str, run_at: Option>| InsertOpts { + queue: queue.to_string(), + priority: PRIORITY, + run_at, + ..Default::default() + }; + + enqueue_job(&pool, &store, &CompleteJob { id: 0 }, opts(hot, None)).await; + let lock_key = enqueue_lane_lock_key(&pool, &store, hot, PRIORITY).await; + + let stop = Arc::new(AtomicBool::new(false)); + let mut holders = Vec::new(); + for _ in 0..2 { + let (held_tx, held_rx) = tokio::sync::oneshot::channel::<()>(); + let pool = pool.clone(); + let stop = Arc::clone(&stop); + holders.push(tokio::spawn(async move { + let mut held_tx = Some(held_tx); + let mut holds = 0u32; + while !stop.load(Ordering::SeqCst) { + let mut tx = pool.begin().await.expect("begin lane holder"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(lock_key) + .execute(tx.as_mut()) + .await + .expect("hold the hot lane's enqueue lock"); + holds += 1; + if let Some(held_tx) = held_tx.take() { + let _ = held_tx.send(()); + } + tokio::time::sleep(HOLD).await; + tx.commit().await.expect("release the hot lane"); + } + holds + })); + // The first holder takes the lock; wait for it so the second queues + // behind it rather than racing it. + if holders.len() == 1 { + held_rx.await.expect("first holder took the lane"); + } + } + + let due = Utc::now() - chrono::Duration::seconds(1); + let mut hot_jobs = Vec::new(); + for n in 1..=3 { + hot_jobs + .push(enqueue_job(&pool, &store, &CompleteJob { id: n }, opts(hot, Some(due))).await); + } + let other_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: 10 }, + opts(other, Some(due)), + ) + .await; + let abandoned_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: 20 }, + opts(abandoned, None), + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, abandoned, 1, Duration::from_secs(3_600)) + .await + .expect("claim the abandoned attempt"); + assert_eq!(claimed.len(), 1); + + let heartbeat_staleness = Duration::from_millis(500); + let worker_queue = || QueueConfig { + max_workers: 4, + poll_interval: Duration::from_millis(25), + ..QueueConfig::default() + }; + let client = Client::builder(pool.clone()) + .queue(hot, worker_queue()) + .queue(other, worker_queue()) + .queue_storage( + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }, + Duration::from_secs(60), + Duration::from_secs(60), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(CompleteWorker) + .promote_interval(Duration::from_millis(50)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .heartbeat_staleness(heartbeat_staleness) + .deadline_rescue_interval(Duration::from_secs(10)) + .callback_rescue_interval(Duration::from_secs(10)) + .build() + .expect("Failed to build busy-lane client"); + client + .start() + .await + .expect("Failed to start busy-lane client"); + + let bound = scaled_timeout(Duration::from_secs(2)); + wait_for_job_state(&store, &pool, other_job, &[JobState::Completed], bound).await; + wait_for_job_state( + &store, + &pool, + abandoned_job, + &[JobState::Retryable, JobState::Available], + bound + heartbeat_staleness, + ) + .await; + + for job_id in &hot_jobs { + wait_for_job_state( + &store, + &pool, + *job_id, + &[JobState::Completed], + scaled_timeout(HOLD * 3), + ) + .await; + } + assert!( + holders.iter().all(|holder| !holder.is_finished()), + "the hot lane must stay busy until its jobs are promoted" + ); + + stop.store(true, Ordering::SeqCst); + for holder in holders { + assert!(holder.await.expect("lane holder") >= 1); + } + client.shutdown(Duration::from_secs(5)).await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_runtime_stale_heartbeat_rescue() { let (_db_guard, pool) = setup_pool(10).await; diff --git a/docs/architecture.md b/docs/architecture.md index 06b967bc3..c4437b0e7 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -150,7 +150,7 @@ sequenceDiagram Enqueue is transactional: if the producer's outer transaction rolls back, the job never becomes visible. Immediate jobs reserve a lane sequence range under a transaction-scoped lane lock, append to `ready_entries_*`, and append a compact `ready_segments` range for the committed lane sequence range; future scheduled or retryable jobs append to `deferred_jobs`. The lane lock is held until commit or rollback so a later producer cannot make lane `N+1` visible while lane `N` can still commit. -`deferred_jobs` is deliberately outside the claim path. The maintenance leader promotes due `scheduled` and `retryable` rows into `ready_entries_*` in batches and notifies the target queues. Cron schedules are just producers for ordinary jobs: when a schedule fires, the cron transaction records the fire and enqueues the job atomically, using `ready_entries_*` for immediate fires or `deferred_jobs` when the enqueue carries a future `run_at`. +`deferred_jobs` is deliberately outside the claim path. The maintenance leader promotes due `scheduled` and `retryable` rows into `ready_entries_*` in batches and notifies the target queues. Re-enqueueing a row takes its lane lock, which producers hold until they commit, so each `(queue, priority)` lane's share of a batch is promoted in its own transaction under a short `lock_timeout`. A lane that times out is handed to a dedicated waiter task, which queues for the lock behind the producers already waiting, so a busy lane is still promoted in turn without holding up other lanes or the leader's rescue, rotation and prune. Cron schedules are just producers for ordinary jobs: when a schedule fires, the cron transaction records the fire and enqueues the job atomically, using `ready_entries_*` for immediate fires or `deferred_jobs` when the enqueue carries a future `run_at`. There are two COPY-shaped producer paths: diff --git a/docs/configuration.md b/docs/configuration.md index 868dfb3cf..71bd3e72c 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -245,6 +245,7 @@ Retry backoff and `Snooze` use the same mechanism: the current attempt closes, a Operational knobs: - `promote_interval` controls how often maintenance checks for due deferred rows. +- A queue whose enqueue lane is kept busy by producers does not delay promotion of other queues or the leader's rescue: promotion waits 50 ms for a lane's lock, then hands the lane to a dedicated waiter that is served in turn. `awa_maintenance_promote_lane_contention_total` counts those hand-offs by queue; a steady `waiter_limit` or `waiter_timed_out` rate means more lanes are saturated than the leader's four waiters cover, or a transaction is holding a lane lock for longer than 30 s. - Per-job `run_at` controls the due time. - Periodic schedules are declared in worker code with `periodic()` and managed through the CLI/UI `cron` surface. diff --git a/docs/grafana/README.md b/docs/grafana/README.md index 823102462..81dcb0f19 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -156,6 +156,7 @@ All metrics use the `awa` OTel meter name and are exported via OTLP to your conf | `awa.maintenance.promote_batches` | Counter | state | Promotion batches | | `awa.maintenance.promote_batch_size` | Histogram | state | Jobs promoted | | `awa.maintenance.promote_duration` | Histogram (s) | state | Promotion time | +| `awa.maintenance.promote.lane_contention` | Counter | state, queue, outcome | Promotion lanes whose lock was busy: `handed_off` to a dedicated waiter, `waiter_limit` when every waiter was busy, `waiter_timed_out` when a waiter's 30 s wait ran out | | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | ### Heartbeat From e4f3d5591fb7d3d0394257fe8528bd470eee042f Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 15:20:23 +1300 Subject: [PATCH 08/16] fix(queue-storage): report Blocked when any prune statement times out on a lock (0.6) Prune and terminal-count rollup set a transaction-local lock_timeout before taking ACCESS EXCLUSIVE on a slot's partitions, but only the LOCK TABLE and TRUNCATE steps mapped a timeout to Blocked. The timeout also applies to the statements after them: the post-lock liveness checks over the lease and claim partitions, the failed-row carry into the current slot's done_entries and terminal-count deltas, the rollup fold and the live-counter cleanup. A wait there surfaced as a 55P03 database error, was logged as a prune error and skipped the prune backoff. Run everything after the timeout is set as one block and settle it in one place: commit on success, and on a lock timeout from any statement roll back and report Blocked. Lease and claim prune use the same helper so the rule holds for every ring. The new tests hold a SHARE lock on each relation the queue prune writes after its child locks, and on queue_terminal_live_counts for the rollup; both fail with 55P03 before this change. The retry-failed outcome test now retries a Blocked prune as the maintenance loop does, since a background session can legitimately block it. --- CHANGELOG.md | 17 + awa-model/src/queue_storage.rs | 1101 +++++++++++------------ awa/tests/queue_storage_runtime_test.rs | 197 +++- 3 files changed, 736 insertions(+), 579 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a17deb..e62d6b908 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,23 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **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 that `LOCK TABLE` 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 writes + `queue_terminal_rollups` 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. - **`admin::job_kind_overviews` on queue storage no longer depends on the planner's row estimates.** Behind the admin UI's kinds page, it counted every job through the `current_jobs` view. Two nodes of that view are diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 930aab715..80647bc5e 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -434,6 +434,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() { @@ -15113,152 +15142,130 @@ impl QueueStorage { return Ok(TerminalDeltaSlotRollup::SkippedActive); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - - let lock_delta = sqlx::query(&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(&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(&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); + } - let grouped_keys: i64 = sqlx::query_scalar(&format!( - r#" - WITH grouped AS MATERIALIZED ( - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - SUM(terminal_delta)::bigint AS delta + let delta_rows: i64 = sqlx::query_scalar(&format!( + r#" + SELECT count(*)::bigint 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(&format!("TRUNCATE TABLE {delta_child}")) - .execute(tx.as_mut()) - .await; + "# + )) + .bind(generation) + .fetch_one(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); } + + let grouped_keys: i64 = sqlx::query_scalar(&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(&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 } /// Prune the oldest sealed queue ring slot. @@ -15411,167 +15418,114 @@ 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_tables = sqlx::query(&format!( - "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE" - )) - .execute(tx.as_mut()) - .await; + sqlx::query(&format!( + "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE" + )) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - if let Err(err) = lock_tables { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); + // 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, &ready_child, 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 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 pending = + queue_prune_has_pending_ready_tx(&mut tx, schema, &ready_child, generation).await?; + if pending { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueuePendingReady, + count: busy_indicator(pending), + }); + } - let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX); - let retention_floor = failed_retention_secs > 0; + 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), + }); + } - // #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 queue_ring_state row is held FOR - // UPDATE above, which serializes against rotate — the current - // slot cannot move while carried rows are re-inserted into it. - let carried_failed_rows = if retention_floor { - let (current_generation,): (i64,) = sqlx::query_as(&format!( - "SELECT generation FROM {schema}.queue_ring_slots WHERE slot = $1" - )) - .bind(state.0) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; + 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 queue_ring_state row is held FOR + // UPDATE above, which serializes against rotate — the current + // slot cannot move while carried rows are re-inserted into it. + let carried_failed_rows = if retention_floor { + let (current_generation,): (i64,) = sqlx::query_as(&format!( + "SELECT generation FROM {schema}.queue_ring_slots WHERE slot = $1" + )) + .bind(state.0) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - // 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(&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(&format!( + // 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(&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) @@ -15579,155 +15533,198 @@ 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(&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(&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 ( - 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 + 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(&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 ) - 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) - "# - )) - .bind(failed_retention_secs) - .bind(retention_floor) - .fetch_all(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; + 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) + "# + )) + .bind(failed_retention_secs) + .bind(retention_floor) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - let truncate = sqlx::query(&format!( - "TRUNCATE TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child}" - )) - .execute(tx.as_mut()) - .await; + sqlx::query(&format!( + "TRUNCATE TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child}" + )) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - match truncate { - Ok(_) => { - if !pruned_terminal_counts.is_empty() { - self.adjust_terminal_rollups_batch(&mut tx, pruned_terminal_counts.into_iter()) - .await?; - } + if !pruned_terminal_counts.is_empty() { + self.adjust_terminal_rollups_batch(&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(&format!( - "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" - )) - .bind(slot) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - tx.commit().await.map_err(map_sqlx_error)?; - 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 { + // #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(&format!( + "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" + )) + .bind(slot) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { + slot, + carried_failed_rows, + }) + } + .await; + + let outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { + slot, + carried_failed_rows, + } = outcome + { + if carried_failed_rows > 0 { + tracing::info!( slot, carried_failed_rows, - }) - } - 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(skip(self, pool), name = "queue_storage.prune_oldest_leases")] @@ -15796,74 +15793,57 @@ impl QueueStorage { return Ok(PruneOutcome::Noop); } - 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_table = sqlx::query(&format!( - "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" - )) - .execute(tx.as_mut()) - .await; + sqlx::query(&format!( + "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" + )) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - if let Err(err) = lock_table { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); - } - return Err(map_sqlx_error(err)); - } + let current_slot: i32 = sqlx::query_scalar(&format!( + r#" + SELECT current_slot + FROM {schema}.lease_ring_state + WHERE singleton = TRUE + "# + )) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - let current_slot: i32 = sqlx::query_scalar(&format!( - r#" - SELECT current_slot - FROM {schema}.lease_ring_state - WHERE singleton = TRUE - "# - )) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; + if current_slot == slot { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::LeaseCurrent, + count: 0, + }); + } - 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?; - 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 { + sqlx::query(&format!("TRUNCATE TABLE {lease_child}")) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { slot, - reason: SkipReason::LeaseActive, - count: busy_indicator(active_leases), - }); + carried_failed_rows: 0, + }) } + .await; - let truncate = sqlx::query(&format!("TRUNCATE TABLE {lease_child}")) - .execute(tx.as_mut()) - .await; - - match truncate { - Ok(_) => { - tx.commit().await.map_err(map_sqlx_error)?; - Ok(PruneOutcome::Pruned { - slot, - carried_failed_rows: 0, - }) - } - 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)) - } - } + settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await } pub async fn vacuum_leases(&self, pool: &PgPool) -> Result<(), AwaError> { @@ -16101,101 +16081,84 @@ 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_tables = sqlx::query(&format!( - "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" - )) - .execute(tx.as_mut()) - .await; + sqlx::query(&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)?; - 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-state lock should already make + // this stable, but keeping the explicit gate documents the + // truncate precondition. + let current_slot: i32 = sqlx::query_scalar(&format!( + r#" + SELECT current_slot FROM {schema}.claim_ring_state WHERE singleton = TRUE + "# + )) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - // After taking ACCESS EXCLUSIVE, recheck that the slot is still - // not the current one. The ring-state lock should already make - // this stable, but keeping the explicit gate documents the - // truncate precondition. - let current_slot: i32 = sqlx::query_scalar(&format!( - r#" - SELECT current_slot FROM {schema}.claim_ring_state WHERE singleton = TRUE - "# - )) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; + if current_slot == slot { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::ClaimCurrent, + count: 0, + }); + } - if current_slot == slot { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::ClaimCurrent, - count: 0, - }); - } + 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 { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::ClaimOpen, + count: busy_indicator(open_claims), + }); + } - 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 { + sqlx::query(&format!( + "TRUNCATE TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child}" + )) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + sqlx::query(&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 + WHERE slot = $1 + "# + )) + .bind(slot) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { slot, - reason: SkipReason::ClaimOpen, - count: busy_indicator(open_claims), - }); + carried_failed_rows: 0, + }) } - - let truncate = sqlx::query(&format!( - "TRUNCATE TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child}" - )) - .execute(tx.as_mut()) .await; - match truncate { - Ok(_) => { - sqlx::query(&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 - WHERE slot = $1 - "# - )) - .bind(slot) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - tx.commit().await.map_err(map_sqlx_error)?; - Ok(PruneOutcome::Pruned { - slot, - carried_failed_rows: 0, - }) - } - 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)) - } - } + settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await } 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 a060369f8..85205ff5e 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -10037,16 +10037,25 @@ async fn test_queue_storage_retry_failed_outcome_surfaces_pruned_rows() { RotateOutcome::Rotated { .. } => {} other => panic!("expected Rotated, got {other:?}"), } - match store - .prune_oldest(&pool, Duration::from_secs(3600)) - .await - .expect("prune_oldest") - { - PruneOutcome::Pruned { - carried_failed_rows: 1, - .. - } => {} - other => panic!("expected Pruned with one carried row, got {other:?}"), + // A background session on a relation prune touches (autovacuum, or a + // hosted Postgres's own workers) can time out its lock wait and report + // Blocked; retry as the maintenance loop would. + let prune_deadline = Instant::now() + Duration::from_secs(5); + loop { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("prune_oldest") + { + PruneOutcome::Pruned { + carried_failed_rows: 1, + .. + } => break, + PruneOutcome::Blocked { slot: 0 } if Instant::now() < prune_deadline => { + tokio::time::sleep(Duration::from_millis(50)).await; + } + other => panic!("expected Pruned with one carried row, got {other:?}"), + } } let outcome = admin::retry_failed_by_queue(&pool, queue) @@ -11844,6 +11853,174 @@ 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 fold 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_rollups"), + format!("{schema}.queue_terminal_live_counts"), + ] { + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(&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(&format!( + "SELECT count(*)::bigint FROM {schema}.queue_terminal_rollups" + )) + .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 whole rollup pass. +#[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; + + 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(&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); + assert_eq!(live_count_sum(&pool, schema, queue).await, 0); + assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 3); + + let outcome = store + .rollup_terminal_count_deltas(&pool, 4) + .await + .expect("rollup after lock released"); + assert_eq!(outcome.rolled_slots, 1); + assert_eq!(outcome.blocked_slots, 0); + assert_eq!(live_count_sum(&pool, schema, queue).await, 3); + assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 0); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_runtime_complete_external() { let (_db_guard, pool) = setup_pool(10).await; From 7e6259e04d9c4d45587a08c7094f6f378d788abf Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 15:26:16 +1300 Subject: [PATCH 09/16] fix(worker): drain finalization on graceful shutdown (0.6) Once a handler returned, the executor removed the attempt from the heartbeat set and finalized it in a detached task that Client::shutdown did not wait for. When the last handler returned as the drain ended, shutdown cancelled the completion batcher before that finalize reached it, and the direct-finalize fallback raced the process exit. The attempt stayed open with no heartbeat, so stale-heartbeat rescue ran the finished job again after heartbeat_staleness. Track finalizers in a TaskTracker and wait for them, within the drain timeout, before stopping the completion batcher, heartbeat and maintenance services. This is the shutdown part of #521 on main. --- CHANGELOG.md | 12 ++ Cargo.lock | 1 + awa-python/Cargo.lock | 1 + awa-worker/Cargo.toml | 2 +- awa-worker/src/client.rs | 18 +- awa-worker/src/dispatcher.rs | 1 + awa-worker/src/executor.rs | 10 +- awa/tests/ci_timing.rs | 186 +++++++++++++++++++ awa/tests/queue_storage_runtime_test.rs | 227 ++++++++++++++++++++++++ docs/test-plan.md | 1 + 10 files changed, 455 insertions(+), 4 deletions(-) create mode 100644 awa/tests/ci_timing.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a17deb..d49c6f89e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,18 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **A graceful shutdown no longer leaves a finished attempt for stale-heartbeat + rescue to run again.** Once a handler returned, its finalize (completion, + retry, cancel or failure, and any ADR-029 follow-up inserts) ran in a task + that `Client::shutdown` did not wait for. When the last handler returned + just as the drain ended, shutdown stopped the completion batcher before + that finalize reached it, and the direct fallback raced the process exit. + The attempt had already left the heartbeat set, so it went stale and was + rescued and executed again `heartbeat_staleness` later; a deploy could + re-run a few jobs on every rolling restart. Shutdown now waits for + pending finalization, within the same drain timeout, before it stops the + completion batcher, heartbeat and maintenance services. A finalize that + outlives the timeout is still left to rescue, as before. - **`admin::job_kind_overviews` on queue storage no longer depends on the planner's row estimates.** Behind the admin UI's kinds page, it counted every job through the `current_jobs` view. Two nodes of that view are diff --git a/Cargo.lock b/Cargo.lock index 30bd194c9..a29ff7c5c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4187,6 +4187,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/awa-python/Cargo.lock b/awa-python/Cargo.lock index fa27a3b79..8922f4d9a 100644 --- a/awa-python/Cargo.lock +++ b/awa-python/Cargo.lock @@ -2355,6 +2355,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/awa-worker/Cargo.toml b/awa-worker/Cargo.toml index 91dd58594..afb66440c 100644 --- a/awa-worker/Cargo.toml +++ b/awa-worker/Cargo.toml @@ -21,7 +21,7 @@ sqlx.workspace = true serde.workspace = true serde_json.workspace = true tokio.workspace = true -tokio-util.workspace = true +tokio-util = { workspace = true, features = ["rt"] } async-trait.workspace = true thiserror.workspace = true chrono.workspace = true diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index e5a498dbb..3cfd9052e 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -27,6 +27,7 @@ use std::time::Duration; use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; use tracing::{info, warn}; use uuid::Uuid; @@ -1085,6 +1086,7 @@ impl ClientBuilder { dispatcher_handles: RwLock::new(Vec::new()), service_handles: RwLock::new(Vec::new()), job_set: Arc::new(Mutex::new(JoinSet::new())), + finalizers: TaskTracker::new(), in_flight: Arc::new(InFlightRegistry::default()), queue_in_flight, dispatcher_alive, @@ -1185,6 +1187,8 @@ pub struct Client { service_handles: RwLock>>, /// JoinSet tracking in-flight job tasks for graceful drain. job_set: Arc>>, + /// Finalization of attempts whose handler has returned, drained on shutdown. + finalizers: TaskTracker, in_flight: InFlightMap, queue_in_flight: Arc>>, dispatcher_alive: Arc>>, @@ -1566,6 +1570,7 @@ impl Client { completion_handle, effective_storage.clone(), self.dlq_policy.clone(), + self.finalizers.clone(), )); // Admin cancellation listener: fires the in-flight cancel flag @@ -1804,8 +1809,9 @@ impl Client { /// 1. Stop dispatchers (no new jobs claimed) /// 2. Signal in-flight jobs to cancel /// 3. Wait for dispatchers to exit - /// 4. Drain in-flight jobs (heartbeat + maintenance still alive!) - /// 5. Stop heartbeat + maintenance + /// 4. Drain in-flight jobs and their finalization (heartbeat, + /// maintenance and the completion batcher still alive!) + /// 5. Stop heartbeat, maintenance and the completion batcher pub async fn shutdown(&self, timeout: Duration) { info!("Initiating graceful shutdown"); @@ -1863,6 +1869,14 @@ impl Client { ); } + self.finalizers.close(); + if tokio::time::timeout_at(drain_deadline, self.finalizers.wait()) + .await + .is_err() + { + warn!("Shutdown drain timeout exceeded waiting for durable finalization"); + } + // Phase 5: Stop background services (heartbeat + maintenance) self.service_cancel.cancel(); let service_handles: Vec<_> = { diff --git a/awa-worker/src/dispatcher.rs b/awa-worker/src/dispatcher.rs index 11e3e9dbd..916edb02f 100644 --- a/awa-worker/src/dispatcher.rs +++ b/awa-worker/src/dispatcher.rs @@ -851,6 +851,7 @@ mod tests { completion, storage.clone(), crate::executor::DlqPolicy::default(), + tokio_util::task::TaskTracker::new(), )); let config = QueueConfig { max_workers: 1, diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index 3d99ec8df..a4fef41c3 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -11,6 +11,7 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::sync::Arc; use std::time::Duration; +use tokio_util::task::TaskTracker; use tracing::{error, info, info_span, warn, Instrument}; /// Result of executing a job handler. @@ -267,6 +268,7 @@ pub struct JobExecutor { completion_batcher: CompletionBatcherHandle, storage: RuntimeStorage, dlq_policy: DlqPolicy, + finalizers: TaskTracker, } impl JobExecutor { @@ -288,6 +290,7 @@ impl JobExecutor { completion_batcher: CompletionBatcherHandle, storage: RuntimeStorage, dlq_policy: DlqPolicy, + finalizers: TaskTracker, ) -> Self { Self { pool, @@ -301,6 +304,7 @@ impl JobExecutor { completion_batcher, storage, dlq_policy, + finalizers, } } @@ -327,6 +331,7 @@ impl JobExecutor { let completion_batcher = self.completion_batcher.clone(); let storage = self.storage.clone(); let dlq_policy = self.dlq_policy.clone(); + let finalizers = self.finalizers.clone(); let job_id = job.id; let job_run_lease = job.run_lease; let job_kind = job.kind.clone(); @@ -444,7 +449,10 @@ impl JobExecutor { return; }; let dlq_enabled = dlq_policy.enabled_for(&job_queue); - tokio::spawn(async move { + // Graceful shutdown waits for this before it stops the completion + // batcher: the attempt has already left `in_flight`, so nothing + // heartbeats it, and a finalize lost at exit leaves it to rescue. + finalizers.spawn(async move { // Every finalize path is one lease-guarded transaction, and a // deadlock abort rolls the whole transaction back, so running // it again is safe. Without the retry the job would sit diff --git a/awa/tests/ci_timing.rs b/awa/tests/ci_timing.rs new file mode 100644 index 000000000..3704a0ba7 --- /dev/null +++ b/awa/tests/ci_timing.rs @@ -0,0 +1,186 @@ +//! Shared CI contention scaling for the nightly chaos and benchmark suites. +//! +//! The nightly suites run on shared GitHub runners whose CPU allocation +//! varies run to run. Three assertion shapes are sensitive to that and +//! flaked repeatedly through 2026-07 (#399, #434): +//! +//! 1. **Wall-clock waits.** A test waits N seconds for a state the +//! runtime reaches in milliseconds when it has a core to itself. +//! Scale with [`scaled_timeout`]. +//! +//! 2. **Aggressive rescue cadences.** A chaos client sets +//! `heartbeat_staleness` a few multiples above `heartbeat_interval` +//! so its own backdating triggers rescue promptly. Under contention +//! a *live* worker's heartbeat can miss that window, so the runtime +//! correctly rescues a healthy attempt — and the test sees a genuine +//! duplicate completion, which reads as a correctness failure. Scale +//! with [`scaled_staleness`]: the interval stays fast so rescue is +//! still exercised, but the staleness window gains the same margin +//! as the waits around it. +//! +//! 3. **Minimum-progress floors.** A gate asserts "at least N of these +//! happened" to catch a stalled subsystem. The floor must sit far +//! below the *observed* operating point, not just below the nominal +//! one. Scale with [`contention_floor`]. +//! +//! Scaling only ever loosens a bound, and only on CI. A local run keeps +//! the strict values, so a real regression still fails fast on a +//! developer machine. Override with `AWA_CHAOS_TIMEOUT_MULTIPLIER`. +#![allow(dead_code)] + +use std::time::Duration; + +/// Upper bound on the override. `Duration::mul_f64` panics on a non-finite +/// or overflowing result, so an unbounded override would turn a typo into a +/// failed nightly rather than a loose bound. +const MAX_MULTIPLIER: f64 = 100.0; + +/// How much slack to give contention-sensitive bounds. `1.0` locally, +/// `3.0` on CI, or the `AWA_CHAOS_TIMEOUT_MULTIPLIER` override (clamped to +/// `1.0 ..= MAX_MULTIPLIER` so the override can only ever loosen, and can +/// never produce a duration `mul_f64` refuses). +pub fn chaos_timeout_multiplier() -> f64 { + let override_var = std::env::var("AWA_CHAOS_TIMEOUT_MULTIPLIER").ok(); + multiplier_from(override_var.as_deref(), std::env::var_os("CI").is_some()) +} + +/// The multiplier decision, with the environment passed in so it can be +/// tested without mutating process-global state. +fn multiplier_from(override_var: Option<&str>, is_ci: bool) -> f64 { + // `"inf"` and `"nan"` parse successfully as f64, and both reach + // `Duration::mul_f64`, which panics rather than saturating. Treat a + // non-finite override as no override at all. + if let Some(parsed) = override_var + .and_then(|raw| raw.parse::().ok()) + .filter(|parsed| parsed.is_finite()) + { + return parsed.clamp(1.0, MAX_MULTIPLIER); + } + + if is_ci { + 3.0 + } else { + 1.0 + } +} + +/// Grow a wait deadline by the contention multiplier. +pub fn scaled_timeout(timeout: Duration) -> Duration { + timeout.mul_f64(chaos_timeout_multiplier()) +} + +/// Grow a heartbeat-staleness window by the contention multiplier. +/// +/// Distinct from [`scaled_timeout`] only in intent: this one is passed to +/// `ClientBuilder::heartbeat_staleness`, where the cost of being too tight +/// is a spurious rescue of a live attempt rather than a timeout. Paired +/// heartbeat/rescue *intervals* are deliberately left unscaled so the +/// rescue path still runs at chaos cadence. +pub fn scaled_staleness(staleness: Duration) -> Duration { + staleness.mul_f64(chaos_timeout_multiplier()) +} + +/// Shrink a minimum-progress floor by the contention multiplier. +/// +/// Use for "this subsystem must have advanced at least N times" gates. +/// The returned floor is at least 1: the regression these gates exist to +/// catch is a fully stalled subsystem, so zero progress must still fail. +pub fn contention_floor(nominal: i64) -> i64 { + let scaled = (nominal as f64 / chaos_timeout_multiplier()).floor() as i64; + scaled.max(1) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn local_runs_keep_strict_bounds() { + assert_eq!(multiplier_from(None, false), 1.0); + } + + #[test] + fn ci_runs_get_margin() { + assert_eq!(multiplier_from(None, true), 3.0); + } + + #[test] + fn override_can_only_loosen() { + assert_eq!(multiplier_from(Some("0.1"), true), 1.0); + assert_eq!(multiplier_from(Some("-5"), true), 1.0); + assert_eq!(multiplier_from(Some("6"), false), 6.0); + } + + #[test] + fn unparseable_override_falls_back_to_the_environment() { + assert_eq!(multiplier_from(Some("banana"), true), 3.0); + assert_eq!(multiplier_from(Some(""), false), 1.0); + } + + #[test] + fn non_finite_override_falls_back_to_the_environment() { + // These all parse as f64 and would otherwise reach `mul_f64`, which + // panics on a non-finite result instead of saturating. + for raw in ["inf", "-inf", "infinity", "NaN", "nan"] { + assert_eq!(multiplier_from(Some(raw), true), 3.0, "raw={raw}"); + assert_eq!(multiplier_from(Some(raw), false), 1.0, "raw={raw}"); + } + } + + #[test] + fn oversized_override_is_clamped() { + assert_eq!(multiplier_from(Some("1e308"), true), MAX_MULTIPLIER); + assert_eq!(multiplier_from(Some("10000"), true), MAX_MULTIPLIER); + } + + #[test] + fn every_multiplier_produces_a_usable_duration() { + // `Duration::mul_f64` panics rather than saturating, so the clamp is + // what keeps a bad override from failing the nightly it was set to + // rescue. + for raw in [ + "inf", "-inf", "NaN", "1e308", "10000", "0", "-1", "banana", "", + ] { + for is_ci in [true, false] { + let m = multiplier_from(Some(raw), is_ci); + assert!(m.is_finite() && (1.0..=MAX_MULTIPLIER).contains(&m)); + // Would panic if the multiplier were unbounded. + let _ = Duration::from_secs(60).mul_f64(m); + } + } + } + + #[test] + fn scaling_a_timeout_grows_it() { + // 250ms staleness against a 50ms heartbeat interval is a 5x margin + // locally; on CI it becomes 15x, which is what stops a contended + // runner's live worker from being rescued as if it had died. + assert_eq!( + Duration::from_millis(250).mul_f64(multiplier_from(None, true)), + Duration::from_millis(750) + ); + } + + #[test] + fn progress_floor_keeps_margin_under_the_observed_rate() { + // The receipt gate's nominal floor is duration_secs / 4 = 45 for a + // 180s run. #399 saw a healthy run produce 41, so the CI floor has + // to sit well under that. + let ci_floor = (45.0_f64 / multiplier_from(None, true)).floor() as i64; + assert_eq!(ci_floor, 15); + assert!( + ci_floor < 41, + "floor must clear the observed operating point" + ); + } + + #[test] + fn progress_floor_never_reaches_zero() { + // Whatever the multiplier, a fully pinned ring (zero rotations) has + // to keep failing — that is the regression the gate exists for. + assert_eq!(contention_floor(45), contention_floor(45).max(1)); + let absurd = (45.0_f64 / multiplier_from(Some("10000"), true)).floor() as i64; + assert_eq!(absurd.max(1), 1); + assert_eq!(contention_floor(0), 1); + } +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index a060369f8..edeacac09 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -3,6 +3,8 @@ //! These tests exercise the full dispatcher/worker/maintenance wiring with the //! queue_storage backend enabled. +mod ci_timing; + use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, BatchOperationSpec, PruneOutcome, QueueStorage, QueueStorageConfig, RotateOutcome, SkipReason, @@ -2009,6 +2011,231 @@ impl Worker for MultiClientTrackingWorker { } } +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct GracefulRestartJob { + id: i64, +} + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct GracefulRestartFollowUp { + id: i64, +} + +/// The first execution reports its attempt and returns `Completed` once +/// shutdown signals cancellation; later executions complete at once. +#[derive(Clone, Default)] +struct GracefulRestartWorker { + executions: Arc, + first_attempt: Arc>>, + entered: Arc, +} + +#[async_trait::async_trait] +impl Worker for GracefulRestartWorker { + fn kind(&self) -> &'static str { + "graceful_restart_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + if self.executions.fetch_add(1, Ordering::SeqCst) > 0 { + return Ok(JobResult::Completed); + } + *self.first_attempt.lock().expect("first attempt") = Some((ctx.job.id, ctx.job.run_lease)); + self.entered.notify_one(); + while !ctx.is_cancelled() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + Ok(JobResult::Completed) + } +} + +/// A runtime on its own tokio runtime and pool, so stopping it after a +/// graceful shutdown drops whatever the client left running, as a process +/// exit does. +struct ProcessHost { + runtime: tokio::runtime::Runtime, + client: Arc, +} + +impl ProcessHost { + async fn start( + options: sqlx::postgres::PgConnectOptions, + build: impl FnOnce(sqlx::PgPool) -> Client + Send + 'static, + ) -> Self { + tokio::task::spawn_blocking(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .expect("build process runtime"); + let client = runtime.block_on(async move { + let pool = PgPoolOptions::new() + .max_connections(6) + .connect_with(options) + .await + .expect("connect process pool"); + let client = build(pool); + client.start().await.expect("start process client"); + client + }); + Self { + runtime, + client: Arc::new(client), + } + }) + .await + .expect("start process host") + } + + async fn exit_gracefully(self, timeout: Duration) { + let client = self.client.clone(); + self.runtime + .spawn(async move { client.shutdown(timeout).await }) + .await + .expect("graceful shutdown"); + tokio::task::spawn_blocking(move || self.runtime.shutdown_background()) + .await + .expect("stop process runtime"); + } +} + +/// A graceful shutdown must finalize every attempt whose handler returned +/// during the drain. Finalization used to run detached from the drain, so a +/// process that exited after `shutdown` returned could drop it: the attempt +/// stayed open with no heartbeat and stale-heartbeat rescue ran the completed +/// job again. A held receipt-attempt lock keeps this attempt's finalize in +/// flight while shutdown drains, which the detached version did not wait for. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_graceful_shutdown_finalizes_returned_attempts() { + // Each host below opens its own pool of 6, and only one runs at a time. + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_graceful_restart"; + let store_config = QueueStorageConfig { + schema: "awa_qs_graceful_restart".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + let job_id = enqueue_job( + &pool, + &store, + &GracefulRestartJob { id: 1 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + + let worker = GracefulRestartWorker::default(); + let options = (*pool.connect_options()).clone(); + let build = |worker: GracefulRestartWorker, store_config: QueueStorageConfig| { + move |pool: sqlx::PgPool| { + Client::builder(pool) + .queue( + queue, + QueueConfig { + max_workers: 2, + poll_interval: Duration::from_millis(25), + deadline_duration: Duration::ZERO, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(worker) + // The follow-up's queue is not served; it only routes the + // completion through a transaction of its own. + .on_completed_enqueue::( + |args, _job| GracefulRestartFollowUp { id: args.id }, + ) + .heartbeat_interval(Duration::from_millis(100)) + .heartbeat_staleness(ci_timing::scaled_staleness(Duration::from_secs(1))) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .promote_interval(Duration::from_millis(25)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build graceful restart client") + } + }; + + let host = + ProcessHost::start(options.clone(), build(worker.clone(), store_config.clone())).await; + tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(10)), + worker.entered.notified(), + ) + .await + .expect("the job should start"); + let (attempt_job_id, run_lease) = worker + .first_attempt + .lock() + .expect("first attempt") + .expect("first attempt recorded"); + assert_eq!(attempt_job_id, job_id); + + let mut lock_holder = pool.acquire().await.expect("lock connection"); + sqlx::query( + "SELECT pg_catalog.pg_advisory_lock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(job_id) + .bind(run_lease) + .execute(&mut *lock_holder) + .await + .expect("hold the attempt's finalize lock"); + + let exit = + tokio::spawn(host.exit_gracefully(ci_timing::scaled_timeout(Duration::from_secs(10)))); + tokio::time::sleep(Duration::from_millis(500)).await; + sqlx::query( + "SELECT pg_catalog.pg_advisory_unlock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(job_id) + .bind(run_lease) + .execute(&mut *lock_holder) + .await + .expect("release the attempt's finalize lock"); + drop(lock_holder); + exit.await.expect("graceful exit"); + + let after_exit = store + .load_job(&pool, job_id) + .await + .expect("load job") + .expect("job row"); + + // A later runtime rescues anything the exit left open. + let rescuer = ProcessHost::start(options, build(worker.clone(), store_config)).await; + let finished = wait_for_job_state( + &store, + &pool, + job_id, + &[JobState::Completed], + ci_timing::scaled_timeout(Duration::from_secs(15)), + ) + .await; + rescuer + .exit_gracefully(ci_timing::scaled_timeout(Duration::from_secs(10))) + .await; + + assert_eq!( + (after_exit.state, after_exit.run_lease), + (JobState::Completed, run_lease), + "graceful shutdown returned before the completed attempt was finalized" + ); + assert_eq!( + (worker.executions.load(Ordering::SeqCst), finished.run_lease), + (1, run_lease), + "the job ran again after a graceful shutdown" + ); +} + /// ADR-023 claim-ring control-plane smoke test. Exercises /// `rotate_claims`, the busy-check, and `prune_oldest_claims` on an /// empty schema: rotate cycles through every slot, prune is a noop diff --git a/docs/test-plan.md b/docs/test-plan.md index 435b18861..618ff69a2 100644 --- a/docs/test-plan.md +++ b/docs/test-plan.md @@ -32,6 +32,7 @@ Tests run against real Postgres 15+ (not managed services) using a dedicated tes | B5 | Deadline rescue signals ctx.is_cancelled() | ✓ | | | B6 | UniqueConflict.constraint field | ✓ | | | B7 | Admin cancel signals in-flight handler cancellation | ✓ | ✓ | +| B8 | Shutdown finalizes attempts whose handler returned during the drain, so a process exit cannot leave them for rescue | ✓ | | | T75 | Priority aging maintenance task promotes long-waiting low-priority jobs | ✓ | | ### Uniqueness (Python) From 4a46e89c9f2da4ff371411bcee3c2eb1bfa768be Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 17:45:29 +1300 Subject: [PATCH 10/16] feat(worker): warn when the prune lock timeout reaches deadlock_timeout or statement_timeout --- CHANGELOG.md | 5 +- awa-model/src/queue_storage.rs | 7 ++- awa-python/python/awa/client.py | 4 +- awa-worker/src/client.rs | 84 ++++++++++++++++++++++++++++++++- awa-worker/src/storage.rs | 82 ++++++++++++++++++++++++++++++++ docs/configuration.md | 4 +- 6 files changed, 178 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ed8acfa6..3579e877c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,7 +15,10 @@ Notable changes between releases. Detailed migration notes for storage transitio `blocked` every tick and the rings stopped recycling. Set it with `ClientBuilder::prune_lock_timeout` (Rust) or `queue_storage_prune_lock_timeout_ms` on `client.start()` (Python); the - default stays 10 ms and zero is rejected. `with_prune_lock_timeout` is now + default stays 10 ms and zero or a value above PostgreSQL's `lock_timeout` + maximum is rejected. Start logs a warning when the value is at or above the + connection's `deadlock_timeout` (prune waiting that long can make deadlock + detection cancel a claim instead) or `statement_timeout`. `with_prune_lock_timeout` is now public, with `set_prune_lock_timeout` and `prune_lock_timeout()`. `awa_maintenance_prune_attempts_total{awa_ring,awa_ring_outcome,awa_ring_reason}` gains `backed_off` (prune skipped by the backoff that follows consecutive diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index c50999925..3401f9cd3 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -2511,8 +2511,11 @@ impl QueueStorage { /// next tick. The default is 10 ms, which keeps prune from queueing behind /// readers and so from stalling claims and enqueues behind its pending /// `ACCESS EXCLUSIVE` request. Raise it when prune is persistently - /// `Blocked` by short-lived readers of the ring partitions. The timeout is - /// applied at millisecond granularity. + /// `Blocked` by short-lived readers of the ring partitions. Keep it well + /// below `deadlock_timeout` (default 1 s) and any `statement_timeout`: a + /// prune that waits as long as `deadlock_timeout` triggers deadlock + /// detection, which may cancel a concurrent claim instead of the prune. The + /// timeout is applied at millisecond granularity. /// /// Returns a validation error for a zero timeout or one above PostgreSQL's /// `lock_timeout` maximum (`i32::MAX` milliseconds, about 24.8 days). diff --git a/awa-python/python/awa/client.py b/awa-python/python/awa/client.py index 57346a6b7..ea3884a24 100644 --- a/awa-python/python/awa/client.py +++ b/awa-python/python/awa/client.py @@ -905,7 +905,9 @@ async def start( ``lock_timeout`` the maintenance leader uses when pruning the queue, lease and claim rings (default 10 ms). A prune that cannot lock its partitions in time is counted as ``blocked`` and retried on the next - tick; raise it when ``blocked`` outcomes persist. Must be > 0. + tick; raise it when ``blocked`` outcomes persist, keeping it well below the + server's ``deadlock_timeout`` and ``statement_timeout`` (a warning is + logged at start otherwise). Must be > 0. Queue dicts also accept ``claimers`` and ``claim_batch_size`` for hot queue-storage queues. ``claimers`` adds dispatcher/claimer loops that diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 1d16bb47c..5cc35083e 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -27,7 +27,7 @@ use std::time::Duration; use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; -use tracing::{info, warn}; +use tracing::{debug, info, warn}; use uuid::Uuid; /// Errors returned when building a worker client. @@ -936,7 +936,11 @@ impl ClientBuilder { /// `ACCESS EXCLUSIVE` request from stalling claims and enqueues behind it. /// Raise it when `blocked` outcomes persist and the rings stop recycling /// because short-lived readers (admin queries, dashboards) repeatedly hold - /// the ring partitions at the moment prune tries. A zero timeout, or one above + /// the ring partitions at the moment prune tries. Keep it well below + /// PostgreSQL's `deadlock_timeout` (default 1 s), so that a waiting prune + /// is abandoned before deadlock detection can pick a claim as the victim, + /// and below any `statement_timeout`; [`Self::start`] logs a warning when + /// it is not. A zero timeout, or one above /// PostgreSQL's `lock_timeout` maximum (`i32::MAX` ms), makes /// [`Self::build`] fail. Only takes effect when queue storage is active. /// May be set before or after [`Self::queue_storage`]; retained across @@ -1495,6 +1499,37 @@ impl Client { reporter.publish_snapshot().await; } + async fn warn_on_unsafe_prune_lock_timeout(&self) { + let Some(runtime) = self.storage.queue_storage() else { + return; + }; + let timeout = runtime.store.prune_lock_timeout(); + let risks = match crate::storage::check_prune_lock_timeout(&self.pool, timeout).await { + Ok(risks) => risks, + Err(err) => { + debug!(error = %err, "Could not read lock settings to check the prune lock timeout"); + return; + } + }; + for risk in risks { + match risk { + crate::storage::PruneLockTimeoutRisk::DeadlockTimeout { limit_ms } => warn!( + prune_lock_timeout_ms = timeout.as_millis() as u64, + deadlock_timeout_ms = limit_ms, + "Prune lock timeout is at or above deadlock_timeout; a waiting prune can \ + make PostgreSQL run deadlock detection and cancel a concurrent claim. \ + Keep the timeout well below deadlock_timeout" + ), + crate::storage::PruneLockTimeoutRisk::StatementTimeout { limit_ms } => warn!( + prune_lock_timeout_ms = timeout.as_millis() as u64, + statement_timeout_ms = limit_ms, + "Prune lock timeout is at or above statement_timeout; prune is cancelled \ + by the statement timeout and fails instead of reporting blocked" + ), + } + } + } + async fn log_transition_startup_status( &self, effective_storage: &RuntimeStorage, @@ -1553,6 +1588,7 @@ impl Client { self.log_transition_startup_status(&effective_storage) .await?; + self.warn_on_unsafe_prune_lock_timeout().await; admin::sync_queue_descriptors( &self.pool, @@ -2909,6 +2945,50 @@ mod tests { ); } + #[tokio::test] + async fn prune_lock_timeout_check_reads_the_session_limits() { + use crate::storage::{check_prune_lock_timeout, PruneLockTimeoutRisk}; + + let url = database_url(); + ensure_database_exists(&url).await; + let pool = PgPoolOptions::new() + .max_connections(1) + .after_connect(|conn, _meta| { + Box::pin(async move { + sqlx::query("SET deadlock_timeout = '200ms'") + .execute(&mut *conn) + .await?; + sqlx::query("SET statement_timeout = '3s'") + .execute(&mut *conn) + .await?; + Ok(()) + }) + }) + .connect(&url) + .await + .expect("connect with session limits"); + + assert!(check_prune_lock_timeout(&pool, Duration::from_millis(100)) + .await + .unwrap() + .is_empty()); + assert_eq!( + check_prune_lock_timeout(&pool, Duration::from_millis(300)) + .await + .unwrap(), + vec![PruneLockTimeoutRisk::DeadlockTimeout { limit_ms: 200 }] + ); + assert_eq!( + check_prune_lock_timeout(&pool, Duration::from_secs(5)) + .await + .unwrap(), + vec![ + PruneLockTimeoutRisk::DeadlockTimeout { limit_ms: 200 }, + PruneLockTimeoutRisk::StatementTimeout { limit_ms: 3000 }, + ] + ); + } + #[tokio::test] async fn out_of_range_prune_lock_timeout_is_rejected_at_build() { for timeout in [Duration::ZERO, Duration::from_millis(i32::MAX as u64 + 1)] { diff --git a/awa-worker/src/storage.rs b/awa-worker/src/storage.rs index e56dc3bfa..828c0a73d 100644 --- a/awa-worker/src/storage.rs +++ b/awa-worker/src/storage.rs @@ -67,3 +67,85 @@ impl RuntimeStorage { self.queue_storage().map(|runtime| runtime.store.clone()) } } + +/// Server settings that a prune `lock_timeout` must stay below. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PruneLockTimeoutRisk { + /// A prune waiting this long runs PostgreSQL's deadlock detector, which + /// may cancel a concurrent claim or enqueue instead of the prune. + DeadlockTimeout { limit_ms: u64 }, + /// `statement_timeout` cancels the statement first, so the prune fails + /// with an error instead of reporting `blocked`. + StatementTimeout { limit_ms: u64 }, +} + +/// Compare a prune lock timeout with the session's `deadlock_timeout` and +/// `statement_timeout` (milliseconds; a zero `statement_timeout` is disabled). +pub(crate) fn prune_lock_timeout_risks( + timeout: Duration, + deadlock_timeout_ms: Option, + statement_timeout_ms: Option, +) -> Vec { + let timeout_ms = timeout.as_millis(); + let mut risks = Vec::new(); + if let Some(limit_ms) = deadlock_timeout_ms.filter(|limit| *limit > 0) { + if timeout_ms >= u128::from(limit_ms) { + risks.push(PruneLockTimeoutRisk::DeadlockTimeout { limit_ms }); + } + } + if let Some(limit_ms) = statement_timeout_ms.filter(|limit| *limit > 0) { + if timeout_ms >= u128::from(limit_ms) { + risks.push(PruneLockTimeoutRisk::StatementTimeout { limit_ms }); + } + } + risks +} + +/// Read the connection's effective `deadlock_timeout` and `statement_timeout` +/// and report the ones the prune lock timeout reaches. +pub(crate) async fn check_prune_lock_timeout( + pool: &sqlx::PgPool, + timeout: Duration, +) -> Result, sqlx::Error> { + let settings: Vec<(String, i64)> = sqlx::query_as( + "SELECT name, setting::bigint FROM pg_settings \ + WHERE name IN ('deadlock_timeout', 'statement_timeout')", + ) + .fetch_all(pool) + .await?; + let setting = |name: &str| { + settings + .iter() + .find(|(setting_name, _)| setting_name == name) + .and_then(|(_, value)| u64::try_from(*value).ok()) + }; + Ok(prune_lock_timeout_risks( + timeout, + setting("deadlock_timeout"), + setting("statement_timeout"), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prune_lock_timeout_risks_flag_timeouts_at_or_above_server_limits() { + let below = Duration::from_millis(999); + let at = Duration::from_millis(1000); + assert!(prune_lock_timeout_risks(below, Some(1000), Some(0)).is_empty()); + assert_eq!( + prune_lock_timeout_risks(at, Some(1000), Some(0)), + vec![PruneLockTimeoutRisk::DeadlockTimeout { limit_ms: 1000 }] + ); + assert_eq!( + prune_lock_timeout_risks(Duration::from_secs(30), Some(1000), Some(5000)), + vec![ + PruneLockTimeoutRisk::DeadlockTimeout { limit_ms: 1000 }, + PruneLockTimeoutRisk::StatementTimeout { limit_ms: 5000 }, + ] + ); + assert!(prune_lock_timeout_risks(Duration::from_secs(30), None, None).is_empty()); + } +} diff --git a/docs/configuration.md b/docs/configuration.md index 4c0007681..36722de60 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -570,7 +570,7 @@ await client.start( | `queue_rotate_interval` | `1000ms` | How often ready/terminal segments rotate | | `lease_rotate_interval` | `1000ms` | How often lease segments rotate | | `claim_rotate_interval` | matches `queue_rotate_interval` | How often the ADR-023 claim-ring rotates. Set with `ClientBuilder::claim_rotate_interval` (Rust) or `queue_storage_claim_rotate_interval_ms` (Python). Test harnesses sometimes set this to a long interval to pin claim-ring layout for deterministic count assertions. | -| `prune_lock_timeout` | `10ms` | The PostgreSQL `lock_timeout` the maintenance leader applies when it takes the exclusive locks to reclaim a queue, lease or claim ring partition (and when it folds terminal-count deltas). A prune that cannot lock its partitions within the timeout rolls back, is counted as `blocked`, and is retried on the next maintenance tick. Set with `ClientBuilder::prune_lock_timeout` (Rust) or `queue_storage_prune_lock_timeout_ms` (Python). Must be greater than zero. The short default keeps prune's pending `ACCESS EXCLUSIVE` request from stalling claims and enqueues queued behind it. See below for when to raise it. | +| `prune_lock_timeout` | `10ms` | The PostgreSQL `lock_timeout` the maintenance leader applies when it takes the exclusive locks to reclaim a queue, lease or claim ring partition (and when it folds terminal-count deltas). A prune that cannot lock its partitions within the timeout rolls back, is counted as `blocked`, and is retried on the next maintenance tick. Set with `ClientBuilder::prune_lock_timeout` (Rust) or `queue_storage_prune_lock_timeout_ms` (Python). Must be greater than zero and at most `2147483647` ms; keep it well below `deadlock_timeout` and `statement_timeout`. The short default keeps prune's pending `ACCESS EXCLUSIVE` request from stalling claims and enqueues queued behind it. See below for when to raise it. | The benchmark harness in [postgresql-job-queue-benchmarking](https://github.com/hardbyte/postgresql-job-queue-benchmarking) reads `QUEUE_SLOT_COUNT`, `LEASE_SLOT_COUNT`, `CLAIM_SLOT_COUNT`, `QUEUE_STRIPE_COUNT`, and `LEASE_CLAIM_RECEIPTS` from the environment. Those env vars are benchmark configuration, not general worker-runtime configuration. @@ -580,7 +580,7 @@ Use the defaults unless you have a reason not to: - Increase `lease_slot_count` if lease churn is high enough that dead tuples in the lease ring stop collapsing promptly. - Increase `claim_slot_count` if the rotation cadence (`claim_rotate_interval`) plus the slot count combine to a partition retention window shorter than your longest in-flight zero-deadline short job; running out of empty slots forces `rotate_claims` to return `SkippedBusy` and the receipt-plane churn falls back onto a smaller working set of partitions. - Increase `queue_stripe_count` only for measured workloads where many small replicas contend on the same logical queue. This is queue-storage configuration, not a per-queue `QueueConfig` field: it gives each logical queue the same number of physical stripes. Striping spreads a hot logical queue over `queue#N` physical coordination paths, but it weakens perfect global ordering and can regress calmer shapes if overused. For single hot-queue workloads, benchmark `2` stripes first; higher values should be treated as workload-specific tuning rather than a general default. -- Increase `prune_lock_timeout` when `increase(awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}[15m])` stays above zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Raise it in steps (for example `100ms`, then `1s`) and keep it well below the time you can tolerate claims and enqueues queueing behind prune; while prune waits, new accesses to the same partition queue behind its lock request. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for up to 32 ticks after consecutive `skipped_active` or `blocked` results, so a persistently blocked ring is retried less often over time, and `error` is a database failure. +- Increase `prune_lock_timeout` when `increase(awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}[15m])` stays above zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Raise it in small steps (for example `50ms`, then `100ms`, then `250ms`) and keep it well below PostgreSQL's `deadlock_timeout` (default `1s`) and below any `statement_timeout` on the pool's connections. A prune that waits as long as `deadlock_timeout` makes PostgreSQL run its deadlock detector, and if claims or enqueues are part of a lock cycle with it, the detector can cancel them rather than the prune; a short `lock_timeout` abandons the prune first, so the prune is always the one that gives way. A `statement_timeout` at or below the lock timeout cancels the prune statement before the lock timeout fires, so it fails with an error instead of reporting `blocked`. Startup logs a warning when the configured timeout is at or above either setting. While prune waits, new accesses to the same partition also queue behind its lock request, so a longer timeout trades claim latency for reclamation. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for up to 32 ticks after consecutive `skipped_active` or `blocked` results, so a persistently blocked ring is retried less often over time, and `error` is a database failure. - Increase rotation intervals to reduce partition churn and metadata activity. - Decrease rotation intervals to tighten dead-tuple bounds at the cost of more frequent rotate/prune work. From ee1786141f43165675822c61676afc0c694e044c Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 17:49:53 +1300 Subject: [PATCH 11/16] fix(queue-storage): skip locked due rows in promotion selection and bound lane waiters by pool size The batch selection was a plain SELECT ... LIMIT and only the per-lane DELETE skipped locked rows. When the oldest due rows were locked by another transaction, such as a lane waiter left running by a former leader, every batch of a tick reselected them and promoted nothing until the lock went away. The selection now takes FOR UPDATE OF deferred SKIP LOCKED. It runs as its own statement, so the row locks end with it, and deferred_jobs is not partitioned, so no run-time pruning meets them. Each lane waiter holds a pool connection while it waits. The leader now runs one waiter per eight pool connections, at least one and at most four, and the configuration docs say so. --- CHANGELOG.md | 8 +- awa-model/src/queue_storage.rs | 7 ++ awa-worker/src/maintenance.rs | 21 ++++- awa/tests/queue_storage_runtime_test.rs | 120 ++++++++++++++++++++++++ docs/configuration.md | 2 +- 5 files changed, 152 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 81e2d5d61..92cc53275 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,9 +18,11 @@ Notable changes between releases. Detailed migration notes for storage transitio `lock_timeout`. A lane that times out keeps its rows deferred and is handed to a dedicated task that waits up to 30 s for the lock, queued behind the producers already waiting, so the lane is still promoted in - turn however busy it stays; the leader skips it meanwhile. At most four - lanes wait this way at once, and a further contended lane is retried on - the next tick. `awa.maintenance.promote.lane_contention` + turn however busy it stays; the leader skips it meanwhile. Each waiter + holds a pool connection, so the leader runs one per eight pool + connections, at least one and at most four; a further contended lane is + retried on the next tick. Batch selection skips rows another transaction + has locked, so locked due rows cannot fill every batch of a tick. `awa.maintenance.promote.lane_contention` (`awa_maintenance_promote_lane_contention_total{awa_job_state,awa_job_queue,awa_promote_outcome}`) counts `handed_off`, `waiter_limit` and `waiter_timed_out`. Promoted rows now take lane sequence numbers in due order within their lane. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index f598d909f..9bc301905 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -14484,6 +14484,12 @@ impl QueueStorage { /// outside the lanes in `skip`, and promote each lane's share in its own /// transaction. /// + /// The selection skips rows another transaction has locked, such as a + /// lane promotion still waiting for its lane, so they cannot fill every + /// batch of a tick. It runs as its own statement, so its row locks end + /// with it; `deferred_jobs` is not partitioned, so no run-time pruning + /// meets those locks. + /// /// Re-enqueueing a row takes its lane's enqueue reservation lock, which /// producers hold until they commit, so a lane under heavy enqueue load /// can keep its promotion waiting. A transaction per lane confines that @@ -14521,6 +14527,7 @@ impl QueueStorage { ) ORDER BY run_at ASC, priority ASC, job_id ASC LIMIT $2 + FOR UPDATE OF deferred SKIP LOCKED "# )) .bind(state) diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index e2a5a5eb3..4623963aa 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -729,9 +729,14 @@ const PROMOTE_LANE_LOCK_TIMEOUT: Duration = Duration::from_millis(50); /// ahead of it, however busy the lane stays. const PROMOTE_LANE_WAITER_LOCK_TIMEOUT: Duration = Duration::from_secs(30); /// Most lanes with a dedicated waiter at once. A waiter holds a pool -/// connection while it waits. +/// connection while it waits, so smaller pools get fewer: one per eight +/// connections, and always one so a busy lane is still promoted. const PROMOTE_MAX_LANE_WAITERS: usize = 4; +fn lane_waiter_limit(pool_max_connections: u32) -> usize { + (pool_max_connections as usize / 8).clamp(1, PROMOTE_MAX_LANE_WAITERS) +} + /// Queue-storage promotion of due scheduled and retryable rows, as the /// maintenance leader runs it for one store. /// @@ -748,11 +753,14 @@ pub struct QueueStoragePromoter { metrics: AwaMetrics, waiting: Arc>>, waiters: std::sync::Mutex>, + max_waiters: usize, } impl QueueStoragePromoter { pub fn new(pool: PgPool, store: Arc, metrics: AwaMetrics) -> Self { + let max_waiters = lane_waiter_limit(pool.options().get_max_connections()); Self { + max_waiters, pool, store, metrics, @@ -825,7 +833,7 @@ impl QueueStoragePromoter { let key = (state, lane); { let mut waiting = self.waiting.lock().expect("waiting lanes poisoned"); - if waiting.len() >= PROMOTE_MAX_LANE_WAITERS { + if waiting.len() >= self.max_waiters { self.metrics.record_promotion_lane_contention( state.as_str(), &key.1.queue, @@ -3677,6 +3685,15 @@ mod tests { } } + #[test] + fn lane_waiters_scale_with_the_pool_and_never_drop_to_zero() { + assert_eq!(lane_waiter_limit(1), 1); + assert_eq!(lane_waiter_limit(10), 1); + assert_eq!(lane_waiter_limit(16), 2); + assert_eq!(lane_waiter_limit(32), 4); + assert_eq!(lane_waiter_limit(200), 4); + } + #[test] fn compute_fire_times_coalesces_missed_existing_fires_by_default() { let last = Utc.with_ymd_and_hms(2026, 5, 7, 12, 0, 0).unwrap(); diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index b0093749c..ac1d65918 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -8276,6 +8276,126 @@ async fn test_queue_storage_runtime_busy_lane_holds_back_neither_promotion_nor_r client.shutdown(Duration::from_secs(5)).await; } +/// Rows another transaction holds locked, as a lane promotion still waiting +/// under a former leader does, must not fill every batch of a promotion tick. +/// With more of the oldest due rows locked than one batch selects, a due job +/// on another queue must still be promoted and run promptly, and the locked +/// rows must be promoted once released. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_locked_due_rows_do_not_starve_promotion() { + const LOCKED: usize = 4_096 + 64; + let (_db_guard, pool) = setup_pool(10).await; + let locked_queue = "qs_locked_due_rows"; + let other = "qs_locked_due_other"; + let schema = "awa_qs_locked_due_rows"; + let store = create_store(&pool, schema).await; + + let oldest = Utc::now() - chrono::Duration::seconds(120); + let params: Vec<_> = (0..LOCKED) + .map(|n| { + insert::params_with( + &CompleteJob { id: n as i64 }, + InsertOpts { + queue: locked_queue.to_string(), + run_at: Some(oldest), + ..Default::default() + }, + ) + .expect("build locked-row params") + }) + .collect(); + for chunk in params.chunks(2_000) { + store + .enqueue_params_batch(&pool, chunk) + .await + .expect("enqueue rows to lock"); + } + let mut holder = pool.begin().await.expect("begin row holder"); + let locked: Vec = sqlx::query_scalar(&format!( + "SELECT job_id FROM {schema}.deferred_jobs WHERE queue = $1 FOR UPDATE" + )) + .bind(locked_queue) + .fetch_all(holder.as_mut()) + .await + .expect("lock the oldest due rows"); + assert_eq!(locked.len(), LOCKED); + + let other_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: -1 }, + InsertOpts { + queue: other.to_string(), + run_at: Some(Utc::now() - chrono::Duration::seconds(1)), + ..Default::default() + }, + ) + .await; + + let client = Client::builder(pool.clone()) + .queue( + other, + QueueConfig { + max_workers: 4, + poll_interval: Duration::from_millis(25), + ..QueueConfig::default() + }, + ) + .queue_storage( + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }, + Duration::from_secs(60), + Duration::from_secs(60), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(CompleteWorker) + .promote_interval(Duration::from_millis(50)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("Failed to build locked-rows client"); + client + .start() + .await + .expect("Failed to start locked-rows client"); + + wait_for_job_state( + &store, + &pool, + other_job, + &[JobState::Completed], + scaled_timeout(Duration::from_secs(2)), + ) + .await; + + holder.rollback().await.expect("release the locked rows"); + let deadline = Instant::now() + scaled_timeout(Duration::from_secs(10)); + loop { + let remaining: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {schema}.deferred_jobs WHERE queue = $1" + )) + .bind(locked_queue) + .fetch_one(&pool) + .await + .expect("count deferred rows"); + if remaining == 0 { + break; + } + assert!( + Instant::now() < deadline, + "{remaining} released rows were never promoted" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + + client.shutdown(Duration::from_secs(5)).await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_runtime_stale_heartbeat_rescue() { let (_db_guard, pool) = setup_pool(10).await; diff --git a/docs/configuration.md b/docs/configuration.md index 71bd3e72c..22d7e7b41 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -245,7 +245,7 @@ Retry backoff and `Snooze` use the same mechanism: the current attempt closes, a Operational knobs: - `promote_interval` controls how often maintenance checks for due deferred rows. -- A queue whose enqueue lane is kept busy by producers does not delay promotion of other queues or the leader's rescue: promotion waits 50 ms for a lane's lock, then hands the lane to a dedicated waiter that is served in turn. `awa_maintenance_promote_lane_contention_total` counts those hand-offs by queue; a steady `waiter_limit` or `waiter_timed_out` rate means more lanes are saturated than the leader's four waiters cover, or a transaction is holding a lane lock for longer than 30 s. +- A queue whose enqueue lane is kept busy by producers does not delay promotion of other queues or the leader's rescue: promotion waits 50 ms for a lane's lock, then hands the lane to a dedicated waiter that is served in turn. Each waiter holds one connection from the worker's pool while it waits, on top of the leader's own. The leader runs at most one waiter per eight pool connections, never fewer than one or more than four, so a 10-connection pool gives up at most one extra connection to a busy lane. `awa_maintenance_promote_lane_contention_total` counts those hand-offs by queue. A steady `waiter_limit` rate means more lanes are saturated than the waiters cover, and a larger pool allows more waiters. `waiter_timed_out` means a transaction held a lane lock for longer than 30 s. - Per-job `run_at` controls the due time. - Periodic schedules are declared in worker code with `periodic()` and managed through the CLI/UI `cron` surface. From 941bd48dd2f3581a99b4d964961dc8fc23f51e7c Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 18:18:30 +1300 Subject: [PATCH 12/16] fix(queue-storage): stop a terminal-count rollup pass at its first blocked slot (0.6) The remaining sealed slots would most likely wait out the same lock timeout, and with a raised prune timeout that holds the maintenance leader for one timeout per slot. The rollup test now seals two slots with pending deltas and expects one blocked slot, then both rolled once the lock is released. --- CHANGELOG.md | 3 +- awa-model/src/queue_storage.rs | 3 ++ awa/tests/queue_storage_runtime_test.rs | 40 ++++++++++++++----------- 3 files changed, 28 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e62d6b908..491090b54 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,7 +22,8 @@ Notable changes between releases. Detailed migration notes for storage transitio 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. + the rollback. A rollup pass stops at its first blocked slot rather than + waiting out the timeout again on each remaining one. - **`admin::job_kind_overviews` on queue storage no longer depends on the planner's row estimates.** Behind the admin UI's kinds page, it counted every job through the `current_jobs` view. Two nodes of that view are diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 80647bc5e..07a4ba7d3 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -14966,7 +14966,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; } } } diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 85205ff5e..fc697d123 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -11973,7 +11973,8 @@ async fn test_queue_storage_prune_lock_timeout_after_child_locks_reports_blocked /// 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 whole rollup pass. +/// 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; @@ -11981,18 +11982,20 @@ async fn test_queue_terminal_count_delta_rollup_lock_timeout_after_child_lock_re let schema = "awa_qs_terminal_delta_rollup_late_lock"; let store = create_store(&pool, schema).await; - 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:?}"), + 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"); @@ -12007,17 +12010,20 @@ async fn test_queue_terminal_count_delta_rollup_lock_timeout_after_child_lock_re 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); + 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, 3); + 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, 1); + assert_eq!(outcome.rolled_slots, 2); assert_eq!(outcome.blocked_slots, 0); - assert_eq!(live_count_sum(&pool, schema, queue).await, 3); + assert_eq!(live_count_sum(&pool, schema, queue).await, 6); assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 0); } From 82ff82d0e78836469c15146f9436dbe7d26d234b Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 18:21:04 +1300 Subject: [PATCH 13/16] fix(queue-storage): repair lane sequences per transaction with a bounded lock wait The definer_lane_provisioning patch reassigned lane sequences inside the patch 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 patch and every repaired lane waited for it to commit. The patch 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. The results are in MigrationReport::lane_sequences. `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 | 14 +- awa-cli/src/main.rs | 64 +++-- .../patches/definer_lane_provisioning.sql | 30 +-- .../v023_install_queue_storage_substrate.sql | 122 +++++---- awa-model/src/migrations.rs | 32 ++- awa-model/src/storage.rs | 237 ++++++++++++++++-- awa/tests/lane_ownership_roles_test.rs | 226 +++++++++++++---- docs/security.md | 27 +- docs/upgrade-0.6.13-to-0.6.14.md | 10 +- 9 files changed, 580 insertions(+), 182 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ff251781..a2b722287 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,14 +19,16 @@ 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. The patch gives existing lane sequences to the schema owner - and adds their missing grants, one sequence at a time and only where the - applying role is allowed to. It never fails on a sequence it cannot - reassign. The new read-only `awa storage lane-sequences` (or + runtime role. Once the patch 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 run. 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 + remain; `--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 before the patch detects that `provision_lane` is missing and creates lanes as before, and older binaries create correctly owned lanes once the patch has refreshed the helpers they call. `awa migrate` applies the diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index 4b994611a..1e347be53 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -438,7 +438,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, }, @@ -681,18 +682,23 @@ async fn main() -> Result<(), Box> { } if report.patches.iter().any(|patch| { patch.name == "definer_lane_provisioning" - && patch.outcome != awa_model::migrations::PatchOutcome::NotApplicable + && patch.outcome == awa_model::migrations::PatchOutcome::AlreadyPresent }) { - for (schema, count) in + // The patch was already in place, so this run repaired + // nothing: report what an earlier repair left behind. + let mut problems = Vec::new(); + for (schema, _) 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." + 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}"); + } } } } @@ -1271,19 +1277,18 @@ async fn main() -> 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?; @@ -1293,15 +1298,14 @@ async fn main() -> 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); } } @@ -1480,6 +1484,18 @@ fn print_migration_report(report: &awa_model::migrations::MigrationReport) { ); } } + if !report.lane_sequences.is_empty() { + let repaired = report + .lane_sequences + .iter() + .filter(|repair| repair.outcome == awa_model::storage::LaneSequenceOutcome::Repaired) + .count(); + // Sequences left behind are logged by the migration runner. + println!( + "Lane sequences repaired: {repaired} of {}", + report.lane_sequences.len() + ); + } match ( &report.storage_routing_before, &report.storage_routing_after, diff --git a/awa-model/migrations/patches/definer_lane_provisioning.sql b/awa-model/migrations/patches/definer_lane_provisioning.sql index 439630db9..2a8ced3f4 100644 --- a/awa-model/migrations/patches/definer_lane_provisioning.sql +++ b/awa-model/migrations/patches/definer_lane_provisioning.sql @@ -8,16 +8,15 @@ -- lane helpers create a missing sequence through them, and each new sequence -- gets the grants of the head table it backs, so a lane is usable by every -- runtime role however its creator's default privileges are set. The block --- below refreshes the helpers on each installed schema, 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 applying role may. It --- never fails on a sequence it cannot reassign; those are reported in a --- warning and by awa.lane_sequence_problems(). +-- below refreshes the helpers on each installed schema and counts existing +-- lane sequences that need repair. It does not reassign them: ALTER SEQUENCE +-- OWNER takes the sequence's ACCESS EXCLUSIVE lock until commit. After the +-- patch commits, `awa migrate` repairs each one 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`. -- --- Live load: function replacement, plus ALTER SEQUENCE OWNER on the lane --- sequences that need it. That holds each such sequence's ACCESS EXCLUSIVE --- lock until commit, so enqueues and claims on those lanes wait for the patch --- to commit; correctly owned lanes are not touched. Reads the catalog only. +-- Live load: function replacement only. Reads the catalog; no table DDL, row +-- scans or sequence locks. -- -- 0.6 delivery: the 0.6 series cannot add a migration version, so this ships as -- an idempotent schema patch. `awa migrate` applies it when the schema is @@ -30,7 +29,6 @@ DO $$ DECLARE v_schema TEXT; - v_leftovers TEXT; v_count BIGINT; BEGIN FOR v_schema IN @@ -61,17 +59,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/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index bcc8fe040..4625ebada 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -418,15 +418,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 @@ -500,55 +503,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; @@ -558,10 +582,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 @@ -571,31 +593,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; CREATE OR REPLACE FUNCTION awa.install_queue_storage_substrate( p_schema TEXT, diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 29a65fc93..d6b4a6734 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -306,8 +306,8 @@ macro_rules! definer_lane_provisioning_pending_sql { "SELECT 'awa.lane_sequence_problems'::text AS target \ WHERE to_regprocedure('awa.lane_sequence_problems(text,text)') IS NULL \ UNION ALL \ - SELECT 'awa.repair_lane_sequences' \ - WHERE to_regprocedure('awa.repair_lane_sequences(text)') IS NULL \ + SELECT 'awa.repair_lane_sequence' \ + WHERE to_regprocedure('awa.repair_lane_sequence(text,text)') IS NULL \ UNION ALL \ SELECT 'awa.refresh_lane_sequence_helpers' \ WHERE EXISTS ( \ @@ -578,6 +578,11 @@ pub struct MigrationReport { /// change routing; both are `None` before the storage tables exist. pub storage_routing_before: Option, pub storage_routing_after: Option, + /// What the run did to each lane sequence that needed repair, when it + /// applied `definer_lane_provisioning`. Each sequence is repaired in its + /// own short transaction after the patch commits; a sequence left behind + /// does not make the run incomplete. + pub lane_sequences: Vec, } impl MigrationReport { @@ -736,6 +741,28 @@ async fn run_inner(conn: &mut PgConnection) -> Result })); } + // ALTER SEQUENCE ... OWNER takes the sequence's ACCESS EXCLUSIVE lock + // until its transaction ends, so the patch only installs the definitions + // and each lane sequence is repaired here, in its own transaction with a + // bounded lock wait, once the patch has committed. + let mut lane_sequences = Vec::new(); + if patches.iter().any(|patch| { + patch.name == "definer_lane_provisioning" && patch.outcome == PatchOutcome::Applied + }) { + match crate::storage::repair_lane_sequences_in_every_schema(conn).await { + Ok(repairs) => { + for next_step in crate::storage::lane_sequence_repair_next_steps(&repairs) { + tracing::warn!("{next_step}"); + } + lane_sequences = repairs; + } + Err(error) => tracing::warn!( + %error, + "Could not repair lane sequences; run `awa storage lane-sequences --repair`" + ), + } + } + // Ensure the admin metadata cache is warm. Since v006 removed the // synchronous triggers on jobs_hot, the cache is only updated by the // maintenance leader. Refreshing here guarantees queue_stats() and @@ -770,6 +797,7 @@ async fn run_inner(conn: &mut PgConnection) -> Result patches, storage_routing_before, storage_routing_after: storage_routing_conn(conn).await?, + lane_sequences, }) } diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index 85d949ee2..21d8f0fa4 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)] @@ -382,30 +382,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. @@ -417,7 +452,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) @@ -426,25 +462,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 n.nspname::text, count(*) \ FROM pg_namespace AS n \ @@ -459,6 +652,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, @@ -593,37 +621,44 @@ 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}" - ); - assert_eq!( - migrations::pending_schema_patches(&cluster.owner) - .await? - .iter() - .map(|patch| patch.name) - .collect::>(), - vec!["definer_lane_provisioning"] - ); + 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}" + ); + } - // `awa migrate` 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. The leftovers do not keep the patch pending. - let report = migrations::run_with_report(&cluster.migrator).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?; + let report = tokio::time::timeout( + Duration::from_secs(120), + migrations::run_with_report(&migrate_pool), + ) + .await??; assert!(report.complete(), "{:?}", report.patches); assert!(migrations::pending_schema_patches(&cluster.owner) .await? .is_empty()); + let outcomes: Vec = report + .lane_sequences + .iter() + .map(|repair| repair.outcome) + .collect(); + assert_eq!(outcomes.len(), 6, "{:?}", report.lane_sequences); + assert_eq!( + outcomes + .iter() + .filter(|outcome| **outcome == LaneSequenceOutcome::Repaired) + .count(), + 4 + ); + migrate_pool.close().await; cluster .assert_lanes_owned_by_owner("migrator_lane", 1) .await?; @@ -631,10 +666,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); @@ -655,21 +690,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. @@ -684,20 +729,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()); @@ -705,7 +744,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!( @@ -713,21 +753,105 @@ 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?; + + // 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 patch completes while the lock is held, and its report names the + // held sequence. + let report = tokio::time::timeout( + Duration::from_secs(120), + migrations::run_with_report(&migrate_pool), + ) + .await??; + assert!(report.complete(), "{:?}", report.patches); + assert_eq!( + report.lane_sequences.len(), + 4, + "{:?}", + report.lane_sequences + ); + for repair in &report.lane_sequences { + let expected = if repair.sequence_name == held_sequence { + LaneSequenceOutcome::Busy + } else { + LaneSequenceOutcome::Repaired + }; + assert_eq!(repair.outcome, expected, "{repair:?}"); + } + 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.md b/docs/security.md index 8c56004f1..8475a9c82 100644 --- a/docs/security.md +++ b/docs/security.md @@ -194,17 +194,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. Applying the patch runs the repair as the applying 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 patch. 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 patch 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 patch or `awa migrate`. 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 +``` ### 5. Configure your processes diff --git a/docs/upgrade-0.6.13-to-0.6.14.md b/docs/upgrade-0.6.13-to-0.6.14.md index 1cccbe370..b18fa715e 100644 --- a/docs/upgrade-0.6.13-to-0.6.14.md +++ b/docs/upgrade-0.6.13-to-0.6.14.md @@ -18,18 +18,18 @@ awa --database-url "$AWA_MIGRATOR_DATABASE_URL" migrate ``` -4. Read the output. The `definer_lane_provisioning` patch reports `applied`. If it could not reassign some lane sequences, `awa migrate` prints a warning; run `awa storage lane-sequences` for the statements to run and the role to run them as (see [Check and repair lane ownership](security.md#check-and-repair-lane-ownership)). The migration exit status does not depend on them. +4. Read the output. The `definer_lane_provisioning` patch reports `applied`, followed by how many lane sequences were repaired. Any left over are listed in warnings: a busy one needs another run of `awa storage lane-sequences --repair` as the migration role (without `PGOPTIONS`), and the others name the roles that can repair them (see [Check and repair lane ownership](security.md#check-and-repair-lane-ownership)). The migration exit status does not depend on them. 5. Roll workers, producers and admin processes to 0.6.14, in any order. ## What `awa migrate` changes -The schema stays at v40. The `definer_lane_provisioning` patch installs `provision_lane` and `provision_lane_sequences`, two `SECURITY DEFINER` functions owned by the schema owner, in `awa` and every other queue-storage schema. It changes the three lane helpers to create a missing sequence through them, and adds the read-only `awa.lane_sequence_problems(schema)` and `awa.repair_lane_sequences(schema)`. Every new lane sequence then belongs to the schema owner and carries the privileges of the head table it backs, whichever role creates the lane: a producer enqueueing to a new queue, priority or shard, the maintenance leader promoting into a shard after `enqueue_shards` was raised, or `awa storage prepare-queue` run as any role. +The schema stays at v40. The `definer_lane_provisioning` patch installs `provision_lane` and `provision_lane_sequences`, two `SECURITY DEFINER` functions owned by the schema owner, in `awa` and every other queue-storage schema. It changes the three lane helpers to create a missing sequence through them, and adds the read-only `awa.lane_sequence_problems(schema)` and `awa.repair_lane_sequence(schema, sequence)`. Every new lane sequence then belongs to the schema owner and carries the privileges of the head table it backs, whichever role creates the lane: a producer enqueueing to a new queue, priority or shard, the maintenance leader promoting into a shard after `enqueue_shards` was raised, or `awa storage prepare-queue` run as any role. -The patch then repairs existing lane sequences. Each one not owned by the schema owner is reassigned if the applying role can `SET ROLE` to both its current owner and the schema owner. Missing grants are added if the applying role can act as the owner. A sequence it cannot change is reported and left as it was. The patch never fails because of one. +Once the patch commits, `awa migrate` repairs existing lane sequences, each in its own short transaction. A sequence not owned by the schema owner is reassigned if the role holds its current owner's privileges and can `SET ROLE` to the schema owner. Missing grants are added if the role holds the owner's privileges. Each sequence is tried as the effective role and then as the login role. Reassigning waits at most 300 ms for the sequence's lock, which an open enqueue or claim on that lane holds, and holds it only for that sequence's transaction. A sequence that is busy, or that neither role may change, is reported and left as it was. The patch: -- replaces functions and changes only the lane sequences that need it. `ALTER SEQUENCE ... OWNER` holds that sequence's lock until the patch commits, so enqueues and claims on those lanes wait for the commit. Correctly owned lanes are not touched; +- only replaces functions; it takes no sequence or table locks; - keeps every lane's sequence position; - leaves storage routing unchanged; - applies as one transaction, so an interrupted run leaves nothing applied; run it again. @@ -38,7 +38,7 @@ The patch: ### External migration runners -`awa migrate --sql --pending` prints the `R__definer_lane_provisioning` patch on a 0.6.13 database; with `--extract-to`, apply `R__definer_lane_provisioning.sql` as a repeatable script after `V40`. Run it as the schema owner, then run `awa storage lane-sequences` to check for leftovers. +`awa migrate --sql --pending` prints the `R__definer_lane_provisioning` patch on a 0.6.13 database; with `--extract-to`, apply `R__definer_lane_provisioning.sql` as a repeatable script after `V40`. Run it as the schema owner. The script does not repair existing lane sequences, so then run `awa storage lane-sequences --repair` as the migration role itself. ## Roll back From a0757054ef6d0ab9ba724da6c20e3b93ffae7ab1 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 20:58:37 +1300 Subject: [PATCH 14/16] docs(security): clarify lane sequence repair outcomes --- awa-model/migrations/patches/definer_lane_provisioning.sql | 2 +- docs/security.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/awa-model/migrations/patches/definer_lane_provisioning.sql b/awa-model/migrations/patches/definer_lane_provisioning.sql index 2a8ced3f4..b75a28cad 100644 --- a/awa-model/migrations/patches/definer_lane_provisioning.sql +++ b/awa-model/migrations/patches/definer_lane_provisioning.sql @@ -65,7 +65,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; diff --git a/docs/security.md b/docs/security.md index 8475a9c82..2b9848ec7 100644 --- a/docs/security.md +++ b/docs/security.md @@ -243,7 +243,7 @@ Since these triggers run with the caller's privileges, the runtime role needs wr The maintenance leader also calls `awa.refresh_admin_metadata()` as a full reconciliation safety net. It runs with invoker privileges and needs DML on the admin metadata tables. `TRUNCATE` in the grant is for queue-storage segment pruning, which truncates retired partitions. -Schema patches (`awa migrate --sql` prints them after the versioned migrations) replace function bodies and create tables under the migrator role, like migrations do; they never change grants. Privileges managed with `GRANT ... ON ALL TABLES` plus the default privileges above cover patched tables without further action. +Schema patches (`awa migrate --sql` prints them after the versioned migrations) replace functions and may create tables or change function ownership and grants. The `definer_lane_provisioning` patch installs provisioning functions owned by the schema owner and grants them `PUBLIC EXECUTE`. After the patch commits, `awa migrate` attempts to repair existing lane-sequence ownership and grants; review its warnings for sequences that still need repair. The default privileges above cover new tables created by the migration role. The runtime also directly upserts into the descriptor catalogs on startup and on each runtime snapshot tick: From 39c6a2c45036e51c2417d010242bc98e1fedc4a5 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 21:05:46 +1300 Subject: [PATCH 15/16] docs(worker): describe prune timeout per lock acquisition --- CHANGELOG.md | 3 ++- awa-model/src/queue_storage.rs | 16 +++++++++++----- awa-python/python/awa/client.py | 9 ++++++--- awa-worker/src/client.rs | 11 +++++++---- docs/configuration.md | 4 ++-- 5 files changed, 28 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3579e877c..5432a7be1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,7 +15,8 @@ Notable changes between releases. Detailed migration notes for storage transitio `blocked` every tick and the rings stopped recycling. Set it with `ClientBuilder::prune_lock_timeout` (Rust) or `queue_storage_prune_lock_timeout_ms` on `client.start()` (Python); the - default stays 10 ms and zero or a value above PostgreSQL's `lock_timeout` + timeout applies separately to each lock acquisition, not to the whole prune; + queue prune can wait on up to eight partition locks. The default stays 10 ms and zero or a value above PostgreSQL's `lock_timeout` maximum is rejected. Start logs a warning when the value is at or above the connection's `deadlock_timeout` (prune waiting that long can make deadlock detection cancel a claim instead) or `statement_timeout`. `with_prune_lock_timeout` is now diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 3401f9cd3..748d4e32c 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -2506,15 +2506,21 @@ impl QueueStorage { /// Set the PostgreSQL `lock_timeout` applied to the exclusive-lock steps of /// ring prune (queue, lease and claim rings) and terminal-delta rollup. /// - /// A prune that cannot take its locks within the timeout rolls back and - /// reports [`PruneOutcome::Blocked`]; the maintenance loop retries on its - /// next tick. The default is 10 ms, which keeps prune from queueing behind - /// readers and so from stalling claims and enqueues behind its pending + /// Each lock acquisition gets its own timeout; this does not bound the + /// whole prune. Queue prune locks up to eight partitions, so its + /// partition-lock + /// waits alone can approach eight times this value. A lock timeout rolls + /// back + /// the prune and reports [`PruneOutcome::Blocked`]; the maintenance loop + /// retries on a later tick. The default is 10 ms, which keeps prune from + /// queueing behind readers and so from stalling claims and enqueues behind + /// its pending /// `ACCESS EXCLUSIVE` request. Raise it when prune is persistently /// `Blocked` by short-lived readers of the ring partitions. Keep it well /// below `deadlock_timeout` (default 1 s) and any `statement_timeout`: a /// prune that waits as long as `deadlock_timeout` triggers deadlock - /// detection, which may cancel a concurrent claim instead of the prune. The + /// detection, which may cancel a concurrent claim instead of the prune. + /// The /// timeout is applied at millisecond granularity. /// /// Returns a validation error for a zero timeout or one above PostgreSQL's diff --git a/awa-python/python/awa/client.py b/awa-python/python/awa/client.py index ea3884a24..dba2a6424 100644 --- a/awa-python/python/awa/client.py +++ b/awa-python/python/awa/client.py @@ -903,9 +903,12 @@ async def start( ``queue_storage_prune_lock_timeout_ms`` sets the PostgreSQL ``lock_timeout`` the maintenance leader uses when pruning the queue, - lease and claim rings (default 10 ms). A prune that cannot lock its - partitions in time is counted as ``blocked`` and retried on the next - tick; raise it when ``blocked`` outcomes persist, keeping it well below the + lease and claim rings (default 10 ms). Each lock acquisition gets its + own timeout; this does not bound the whole prune. Queue prune locks up + to eight partitions, so its partition-lock waits alone can approach + eight times this value. A lock timeout rolls back the prune and counts + it as ``blocked``; the maintenance loop retries on a later tick. Raise + it when ``blocked`` outcomes persist, keeping it well below the server's ``deadlock_timeout`` and ``statement_timeout`` (a warning is logged at start otherwise). Must be > 0. diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 5cc35083e..33f375695 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -930,10 +930,13 @@ impl ClientBuilder { /// the exclusive-lock steps of queue, lease and claim ring prune and of /// terminal-delta rollup (default 10 ms). /// - /// A prune that cannot take its locks within the timeout rolls back and is - /// recorded as `blocked` in `awa.maintenance.prune.attempts`; it retries on - /// the next maintenance tick. The short default keeps the pending - /// `ACCESS EXCLUSIVE` request from stalling claims and enqueues behind it. + /// Each lock acquisition gets its own timeout; this does not bound the + /// whole prune. Queue prune locks up to eight partitions, so its + /// partition-lock waits alone can approach eight times this value. A lock + /// timeout rolls back the prune and records `blocked` in + /// `awa.maintenance.prune.attempts`; the maintenance loop retries on a + /// later tick. The short default keeps the pending `ACCESS EXCLUSIVE` + /// request from stalling claims and enqueues behind it. /// Raise it when `blocked` outcomes persist and the rings stop recycling /// because short-lived readers (admin queries, dashboards) repeatedly hold /// the ring partitions at the moment prune tries. Keep it well below diff --git a/docs/configuration.md b/docs/configuration.md index 36722de60..dd308b95a 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -570,7 +570,7 @@ await client.start( | `queue_rotate_interval` | `1000ms` | How often ready/terminal segments rotate | | `lease_rotate_interval` | `1000ms` | How often lease segments rotate | | `claim_rotate_interval` | matches `queue_rotate_interval` | How often the ADR-023 claim-ring rotates. Set with `ClientBuilder::claim_rotate_interval` (Rust) or `queue_storage_claim_rotate_interval_ms` (Python). Test harnesses sometimes set this to a long interval to pin claim-ring layout for deterministic count assertions. | -| `prune_lock_timeout` | `10ms` | The PostgreSQL `lock_timeout` the maintenance leader applies when it takes the exclusive locks to reclaim a queue, lease or claim ring partition (and when it folds terminal-count deltas). A prune that cannot lock its partitions within the timeout rolls back, is counted as `blocked`, and is retried on the next maintenance tick. Set with `ClientBuilder::prune_lock_timeout` (Rust) or `queue_storage_prune_lock_timeout_ms` (Python). Must be greater than zero and at most `2147483647` ms; keep it well below `deadlock_timeout` and `statement_timeout`. The short default keeps prune's pending `ACCESS EXCLUSIVE` request from stalling claims and enqueues queued behind it. See below for when to raise it. | +| `prune_lock_timeout` | `10ms` | The PostgreSQL `lock_timeout` the maintenance leader applies when it takes the exclusive locks to reclaim a queue, lease or claim ring partition (and when it folds terminal-count deltas). The timeout applies separately to each lock acquisition, not to the whole prune. Queue prune locks up to eight partitions, so its partition-lock waits alone can approach eight times this value. A lock timeout rolls the prune back, is counted as `blocked`, and is retried on a later maintenance tick. Set with `ClientBuilder::prune_lock_timeout` (Rust) or `queue_storage_prune_lock_timeout_ms` (Python). Must be greater than zero and at most `2147483647` ms; keep it well below `deadlock_timeout` and `statement_timeout`. The short default keeps prune's pending `ACCESS EXCLUSIVE` request from stalling claims and enqueues queued behind it. See below for when to raise it. | The benchmark harness in [postgresql-job-queue-benchmarking](https://github.com/hardbyte/postgresql-job-queue-benchmarking) reads `QUEUE_SLOT_COUNT`, `LEASE_SLOT_COUNT`, `CLAIM_SLOT_COUNT`, `QUEUE_STRIPE_COUNT`, and `LEASE_CLAIM_RECEIPTS` from the environment. Those env vars are benchmark configuration, not general worker-runtime configuration. @@ -580,7 +580,7 @@ Use the defaults unless you have a reason not to: - Increase `lease_slot_count` if lease churn is high enough that dead tuples in the lease ring stop collapsing promptly. - Increase `claim_slot_count` if the rotation cadence (`claim_rotate_interval`) plus the slot count combine to a partition retention window shorter than your longest in-flight zero-deadline short job; running out of empty slots forces `rotate_claims` to return `SkippedBusy` and the receipt-plane churn falls back onto a smaller working set of partitions. - Increase `queue_stripe_count` only for measured workloads where many small replicas contend on the same logical queue. This is queue-storage configuration, not a per-queue `QueueConfig` field: it gives each logical queue the same number of physical stripes. Striping spreads a hot logical queue over `queue#N` physical coordination paths, but it weakens perfect global ordering and can regress calmer shapes if overused. For single hot-queue workloads, benchmark `2` stripes first; higher values should be treated as workload-specific tuning rather than a general default. -- Increase `prune_lock_timeout` when `increase(awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}[15m])` stays above zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Raise it in small steps (for example `50ms`, then `100ms`, then `250ms`) and keep it well below PostgreSQL's `deadlock_timeout` (default `1s`) and below any `statement_timeout` on the pool's connections. A prune that waits as long as `deadlock_timeout` makes PostgreSQL run its deadlock detector, and if claims or enqueues are part of a lock cycle with it, the detector can cancel them rather than the prune; a short `lock_timeout` abandons the prune first, so the prune is always the one that gives way. A `statement_timeout` at or below the lock timeout cancels the prune statement before the lock timeout fires, so it fails with an error instead of reporting `blocked`. Startup logs a warning when the configured timeout is at or above either setting. While prune waits, new accesses to the same partition also queue behind its lock request, so a longer timeout trades claim latency for reclamation. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for up to 32 ticks after consecutive `skipped_active` or `blocked` results, so a persistently blocked ring is retried less often over time, and `error` is a database failure. +- Increase `prune_lock_timeout` when `increase(awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}[15m])` stays above zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Raise it in small steps (for example `50ms`, then `100ms`, then `250ms`) and keep it well below PostgreSQL's `deadlock_timeout` (default `1s`) and below any `statement_timeout` on the pool's connections. A single lock wait that reaches `deadlock_timeout` makes PostgreSQL run its deadlock detector, and if claims or enqueues are part of a lock cycle with it, the detector can cancel them rather than the prune; a short `lock_timeout` abandons the prune first, so the prune is always the one that gives way. A `statement_timeout` at or below the lock timeout cancels the prune statement before the lock timeout fires, so it fails with an error instead of reporting `blocked`. Startup logs a warning when the configured timeout is at or above either setting. While prune waits, new accesses to the same partition also queue behind its lock request, so a longer timeout trades claim latency for reclamation. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for up to 32 ticks after consecutive `skipped_active` or `blocked` results, so a persistently blocked ring is retried less often over time, and `error` is a database failure. - Increase rotation intervals to reduce partition churn and metadata activity. - Decrease rotation intervals to tighten dead-tuple bounds at the cost of more frequent rotate/prune work. From 9dcb8c13ca6a4a465122f32b52866287968d290a Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 22:31:10 +1300 Subject: [PATCH 16/16] fix(queue-storage): ignore retired ready rows in point reads --- CHANGELOG.md | 5 ++ awa-model/src/queue_storage.rs | 13 ++- awa/tests/queue_storage_runtime_test.rs | 111 ++++++++++++++++++++++++ 3 files changed, 127 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a1b0a6839..a59672bac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **Point reads show the live queue and priority after a lane move.** Retained + ready rows from before a priority or queue change could win a read's tie + after the ring wrapped, returning the old fields even though the batch + operation succeeded. Point reads now exclude tombstoned source rows. + - **Callback completion, `get_job`, lifecycle events and the admin reads find running jobs that priority aging promoted on queue storage** (backport of #600). Claims and leases record a job's aged effective priority while its diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 3b0b1503f..fa4f8fb22 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -11695,8 +11695,17 @@ impl QueueStorage { unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload - FROM {schema}.ready_entries - WHERE job_id = $1 + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC "#, )) diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 376548e7f..e03eece6b 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -15547,3 +15547,114 @@ async fn test_aged_materialized_lease_is_visible_to_admin_reads() { .expect("cancel by unique key finds the aged running job"); assert_eq!(cancelled.id, job_id); } + +/// A lane move leaves immutable source rows behind. Wrap the ring so the +/// destination is in an earlier partition: point reads must use live evidence, +/// rather than depending on which equal-timestamp ready row the plan emits last. +async fn ready_point_read_wrap_fixture() -> (TestDbGuard, sqlx::PgPool, QueueStorage, i64) { + let (guard, pool) = setup_pool(4).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: "awa_qs_point_read_wrap".to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + ..Default::default() + }, + ) + .await; + // A fresh empty ring can be positioned just before wrap without retaining + // unrelated jobs. This is the state reached after three rotations/prunes. + let schema = store.schema(); + sqlx::query(&format!( + "UPDATE {schema}.queue_ring_state SET current_slot = 3, generation = 3 WHERE singleton" + )) + .execute(&pool) + .await + .expect("position empty ring before wrap"); + sqlx::query(&format!( + "UPDATE {schema}.queue_ring_slots SET generation = 3 WHERE slot = 3" + )) + .execute(&pool) + .await + .expect("set current slot generation"); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: "point_read_source".to_string(), + priority: 4, + ..Default::default() + }, + ) + .await; + assert!(matches!( + store.rotate(&pool).await.expect("wrap to slot zero"), + RotateOutcome::Rotated { slot: 0, .. } + )); + (guard, pool, store, job_id) +} + +async fn assert_live_ready_row( + pool: &sqlx::PgPool, + store: &QueueStorage, + job_id: i64, + expected_queue: &str, + expected_priority: i16, +) { + let schema = store.schema(); + let live: Vec<(String, i16, i32)> = sqlx::query_as(&format!( + "SELECT ready.queue, ready.priority, ready.ready_slot + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + )" + )) + .bind(job_id) + .fetch_all(pool) + .await + .expect("read live ready evidence"); + assert_eq!( + live, + vec![(expected_queue.to_string(), expected_priority, 0)] + ); +} + +#[tokio::test] +async fn test_point_read_uses_live_ready_priority_after_ring_wrap() { + let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; + assert!(store + .set_priority(&pool, job_id, 1) + .await + .expect("set priority")); + assert_live_ready_row(&pool, &store, job_id, "point_read_source", 1).await; + let job = admin::get_job(&pool, job_id).await.expect("read moved job"); + assert_eq!( + job.priority, 1, + "point read must ignore the tombstoned source" + ); +} + +#[tokio::test] +async fn test_point_read_uses_live_ready_queue_after_ring_wrap() { + let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; + assert!(store + .move_queue(&pool, job_id, "point_read_destination", None) + .await + .expect("move queue")); + assert_live_ready_row(&pool, &store, job_id, "point_read_destination", 4).await; + let job = admin::get_job(&pool, job_id).await.expect("read moved job"); + assert_eq!( + job.queue, "point_read_destination", + "point read must ignore the tombstoned source" + ); +}