diff --git a/CHANGELOG.md b/CHANGELOG.md index 39b44fcb2..e0fd687e5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,21 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **`admin::cancel_by_unique_key` and `cancel_by_unique_key_tx` on queue + storage find running jobs and no longer read retained history** + ([#585](https://github.com/hardbyte/awa/issues/585)). The candidate lookup + built every available ready row before filtering by key, reading every + retained ready row, completed jobs included. A missing key cost as much + as a full scan, and the transaction variant held the caller's + transaction while it ran. Given a queue, it now starts from that queue's + ready rows holding the key and checks only those jobs for being available, + leased or on an open receipt claim; with no such row, only the + `deferred_jobs` key index is read. Its cost follows that queue's own rows, + not other queues' backlogs or the claim ring, which a long-running job can + keep from pruning. Without a queue, available rows are walked lane by lane + from the claim cursor and open receipt claims are read in full. It also + finds jobs running on receipt claims not yet materialised into `leases`, + which it missed, so they couldn't be cancelled by key. - **The admin UI's runtime page no longer shows the storage transition card once the transition is finalized.** It stayed up on every cluster running queue storage; it now appears only while an engine is prepared or a mixed diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index 5cbeacd30..0a66610e2 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -897,34 +897,86 @@ async fn cancel_in_tx<'a>( Ok(Some(job)) } -/// Queue-storage candidate lookup for a unique key. Shared by -/// [`cancel_by_unique_key`] and [`cancel_by_unique_key_tx`] so the two cannot -/// drift. -fn unique_key_candidate_sql(schema: &str) -> String { +/// Queue-storage candidate lookup for a unique key: the oldest nonterminal +/// job holding key `$1`. Shared by [`cancel_by_unique_key`] and +/// [`cancel_by_unique_key_tx`] so the two cannot drift. +/// +/// Each source is filtered by the key where it is read, so a missing key +/// costs no more than reading the live work: +/// - available rows are walked lane by lane through `ready_segments` from the +/// claim cursor, the way a claim finds them, and probed along the ready key +/// within each segment's ring slot. The ready key leads with the ring slot, +/// so a lane range without it scans every retained ready row, completed +/// jobs included; +/// - scheduled and retryable rows use the `deferred_jobs` unique-key index; +/// - running and waiting jobs are read from `leases` and from open receipt +/// claims that have not been materialised into `leases`. +/// +/// Without a queue, the open receipt claims are read in full before the key +/// filter, so this lookup costs what the claim ring retains. With `scoped` +/// see [`scoped_unique_key_candidate_sql`], which starts from the key. +fn unique_key_candidate_sql(schema: &str, scoped: bool) -> String { + if scoped { + return scoped_unique_key_candidate_sql(schema); + } + let open_running = crate::queue_storage::open_receipt_running_claims_sql(schema); format!( r#" - WITH current_available AS ( - SELECT ready.job_id, ready.unique_key - FROM {schema}.ready_entries AS ready - JOIN {schema}.queue_claim_heads AS claims - ON claims.queue = ready.queue - AND claims.priority = ready.priority - AND claims.enqueue_shard = ready.enqueue_shard - WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) - 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 - ) + WITH heads AS ( + SELECT + claims.queue, + claims.priority, + claims.enqueue_shard, + {schema}.sequence_next_value(claims.seq_name) AS claim_seq, + {schema}.sequence_next_value(enqueues.seq_name) AS next_seq + FROM {schema}.queue_claim_heads AS claims + JOIN {schema}.queue_enqueue_heads AS enqueues + ON enqueues.queue = claims.queue + AND enqueues.priority = claims.priority + AND enqueues.enqueue_shard = claims.enqueue_shard ), candidates AS ( - SELECT job_id - FROM current_available - WHERE unique_key = $1 + SELECT ready.job_id + FROM heads + CROSS JOIN LATERAL ( + SELECT + segment.ready_slot, + segment.ready_generation, + GREATEST(segment.first_lane_seq, heads.claim_seq) AS from_seq, + LEAST(segment.next_lane_seq, heads.next_seq) AS to_seq + FROM {schema}.ready_segments AS segment + WHERE segment.queue = heads.queue + AND segment.priority = heads.priority + AND segment.enqueue_shard = heads.enqueue_shard + AND segment.next_lane_seq > heads.claim_seq + AND segment.first_lane_seq < heads.next_seq + ORDER BY segment.next_lane_seq + ) AS segment + -- OFFSET 0 keeps this a per-segment probe along the ready key; + -- flattened, the key filter can lead the planner to scan every + -- retained ready row instead. + CROSS JOIN LATERAL ( + SELECT ready.* + FROM {schema}.ready_entries AS ready + WHERE ready.ready_slot = segment.ready_slot + AND ready.ready_generation = segment.ready_generation + AND ready.queue = heads.queue + AND ready.priority = heads.priority + AND ready.enqueue_shard = heads.enqueue_shard + AND ready.lane_seq >= segment.from_seq + AND ready.lane_seq < segment.to_seq + AND ready.unique_key = $1 + OFFSET 0 + ) AS ready + WHERE 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 + ) UNION ALL SELECT job_id FROM {schema}.deferred_jobs @@ -934,24 +986,207 @@ fn unique_key_candidate_sql(schema: &str) -> String { -- reach it through the lease's lane identity. SELECT leases.job_id FROM {schema}.leases AS leases - JOIN {schema}.ready_entries AS ready - ON ready.ready_slot = leases.ready_slot - AND ready.ready_generation = leases.ready_generation - AND ready.queue = leases.queue - AND ready.priority = leases.priority - AND ready.enqueue_shard = leases.enqueue_shard - AND ready.lane_seq = leases.lane_seq - WHERE ready.unique_key = $1 + JOIN LATERAL ( + SELECT 1 + FROM {schema}.ready_entries AS ready + WHERE ready.ready_slot = leases.ready_slot + AND ready.ready_generation = leases.ready_generation + AND ready.queue = leases.queue + AND ready.priority = leases.priority + AND ready.enqueue_shard = leases.enqueue_shard + AND ready.lane_seq = leases.lane_seq + AND ready.unique_key = $1 + OFFSET 0 + ) AS ready ON TRUE + UNION ALL + -- Running receipt claims not yet materialised into `leases`. + -- This leg reads every claim still in the claim ring before the + -- key filter, so its cost follows claim-ring retention (bounded + -- by claim rotation and pruning), not completed-job history. + SELECT open_claims.job_id + FROM ({open_running}) AS open_claims + JOIN LATERAL ( + SELECT 1 + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = open_claims.job_id + AND ready.unique_key = $1 + LIMIT 1 + ) AS ready ON TRUE ) SELECT job_id FROM candidates ORDER BY job_id ASC LIMIT 1 - "#, - schema = schema + "# + ) +} + +/// The ready rows of the physical queues in `$2` that hold the key `$1`, +/// probed one ring slot at a time: the ready key leads with `ready_slot`, so +/// a filter on `queue` alone scans every retained ready row. The slot count +/// is the schema's own; the admin store only knows the default. +fn keyed_ready_rows_sql(schema: &str) -> String { + format!( + r#" + SELECT + ready.job_id, + ready.ready_slot, + ready.ready_generation, + ready.queue, + ready.priority, + ready.enqueue_shard, + ready.lane_seq + FROM {schema}.queue_ring_state AS ring + CROSS JOIN LATERAL generate_series(0, ring.slot_count - 1) AS slot(ready_slot) + CROSS JOIN LATERAL ( + SELECT ready.* + FROM {schema}.ready_entries AS ready + WHERE ready.ready_slot = slot.ready_slot + AND ready.queue = ANY($2) + AND ready.unique_key = $1 + OFFSET 0 + ) AS ready + "# + ) +} + +/// The queue-scoped candidate lookup: starts from the queue's ready rows that +/// hold the key, read through the lane index, then keeps the ones still live. +/// Its cost follows that queue's retained rows, not other queues' backlogs or +/// the claim ring, which a large backlog elsewhere can keep from pruning. +/// +/// - available: at or past the lane's claim cursor and not tombstoned; +/// - scheduled and retryable: the `deferred_jobs` unique-key index; +/// - running or waiting: in `leases` or on an open receipt claim, both +/// checked by `job_id` for just these jobs. +/// +/// `$2` holds the physical queues of the logical queue. +fn scoped_unique_key_candidate_sql(schema: &str) -> String { + let open_running = crate::queue_storage::open_receipt_running_claims_sql_for( + schema, + Some("ARRAY(SELECT keyed.job_id FROM keyed)"), + ); + let keyed = keyed_ready_rows_sql(schema); + format!( + r#" + WITH keyed AS MATERIALIZED ( + {keyed} + ), + candidates AS ( + SELECT keyed.job_id + FROM keyed + JOIN {schema}.queue_claim_heads AS claims + ON claims.queue = keyed.queue + AND claims.priority = keyed.priority + AND claims.enqueue_shard = keyed.enqueue_shard + WHERE keyed.lane_seq >= {schema}.sequence_next_value(claims.seq_name) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = keyed.ready_slot + AND tomb.ready_generation = keyed.ready_generation + AND tomb.queue = keyed.queue + AND tomb.priority = keyed.priority + AND tomb.enqueue_shard = keyed.enqueue_shard + AND tomb.lane_seq = keyed.lane_seq + ) + UNION ALL + SELECT job_id + FROM {schema}.deferred_jobs + WHERE unique_key = $1 + UNION ALL + -- By job id rather than lane: a job moved to another queue + -- while available is leased on its new lane. + SELECT keyed.job_id + FROM keyed + WHERE EXISTS ( + SELECT 1 FROM {schema}.leases AS lease + WHERE lease.job_id = keyed.job_id + ) + UNION ALL + SELECT open_claims.job_id + FROM ({open_running}) AS open_claims + ) + SELECT job_id + FROM candidates + ORDER BY job_id ASC + LIMIT 1 + "# ) } +/// The physical queues that hold a logical queue's lanes, read from the +/// schema's `queue_lanes`: the queue itself and its `queue#N` stripes. The +/// admin store is rebuilt from the schema name alone, so its configured +/// stripe count is the default, not the one the runtime enqueues with. +async fn persisted_physical_queues( + conn: &mut sqlx::PgConnection, + schema: &str, + queue: &str, +) -> Result, AwaError> { + let mut queues: Vec = sqlx::query_scalar(audited_sql(format!( + r#" + SELECT DISTINCT lanes.queue + FROM {schema}.queue_lanes AS lanes + WHERE left(lanes.queue, char_length($1) + 1) = $1 || '#' + AND substr(lanes.queue, char_length($1) + 2) ~ '^[0-9]+$' + "# + ))) + .bind(queue) + .fetch_all(&mut *conn) + .await?; + queues.push(queue.to_string()); + Ok(queues) +} + +/// The oldest nonterminal queue-storage job holding `unique_key`. +/// +/// A queue is part of a key only when it was inserted with `by_queue`, so a +/// job holding a queue-scoped key was enqueued there: with a queue, its ready +/// rows are read from that queue alone. A job moved to another queue is +/// found through its original ready row while that row is retained, except +/// while it is available in the new queue; cancel it by id otherwise. +async fn queue_storage_unique_key_candidate( + conn: &mut sqlx::PgConnection, + store: &QueueStorage, + unique_key: &[u8], + queue: Option<&str>, +) -> Result, AwaError> { + let schema = store.schema(); + let Some(queue) = queue else { + let sql = unique_key_candidate_sql(schema, false); + return Ok(sqlx::query_scalar(audited_sql(sql)) + .bind(unique_key) + .fetch_optional(&mut *conn) + .await?); + }; + let queues = persisted_physical_queues(&mut *conn, schema, queue).await?; + // Available, running and waiting jobs all have a ready row in the queue. + // Without one, only a scheduled or retryable job can hold the key, and + // the deferred index answers that alone; the full lookup spans enough + // partitioned tables that planning it costs several milliseconds. + // One statement, so a job promoted from deferred_jobs to a ready row + // between the two reads cannot be missed by both. + let keyed = keyed_ready_rows_sql(schema); + let (has_ready_rows, deferred): (bool, Option) = sqlx::query_as(audited_sql(format!( + "SELECT EXISTS ({keyed}), \ + (SELECT job_id FROM {schema}.deferred_jobs WHERE unique_key = $1 \ + ORDER BY job_id LIMIT 1)" + ))) + .bind(unique_key) + .bind(&queues) + .fetch_one(&mut *conn) + .await?; + if !has_ready_rows { + return Ok(deferred); + } + let sql = unique_key_candidate_sql(schema, true); + Ok(sqlx::query_scalar(audited_sql(sql)) + .bind(unique_key) + .bind(&queues) + .fetch_optional(&mut *conn) + .await?) +} + /// Canonical (non queue-storage) candidate lookup for a unique key. const UNIQUE_KEY_CANDIDATE_SQL_CANONICAL: &str = r#" WITH candidates AS ( @@ -996,10 +1231,17 @@ const UNIQUE_KEY_CANDIDATE_SQL_CANONICAL: &str = r#" /// in-flight cancellation notification. If the selected job completes before /// `cancel(...)` locks it, the cancel becomes a no-op and returns `None`. /// -/// The lookup scans `unique_key` on both physical tables without a dedicated -/// index. This is acceptable for low-volume use cases. For high-volume tables, -/// consider adding a partial index on `unique_key WHERE unique_key IS NOT NULL` -/// or routing through `job_unique_claims` (which is already indexed). +/// On the canonical engine the lookup scans `unique_key` on both physical +/// tables without a dedicated index. On queue storage, scheduled and +/// retryable rows are found by their unique-key index. With a `queue`, the +/// lookup reads that queue's ready rows in each ring slot, retained ones +/// included, and checks only the rows holding the key for being available, +/// leased or on an open receipt claim; its cost follows that queue's own +/// retained rows, not other queues' backlogs. Without a `queue`, available +/// rows are walked lane by lane from the claim cursor and open receipt claims +/// are read in full, so the cost follows what the claim ring retains. A job +/// moved to another queue while available is not found by a lookup that +/// names its original queue. pub async fn cancel_by_unique_key( pool: &PgPool, kind: &str, @@ -1010,14 +1252,16 @@ pub async fn cancel_by_unique_key( let unique_key = crate::unique::compute_unique_key(kind, queue, args, period_bucket); if let Some(store) = active_queue_storage(pool).await? { - let sql = unique_key_candidate_sql(store.schema()); - let candidate: Option = sqlx::query_scalar(audited_sql(sql)) - .bind(&unique_key) - .fetch_optional(pool) - .await?; + let mut conn = pool.acquire().await?; + let candidate = + queue_storage_unique_key_candidate(&mut conn, &store, &unique_key, queue).await?; + drop(conn); return match candidate { - Some(job_id) => cancel(pool, job_id).await, + Some(job_id) => match cancel(pool, job_id).await { + Err(AwaError::JobNotFound { .. }) => Ok(None), + other => other, + }, None => Ok(None), }; } @@ -1077,14 +1321,14 @@ async fn cancel_by_unique_key_in_tx<'a>( let unique_key = crate::unique::compute_unique_key(kind, queue, args, period_bucket); if let Some(store) = active_queue_storage_in_tx(tx).await? { - let sql = unique_key_candidate_sql(store.schema()); - let candidate: Option = sqlx::query_scalar(audited_sql(sql)) - .bind(&unique_key) - .fetch_optional(tx.as_mut()) - .await?; + let candidate = + queue_storage_unique_key_candidate(tx.as_mut(), &store, &unique_key, queue).await?; return match candidate { - Some(job_id) => cancel_in_tx(tx, job_id).await, + Some(job_id) => match cancel_in_tx(tx, job_id).await { + Err(AwaError::JobNotFound { .. }) => Ok(None), + other => other, + }, None => Ok(None), }; } diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 746a596e7..cdae8148a 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -1298,15 +1298,32 @@ pub fn queue_signals_sql(schema: &str, reads: QueueSignalReads) -> String { /// is reported by the caller's own `leases` leg, so it is excluded here to /// avoid double counting. pub fn open_receipt_running_claims_sql(schema: &str) -> String { + open_receipt_running_claims_sql_for(schema, None) +} + +/// [`open_receipt_running_claims_sql`] restricted to the jobs in `job_ids`, +/// an SQL expression yielding `bigint[]`. The restriction applies before the +/// closure and materialisation checks. Row-local claims are found through the +/// `job_id` index. Compact claim batches have no index on `job_ids`, so every +/// retained batch header is read, but only batches holding one of the jobs +/// are unnested and checked. +pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option<&str>) -> String { let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let row_closed = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); + let (row_filter, batch_filter) = match job_ids { + Some(ids) => ( + format!("claims.job_id = ANY({ids}) AND "), + format!("batches.job_ids && {ids} AND items.job_id = ANY({ids}) AND "), + ), + None => (String::new(), String::new()), + }; format!( r#" SELECT claims.job_id, claims.run_lease, claims.queue FROM {schema}.lease_claims AS claims - WHERE NOT {row_closed} + WHERE {row_filter}NOT {row_closed} AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id @@ -1322,7 +1339,7 @@ pub fn open_receipt_running_claims_sql(schema: &str) -> String { ) AS items(job_id, run_lease, receipt_id) -- Skips, before unnesting, a batch whose receipts closure batches -- have all closed: most retained batches, once their jobs finish. - WHERE NOT ( + WHERE {batch_filter}NOT ( batches.receipt_ranges <@ COALESCE(( SELECT range_agg(closed) FROM {schema}.lease_claim_closure_batches AS closure_batches @@ -3600,7 +3617,7 @@ impl QueueStorage { format!("{queue}{QUEUE_STRIPE_DELIMITER}{stripe}") } - fn physical_queues_for_logical(&self, queue: &str) -> Vec { + pub(crate) fn physical_queues_for_logical(&self, queue: &str) -> Vec { if !self.uses_queue_striping() || self.is_physical_stripe_queue(queue) { return vec![queue.to_string()]; } diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index d1cf9fe2b..caedc9c23 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -19476,3 +19476,585 @@ async fn test_brin_fsm_baseline_block_size_model() { } pool.close().await; } + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct UniqueCancelJob { + key: String, +} + +fn unique_cancel_job(key: &str) -> UniqueCancelJob { + UniqueCancelJob { key: key.into() } +} + +fn unique_cancel_opts(queue: &str, run_at: Option>) -> InsertOpts { + InsertOpts { + queue: queue.into(), + run_at, + unique: Some(UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + } +} + +fn unique_cancel_store_config(schema: &str) -> QueueStorageConfig { + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + } +} + +async fn cancel_unique(pool: &sqlx::PgPool, queue: Option<&str>, key: &str) -> Option { + admin::cancel_by_unique_key( + pool, + UniqueCancelJob::kind(), + queue, + Some(&serde_json::to_value(unique_cancel_job(key)).unwrap()), + None, + ) + .await + .expect("cancel by unique key") +} + +/// `cancel_by_unique_key` finds a job in every live state on queue storage: +/// scheduled, available, running on a receipt claim, and waiting on a +/// callback. When two live jobs share a key, the oldest is cancelled first. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_queue_storage_cancel_by_unique_key_finds_each_live_state() { + let (_db_guard, pool) = setup_pool(8).await; + let queue = "qs_cancel_by_key"; + let store = + create_store_with_config(&pool, unique_cancel_store_config("awa_qs_cancel_by_key")).await; + + let running = enqueue_job( + &pool, + &store, + &unique_cancel_job("running"), + unique_cancel_opts(queue, None), + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim the running job"); + assert_eq!(claimed[0].job.id, running); + + let waiting = enqueue_job( + &pool, + &store, + &unique_cancel_job("waiting"), + unique_cancel_opts(queue, None), + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim the waiting job"); + assert_eq!(claimed[0].job.id, waiting); + let callback_id = store + .register_callback( + &pool, + waiting, + claimed[0].job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("register the callback"); + assert!(store + .enter_callback_wait(&pool, waiting, claimed[0].job.run_lease, callback_id) + .await + .expect("enter the callback wait")); + assert_eq!( + admin::get_job(&pool, waiting).await.unwrap().state, + JobState::WaitingExternal + ); + // Two live jobs sharing a key: uniqueness enforced only among failed jobs. + let shared_opts = || InsertOpts { + unique: Some(UniqueOpts { + by_queue: true, + by_args: true, + states: 1 << JobState::Failed.bit_position(), + ..Default::default() + }), + ..unique_cancel_opts(queue, None) + }; + let shared_older = + enqueue_job(&pool, &store, &unique_cancel_job("shared"), shared_opts()).await; + let shared_newer = + enqueue_job(&pool, &store, &unique_cancel_job("shared"), shared_opts()).await; + + let available = enqueue_job( + &pool, + &store, + &unique_cancel_job("available"), + unique_cancel_opts(queue, None), + ) + .await; + let scheduled = enqueue_job( + &pool, + &store, + &unique_cancel_job("scheduled"), + unique_cancel_opts(queue, Some(Utc::now() + chrono::Duration::hours(1))), + ) + .await; + + assert!(cancel_unique(&pool, Some(queue), "absent").await.is_none()); + for (key, job_id) in [ + ("scheduled", scheduled), + ("available", available), + ("running", running), + ("waiting", waiting), + ("shared", shared_older), + ] { + let cancelled = cancel_unique(&pool, Some(queue), key) + .await + .unwrap_or_else(|| panic!("a {key} job should be found")); + assert_eq!(cancelled.id, job_id, "{key}"); + assert_eq!( + admin::get_job(&pool, job_id).await.unwrap().state, + JobState::Cancelled, + "{key}" + ); + } + assert_eq!( + admin::get_job(&pool, shared_newer).await.unwrap().state, + JobState::Available, + "only the oldest job holding the key is cancelled" + ); + assert!( + cancel_unique(&pool, None, "shared").await.is_none(), + "the key includes the queue, so a lookup without it matches nothing" + ); + assert_eq!( + cancel_unique(&pool, Some(queue), "shared") + .await + .map(|job| job.id), + Some(shared_newer), + "the next lookup finds the newer job" + ); +} + +/// The admin store is rebuilt from the schema name alone, with the default +/// ring size. A schema installed with more ready-ring slots still has its +/// higher slots searched. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_queue_storage_cancel_by_unique_key_reads_every_installed_slot() { + let (_db_guard, pool) = setup_pool(8).await; + let queue = "qs_cancel_by_key_wide"; + let schema = "awa_qs_cancel_by_key_wide"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + queue_slot_count: 20, + ..unique_cancel_store_config(schema) + }, + ) + .await; + // A job in another queue keeps the ring from idling, so each rotation + // advances the cursor. + enqueue_job( + &pool, + &store, + &unique_cancel_job("anchor"), + unique_cancel_opts("qs_cancel_by_key_anchor", None), + ) + .await; + for _ in 0..16 { + let outcome = store.rotate(&pool).await.expect("rotate the ready ring"); + assert!( + matches!(outcome, RotateOutcome::Rotated { .. }), + "{outcome:?}" + ); + } + assert_eq!(ring_cursor(&pool, schema, "queue").await.0, 16); + + let running = enqueue_job( + &pool, + &store, + &unique_cancel_job("running"), + unique_cancel_opts(queue, None), + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim the running job"); + assert_eq!(claimed[0].job.id, running); + let available = enqueue_job( + &pool, + &store, + &unique_cancel_job("available"), + unique_cancel_opts(queue, None), + ) + .await; + let slots: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT DISTINCT ready_slot FROM {schema}.ready_entries WHERE queue = $1" + ))) + .bind(queue) + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(slots, vec![16]); + + for (key, job_id) in [("available", available), ("running", running)] { + let cancelled = cancel_unique(&pool, Some(queue), key) + .await + .unwrap_or_else(|| panic!("the {key} job in slot 16 should be found")); + assert_eq!(cancelled.id, job_id, "{key}"); + } +} + +/// A striped queue's jobs live in `queue#N` lanes. The admin store's default +/// stripe count is 1, so a lookup by the logical queue name must find the +/// stripes from the lanes the schema holds. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_queue_storage_cancel_by_unique_key_finds_striped_lanes() { + let (_db_guard, pool) = setup_pool(8).await; + let queue = "qs_cancel_by_key_striped"; + let schema = "awa_qs_cancel_by_key_striped"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + queue_stripe_count: 2, + ..unique_cancel_store_config(schema) + }, + ) + .await; + + let running = enqueue_job( + &pool, + &store, + &unique_cancel_job("running"), + unique_cancel_opts(queue, None), + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim the running job"); + assert_eq!(claimed[0].job.id, running); + let available = enqueue_job( + &pool, + &store, + &unique_cancel_job("available"), + unique_cancel_opts(queue, None), + ) + .await; + let scheduled = enqueue_job( + &pool, + &store, + &unique_cancel_job("scheduled"), + unique_cancel_opts(queue, Some(Utc::now() + chrono::Duration::hours(1))), + ) + .await; + let unstriped: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.ready_entries WHERE queue NOT LIKE $1" + ))) + .bind(format!("{queue}#%")) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(unstriped, 0, "every ready row is in a stripe lane"); + + for (key, job_id) in [ + ("scheduled", scheduled), + ("available", available), + ("running", running), + ] { + let cancelled = cancel_unique(&pool, Some(queue), key) + .await + .unwrap_or_else(|| panic!("a striped {key} job should be found")); + assert_eq!(cancelled.id, job_id, "{key}"); + } +} + +/// A job moved to another queue while available keeps the key its original +/// queue gave it. Once it runs or waits in the new queue, a lookup naming +/// the original queue still finds it, through its retained ready row there. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_queue_storage_cancel_by_unique_key_finds_a_moved_running_job() { + let (_db_guard, pool) = setup_pool(8).await; + let origin = "qs_cancel_by_key_origin"; + let target = "qs_cancel_by_key_target"; + let store = create_store_with_config( + &pool, + unique_cancel_store_config("awa_qs_cancel_by_key_moved"), + ) + .await; + + let running = enqueue_job( + &pool, + &store, + &unique_cancel_job("moved-running"), + unique_cancel_opts(origin, None), + ) + .await; + let waiting = enqueue_job( + &pool, + &store, + &unique_cancel_job("moved-waiting"), + unique_cancel_opts(origin, None), + ) + .await; + for job_id in [running, waiting] { + assert!(store.move_queue(&pool, job_id, target, None).await.unwrap()); + } + let claimed = store + .claim_runtime_batch(&pool, target, 2, Duration::from_secs(300)) + .await + .expect("claim the moved jobs"); + assert_eq!(claimed.len(), 2); + let waiting_claim = claimed + .iter() + .find(|claim| claim.job.id == waiting) + .expect("waiting job claimed"); + let callback_id = store + .register_callback( + &pool, + waiting, + waiting_claim.job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("register the callback"); + assert!(store + .enter_callback_wait(&pool, waiting, waiting_claim.job.run_lease, callback_id) + .await + .expect("enter the callback wait")); + + for (key, job_id) in [("moved-running", running), ("moved-waiting", waiting)] { + let cancelled = cancel_unique(&pool, Some(origin), key) + .await + .unwrap_or_else(|| panic!("the {key} job should be found by its original queue")); + assert_eq!(cancelled.id, job_id, "{key}"); + } +} + +/// The transaction variant cancels only when the caller commits, and a +/// running job's cancellation notice is sent on commit. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_queue_storage_cancel_by_unique_key_tx_follows_the_callers_transaction() { + let (_db_guard, pool) = setup_pool(8).await; + let queue = "qs_cancel_by_key_tx"; + let store = + create_store_with_config(&pool, unique_cancel_store_config("awa_qs_cancel_by_key_tx")) + .await; + let job_id = enqueue_job( + &pool, + &store, + &unique_cancel_job("tx"), + unique_cancel_opts(queue, None), + ) + .await; + store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim the job"); + let args = serde_json::to_value(unique_cancel_job("tx")).unwrap(); + + let mut listener = sqlx::postgres::PgListener::connect_with(&pool) + .await + .expect("connect listener"); + listener.listen("awa:cancel").await.expect("listen"); + + let mut tx = pool.begin().await.unwrap(); + let cancelled = admin::cancel_by_unique_key_tx( + &mut tx, + UniqueCancelJob::kind(), + Some(queue), + Some(&args), + None, + ) + .await + .unwrap(); + assert_eq!(cancelled.map(|job| job.id), Some(job_id)); + tx.rollback().await.unwrap(); + assert_eq!( + admin::get_job(&pool, job_id).await.unwrap().state, + JobState::Running, + "a rolled-back cancel leaves the job running" + ); + + let mut tx = pool.begin().await.unwrap(); + admin::cancel_by_unique_key_tx( + &mut tx, + UniqueCancelJob::kind(), + Some(queue), + Some(&args), + None, + ) + .await + .unwrap() + .expect("the job is still found"); + tx.commit().await.unwrap(); + let notice = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let notice = listener.recv().await.expect("receive notification"); + let payload: serde_json::Value = serde_json::from_str(notice.payload()).unwrap(); + if payload["job_id"] == job_id { + return payload; + } + } + }) + .await + .expect("the cancellation notice arrives after commit"); + assert_eq!(notice["job_id"], job_id); + assert_eq!( + admin::get_job(&pool, job_id).await.unwrap().state, + JobState::Cancelled + ); +} + +/// Rows the lookup reads on one connection, from that backend's statistics. +async fn unique_key_lookup_rows( + conn: &mut sqlx::PgConnection, + schema: &str, + queue: Option<&str>, + key: &str, +) -> i64 { + async fn rows_read(conn: &mut sqlx::PgConnection, schema: &str) -> i64 { + sqlx::query_scalar( + "SELECT COALESCE(sum(COALESCE(seq_tup_read, 0) + COALESCE(idx_tup_fetch, 0)), 0)::bigint \ + FROM pg_stat_user_tables WHERE schemaname = $1", + ) + .bind(schema) + .fetch_one(&mut *conn) + .await + .unwrap() + } + let before = rows_read(conn, schema).await; + let found = admin::cancel_by_unique_key_tx( + &mut *conn, + UniqueCancelJob::kind(), + queue, + Some(&serde_json::to_value(unique_cancel_job(key)).unwrap()), + None, + ) + .await + .unwrap(); + assert!(found.is_none()); + sqlx::query("SELECT pg_stat_force_next_flush()") + .execute(&mut *conn) + .await + .unwrap(); + sqlx::query("SELECT 1").execute(&mut *conn).await.unwrap(); + rows_read(conn, schema).await - before +} + +/// A unique-key lookup reads the live work, not retained history: rows +/// read stay flat as completed jobs accumulate, and with a queue they stay +/// flat as other queues' backlogs grow and as the claim ring retains more +/// claims, which a long-running job elsewhere can keep from pruning. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_queue_storage_cancel_by_unique_key_reads_live_work_only() { + let (_db_guard, pool) = setup_pool(8).await; + let schema = "awa_qs_cancel_by_key_cost"; + let queue = "qs_cancel_by_key_cost"; + let other = "qs_cancel_by_key_cost_other"; + let store = create_store_with_config(&pool, unique_cancel_store_config(schema)).await; + // A finished job still holding a key: its retained ready row sends a + // lookup for that key through the full scoped query. + let finished = enqueue_job( + &pool, + &store, + &unique_cancel_job("finished"), + unique_cancel_opts(queue, None), + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .unwrap(); + assert_eq!(claimed[0].job.id, finished); + store.complete_runtime_batch(&pool, &claimed).await.unwrap(); + store.enqueue_batch(&pool, queue, 2, 200).await.unwrap(); + + let complete = |count: i64| { + let (pool, store) = (pool.clone(), &store); + async move { + let mut remaining = count; + while remaining > 0 { + let n = remaining.min(500); + store.enqueue_batch(&pool, other, 2, n).await.unwrap(); + let claims = store + .claim_runtime_batch(&pool, other, n, Duration::ZERO) + .await + .unwrap(); + store.complete_runtime_batch(&pool, &claims).await.unwrap(); + remaining -= n; + } + for _ in 0..=store.claim_slot_count() { + store.rotate_claims(&pool).await.unwrap(); + store.prune_oldest_claims(&pool).await.unwrap(); + } + } + }; + complete(500).await; + + let mut conn = pool.acquire().await.unwrap(); + sqlx::raw_sql("SET max_parallel_workers_per_gather = 0; SET stats_fetch_consistency = 'none'") + .execute(&mut *conn) + .await + .unwrap(); + sqlx::query("ANALYZE").execute(&mut *conn).await.unwrap(); + let scoped = unique_key_lookup_rows(&mut conn, schema, Some(queue), "absent").await; + let unscoped = unique_key_lookup_rows(&mut conn, schema, None, "absent").await; + + // Ten times the retained history. + complete(4_500).await; + sqlx::query("ANALYZE").execute(&mut *conn).await.unwrap(); + let flat = |before: i64| before + before / 2 + 64; + let scoped_history = unique_key_lookup_rows(&mut conn, schema, Some(queue), "absent").await; + let unscoped_history = unique_key_lookup_rows(&mut conn, schema, None, "absent").await; + assert!( + scoped_history <= flat(scoped), + "scoped lookup read {scoped} rows, then {scoped_history} with more history" + ); + assert!( + unscoped_history <= flat(unscoped), + "unscoped lookup read {unscoped} rows, then {unscoped_history} with more history" + ); + + // Another queue's backlog: a lookup given its queue doesn't read it. + for _ in 0..10 { + store.enqueue_batch(&pool, other, 2, 500).await.unwrap(); + } + sqlx::query("ANALYZE").execute(&mut *conn).await.unwrap(); + let scoped_backlog = unique_key_lookup_rows(&mut conn, schema, Some(queue), "absent").await; + assert!( + scoped_backlog <= flat(scoped), + "scoped lookup read {scoped} rows, then {scoped_backlog} beside another queue's backlog" + ); + + // Claims retained without pruning: a scoped lookup reads none of them. + // A deadline makes them row-local claims, as jobs with a deadline get. + let scoped_finished = unique_key_lookup_rows(&mut conn, schema, Some(queue), "finished").await; + let mut retained = 0; + while retained < 5_000 { + let claims = store + .claim_runtime_batch(&pool, other, 500, Duration::from_secs(300)) + .await + .unwrap(); + store.complete_runtime_batch(&pool, &claims).await.unwrap(); + retained += 500; + } + sqlx::query("ANALYZE").execute(&mut *conn).await.unwrap(); + let scoped_claims = unique_key_lookup_rows(&mut conn, schema, Some(queue), "absent").await; + let scoped_finished_claims = + unique_key_lookup_rows(&mut conn, schema, Some(queue), "finished").await; + assert!( + scoped_claims <= flat(scoped), + "scoped lookup read {scoped} rows, then {scoped_claims} with retained claims" + ); + assert!( + scoped_finished_claims <= flat(scoped_finished), + "scoped lookup of a finished job's key read {scoped_finished} rows, then \ + {scoped_finished_claims} with retained claims" + ); +} diff --git a/skills/awa-jobs/SKILL.md b/skills/awa-jobs/SKILL.md index 037fdf85e..3895bf63a 100644 --- a/skills/awa-jobs/SKILL.md +++ b/skills/awa-jobs/SKILL.md @@ -129,7 +129,13 @@ depends on the API: a single-row `insert` / `insert_with` and the row-by-row corresponding error) on a duplicate — you must handle it, do not assume a silent no-op. Only the COPY bulk path (`insert_many_copy`) silently skips duplicates. Cancel without storing the id via `cancel_by_unique_key(kind, queue=, args=, -period_bucket=)` — the components must match those used at insert. +period_bucket=)` — the components must match those used at insert. It cancels +the oldest live job holding the key. On queue storage, a `by_queue` key is +looked up through that queue's ready rows, so a job moved to another queue +can't be cancelled by key while it is available there. A moved running or +waiting job can't be cancelled by key once its original ready row is pruned; +cancel it by id. Scheduled and retryable jobs are found by key wherever they +are. Args-based uniqueness relies on canonical (sorted-key) JSON, pinned by cross-language golden tests; the default state set includes `completed`, so a still-retained completed job suppresses a re-insert with the same args. Scope