fix(queue-storage): make an explicit progress flush wait for the attempt lock

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 Notable changes between releases. Detailed migration notes for storage transitions live in [`docs/upgrade-0.5-to-0.6.md`](docs/upgrade-0.5-to-0.6.md).
 
+## [Unreleased]
+
+### Fixed
+
+- `JobContext::flush_progress` (Python `ctx.flush_progress()`) on a queue-storage attempt with receipts, the default, no longer returns success without writing the progress. A heartbeat leaves out an attempt whose lock another transaction holds, and since 0.6.16 an explicit flush shared that path, so a flush that overlapped a heartbeat or a completion of the same attempt was dropped and the worker treated the progress as persisted. An explicit flush now waits for the attempt's lock; heartbeats still skip.
+
 ## [0.6.17] — 2026-11-12
 
 Recommended for everyone on 0.6.16 using queue storage. Claims commit asynchronously and reuse the runtime's claimer lease, which raised throughput by 19 to 21% at 4, 16 and 64 workers and cut end-to-end p99 at 800 jobs/s from 31 ms to 24 ms in a three-round comparison against 0.6.16 on PostgreSQL 18 with `synchronous_commit = on`. Promotion and rescue no longer wait behind the leader's other maintenance work, a runtime can opt out of maintenance leadership, and several admin reads and the admin UI are fixed. See the [upgrade guide](docs/upgrade-0.6.16-to-0.6.17.md).

awa-model/src/queue_storage.rs

@@ -11214,10 +11214,14 @@ impl QueueStorage {
         Ok(updated as usize)
     }
 
+    /// With `skip_locked`, attempts whose lock another transaction holds are
+    /// left out, as a heartbeat batch wants. Without it the call waits for
+    /// each attempt's lock, so an explicit flush is never dropped.
     async fn upsert_attempt_state_progress_from_receipts_tx<'a>(
         &self,
         tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
         jobs: &[(i64, i64, serde_json::Value)],
