diff --git a/CHANGELOG.md b/CHANGELOG.md index a9228e5a2..79c8d15a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,9 @@ Notable changes between releases. Detailed migration notes for storage transitio - Heartbeat, deadline and callback rescue of lease-backed attempts reach the stale attempts behind a full batch of attempts whose retained ready bodies are ambiguous. Before, every pass selected the same ambiguous attempts first, left them running, and rescued nothing else. A pass now sets all of a batch's ambiguous attempts aside at once. - Admin job lists and kind counts no longer read other claim slots' receipt closure partitions when they read one slot's open claims. - A `Completed` lifecycle event, and the job a `Completed` follow-up enqueue spec receives, no longer carry the finished attempt's `heartbeat_at` and `deadline_at`; both are `None`, as a read of the completed job returns them. +- Worker finalization is bounded to the larger of the runtime's worker capacity and twice the completion batch size times the completion shard count. A finished handler keeps its worker slot while it waits for finalizer capacity, so when durable completion falls behind the runtime stops claiming instead of holding an unbounded backlog of finished attempts and their payloads in memory. Attempt heartbeats continue until finalization settles. +- Retried queue-storage jobs and jobs carrying metadata complete by their compact claim batch and item index, instead of reading every retained claim batch. +- Lifecycle hooks (`on_event`, `on_event_kind`) no longer accumulate without bound when they are slower than job throughput. Each attempt's `Started` and outcome hooks now hold its finalizer slot until they return, so slow hooks slow claiming instead of growing worker memory, and a worker-finalized outcome hook runs after that attempt's `Started` hook has returned. Maintenance's `Rescued` and duplicate `Cancelled` events allow 1,024 outstanding hook tasks per worker and drop further events with a warning. ## [0.6.14] — 2026-11-08 diff --git a/README.md b/README.md index 3bcd74cb1..4fd0e2f89 100644 --- a/README.md +++ b/README.md @@ -295,7 +295,7 @@ tx.commit().await?; awa::admin::cancel_by_unique_key(&pool, "send_email", None, Some(&serde_json::json!({"to": "alice@example.com", "subject": "Welcome"})), None).await?; ``` -Lifecycle hooks receive `Started` after a claim commits as the worker handler is about to run, plus outcome events (`Completed`, `Retried`, `Exhausted`, `Cancelled`) after the corresponding state transition commits. Hooks run in detached tasks, so they do not block job execution. +Lifecycle hooks receive `Started` after a claim commits as the worker handler is about to run, plus outcome events (`Completed`, `Retried`, `Exhausted`, `Cancelled`) after the corresponding state transition commits. Hooks run alongside job execution and never delay a job's state transition, but each attempt's hooks hold a finalizer slot until they return, so hooks slower than job throughput slow claiming. Cancellation is cooperative for running handlers: diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 3559aa4e0..7ae6cb23a 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -2150,6 +2150,62 @@ mod ready_segment_tests { } } + #[sqlx::test(migrations = false)] + async fn enqueue_batch_writes_more_segments_than_the_old_bind_limit(pool: sqlx::PgPool) { + use super::{QueueStorageConfig, RuntimeReadyRow}; + + crate::migrations::run(&pool).await.unwrap(); + let store = QueueStorage::new(QueueStorageConfig { + schema: "awa_qs_segment_bind_limit".into(), + ..Default::default() + }) + .unwrap(); + store.install(&pool).await.unwrap(); + + // Distinct microsecond timestamps guarantee one segment per row. + // Reading the clock for every row can coalesce adjacent timestamps + // and fail to exercise the old 8,191-segment bind limit. + const JOBS: i64 = 9_000; + let first_run_at = Utc::now() - Duration::seconds(1); + let rows = (0..JOBS) + .map(|seq| RuntimeReadyRow { + kind: "segment_test".into(), + queue: "segment-q".into(), + args: serde_json::json!({"seq": seq}), + priority: 2, + attempt: 0, + run_lease: 0, + max_attempts: 25, + run_at: first_run_at + Duration::microseconds(seq), + attempted_at: None, + created_at: first_run_at, + unique_key: None, + unique_states: None, + payload: serde_json::json!({}), + ordering_key: None, + }) + .collect(); + assert_eq!( + store.enqueue_runtime_rows(&pool, rows).await.unwrap(), + JOBS as usize + ); + let counts: (i64, i64) = sqlx::query_as(&format!( + "SELECT (SELECT count(*) FROM {}.ready_entries), \ + (SELECT count(*) FROM {}.ready_segments)", + store.schema(), + store.schema() + )) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(counts, (JOBS, JOBS)); + let claimed = store + .claim_runtime_batch(&pool, "segment-q", 100, std::time::Duration::from_secs(30)) + .await + .unwrap(); + assert_eq!(claimed.len(), 100); + } + #[test] fn ready_segments_split_on_run_at_boundaries() { let first = Utc @@ -7010,14 +7066,22 @@ impl QueueStorage { .expect("receipt-backed slow completion requires receipt_id") }) .collect(); + let receipt_claim_batch_ids: Vec> = receipt_claimed + .iter() + .map(|entry| entry.claim.claim_batch_id) + .collect(); + let receipt_claim_batch_indices: Vec> = receipt_claimed + .iter() + .map(|entry| entry.claim.claim_batch_index) + .collect(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let updated: Vec<(i64, i64)> = sqlx::query_as(&format!( r#" - WITH completed(claim_slot, job_id, run_lease, receipt_id) AS ( - SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[]) + WITH completed(claim_slot, job_id, run_lease, receipt_id, claim_batch_id, claim_batch_index) AS ( + SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[], $5::bigint[], $6::int[]) ), locked_row_claims AS ( SELECT claims.claim_slot, claims.job_id, claims.run_lease @@ -7044,18 +7108,22 @@ impl QueueStorage { items.job_id, items.run_lease, items.receipt_id - FROM {schema}.lease_claim_batches AS claim_batches - CROSS JOIN LATERAL unnest( - claim_batches.job_ids, - claim_batches.run_leases, - claim_batches.receipt_ids - ) AS items(job_id, run_lease, receipt_id) - JOIN completed - ON completed.claim_slot = claim_batches.claim_slot - AND completed.job_id = items.job_id - AND completed.run_lease = items.run_lease - AND completed.receipt_id = items.receipt_id - WHERE claim_batches.job_ids && $2::bigint[] + FROM completed + JOIN {schema}.lease_claim_batches AS claim_batches + ON claim_batches.claim_slot = completed.claim_slot + AND claim_batches.batch_id = completed.claim_batch_id + CROSS JOIN LATERAL ( + SELECT + claim_batches.job_ids[completed.claim_batch_index] AS job_id, + claim_batches.run_leases[completed.claim_batch_index] AS run_lease, + claim_batches.receipt_ids[completed.claim_batch_index] AS receipt_id + ) AS items + WHERE completed.claim_batch_id IS NOT NULL + AND completed.claim_batch_index IS NOT NULL + AND completed.claim_batch_index BETWEEN 1 AND claim_batches.claimed_count + AND completed.job_id = items.job_id + AND completed.run_lease = items.run_lease + AND completed.receipt_id = items.receipt_id AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures @@ -7093,7 +7161,6 @@ impl QueueStorage { WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) - FOR UPDATE OF claim_batches ), locked_claims AS ( SELECT claim_slot, job_id, run_lease FROM locked_row_claims @@ -7185,6 +7252,8 @@ impl QueueStorage { .bind(&receipt_job_ids) .bind(&receipt_run_leases) .bind(&receipt_receipt_ids) + .bind(&receipt_claim_batch_ids) + .bind(&receipt_claim_batch_indices) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 87ae548ed..902a91750 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -321,10 +321,11 @@ impl ClientBuilder { /// `Started` dispatch is scheduled after the claim commits and before the /// worker handler is invoked. Outcome handlers run only after the /// corresponding DB state transition commits. Hooks are best-effort - /// in-process notifications, not a durable workflow mechanism; because - /// they run in detached tasks, a very short job can complete before a - /// `Started` handler finishes. Capture any shared dependencies you need in - /// the closure environment. + /// in-process notifications, not a durable workflow mechanism. `Started` + /// runs alongside the handler, so a very short job can complete before a + /// `Started` handler finishes. Each attempt's hooks hold a finalizer slot + /// until they return, so hooks slower than job throughput slow claiming. + /// Capture any shared dependencies you need in the closure environment. pub fn on_event(mut self, handler: F) -> Self where T: JobArgs + DeserializeOwned + Send + Sync + 'static, @@ -1640,6 +1641,8 @@ impl Client { effective_storage.clone(), self.dlq_policy.clone(), self.finalizers.clone(), + self.service_cancel.clone(), + runtime_worker_capacity, )); // Admin cancellation listener: fires the in-flight cancel flag diff --git a/awa-worker/src/completion.rs b/awa-worker/src/completion.rs index 50fe63e38..0872bc392 100644 --- a/awa-worker/src/completion.rs +++ b/awa-worker/src/completion.rs @@ -105,9 +105,19 @@ fn completion_sort_key(request: &CompletionRequest) -> (i32, i64, RunLease) { #[derive(Clone)] pub(crate) struct CompletionBatcherHandle { shards: Vec>, + pipeline_capacity: usize, } impl CompletionBatcherHandle { + /// How many finished attempts may wait for finalization at once: enough + /// to keep every shard's in-flight and next batch full, and at least one + /// per worker. + pub fn finalizer_capacity(&self, runtime_worker_capacity: u32) -> usize { + self.pipeline_capacity + .max(runtime_worker_capacity.max(1) as usize) + .min(tokio::sync::Semaphore::MAX_PERMITS) + } + pub async fn complete(&self, job_id: i64, run_lease: RunLease) -> Result { self.complete_inner(job_id, run_lease, None, None).await } @@ -185,7 +195,14 @@ impl CompletionBatcher { }); } - (Self { workers }, CompletionBatcherHandle { shards }) + ( + Self { workers }, + CompletionBatcherHandle { + shards, + // Each shard needs an in-flight batch and a queued batch. + pipeline_capacity: batch_size.saturating_mul(shard_count).saturating_mul(2), + }, + ) } pub fn spawn(self) -> Vec> { diff --git a/awa-worker/src/dispatcher.rs b/awa-worker/src/dispatcher.rs index 916edb02f..44d7d5fd7 100644 --- a/awa-worker/src/dispatcher.rs +++ b/awa-worker/src/dispatcher.rs @@ -852,6 +852,8 @@ mod tests { storage.clone(), crate::executor::DlqPolicy::default(), tokio_util::task::TaskTracker::new(), + CancellationToken::new(), + 1, )); let config = QueueConfig { max_workers: 1, diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index d1c5db277..59ccd581a 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -11,6 +11,8 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::sync::Arc; use std::time::Duration; +use tokio::sync::Semaphore; +use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; use tracing::{debug, error, info, info_span, warn, Instrument}; @@ -269,6 +271,8 @@ pub struct JobExecutor { storage: RuntimeStorage, dlq_policy: DlqPolicy, finalizers: TaskTracker, + finalize_cancel: CancellationToken, + finalize_capacity: Arc, } impl JobExecutor { @@ -291,7 +295,12 @@ impl JobExecutor { storage: RuntimeStorage, dlq_policy: DlqPolicy, finalizers: TaskTracker, + finalize_cancel: CancellationToken, + runtime_worker_capacity: u32, ) -> Self { + let finalize_capacity = Arc::new(Semaphore::new( + completion_batcher.finalizer_capacity(runtime_worker_capacity), + )); Self { pool, workers, @@ -305,6 +314,8 @@ impl JobExecutor { storage, dlq_policy, finalizers, + finalize_cancel, + finalize_capacity, } } @@ -332,6 +343,8 @@ impl JobExecutor { let storage = self.storage.clone(); let dlq_policy = self.dlq_policy.clone(); let finalizers = self.finalizers.clone(); + let finalize_cancel = self.finalize_cancel.clone(); + let finalize_capacity = self.finalize_capacity.clone(); let job_id = job.id; let job_run_lease = job.run_lease; let job_kind = job.kind.clone(); @@ -379,20 +392,21 @@ impl JobExecutor { progress_state.clone(), ); + let mut started_event = None; let result = match workers.get(&job.kind) { Some(worker) => { if has_lifecycle_handlers { let started_handlers = lifecycle_handlers.clone(); let started_kind = job_kind.clone(); let started_job = job.clone(); - tokio::spawn(async move { + started_event = Some(tokio::spawn(async move { dispatch_lifecycle_event( &started_handlers, &started_kind, UntypedJobEvent::Started { job: started_job }, ) .await; - }); + })); } let perform = CatchUnwind(worker.perform(&ctx)); let outcome = match abandon_after { @@ -439,6 +453,20 @@ impl JobExecutor { guard.clone_latest() }; + // Keep worker capacity and heartbeats until a bounded finalizer slot + // is available. Otherwise returning handlers can create an unbounded + // set of tasks retaining job payloads while completion is stalled. + let finalize_permit = match result { + Some(_) => tokio::select! { + biased; + _ = finalize_cancel.cancelled() => None, + permit = finalize_capacity.acquire_owned() => { + Some(permit.expect("finalizer capacity is never closed")) + } + }, + None => None, + }; + if heartbeat_guard.is_none() { in_flight.remove((job_id, job_run_lease)); } @@ -447,7 +475,7 @@ impl JobExecutor { } metrics.record_in_flight_change(&job_queue, -1); - let Some(result) = result else { + let (Some(result), Some(finalize_permit)) = (result, finalize_permit) else { return; }; let dlq_enabled = dlq_policy.enabled_for(&job_queue); @@ -547,12 +575,19 @@ impl JobExecutor { } drop(heartbeat_guard); + // Lifecycle hooks run under the finalizer permit, so a hook slower + // than job throughput holds back claiming instead of accumulating + // tasks that each retain a copy of the job. + if let Some(started_event) = started_event { + let _ = started_event.await; + } if let Ok(CompletionOutcome::Applied { event: Some(event), .. }) = outcome { dispatch_lifecycle_event(&lifecycle_handlers, &job_kind, event).await; } + drop(finalize_permit); }); } .instrument(span) @@ -3444,6 +3479,8 @@ mod tests { storage, DlqPolicy::default(), finalizers.clone(), + CancellationToken::new(), + 1, ); executor .execute_task( @@ -3505,6 +3542,99 @@ mod tests { ); } + #[tokio::test(start_paused = true)] + async fn saturated_finalizer_capacity_retains_worker_until_shutdown() { + for storage in [ + RuntimeStorage::Canonical, + RuntimeStorage::QueueStorage( + QueueStorageRuntime::new( + awa_model::QueueStorageConfig::default(), + Duration::from_secs(1), + Duration::from_secs(1), + ) + .unwrap(), + ), + ] { + let now = chrono::Utc::now(); + let job: JobRow = serde_json::from_value(serde_json::json!({ + "id": 1, "kind": "finalize_parent", "queue": "test", "args": {}, + "state": "running", "priority": 2, "attempt": 1, "run_lease": 1, + "max_attempts": 5, "run_at": now, "attempted_at": now, + "created_at": now, "metadata": {}, "tags": [], + })) + .unwrap(); + let key = (job.id, job.run_lease); + let pool = sqlx::postgres::PgPoolOptions::new() + .connect_lazy("postgres://postgres:test@localhost/awa_test") + .unwrap(); + let metrics = test_metrics(); + let finalizers = TaskTracker::new(); + let finalize_cancel = CancellationToken::new(); + let in_flight = Arc::new(crate::runtime::InFlightRegistry::default()); + let count = Arc::new(AtomicU32::new(0)); + let (_unstarted, handle) = crate::completion::CompletionBatcher::new( + pool.clone(), + CancellationToken::new(), + metrics.clone(), + storage.clone(), + 1, + ); + let executor = JobExecutor::new( + pool.clone(), + Arc::new(HashMap::from([( + "finalize_parent".into(), + Box::new(FinalizeWorker) as BoxedWorker, + )])), + Arc::new(HashMap::new()), + Arc::new(HashMap::new()), + in_flight.clone(), + Arc::new(HashMap::from([("test".into(), count.clone())])), + Arc::new(HashMap::new()), + metrics, + handle, + storage, + DlqPolicy::default(), + finalizers.clone(), + finalize_cancel.clone(), + 1, + ); + // Occupy the entire budget without starting database work. + let capacity = executor.finalize_capacity.clone(); + let held = capacity + .clone() + .acquire_many_owned(capacity.available_permits() as u32) + .await + .unwrap(); + let task = tokio::spawn(executor.execute_task( + DispatchedJob { + job, + queue_storage_claim: None, + queue_storage_unique_states: None, + abandon_after: None, + }, + Arc::new(AtomicBool::new(false)), + )); + tokio::task::yield_now().await; + assert!(!task.is_finished()); + assert_eq!(count.load(Ordering::SeqCst), 1); + assert!(in_flight.get_cancel(key).is_some()); + assert!(finalizers.is_empty()); + + // Cancelling background services at the drain deadline must release + // the capacity waiter without spawning another finalizer. + finalize_cancel.cancel(); + tokio::time::timeout(Duration::from_secs(1), task) + .await + .unwrap() + .unwrap(); + assert_eq!(count.load(Ordering::SeqCst), 0); + assert!(in_flight.get_cancel(key).is_none()); + assert!(finalizers.is_empty()); + drop(held); + pool.close().await; + } + } + #[tokio::test] async fn cancelled_finalization_unregisters_only_its_exact_lease() { let in_flight = Arc::new(crate::runtime::InFlightRegistry::default()); diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index 4b9765602..48f1ea066 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -16,7 +16,7 @@ use croner::Cron; use sqlx::pool::PoolConnection; use sqlx::{PgPool, Postgres}; use std::collections::{HashMap, HashSet}; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::task::{JoinHandle, JoinSet}; @@ -24,6 +24,9 @@ use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, warn}; use uuid::Uuid; +/// Detached lifecycle hook tasks maintenance may have outstanding at once. +const MAINTENANCE_LIFECYCLE_EVENT_CAPACITY: usize = 1024; + /// The canonical unique-claim conflict raised by `awa.sync_job_unique_claims` /// (SQLSTATE 23505, constraint reported as `idx_awa_jobs_unique`). See #388. fn is_unique_claim_conflict(err: &awa_model::AwaError) -> bool { @@ -675,6 +678,10 @@ pub struct MaintenanceService { /// In-process `Rescued` lifecycle hooks (ADR-015), shared with the /// executor. Best-effort fire-and-forget per rescued job. lifecycle_handlers: Arc>>, + /// Bounds the detached lifecycle hook tasks maintenance has outstanding; + /// events beyond it are dropped and counted. + lifecycle_event_capacity: Arc, + dropped_lifecycle_events: Arc, /// In-flight job cancellation flags — used to signal deadline/heartbeat rescue /// to running handlers on this worker instance. in_flight: InFlightMap, @@ -978,6 +985,10 @@ impl MaintenanceService { standby_queue_storage: None, enqueue_specs, lifecycle_handlers, + lifecycle_event_capacity: Arc::new(tokio::sync::Semaphore::new( + MAINTENANCE_LIFECYCLE_EVENT_CAPACITY, + )), + dropped_lifecycle_events: Arc::new(AtomicU64::new(0)), heartbeat_rescue_interval: Duration::from_secs(30), deadline_rescue_interval: DEFAULT_DEADLINE_RESCUE_INTERVAL, callback_rescue_interval: Duration::from_secs(30), @@ -2412,9 +2423,7 @@ impl MaintenanceService { job: job.clone(), reason, }; - tokio::spawn(async move { - crate::executor::dispatch_lifecycle_event(&handlers, &kind, event).await; - }); + self.spawn_lifecycle_event(handlers, kind, event); } } @@ -2501,8 +2510,41 @@ impl MaintenanceService { job: job.clone(), reason, }; + self.spawn_lifecycle_event(handlers, kind, event); + } + + /// Runs `kind`'s lifecycle hooks for `event` on a detached task, so a slow + /// hook does not hold up maintenance. At most + /// [`MAINTENANCE_LIFECYCLE_EVENT_CAPACITY`] such tasks are outstanding; an + /// event beyond that is dropped rather than retaining its job until the + /// hooks catch up. + fn spawn_lifecycle_event( + &self, + handlers: Arc>>, + kind: String, + event: crate::events::UntypedJobEvent, + ) { + if !handlers.contains_key(&kind) { + return; + } + let Ok(permit) = self.lifecycle_event_capacity.clone().try_acquire_owned() else { + let dropped = self + .dropped_lifecycle_events + .fetch_add(1, Ordering::Relaxed) + + 1; + if dropped % 1000 == 1 { + warn!( + kind, + dropped, + capacity = MAINTENANCE_LIFECYCLE_EVENT_CAPACITY, + "Dropped a maintenance lifecycle event: earlier hooks have not returned" + ); + } + return; + }; tokio::spawn(async move { crate::executor::dispatch_lifecycle_event(&handlers, &kind, event).await; + drop(permit); }); } diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index c269f3acd..545898877 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -20,7 +20,7 @@ use opentelemetry_sdk::metrics::{InMemoryMetricExporter, SdkMeterProvider}; use serde::{Deserialize, Serialize}; use sqlx::postgres::PgPoolOptions; use std::collections::{HashMap, HashSet}; -use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::{Mutex, Notify}; @@ -2240,6 +2240,289 @@ async fn test_queue_storage_graceful_shutdown_finalizes_returned_attempts() { ); } +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct FinalizeBacklogJob { + n: i64, +} + +/// Stalling durable completion must backpressure claiming rather than grow +/// detached tasks holding job payloads. Release the locks and verify that +/// every job completes once. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_stops_claiming_while_finalization_is_stalled() { + assert_claiming_stops_while_finalization_is_stalled( + "qs_finalize_backlog", + "awa_qs_finalize_backlog", + true, + ) + .await; +} + +/// The same bound holds when completion deletes materialized leases and +/// writes terminal rows instead of closing receipt claims. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_lease_runtime_stops_claiming_while_finalization_is_stalled() { + assert_claiming_stops_while_finalization_is_stalled( + "qs_finalize_backlog_leases", + "awa_qs_finalize_backlog_leases", + false, + ) + .await; +} + +/// Lifecycle hooks run under the finalizer bound: while a hook does not +/// return, claiming stops at that bound instead of accumulating hook tasks, +/// and every job completes once the hooks return. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_slow_lifecycle_hook_holds_back_claiming() { + const FINALIZER_CAPACITY: usize = 1_024; + const WORKERS: u32 = 64; + const JOBS: i64 = 1_600; + let queue = "qs_slow_lifecycle_hook"; + let (_db_guard, pool) = setup_pool(10).await; + let store_config = QueueStorageConfig { + schema: "awa_qs_slow_lifecycle_hook".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + for first in (0..JOBS).step_by(500) { + let params: Vec<_> = (first..(first + 500).min(JOBS)) + .map(|n| { + insert::params_with( + &FinalizeBacklogJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue backlog"); + } + + let (open_hooks, hooks_open) = tokio::sync::watch::channel(false); + let runs = Arc::new(AtomicUsize::new(0)); + let recorded = runs.clone(); + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: WORKERS, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: WORKERS as usize, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register::(move |_, _| { + recorded.fetch_add(1, Ordering::SeqCst); + async { Ok(JobResult::Completed) } + }) + .on_event::(move |_| { + let mut hooks_open = hooks_open.clone(); + async move { + let _ = hooks_open.wait_for(|open| *open).await; + } + }) + .heartbeat_staleness(Duration::from_secs(300)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build slow hook client"); + client.start().await.expect("start slow hook client"); + + let limit = FINALIZER_CAPACITY + WORKERS as usize; + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); + let mut last = (0, Instant::now()); + let started_while_blocked = loop { + let started = runs.load(Ordering::SeqCst); + if started > limit { + break started; + } + if started != last.0 { + last = (started, Instant::now()); + } else if started > 0 && last.1.elapsed() >= Duration::from_secs(5) { + break started; + } + assert!( + Instant::now() < deadline, + "claiming neither stopped nor passed {limit}" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + }; + + open_hooks.send(true).expect("open the hooks"); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(60)); + while completed_terminal_count(&pool, &store, queue).await < JOBS { + assert!(Instant::now() < deadline, "the backlog did not drain"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + client.shutdown(Duration::from_secs(10)).await; + + assert!( + started_while_blocked <= limit, + "{started_while_blocked} jobs started while lifecycle hooks were blocked; \ + at most {limit} may occupy finalizer slots or hold a worker" + ); + assert_eq!(runs.load(Ordering::SeqCst) as i64, JOBS); +} + +async fn assert_claiming_stops_while_finalization_is_stalled( + queue: &str, + schema: &str, + lease_claim_receipts: bool, +) { + // The default completion pipeline: one shard of 512-job batches, one + // batch in flight and one queued. + const FINALIZER_CAPACITY: usize = 1_024; + const WORKERS: u32 = 64; + const JOBS: i64 = 1_600; + let (_db_guard, pool) = setup_pool(10).await; + let store_config = QueueStorageConfig { + schema: schema.to_string(), + lease_claim_receipts, + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + for first in (0..JOBS).step_by(500) { + let params: Vec<_> = (first..(first + 500).min(JOBS)) + .map(|n| { + insert::params_with( + &FinalizeBacklogJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue backlog"); + } + + let mut stall = pool.begin().await.expect("begin stall transaction"); + if lease_claim_receipts { + let first_jobs: Vec = sqlx::query_scalar(&format!( + "SELECT job_id FROM {}.ready_entries WHERE queue = $1 ORDER BY lane_seq LIMIT $2", + store.schema() + )) + .bind(queue) + .bind(i64::from(WORKERS) * 2) + .fetch_all(&pool) + .await + .expect("first jobs"); + sqlx::query( + "SELECT pg_catalog.pg_advisory_xact_lock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', job_id, 1), 0)) + FROM unnest($1::bigint[]) AS job_id", + ) + .bind(&first_jobs) + .execute(stall.as_mut()) + .await + .expect("stall the first jobs' finalization"); + } else { + sqlx::query(&format!( + "LOCK TABLE {}.done_entries IN SHARE MODE", + store.schema() + )) + .execute(stall.as_mut()) + .await + .expect("stall terminal writes"); + } + + let runs: Arc>> = Arc::default(); + let recorded = runs.clone(); + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: WORKERS, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: WORKERS as usize, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register::(move |_, ctx| { + *recorded + .lock() + .expect("runs") + .entry(ctx.job.id) + .or_insert(0) += 1; + async { Ok(JobResult::Completed) } + }) + .heartbeat_staleness(Duration::from_secs(300)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build finalize backlog client"); + client.start().await.expect("start finalize backlog client"); + + // Wait for claiming to stop, or to run past the bound. A fixed runtime + // can never pass it, so a pause in claiming only weakens the check. + let limit = FINALIZER_CAPACITY + WORKERS as usize; + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); + let mut last = (0, Instant::now()); + let started_while_stalled = loop { + let started = runs.lock().expect("runs").len(); + if started > limit { + break started; + } + if started != last.0 { + last = (started, Instant::now()); + } else if started > 0 && last.1.elapsed() >= Duration::from_secs(5) { + break started; + } + assert!( + Instant::now() < deadline, + "claiming neither stopped nor passed {limit}" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + }; + + stall + .rollback() + .await + .expect("release the stalled finalization"); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(60)); + while completed_terminal_count(&pool, &store, queue).await < JOBS { + assert!(Instant::now() < deadline, "the backlog did not drain"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + client.shutdown(Duration::from_secs(10)).await; + + assert!( + started_while_stalled <= limit, + "{started_while_stalled} jobs started while finalization was stalled; \ + at most {limit} may occupy finalizer slots or hold a worker" + ); + let runs = runs.lock().expect("runs"); + assert_eq!(runs.len() as i64, JOBS); + assert!( + runs.values().all(|count| *count == 1), + "a job ran more than once" + ); +} + /// 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 @@ -4613,46 +4896,6 @@ async fn test_queue_storage_promote_due_preserves_row_fields_across_lanes() { .expect("a different key must still enqueue"); } -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn test_queue_storage_enqueue_batch_writing_more_segments_than_the_old_bind_limit() { - // `enqueue_batch` stamps each row with its own run_at, and segments split - // on every run_at change, so one batch yields about one segment per row: - // more than the 8,191 rows an 8-bind-per-row statement could carry. - const JOBS: usize = 9_000; - let (_db_guard, pool) = setup_pool(10).await; - let queue = "qs_bind_limit_segments"; - let schema = "awa_qs_bind_limit_segments"; - let store = create_store(&pool, schema).await; - - let inserted = store - .enqueue_batch(&pool, queue, 2, JOBS as i64) - .await - .expect("a batch producing one segment per row should insert"); - assert_eq!(inserted, JOBS as i64); - assert_eq!( - table_row_count(&pool, schema, "ready_entries").await, - JOBS as i64 - ); - - let segments: i64 = sqlx::query_scalar(&format!( - "SELECT count(*) FROM {schema}.ready_segments WHERE queue = $1" - )) - .bind(queue) - .fetch_one(&pool) - .await - .expect("Failed to count ready segments"); - assert!( - segments > 8_191, - "expected more than 8,191 segments from one batch, got {segments}" - ); - - let claimed = store - .claim_runtime_batch(&pool, queue, 100, Duration::from_secs(30)) - .await - .expect("Failed to claim jobs"); - assert_eq!(claimed.len(), 100); -} - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_late_completion_after_cancel_is_noop() { let (_db_guard, pool) = setup_pool(10).await; @@ -20616,6 +20859,80 @@ async fn test_busy_retained_head_does_not_block_materialized_heartbeat_rescue() historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Heartbeat, true).await; } +/// Metadata completions validate the exact compact item and do not lock the +/// immutable batch shared with other attempts. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_metadata_completion_validates_and_does_not_lock_compact_batch() { + let (_db_guard, pool) = setup_pool(10).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: "awa_qs_metadata_completion_identity".into(), + ..Default::default() + }, + ) + .await; + let queue = "metadata_completion_identity"; + for id in 0..2 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + tags: vec!["retained".into()], + metadata: serde_json::json!({"source": "recovery"}), + ..Default::default() + }, + ) + .await; + } + let jobs = store + .claim_runtime_batch(&pool, queue, 2, Duration::ZERO) + .await + .expect("claim compact jobs with metadata"); + assert_eq!(jobs.len(), 2); + assert_eq!(jobs[0].claim.claim_batch_id, jobs[1].claim.claim_batch_id); + let mut invalid = jobs[0].clone(); + invalid.claim.claim_batch_index = jobs[1].claim.claim_batch_index; + assert!(store + .complete_runtime_batch(&pool, &[invalid]) + .await + .expect("reject a different batch item") + .is_empty()); + + let mut holder = pool.begin().await.expect("begin batch holder"); + sqlx::query(&format!( + "SELECT batch_id FROM {}.lease_claim_batches \ + WHERE claim_slot = $1 AND batch_id = $2 FOR UPDATE", + store.schema() + )) + .bind(jobs[0].claim.claim_slot) + .bind(jobs[0].claim.claim_batch_id.expect("compact batch")) + .execute(holder.as_mut()) + .await + .expect("hold the immutable batch row"); + let completed = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.complete_runtime_batch(&pool, &jobs), + ) + .await + .expect("completion must not wait on an unrelated batch-row lock") + .expect("complete metadata jobs"); + assert_eq!(completed.len(), 2); + holder.rollback().await.expect("release batch row"); + for job in jobs { + let durable = store + .load_job(&pool, job.job.id) + .await + .expect("read completed job") + .expect("completed job"); + assert_eq!(durable.state, JobState::Completed); + assert_eq!(durable.tags, vec!["retained".to_string()]); + assert_eq!(durable.metadata, serde_json::json!({"source": "recovery"})); + } +} + /// Claim slots of `schema` whose `lease_claim_batches` partition carries the /// GIN index on `job_ids`. async fn claim_batch_job_index_slots(pool: &sqlx::PgPool, schema: &str) -> Vec { diff --git a/docs/lifecycle-hooks.md b/docs/lifecycle-hooks.md index 43304f67c..46ab87e66 100644 --- a/docs/lifecycle-hooks.md +++ b/docs/lifecycle-hooks.md @@ -75,9 +75,11 @@ A callback that _resumes_ the job to `running` emits no event itself; the job re Hooks are **best-effort, in-process notifications** — not a durable workflow mechanism: - They observe **committed** state. An outcome hook carries the post-transition `JobRow`; a stale/rejected finalization emits no outcome event (though a `Started` may already have fired for that attempt). -- They do not block executor progress or queue capacity, and `shutdown()` does not wait for them. They can be lost on process crash or shutdown. +- They do not delay a job's state transition, and `shutdown()` does not wait for them. They can be lost on process crash or shutdown. +- Each attempt's hooks hold one of the worker's finalizer slots until they return, the same bounded slots that wait for completions to commit. Hooks slower than job throughput therefore slow claiming instead of piling up in memory. Keep hooks fast, and hand slow work to a job or follow-up instead. +- Maintenance dispatches `Rescued` events, and `Cancelled` for duplicates it supersedes, on detached tasks with at most 1,024 outstanding per worker. Events beyond that are dropped with a warning rather than holding up rescue. - **Dispatch is not atomic with the state transition.** The hook runs _after_ the job's transition commits, as a separate in-process step — it is not a transactional outbox. If the process dies between the commit and the dispatch, the job state is durable but the event is lost. The same applies to callback resolution: `complete_external` / `resolve_callback` commit the outcome first, then dispatch. -- `Started` is dispatched detached, so a very fast job may finish before its `Started` hook completes — do not rely on strict started-before-outcome ordering for short jobs. +- `Started` runs alongside the handler, so a very fast job's state transition may commit before its `Started` hook completes. For outcomes the worker finalizes itself, the outcome hook runs only after the `Started` hook has returned. - A panicking hook is logged; later hooks for the same kind still run. If a side effect must fire exactly once or survive a crash, do not drive it from a hook at all — the dispatch may never run, and enqueueing the follow-up work _from_ the hook just inherits that same unreliability. See the next section for the durable mechanism.