fix(queue-storage): keep a late heartbeat from replacing parked progress

.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

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<bool, AwaError> {
+        // 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<i64>,
+    ) -> 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<DeletedLeaseRow> = sqlx::query_as(audited_sql(format!(
             r#"
@@ -17062,6 +17108,8 @@ impl QueueStorage {
         error_entry: serde_json::Value,
         run_lease: Option<i64>,
     ) -> Result<JobRow, AwaError> {
+        self.lock_callback_receipt_attempts_tx(tx, callback_id, run_lease)
+            .await?;
         let schema = self.schema();
         let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(audited_sql(format!(
             r#"
@@ -17161,6 +17209,8 @@ impl QueueStorage {
         callback_id: Uuid,
         run_lease: Option<i64>,
     ) -> Result<JobRow, AwaError> {
+        self.lock_callback_receipt_attempts_tx(tx, callback_id, run_lease)
+            .await?;
         let schema = self.schema();
         let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(audited_sql(format!(
             r#"

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<i64> = 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;