From d064c91ed9df658e5a0d6cdcff6c40b12df6cfbe Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Mon, 2 Nov 2026 12:02:43 +0000 Subject: [PATCH 1/6] Find unique-key cancel candidates without scanning retained ready rows (#585) On queue storage, cancel_by_unique_key and cancel_by_unique_key_tx built every available ready row before filtering by key. The range didn't lead with the ready key's ring slot, so it read every retained ready row, completed jobs included, and a missing key cost a full scan while the tx variant held the caller's transaction. The key is now filtered where each source is read. Available rows are walked lane by lane through ready_segments from the claim cursor, as claims find them, and with a queue only that queue's lanes are walked. Scheduled and retryable rows use the deferred unique-key index, and running and waiting jobs come from leases and open receipt claims. The lookup also missed running receipt claims not yet materialised into leases, so a running job couldn't be cancelled by key; it now finds them. Queue-storage tests cover each live state, oldest-first among jobs sharing a key, the transaction variant's rollback and commit-time notice, and rows read staying flat as retained completions and other queues' backlogs grow. Against the previous lookup they fail: the running job isn't found, and a scoped missing-key lookup reads 726 rows, then 5,206. --- CHANGELOG.md | 11 + awa-model/src/admin.rs | 186 +++++++++---- awa-model/src/queue_storage.rs | 2 +- awa/tests/queue_storage_runtime_test.rs | 335 ++++++++++++++++++++++++ 4 files changed, 485 insertions(+), 49 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e5449bcfe..647e000f7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,17 @@ 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. It now filters by key where each source is + read, walks available rows lane by lane from the claim cursor (only the + given queue's lanes when a queue is passed), and also checks running + receipt claims, which it missed: a job running on a claim not yet + materialised into `leases` couldn't be cancelled by key. - **A request the admin UI's client abandons no longer holds a database connection for good.** When a browser closes a tab or a proxy times out, the server drops the request but Postgres keeps running its query. If that diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index 53618de00..426a4fe6e 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -866,34 +866,90 @@ 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`. +/// +/// With `scoped`, the available-row walk covers only the lanes of the +/// physical queues in `$2`, so a lookup in one queue doesn't read every +/// queue's backlog. The other sources are indexed or small and are read for +/// every queue. +fn unique_key_candidate_sql(schema: &str, scoped: bool) -> String { + let heads_scope = if scoped { + "WHERE claims.queue = ANY($2)" + } else { + "" + }; + 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 + {heads_scope} ), 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 @@ -903,24 +959,59 @@ 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`. + 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 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: available rows are +/// read from that queue's lanes alone. A job moved to another queue while +/// available is therefore not found by a queue-scoped lookup; scheduled, +/// retryable, running and waiting jobs are found wherever they are. +async fn queue_storage_unique_key_candidate( + conn: &mut sqlx::PgConnection, + store: &QueueStorage, + unique_key: &[u8], + queue: Option<&str>, +) -> Result, AwaError> { + let sql = unique_key_candidate_sql(store.schema(), queue.is_some()); + let mut query = sqlx::query_scalar(audited_sql(sql)).bind(unique_key); + if let Some(queue) = queue { + query = query.bind(store.physical_queues_for_logical(queue)); + } + Ok(query.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 ( @@ -965,10 +1056,13 @@ 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 it reads the live work: +/// available rows along each lane from its claim cursor (only the lanes of +/// `queue` when one is given), scheduled and retryable rows by their +/// unique-key index, and running and waiting jobs, so its cost doesn't grow +/// with retained history. 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, @@ -979,11 +1073,10 @@ 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, @@ -1046,11 +1139,8 @@ 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, diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 71c1c24b1..2a020fbf3 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -3485,7 +3485,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 b86f953e5..ebab808c1 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -19392,3 +19392,338 @@ 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); + store + .register_callback( + &pool, + waiting, + claimed[0].job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("park the job on a callback"); + // 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 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>, +) -> 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("absent")).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. +#[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; + 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)).await; + let unscoped = unique_key_lookup_rows(&mut conn, schema, None).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)).await; + let unscoped_history = unique_key_lookup_rows(&mut conn, schema, None).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)).await; + assert!( + scoped_backlog <= flat(scoped), + "scoped lookup read {scoped} rows, then {scoped_backlog} beside another queue's backlog" + ); +} From cab717e51b6e6045df2e55cdd2a134069be2bf75 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Mon, 2 Nov 2026 12:05:34 +0000 Subject: [PATCH 2/6] Note the moved-job exception for cancel-by-key in the jobs skill --- skills/awa-jobs/SKILL.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/skills/awa-jobs/SKILL.md b/skills/awa-jobs/SKILL.md index 037fdf85e..8297f252a 100644 --- a/skills/awa-jobs/SKILL.md +++ b/skills/awa-jobs/SKILL.md @@ -129,7 +129,10 @@ 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 in that queue's available rows only, so a job moved to another queue +while available can't be cancelled by key; cancel it by id. 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 From 9cac59cee5c269938770dbe2244328b54f4fc1c0 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 12:24:30 +0000 Subject: [PATCH 3/6] Cover a waiting_external job in the cancel-by-key test register_callback alone leaves the job running; enter the callback wait so the test cancels a job in waiting_external. Note that the open-receipt-claim leg of the candidate lookup costs claim-ring retention, not job history. --- awa-model/src/admin.rs | 3 +++ awa/tests/queue_storage_runtime_test.rs | 12 ++++++++++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index 426a4fe6e..935739eaf 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -973,6 +973,9 @@ fn unique_key_candidate_sql(schema: &str, scoped: bool) -> String { ) 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 ( diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index ebab808c1..003263d76 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -19473,7 +19473,7 @@ async fn test_queue_storage_cancel_by_unique_key_finds_each_live_state() { .await .expect("claim the waiting job"); assert_eq!(claimed[0].job.id, waiting); - store + let callback_id = store .register_callback( &pool, waiting, @@ -19481,7 +19481,15 @@ async fn test_queue_storage_cancel_by_unique_key_finds_each_live_state() { Duration::from_secs(3600), ) .await - .expect("park the job on a callback"); + .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 { From ef5123427854d1d9ecdfd656a70682431624bea3 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 01:31:34 +0000 Subject: [PATCH 4/6] fix(admin): start a queue-scoped unique-key lookup from the key Given a queue, the lookup now reads that queue's ready rows holding the key, one ring slot at a time along the ready key, and checks only those jobs: available past the claim cursor and not tombstoned, in leases, or on an open receipt claim found by job_id. With no such row only the deferred_jobs key index can hold the key, so that query alone answers. The previous lookup read every open receipt claim before filtering by key. A long-running job keeps the claim ring from pruning; on the 0.6 line, 120,000 retained row-local claims made each lookup take about 9.7 s. With this change the same database answers in 1.7 ms for a missing key and 11-20 ms to cancel a scheduled or available job. The cost test also measures a lookup for a finished job's key and keeps row-local claims retained. --- CHANGELOG.md | 14 ++- awa-model/src/admin.rs | 160 +++++++++++++++++++++--- awa-model/src/queue_storage.rs | 19 ++- awa/tests/queue_storage_runtime_test.rs | 57 +++++++-- skills/awa-jobs/SKILL.md | 5 +- 5 files changed, 219 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 647e000f7..67befedea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,11 +12,15 @@ Notable changes between releases. Detailed migration notes for storage transitio 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. It now filters by key where each source is - read, walks available rows lane by lane from the claim cursor (only the - given queue's lanes when a queue is passed), and also checks running - receipt claims, which it missed: a job running on a claim not yet - materialised into `leases` couldn't be cancelled by key. + 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. - **A request the admin UI's client abandons no longer holds a database connection for good.** When a browser closes a tab or a proxy times out, the server drops the request but Postgres keeps running its query. If that diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index 935739eaf..f28b7068f 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -881,16 +881,14 @@ async fn cancel_in_tx<'a>( /// - running and waiting jobs are read from `leases` and from open receipt /// claims that have not been materialised into `leases`. /// -/// With `scoped`, the available-row walk covers only the lanes of the -/// physical queues in `$2`, so a lookup in one queue doesn't read every -/// queue's backlog. The other sources are indexed or small and are read for -/// every queue. -fn unique_key_candidate_sql(schema: &str, scoped: bool) -> String { - let heads_scope = if scoped { - "WHERE claims.queue = ANY($2)" - } else { - "" - }; +/// 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` +/// (the store's ready-ring slot count) see +/// [`scoped_unique_key_candidate_sql`], which starts from the key. +fn unique_key_candidate_sql(schema: &str, scoped: Option) -> String { + if let Some(queue_slots) = scoped { + return scoped_unique_key_candidate_sql(schema, queue_slots); + } let open_running = crate::queue_storage::open_receipt_running_claims_sql(schema); format!( r#" @@ -906,7 +904,6 @@ fn unique_key_candidate_sql(schema: &str, scoped: bool) -> String { ON enqueues.queue = claims.queue AND enqueues.priority = claims.priority AND enqueues.enqueue_shard = claims.enqueue_shard - {heads_scope} ), candidates AS ( SELECT ready.job_id @@ -994,25 +991,148 @@ fn unique_key_candidate_sql(schema: &str, scoped: bool) -> String { ) } +/// 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. +fn keyed_ready_rows_sql(schema: &str, queue_slots: usize) -> String { + format!( + r#" + SELECT + ready.job_id, + ready.ready_slot, + ready.ready_generation, + ready.queue, + ready.priority, + ready.enqueue_shard, + ready.lane_seq + FROM generate_series(0, {last_slot}) 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 + "#, + last_slot = queue_slots.saturating_sub(1) + ) +} + +/// 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 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, queue_slots: usize) -> 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, queue_slots); + 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 + SELECT keyed.job_id + FROM keyed + WHERE EXISTS ( + SELECT 1 FROM {schema}.leases AS lease + WHERE lease.ready_slot = keyed.ready_slot + AND lease.ready_generation = keyed.ready_generation + AND lease.queue = keyed.queue + AND lease.priority = keyed.priority + AND lease.enqueue_shard = keyed.enqueue_shard + AND lease.lane_seq = keyed.lane_seq + ) + 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 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: available rows are -/// read from that queue's lanes alone. A job moved to another queue while -/// available is therefore not found by a queue-scoped lookup; scheduled, -/// retryable, running and waiting jobs are found wherever they are. +/// 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 while +/// available is therefore not found by a queue-scoped lookup unless it is +/// scheduled or retryable; cancel it by id. async fn queue_storage_unique_key_candidate( conn: &mut sqlx::PgConnection, store: &QueueStorage, unique_key: &[u8], queue: Option<&str>, ) -> Result, AwaError> { - let sql = unique_key_candidate_sql(store.schema(), queue.is_some()); - let mut query = sqlx::query_scalar(audited_sql(sql)).bind(unique_key); - if let Some(queue) = queue { - query = query.bind(store.physical_queues_for_logical(queue)); + let schema = store.schema(); + let Some(queue) = queue else { + let sql = unique_key_candidate_sql(schema, None); + return Ok(sqlx::query_scalar(audited_sql(sql)) + .bind(unique_key) + .fetch_optional(&mut *conn) + .await?); + }; + let queues = store.physical_queues_for_logical(queue); + // 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. + let keyed = keyed_ready_rows_sql(schema, store.queue_slot_count()); + let has_ready_rows: bool = sqlx::query_scalar(audited_sql(format!("SELECT EXISTS ({keyed})"))) + .bind(unique_key) + .bind(&queues) + .fetch_one(&mut *conn) + .await?; + if !has_ready_rows { + return Ok(sqlx::query_scalar(audited_sql(format!( + "SELECT job_id FROM {schema}.deferred_jobs WHERE unique_key = $1 \ + ORDER BY job_id LIMIT 1" + ))) + .bind(unique_key) + .fetch_optional(&mut *conn) + .await?); } - Ok(query.fetch_optional(&mut *conn).await?) + let sql = unique_key_candidate_sql(schema, Some(store.queue_slot_count())); + 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. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 2a020fbf3..c1c6f32eb 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -1194,15 +1194,30 @@ 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, through the `job_id` indexes, so the +/// cost follows the number of jobs asked about rather than the claim ring. +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 @@ -1216,7 +1231,7 @@ pub fn open_receipt_running_claims_sql(schema: &str) -> String { batches.run_leases, batches.receipt_ids ) AS items(job_id, run_lease, receipt_id) - WHERE NOT EXISTS ( + WHERE {batch_filter}NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = batches.claim_slot AND closures.job_id = items.job_id diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 003263d76..b3b364073 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -19636,6 +19636,7 @@ 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( @@ -19652,7 +19653,7 @@ async fn unique_key_lookup_rows( &mut *conn, UniqueCancelJob::kind(), queue, - Some(&serde_json::to_value(unique_cancel_job("absent")).unwrap()), + Some(&serde_json::to_value(unique_cancel_job(key)).unwrap()), None, ) .await @@ -19668,7 +19669,8 @@ async fn unique_key_lookup_rows( /// 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. +/// 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; @@ -19676,6 +19678,21 @@ async fn test_queue_storage_cancel_by_unique_key_reads_live_work_only() { 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| { @@ -19706,15 +19723,15 @@ async fn test_queue_storage_cancel_by_unique_key_reads_live_work_only() { .await .unwrap(); sqlx::query("ANALYZE").execute(&mut *conn).await.unwrap(); - let scoped = unique_key_lookup_rows(&mut conn, schema, Some(queue)).await; - let unscoped = unique_key_lookup_rows(&mut conn, schema, None).await; + 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)).await; - let unscoped_history = unique_key_lookup_rows(&mut conn, schema, None).await; + 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" @@ -19729,9 +19746,35 @@ async fn test_queue_storage_cancel_by_unique_key_reads_live_work_only() { 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)).await; + 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 8297f252a..b49e4de17 100644 --- a/skills/awa-jobs/SKILL.md +++ b/skills/awa-jobs/SKILL.md @@ -131,8 +131,9 @@ 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. It cancels the oldest live job holding the key. On queue storage, a `by_queue` key is -looked up in that queue's available rows only, so a job moved to another queue -while available can't be cancelled by key; cancel it by id. +looked up through that queue's ready rows, so a job moved to another queue +while available can't be cancelled by key unless it is scheduled or retryable; +cancel it by id. 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 From 157332abd7da212eca2698178dd5a461dd0a2569 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 01:59:15 +0000 Subject: [PATCH 5/6] fix(admin): find a moved running job by key, in one snapshot Review follow-ups on the queue-scoped unique-key lookup: - Match leases by job_id, as open receipt claims already are. A job moved to another queue while available is leased on its new lane, so matching by its original ready row's lane missed it once it ran or waited there. - Read the ready-row check and the deferred_jobs lookup in one statement, so a job promoted from deferred_jobs between two reads is not missed by both. - Say what the compact-batch leg of the job-restricted open-claims query reads: every retained batch header, unnesting only batches that hold one of the jobs. New test: two jobs moved to another queue, one running on a receipt claim and one waiting on a callback, are both found by a lookup naming their original queue. --- awa-model/src/admin.rs | 43 ++++++++--------- awa-model/src/queue_storage.rs | 6 ++- awa/tests/queue_storage_runtime_test.rs | 62 +++++++++++++++++++++++++ skills/awa-jobs/SKILL.md | 4 +- 4 files changed, 88 insertions(+), 27 deletions(-) diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index f28b7068f..c50c27de8 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -1026,8 +1026,8 @@ fn keyed_ready_rows_sql(schema: &str, queue_slots: usize) -> String { /// /// - 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 checked by -/// `job_id` for just these jobs. +/// - 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, queue_slots: usize) -> String { @@ -1063,16 +1063,13 @@ fn scoped_unique_key_candidate_sql(schema: &str, queue_slots: usize) -> String { 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.ready_slot = keyed.ready_slot - AND lease.ready_generation = keyed.ready_generation - AND lease.queue = keyed.queue - AND lease.priority = keyed.priority - AND lease.enqueue_shard = keyed.enqueue_shard - AND lease.lane_seq = keyed.lane_seq + WHERE lease.job_id = keyed.job_id ) UNION ALL SELECT open_claims.job_id @@ -1090,9 +1087,9 @@ fn scoped_unique_key_candidate_sql(schema: &str, queue_slots: usize) -> String { /// /// 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 while -/// available is therefore not found by a queue-scoped lookup unless it is -/// scheduled or retryable; cancel it by id. +/// 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, @@ -1112,20 +1109,20 @@ async fn queue_storage_unique_key_candidate( // 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, store.queue_slot_count()); - let has_ready_rows: bool = sqlx::query_scalar(audited_sql(format!("SELECT EXISTS ({keyed})"))) - .bind(unique_key) - .bind(&queues) - .fetch_one(&mut *conn) - .await?; + 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(sqlx::query_scalar(audited_sql(format!( - "SELECT job_id FROM {schema}.deferred_jobs WHERE unique_key = $1 \ - ORDER BY job_id LIMIT 1" - ))) - .bind(unique_key) - .fetch_optional(&mut *conn) - .await?); + return Ok(deferred); } let sql = unique_key_candidate_sql(schema, Some(store.queue_slot_count())); Ok(sqlx::query_scalar(audited_sql(sql)) diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index c1c6f32eb..13f3aee75 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -1199,8 +1199,10 @@ pub fn open_receipt_running_claims_sql(schema: &str) -> String { /// [`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, through the `job_id` indexes, so the -/// cost follows the number of jobs asked about rather than the claim ring. +/// 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"); diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index b3b364073..dae03cc34 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -19556,6 +19556,68 @@ async fn test_queue_storage_cancel_by_unique_key_finds_each_live_state() { ); } +/// 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)] diff --git a/skills/awa-jobs/SKILL.md b/skills/awa-jobs/SKILL.md index b49e4de17..414ae47dd 100644 --- a/skills/awa-jobs/SKILL.md +++ b/skills/awa-jobs/SKILL.md @@ -132,8 +132,8 @@ Cancel without storing the id via `cancel_by_unique_key(kind, queue=, args=, period_bucket=)` — the components must match those used at insert. It cancels the oldest live job holding the key. On queue storage, a `by_queue` key is looked up through that queue's ready rows, so a job moved to another queue -while available can't be cancelled by key unless it is scheduled or retryable; -cancel it by id. +can't be cancelled by key while it is available there, or once its original +ready row is pruned; cancel it by id. 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 From 7254e5e5c50820a729c519ae1ad12fdb294002b5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 03:25:26 +0000 Subject: [PATCH 6/6] fix(admin): scope unique-key lookups by the schema's own ring and lanes The admin store is rebuilt from the active schema name with the default config, so the queue-scoped lookup probed only the default 16 ready-ring slots and bound only the logical queue name. On a schema installed with more slots it missed available, running and waiting jobs in the higher slots, and on a striped queue it missed jobs in the queue#N lanes. The slot probe now takes its range from the schema's queue_ring_state, and the physical queues come from the schema's queue_lanes. A job that finishes between the lookup and the cancel returns None on queue storage too, as it does on the canonical engine. Two regression tests cover a 20-slot ring with rows in slot 16 and a two-stripe queue; both fail on the previous lookup. The cancel_by_unique_key doc and the jobs skill now say what the lookup reads. --- awa-model/src/admin.rs | 82 ++++++++++----- awa/tests/queue_storage_runtime_test.rs | 134 ++++++++++++++++++++++++ skills/awa-jobs/SKILL.md | 6 +- 3 files changed, 196 insertions(+), 26 deletions(-) diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index 1d698c979..0a66610e2 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -914,11 +914,10 @@ async fn cancel_in_tx<'a>( /// /// 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` -/// (the store's ready-ring slot count) see -/// [`scoped_unique_key_candidate_sql`], which starts from the key. -fn unique_key_candidate_sql(schema: &str, scoped: Option) -> String { - if let Some(queue_slots) = scoped { - return scoped_unique_key_candidate_sql(schema, queue_slots); +/// 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!( @@ -1024,8 +1023,9 @@ fn unique_key_candidate_sql(schema: &str, scoped: Option) -> String { /// 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. -fn keyed_ready_rows_sql(schema: &str, queue_slots: usize) -> String { +/// 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 @@ -1036,7 +1036,8 @@ fn keyed_ready_rows_sql(schema: &str, queue_slots: usize) -> String { ready.priority, ready.enqueue_shard, ready.lane_seq - FROM generate_series(0, {last_slot}) AS slot(ready_slot) + 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 @@ -1045,8 +1046,7 @@ fn keyed_ready_rows_sql(schema: &str, queue_slots: usize) -> String { AND ready.unique_key = $1 OFFSET 0 ) AS ready - "#, - last_slot = queue_slots.saturating_sub(1) + "# ) } @@ -1061,12 +1061,12 @@ fn keyed_ready_rows_sql(schema: &str, queue_slots: usize) -> String { /// 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, queue_slots: usize) -> String { +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, queue_slots); + let keyed = keyed_ready_rows_sql(schema); format!( r#" WITH keyed AS MATERIALIZED ( @@ -1114,6 +1114,30 @@ fn scoped_unique_key_candidate_sql(schema: &str, queue_slots: usize) -> String { ) } +/// 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 @@ -1129,20 +1153,20 @@ async fn queue_storage_unique_key_candidate( ) -> Result, AwaError> { let schema = store.schema(); let Some(queue) = queue else { - let sql = unique_key_candidate_sql(schema, None); + 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 = store.physical_queues_for_logical(queue); + 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, store.queue_slot_count()); + 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 \ @@ -1155,7 +1179,7 @@ async fn queue_storage_unique_key_candidate( if !has_ready_rows { return Ok(deferred); } - let sql = unique_key_candidate_sql(schema, Some(store.queue_slot_count())); + let sql = unique_key_candidate_sql(schema, true); Ok(sqlx::query_scalar(audited_sql(sql)) .bind(unique_key) .bind(&queues) @@ -1208,12 +1232,16 @@ const UNIQUE_KEY_CANDIDATE_SQL_CANONICAL: &str = r#" /// `cancel(...)` locks it, the cancel becomes a no-op and returns `None`. /// /// On the canonical engine the lookup scans `unique_key` on both physical -/// tables without a dedicated index. On queue storage it reads the live work: -/// available rows along each lane from its claim cursor (only the lanes of -/// `queue` when one is given), scheduled and retryable rows by their -/// unique-key index, and running and waiting jobs, so its cost doesn't grow -/// with retained history. A job moved to another queue while available is -/// not found by a lookup that names its original queue. +/// 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, @@ -1230,7 +1258,10 @@ pub async fn cancel_by_unique_key( 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), }; } @@ -1294,7 +1325,10 @@ async fn cancel_by_unique_key_in_tx<'a>( 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/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index c1d1b3ccd..caedc9c23 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -19640,6 +19640,140 @@ async fn test_queue_storage_cancel_by_unique_key_finds_each_live_state() { ); } +/// 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. diff --git a/skills/awa-jobs/SKILL.md b/skills/awa-jobs/SKILL.md index 414ae47dd..3895bf63a 100644 --- a/skills/awa-jobs/SKILL.md +++ b/skills/awa-jobs/SKILL.md @@ -132,8 +132,10 @@ Cancel without storing the id via `cancel_by_unique_key(kind, queue=, args=, period_bucket=)` — the components must match those used at insert. It cancels the oldest live job holding the key. On queue storage, a `by_queue` key is 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, or once its original -ready row is pruned; cancel it by id. +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