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 diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 4b8d285b3..0b09d4fa2 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?; @@ -16943,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 @@ -16959,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#" @@ -17062,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#" @@ -17161,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 656c90753..59fa1f36b 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,166 @@ 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)); +} + +/// 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;