From 95b0c99287a041bf3cd52f2a72eb2666b638ab38 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Thu, 12 Nov 2026 08:40:01 +0000 Subject: [PATCH 1/3] fix(queue-storage): keep a late heartbeat from replacing parked progress A heartbeat that read a progress snapshot before its attempt parked for a callback could write that snapshot after the park committed, replacing the final progress the park stored. Parking a receipt attempt now takes the attempt lock, which heartbeat writes take too, and the receipt progress upsert leaves out attempts that are waiting for a callback. The explicit flush test now waits until the flush is seen waiting on a lock before releasing it. --- awa-model/src/queue_storage.rs | 15 ++++++ awa/tests/queue_storage_runtime_test.rs | 72 ++++++++++++++++++++++++- 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 4b8d285b3..92391b816 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -12798,6 +12798,7 @@ impl QueueStorage { self.lock_receipt_attempts_tx(tx, &receipt_jobs).await?; let schema = self.schema(); + let leases = self.leases_table(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = @@ -12885,6 +12886,13 @@ impl QueueStorage { claim_refs.progress, clock_timestamp() FROM claim_refs + WHERE NOT EXISTS ( + SELECT 1 + FROM {leases} AS parked + WHERE parked.job_id = claim_refs.job_id + AND parked.run_lease = claim_refs.run_lease + AND parked.state = 'waiting_external' + ) ON CONFLICT (job_id, run_lease) DO UPDATE SET heartbeat_at = clock_timestamp(), @@ -16592,6 +16600,13 @@ impl QueueStorage { callback_id: Uuid, progress: Option<&serde_json::Value>, ) -> Result { + // Heartbeats write receipt progress under the attempt lock and skip + // parked attempts, so holding it here keeps an older heartbeat + // snapshot from replacing the progress stored with the park. + if self.lease_claim_receipts() { + self.lock_receipt_attempts_tx(tx, &[(job_id, run_lease)]) + .await?; + } let entered = self .enter_callback_wait_in_tx(tx, job_id, run_lease, callback_id) .await?; diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 656c90753..f2597435c 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -9914,7 +9914,7 @@ async fn test_queue_storage_flush_progress_waits_for_a_locked_receipt_attempt() .await }) }; - tokio::time::sleep(Duration::from_millis(200)).await; + wait_for_lock_waiter(&pool, "pg_advisory_xact_lock").await; assert!( !flush.is_finished(), "an explicit flush must wait for the attempt lock rather than skip the attempt" @@ -18436,6 +18436,76 @@ async fn test_queue_storage_legacy_callback_wait_stores_the_final_progress() { callback_wait_stores_the_final_progress(false).await; } +/// A heartbeat that read an older progress snapshot before the attempt was +/// parked must not replace the progress stored with the park. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_heartbeat_does_not_overwrite_parked_progress() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_parked_progress_heartbeat"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: "awa_qs_parked_progress_heartbeat".to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: true, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CallbackJob { id: 6 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claim = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim") + .into_iter() + .next() + .expect("claimed job"); + let run_lease = claim.job.run_lease; + let callback_id = store + .register_callback(&pool, job_id, run_lease, Duration::from_secs(300)) + .await + .expect("register callback"); + + let parked_progress = serde_json::json!({"percent": 60}); + assert!(store + .enter_callback_wait_with_progress( + &pool, + job_id, + run_lease, + callback_id, + Some(&parked_progress), + ) + .await + .expect("park the attempt")); + + let updated = store + .heartbeat_progress_batch( + &pool, + &[(job_id, run_lease, serde_json::json!({"percent": 40}))], + ) + .await + .expect("late heartbeat"); + assert_eq!(updated, 0, "a heartbeat must skip a parked attempt"); + + let parked = store + .load_job(&pool, job_id) + .await + .expect("load parked job") + .expect("parked job exists"); + assert_eq!(parked.state, JobState::WaitingExternal); + assert_eq!(parked.progress, Some(parked_progress)); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_runtime_complete_external() { let (_db_guard, pool) = setup_pool(10).await; From d325f99c35f7840e634a25d2f075887591597af5 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Thu, 12 Nov 2026 08:50:31 +0000 Subject: [PATCH 2/3] fix(queue-storage): lock the attempt before the lease in external finalizers External completion, failure and retry deleted the lease row and then waited for the attempt lock, the reverse of the order every other receipt path and the callback park use, so a fast callback could deadlock with a parking worker. They now take the attempt locks for the callback's leases first. --- awa-model/src/queue_storage.rs | 35 ++++++++++ awa/tests/queue_storage_runtime_test.rs | 90 +++++++++++++++++++++++++ 2 files changed, 125 insertions(+) diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 92391b816..0b09d4fa2 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -16958,6 +16958,35 @@ impl QueueStorage { .await } + /// Take the attempt locks for the leases parked on `callback_id` before + /// an external finalizer deletes them, so it locks the attempt before the + /// lease row like the other receipt paths. + async fn lock_callback_receipt_attempts_tx( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + callback_id: Uuid, + run_lease: Option, + ) -> Result<(), AwaError> { + if !self.lease_claim_receipts() { + return Ok(()); + } + let pairs: Vec<(i64, i64)> = sqlx::query_as(audited_sql(format!( + r#" + SELECT job_id, run_lease + FROM {} + WHERE callback_id = $1 + AND ($2::bigint IS NULL OR run_lease = $2) + "#, + self.leases_table() + ))) + .bind(callback_id) + .bind(run_lease) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.lock_receipt_attempts_tx(tx, &pairs).await + } + /// Complete a job waiting on `callback_id` and store `output` as its /// output (ADR-048), if this schema has the output columns. The caller /// validates the output. Returns the completed row, whose `output` is @@ -16974,6 +17003,8 @@ impl QueueStorage { Some(output) if self.supports_job_output(tx.as_mut()).await? => Some(output), _ => None, }; + self.lock_callback_receipt_attempts_tx(tx, callback_id, run_lease) + .await?; let schema = self.schema(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" @@ -17077,6 +17108,8 @@ impl QueueStorage { error_entry: serde_json::Value, run_lease: Option, ) -> Result { + self.lock_callback_receipt_attempts_tx(tx, callback_id, run_lease) + .await?; let schema = self.schema(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" @@ -17176,6 +17209,8 @@ impl QueueStorage { callback_id: Uuid, run_lease: Option, ) -> Result { + self.lock_callback_receipt_attempts_tx(tx, callback_id, run_lease) + .await?; let schema = self.schema(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index f2597435c..59fa1f36b 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -18506,6 +18506,96 @@ async fn test_queue_storage_heartbeat_does_not_overwrite_parked_progress() { assert_eq!(parked.progress, Some(parked_progress)); } +/// An external completion waits for the attempt lock before it locks the +/// lease row, so a worker parking the same attempt cannot deadlock with it. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_external_completion_locks_the_attempt_before_the_lease() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_external_completion_lock_order"; + let schema = "awa_qs_external_completion_lock_order"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: true, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CallbackJob { id: 7 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claim = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim") + .into_iter() + .next() + .expect("claimed job"); + let run_lease = claim.job.run_lease; + let callback_id = store + .register_callback(&pool, job_id, run_lease, Duration::from_secs(300)) + .await + .expect("register callback"); + + let mut attempt_lock = pool.begin().await.expect("begin attempt lock tx"); + sqlx::query( + r#" + SELECT pg_catalog.pg_advisory_xact_lock( + pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), + 0 + ) + ) + "#, + ) + .bind(job_id) + .bind(run_lease) + .execute(attempt_lock.as_mut()) + .await + .expect("hold the attempt lock"); + + let store = Arc::new(store); + let completion = { + let store = store.clone(); + let pool = pool.clone(); + tokio::spawn(async move { + store + .complete_external(&pool, callback_id, None, Some(run_lease), false) + .await + }) + }; + wait_for_lock_waiter(&pool, "pg_advisory_xact_lock").await; + + let lease_rows: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT job_id FROM {schema}.leases WHERE job_id = $1 FOR UPDATE NOWAIT" + ))) + .bind(job_id) + .fetch_all(attempt_lock.as_mut()) + .await + .expect("the lease row stays unlocked while the completion waits for the attempt"); + assert_eq!(lease_rows, vec![job_id]); + + attempt_lock + .commit() + .await + .expect("release the attempt lock"); + let completed = completion + .await + .expect("completion task") + .expect("complete external after the lock is released"); + assert_eq!(completed.state, JobState::Completed); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_runtime_complete_external() { let (_db_guard, pool) = setup_pool(10).await; From 0a0475ae8fc264d2a670719c32ffa3e7a6626eac Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Thu, 12 Nov 2026 09:01:38 +0000 Subject: [PATCH 3/3] ci: install psycopg with its binary extra psycopg 3.3.9 has no Python 3.15 binary wheel yet, so installing psycopg and psycopg-binary separately paired psycopg 3.3.9 with psycopg-binary 3.3.8, which fails to import. The extra pins matching versions. --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 83079c307..b1dfa8349 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -477,7 +477,7 @@ jobs: - name: Create venv and install deps run: | uv venv .venv --python "${{ matrix.python }}" - uv pip install maturin pytest pytest-asyncio asyncpg django greenlet psycopg psycopg-binary sqlalchemy + uv pip install maturin pytest pytest-asyncio asyncpg django greenlet "psycopg[binary]" sqlalchemy - name: Build extension (maturin develop) run: .venv/bin/maturin develop - name: Wait for Postgres