fix(worker): heartbeat attempts until durable finalization

CHANGELOG.md

@@ -6,6 +6,12 @@ Notable changes between releases. Detailed migration notes for storage transitio
 
 ### Fixed
 
+- Finished attempts keep heartbeating until their durable finalization finishes,
+  including while completion writes are stalled. Compact receipt heartbeats
+  address claims by batch identity and skip attempts a closer already holds,
+  so one stalled completion cannot block the rest of a heartbeat pass. Late
+  heartbeats cannot recreate state for a completed attempt.
+
 - **Batch enqueues take their lane locks in one order, so they no longer
   deadlock with each other.** On queue storage every enqueue holds a
   transaction-scoped advisory lock on its `(queue, priority, shard)` lane

awa-model/src/queue_storage.rs

@@ -11783,6 +11783,47 @@ impl QueueStorage {
         Ok(())
     }
 
+    /// Take the closers' attempt locks without waiting, returning the attempts
+    /// locked. A heartbeat must not queue behind a closer: an attempt whose lock
+    /// is held is being finalized or rescued right now, and rescue skips it for
+    /// as long as the lock is held. Taking the locks in a statement of their own
+    /// means the statement that then decides what is still open starts after
+    /// every closer that held a lock first has committed (#614), so a heartbeat
+    /// cannot recreate state for an attempt that has just been closed.
+    async fn try_lock_receipt_attempt_keys_tx<'a>(
+        &self,
+        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
+        jobs: &[(i64, i64)],
+    ) -> Result<Vec<(i64, i64)>, AwaError> {
+        if jobs.is_empty() {
+            return Ok(Vec::new());
+        }
+        let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
+        let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
+        sqlx::query_as::<_, (i64, i64)>(
+            r#"
+            WITH attempts AS MATERIALIZED (
+                SELECT DISTINCT
+                    input.job_id,
+                    input.run_lease,
+                    pg_catalog.hashtextextended(
+                        format('awa.receipt.complete:%s:%s', input.job_id, input.run_lease),
+                        0
+                    ) AS lock_key
+                FROM unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease)
+            )
+            SELECT job_id, run_lease
+            FROM attempts
+            WHERE pg_catalog.pg_try_advisory_xact_lock(lock_key)
+            "#,
+        )
+        .bind(&job_ids)
+        .bind(&run_leases)
+        .fetch_all(tx.as_mut())
+        .await
+        .map_err(map_sqlx_error)
+    }
+
     async fn ensure_running_leases_from_receipts_tx<'a>(
         &self,
         tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
@@ -12086,7 +12127,6 @@ impl QueueStorage {
                       WHERE dlq.job_id = items.job_id
                         AND dlq.run_lease = items.run_lease
                   )