+        skip_locked: bool,
     ) -> Result<usize, AwaError> {
         if jobs.is_empty() {
             return Ok(0);
@@ -11227,14 +11231,23 @@ impl QueueStorage {
             .iter()
             .map(|(job_id, run_lease, _)| (*job_id, *run_lease))
             .collect();
-        let locked: HashSet<_> = self
-            .try_lock_receipt_attempt_keys_tx(tx, &receipt_jobs)
-            .await?
-            .into_iter()
-            .collect();
+        let locked: Option<HashSet<_>> = if skip_locked {
+            Some(
+                self.try_lock_receipt_attempt_keys_tx(tx, &receipt_jobs)
+                    .await?
+                    .into_iter()
+                    .collect(),
+            )
+        } else {
+            None
+        };
         let jobs: Vec<_> = jobs
             .iter()
-            .filter(|(job_id, run_lease, _)| locked.contains(&(*job_id, *run_lease)))
+            .filter(|(job_id, run_lease, _)| {
+                locked
+                    .as_ref()
+                    .is_none_or(|locked| locked.contains(&(*job_id, *run_lease)))
+            })
             .collect();
         if jobs.is_empty() {
             return Ok(0);
@@ -15069,6 +15082,7 @@ impl QueueStorage {
             self.upsert_attempt_state_progress_from_receipts_tx(
                 &mut tx,
                 &[(job_id, run_lease, progress.clone())],
+                false,
             )
             .await?;
         }
@@ -15312,7 +15326,7 @@ impl QueueStorage {
             // anti-joined against durable closure evidence so it picks
             // up every open claim regardless of whether
             // `materialize_claims` has fanned it out to `leases` yet.
-            self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs)
+            self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs, true)
                 .await?
         } else {
             // Prelock independently so a finalizer cannot stall the entire

awa-worker/src/runtime_control.rs

@@ -780,14 +780,14 @@ mod tests {
         assert!(!control.allowed());
         tokio::time::timeout(Duration::from_secs(3), async {
             loop {
-                let value: Option<serde_json::Value> = sqlx::query_scalar(
+                let value: Option<Option<serde_json::Value>> = sqlx::query_scalar(
                     "SELECT observation FROM awa.runtime_control_observations WHERE instance_id=$1",
                 )
                 .bind(id)
-                .fetch_one(&admin)
+                .fetch_optional(&admin)
                 .await
                 .unwrap();
-                if value.is_some_and(|v| {
+                if value.flatten().is_some_and(|v| {
                     v["observed_revision"] == halt.revision && v["state"] == "halted"
                 }) {
                     break;

awa/tests/bulk_rescue_test.rs

@@ -801,13 +801,20 @@ async fn rescue_drains_an_abandoned_backlog_without_rescuing_live_claims() {
             .expect("recovery client")
     };
     recovery.start().await.expect("start recovery client");
+    // Completed rows leave the ring once queue prune wraps past them, so the
+    // handler counts the recovery runs.
     wait_until(
         scaled_timeout(Duration::from_secs(120)),
-        "the retried jobs to complete",
+        "every retried job to run",
+        || async { recovered.lock().expect("recovered").len() as i64 >= ABANDONED - exhausted },
+    )
+    .await;
+    wait_until(
+        scaled_timeout(Duration::from_secs(30)),
+        "the retried jobs to finish",
         || async {
-            count(&pool, format!(
-                "SELECT count(*) FROM awa.terminal_jobs WHERE queue = '{ABANDONED_QUEUE}' AND state = 'completed'"
-            )).await == ABANDONED - exhausted
+            running(&pool, &store, ABANDONED_QUEUE).await == 0
+                && waiting_retry_counts(&pool, ABANDONED_QUEUE).await == (0, 0)
         },
     )
     .await;

awa/tests/queue_storage_runtime_test.rs

@@ -6623,6 +6623,97 @@ async fn test_queue_storage_receipt_rescue_cursor_sweeps_past_fresh_claims_and_w
     assert_eq!(open_receipt_claim_count(&pool, &store).await, 0);
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn test_queue_storage_flush_progress_waits_for_a_locked_receipt_attempt() {
+    let (_db_guard, pool) = setup_pool(10).await;
+    let queue = "qs_flush_progress_locked";
+    let schema = "awa_qs_flush_progress_locked";
+    let store = Arc::new(
+        create_store_with_config(
+            &pool,
+            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,
+            },
+        )
+        .await,
+    );
+
+    let job_id = enqueue_job(
+        &pool,
+        &store,
+        &CompleteJob { id: 78 },
+        InsertOpts {
+            queue: queue.to_string(),
+            ..Default::default()
+        },
+    )
+    .await;
+    let claim = store
+        .claim_runtime_batch(&pool, queue, 1, Duration::ZERO)
+        .await
+        .expect("claim receipt-backed job")
+        .into_iter()
+        .next()
+        .expect("claimed job");
+    assert!(claim.claim.lease_claim_receipt);
+    let run_lease = claim.job.run_lease;
+
+    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 flush = {
+        let store = store.clone();
+        let pool = pool.clone();
+        tokio::spawn(async move {
+            store
+                .flush_progress(&pool, job_id, run_lease, serde_json::json!({"step": 3}))
+                .await
+        })
+    };
+    tokio::time::sleep(Duration::from_millis(200)).await;
+    assert!(
+        !flush.is_finished(),
+        "an explicit flush must wait for the attempt lock rather than skip the attempt"
+    );
+    attempt_lock
+        .commit()
+        .await
+        .expect("release the attempt lock");
+    flush
+        .await
+        .expect("flush task")
+        .expect("flush progress after the lock is released");
+
+    let progress: Option<serde_json::Value> = sqlx::query_scalar(&format!(
+        "SELECT progress FROM {schema}.attempt_state WHERE job_id = $1 AND run_lease = $2"
+    ))
+    .bind(job_id)
+    .bind(run_lease)
+    .fetch_optional(&pool)
+    .await
+    .expect("read attempt state");
+    assert_eq!(progress, Some(serde_json::json!({"step": 3})));
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
 async fn test_queue_storage_receipt_rescue_skips_attempt_locked_for_completion() {
     let (_db_guard, pool) = setup_pool(10).await;