-                FOR UPDATE OF claim_batches
             ),
             claim_refs AS (
                 SELECT * FROM row_claim_refs
@@ -12135,6 +12175,19 @@ 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 jobs: Vec<_> = jobs
+            .iter()
+            .filter(|(job_id, run_lease, _)| locked.contains(&(*job_id, *run_lease)))
+            .collect();
+        if jobs.is_empty() {
+            return Ok(0);
+        }
+        let receipt_jobs: Vec<_> = jobs.iter().map(|(id, lease, _)| (*id, *lease)).collect();
         self.lock_receipt_attempts_tx(tx, &receipt_jobs).await?;
 
         let schema = self.schema();
@@ -12203,7 +12256,6 @@ impl QueueStorage {
                       WHERE dlq.job_id = items.job_id
                         AND dlq.run_lease = items.run_lease
                   )
-                FOR UPDATE OF claim_batches
             ),
             claim_refs AS (
                 SELECT * FROM row_claim_refs
@@ -15628,7 +15680,8 @@ impl QueueStorage {
     /// Heartbeat compact receipt claims by their batch identity. Each claim
     /// batch is addressed by primary key and the attempt by item index, with
     /// index-only closure probes and no batch row lock, so the cost is
-    /// O(in-flight attempts) rather than O(open claim batches).
+    /// O(in-flight attempts) rather than O(open claim batches). Attempts that
+    /// a closer holds right now are skipped rather than waited for.
     pub async fn heartbeat_claims_batch(
         &self,
         pool: &PgPool,
@@ -15643,7 +15696,19 @@ impl QueueStorage {
             .iter()
             .map(|claim| (claim.job_id, claim.run_lease))
             .collect();
-        self.lock_receipt_attempt_keys_tx(&mut tx, &pairs).await?;
+        let locked: HashSet<(i64, i64)> = self
+            .try_lock_receipt_attempt_keys_tx(&mut tx, &pairs)
+            .await?
+            .into_iter()
+            .collect();
+        let claims: Vec<&ClaimHeartbeatRef> = claims
+            .iter()
+            .filter(|claim| locked.contains(&(claim.job_id, claim.run_lease)))
+            .collect();
+        if claims.is_empty() {
+            tx.commit().await.map_err(map_sqlx_error)?;
+            return Ok(0);
+        }
         let job_ids: Vec<i64> = claims.iter().map(|claim| claim.job_id).collect();
         let run_leases: Vec<i64> = claims.iter().map(|claim| claim.run_lease).collect();
         let claim_slots: Vec<i32> = claims.iter().map(|claim| claim.claim_slot).collect();
@@ -15715,6 +15780,38 @@ impl QueueStorage {
         Ok(updated as usize)
     }
 
+    /// Lock legacy lease rows without waiting behind a finalizer. Keep this
+    /// separate from the updating statement: the latter must use a fresh
+    /// snapshot and must not mix partition pruning with row locking (PG18).
+    async fn lock_live_leases_for_heartbeat_tx(
+        &self,
+        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
+        jobs: &[(i64, i64)],
+    ) -> Result<Vec<(i64, i64)>, AwaError> {
+        if jobs.is_empty() {
+            return Ok(Vec::new());
+        }
+        let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
+        let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
+        sqlx::query_as::<_, (i64, i64)>(audited_sql(format!(
+            r#"
+            SELECT lease.job_id, lease.run_lease
+            FROM {table} AS lease
+            JOIN unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease)
+              ON lease.job_id = input.job_id AND lease.run_lease = input.run_lease
+            WHERE lease.state = 'running'
+            ORDER BY lease.job_id, lease.run_lease
+            FOR UPDATE OF lease SKIP LOCKED
+            "#,
+            table = self.leases_table(),
+        )))
+        .bind(&job_ids)
+        .bind(&run_leases)
+        .fetch_all(tx.as_mut())
+        .await
+        .map_err(map_sqlx_error)
+    }
+
     pub async fn heartbeat_batch(
         &self,
         pool: &PgPool,
@@ -15735,14 +15832,18 @@ impl QueueStorage {
             // dead entries per write under sustained churn). Compat
             // reads in `LeaseJobRow` SELECTs and the `awa.jobs` view
             // COALESCE attempt_state.heartbeat_at first.
+            let locked = self.try_lock_receipt_attempt_keys_tx(&mut tx, jobs).await?;
             updated += self
-                .upsert_attempt_state_from_receipts_tx(&mut tx, jobs)
+                .upsert_attempt_state_from_receipts_tx(&mut tx, &locked)
                 .await?;
         } else {
             // Legacy non-receipts mode (custom schemas with
             // `lease_claim_receipts=FALSE`): the `leases.heartbeat_at`
             // write is still the heartbeat home, since there is no
             // upsert_attempt_state path firing.
+            let jobs = self
+                .lock_live_leases_for_heartbeat_tx(&mut tx, jobs)
+                .await?;
             let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
             let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
             let result = sqlx::query(audited_sql(format!(
@@ -15793,9 +15894,19 @@ impl QueueStorage {
             self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs)
                 .await?
         } else {
-            // Legacy non-receipts mode keeps the old CTE shape that
-            // updates leases.heartbeat_at + leases.progress
-            // (via attempt_state upsert) in a single round-trip.
+            // Prelock independently so a finalizer cannot stall the entire
+            // heartbeat pass. Skipped progress snapshots remain unacknowledged
+            // in the worker and are retried on the next pass.
+            let keys: Vec<_> = jobs.iter().map(|(id, lease, _)| (*id, *lease)).collect();
+            let locked: HashSet<_> = self
+                .lock_live_leases_for_heartbeat_tx(&mut tx, &keys)
+                .await?
+                .into_iter()
+                .collect();
+            let jobs: Vec<_> = jobs
+                .iter()
+                .filter(|(id, lease, _)| locked.contains(&(*id, *lease)))
+                .collect();
             let schema = self.schema();
             let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _, _)| *job_id).collect();
             let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect();

awa-worker/src/executor.rs

@@ -409,6 +409,7 @@ impl JobExecutor {
                         claim,
                     )
                 }),
+                finalizing: false,
             };
             in_flight.insert((job_id, job_run_lease), in_flight_state);
             if let Some(counter) = queue_in_flight.get(&job_queue) {
@@ -514,7 +515,12 @@ impl JobExecutor {
                 }
             }
 
-            // Keep worker/heartbeat ownership while the bounded finalizer queue is full.
+            // Keep heartbeating through the finalization queue and its durable write.
+            in_flight.mark_finalizing((job_id, job_run_lease));
+            let heartbeat = HeartbeatUntilFinalized {
+                in_flight: in_flight.clone(),
+                key: (job_id, job_run_lease),
+            };
             let finalize_permit = match result {
                 Some(_) => tokio::select! {
                     permit = finalize_capacity.acquire_owned() => Some(permit.expect("finalizer capacity is never closed")),
@@ -523,7 +529,6 @@ impl JobExecutor {
                 None => None,
             };
 
-            in_flight.remove((job_id, job_run_lease));
             if let Some(counter) = queue_in_flight.get(&job_queue) {
                 counter.fetch_sub(1, Ordering::SeqCst);
             }
@@ -543,6 +548,7 @@ impl JobExecutor {
             job.output = sealed_output;
             finalizers.spawn(
                 async move {
+                    let _heartbeat = heartbeat;
                     let event_kind = job_kind.clone();
                     let abandon_output = had_output.then(|| (metrics.clone(), job_queue.clone()));
                     let finalize = async move {
@@ -3369,6 +3375,18 @@ async fn retry_backoff_canonical_with_followups(
     Ok(Some(updated_job))
 }
 
+/// Keep finalizing attempts live until their durable write or abandonment.
+struct HeartbeatUntilFinalized {
+    in_flight: InFlightMap,
+    key: (i64, i64),
+}
+
+impl Drop for HeartbeatUntilFinalized {
+    fn drop(&mut self) {
+        self.in_flight.remove(self.key);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

awa-worker/src/heartbeat.rs

@@ -1,5 +1,6 @@
 use crate::runtime::{InFlightMap, RunLease};
 use crate::storage::{QueueStorageRuntime, RuntimeStorage};
+use awa_model::queue_storage::ClaimHeartbeatRef;
 use sqlx::PgPool;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
@@ -140,7 +141,7 @@ impl HeartbeatService {
                     .iter()
                     .filter(|(key, _)| !progress_keys.contains(key))
                     .partition(|(_, claim)| claim.is_some());
-                let claimed: Vec<awa_model::queue_storage::ClaimHeartbeatRef> =
+                let claimed: Vec<ClaimHeartbeatRef> =
                     claimed.into_iter().filter_map(|(_, claim)| claim).collect();
                 let plain: Vec<(i64, RunLease)> = plain.into_iter().map(|(key, _)| key).collect();
                 self.heartbeat_once_queue_storage(runtime, &claimed, &plain, &pending_progress)
@@ -170,7 +171,7 @@ impl HeartbeatService {
                      AND jobs.run_lease = inflight.run_lease
                     WHERE jobs.state = 'running'
                     ORDER BY jobs.id
-                    FOR UPDATE OF jobs
+                    FOR UPDATE OF jobs SKIP LOCKED
                 )
                 UPDATE awa.jobs_hot AS jobs
                 SET heartbeat_at = now()
@@ -217,7 +218,7 @@ impl HeartbeatService {
                          AND jobs.run_lease = inflight.run_lease
                         WHERE jobs.state = 'running'
                         ORDER BY jobs.id
-                        FOR UPDATE OF jobs
+                        FOR UPDATE OF jobs SKIP LOCKED
                     )
                     UPDATE awa.jobs_hot AS jobs
                     SET heartbeat_at = now(),
@@ -243,7 +244,13 @@ impl HeartbeatService {
                             .iter()
                             .map(|(id, lease, gen, _)| (*id, *lease, *gen))
                             .collect();
-                        self.in_flight.ack_progress(&acked);
+                        if result.rows_affected() as usize == chunk.len() {
+                            self.in_flight.ack_progress(&acked);
+                        } else {
+                            // Re-send successful snapshots too: replay is safe,
+                            // whereas acknowledging a skipped write loses progress.
+                            self.in_flight.clear_in_flight_progress(&acked);
+                        }
                     }
                     Err(err) => {
                         warn!(error = %err, "Failed to send heartbeat+progress batch");
@@ -261,7 +268,7 @@ impl HeartbeatService {
     async fn heartbeat_once_queue_storage(
         &self,
         runtime: &QueueStorageRuntime,
-        claimed: &[awa_model::queue_storage::ClaimHeartbeatRef],
+        claimed: &[ClaimHeartbeatRef],
         heartbeat_only: &[(i64, RunLease)],
         pending_progress: &[(i64, i64, u64, serde_json::Value)],
     ) {
@@ -314,7 +321,11 @@ impl HeartbeatService {
                             .iter()
                             .map(|(id, lease, gen, _)| (*id, *lease, *gen))
                             .collect();
-                        self.in_flight.ack_progress(&acked);
+                        if updated == chunk.len() {
+                            self.in_flight.ack_progress(&acked);
+                        } else {
+                            self.in_flight.clear_in_flight_progress(&acked);
+                        }
                     }
                     Err(err) => {
                         warn!(error = %err, "Failed to send heartbeat+progress batch");

awa-worker/src/runtime.rs

@@ -23,6 +23,9 @@ pub(crate) struct InFlightState {
     /// queue storage with one; lets the heartbeat address the claim batch
     /// directly.
     pub claim: Option<ClaimHeartbeatRef>,
+    /// The handler has returned and the attempt waits for its finalize. It is
+    /// still heartbeated, but its progress is settled by the finalize.
+    pub finalizing: bool,
 }
 
 /// Mutable progress buffer shared between handler and heartbeat service.
@@ -163,7 +166,7 @@ impl InFlightRegistry {
         guard.get(&key).map(|s| s.cancel.clone())
     }
 
-    /// Every in-flight attempt with its compact claim identity, if any.
+    /// Every attempt to heartbeat, with its compact claim identity, if any.
     pub fn claim_refs(&self) -> Vec<(InFlightKey, Option<ClaimHeartbeatRef>)> {
         let mut refs = Vec::new();
         for shard in &*self.shards {
@@ -173,6 +176,18 @@ impl InFlightRegistry {
         refs
     }
 
+    /// The handler has returned: keep heartbeating the attempt until its
+    /// finalize is done, and stop flushing its progress.
+    pub fn mark_finalizing(&self, key: InFlightKey) {
+        let mut guard = self
+            .shard_for(&key)
+            .write()
+            .expect("in_flight shard poisoned");
+        if let Some(state) = guard.get_mut(&key) {
+            state.finalizing = true;
+        }
+    }
+
     pub fn flags(&self) -> Vec<CancelFlag> {
         let mut flags = Vec::new();
         for shard in &*self.shards {
@@ -192,6 +207,9 @@ impl InFlightRegistry {
         for shard in &*self.shards {
             let guard = shard.read().expect("in_flight shard poisoned");
             for (&(job_id, run_lease), state) in guard.iter() {
+                if state.finalizing {
+                    continue;
+                }
                 let mut progress = state.progress.lock().expect("progress lock poisoned");
                 if progress.has_pending() && progress.in_flight.is_none() {
                     if let Some(ref value) = progress.latest {
@@ -278,6 +296,7 @@ mod tests {
                 cancel: Arc::new(AtomicBool::new(false)),
                 progress: progress.clone(),
                 claim: None,
+                finalizing: false,
             },
         );