fix(worker): heartbeat finalizing attempts until completion commits (0.6)

CHANGELOG.md

@@ -30,6 +30,29 @@ Notable changes between releases. Detailed migration notes for storage transitio
 
 ### Fixed
 
+- Queue-storage completion of retried jobs or jobs carrying metadata now
+  addresses compact claims by their batch identity and item index, as on main.
+  It no longer expands every retained claim batch for each completion flush,
+  so bounded finalization can drain large recovery backlogs without that scan.
+
+- 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.
+
+- **A completion backlog no longer gets finished jobs rescued and run
+  again.** An attempt stops heartbeating when its handler returns, and its
+  finalize then waits for the completion batcher. The runtime released the
+  worker slot at the same moment and kept claiming, so when finalization was
+  slower than claiming (for example, retried jobs, whose completions take
+  the slower per-row path) the queue of finished, unheartbeated attempts grew
+  without bound. Attempts that waited longer than `heartbeat_staleness` were
+  rescued and executed again, some of them twice. The bound on attempts waiting
+  for finalization is `max(worker count, 2 × batch size × completion shards)`.
+  Past that, a finished attempt keeps its
+  worker slot and its heartbeat until there is room, so the dispatcher stops
+  claiming until finalization catches up.
 - **Ring prune and terminal-count rollup report `Blocked` when any of their
   statements times out on a lock.** Prune sets a short `lock_timeout` (10 ms
   by default) before it takes `ACCESS EXCLUSIVE` on a slot's partitions, but

awa-worker/src/client.rs

@@ -1640,6 +1640,8 @@ impl Client {
             effective_storage.clone(),
             self.dlq_policy.clone(),
             self.finalizers.clone(),
+            self.service_cancel.clone(),
+            runtime_worker_capacity,
         ));
 
         // Admin cancellation listener: fires the in-flight cancel flag

awa-worker/src/completion.rs

@@ -105,9 +105,19 @@ fn completion_sort_key(request: &CompletionRequest) -> (i32, i64, RunLease) {
 #[derive(Clone)]
 pub(crate) struct CompletionBatcherHandle {
     shards: Vec<mpsc::Sender<CompletionRequest>>,
+    pipeline_capacity: usize,
 }
 
 impl CompletionBatcherHandle {
+    /// How many finished attempts may wait for finalization at once: enough
+    /// to keep every shard's in-flight and next batch full, and at least one
+    /// per worker.
+    pub fn finalizer_capacity(&self, runtime_worker_capacity: u32) -> usize {
+        self.pipeline_capacity
+            .max(runtime_worker_capacity.max(1) as usize)
+            .min(tokio::sync::Semaphore::MAX_PERMITS)
+    }
+
     pub async fn complete(&self, job_id: i64, run_lease: RunLease) -> Result<bool, AwaError> {
         self.complete_inner(job_id, run_lease, None, None).await
     }
@@ -185,7 +195,14 @@ impl CompletionBatcher {
             });
         }
 
-        (Self { workers }, CompletionBatcherHandle { shards })
+        (
+            Self { workers },
+            CompletionBatcherHandle {
+                shards,
+                // Each shard needs an in-flight batch and a queued batch.
+                pipeline_capacity: batch_size.saturating_mul(shard_count).saturating_mul(2),
+            },
+        )
     }
 
     pub fn spawn(self) -> Vec<tokio::task::JoinHandle<()>> {

awa-worker/src/dispatcher.rs

@@ -852,6 +852,8 @@ mod tests {
             storage.clone(),
             crate::executor::DlqPolicy::default(),
             tokio_util::task::TaskTracker::new(),
+            CancellationToken::new(),
+            1,
         ));
         let config = QueueConfig {
             max_workers: 1,

awa-worker/src/executor.rs

@@ -11,6 +11,8 @@ use std::collections::HashMap;
 use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
+use tokio::sync::Semaphore;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 use tracing::{error, info, info_span, warn, Instrument};
 
@@ -269,6 +271,8 @@ pub struct JobExecutor {
     storage: RuntimeStorage,
     dlq_policy: DlqPolicy,
     finalizers: TaskTracker,
+    finalize_cancel: CancellationToken,
+    finalize_capacity: Arc<Semaphore>,
 }
 
 impl JobExecutor {
@@ -291,7 +295,12 @@ impl JobExecutor {
         storage: RuntimeStorage,
         dlq_policy: DlqPolicy,
         finalizers: TaskTracker,
+        finalize_cancel: CancellationToken,
+        runtime_worker_capacity: u32,
     ) -> Self {
+        let finalize_capacity = Arc::new(Semaphore::new(
+            completion_batcher.finalizer_capacity(runtime_worker_capacity),
+        ));
         Self {
             pool,
             workers,
@@ -305,6 +314,8 @@ impl JobExecutor {
             storage,
             dlq_policy,
             finalizers,
+            finalize_cancel,
+            finalize_capacity,
         }
     }
 
@@ -332,6 +343,8 @@ impl JobExecutor {
         let storage = self.storage.clone();
         let dlq_policy = self.dlq_policy.clone();
         let finalizers = self.finalizers.clone();
+        let finalize_cancel = self.finalize_cancel.clone();
+        let finalize_capacity = self.finalize_capacity.clone();
         let job_id = job.id;
         let job_run_lease = job.run_lease;
         let job_kind = job.kind.clone();
@@ -357,6 +370,14 @@ impl JobExecutor {
             let in_flight_state = InFlightState {
                 cancel: cancel.clone(),
                 progress: progress_state.clone(),
+                claim: queue_storage_claim.as_ref().and_then(|claim| {
+                    awa_model::queue_storage::ClaimHeartbeatRef::from_claim(
+                        job_id,
+                        job_run_lease,
+                        claim,
+                    )
+                }),
+                finalizing: false,
             };
             in_flight.insert((job_id, job_run_lease), in_flight_state);
             if let Some(counter) = queue_in_flight.get(&job_queue) {
@@ -435,23 +456,38 @@ impl JobExecutor {
                 guard.clone_latest()
             };
 
-            // Remove from in-flight immediately after the handler returns and
-            // the progress snapshot is captured. This keeps local worker
-            // capacity tied to active handler execution, not to the tail
-            // latency of durable completion bookkeeping.
-            in_flight.remove((job_id, job_run_lease));
+            // The attempt stays in `in_flight`, and so keeps its heartbeat,
+            // until its finalize is over: a finalize that waits behind a
+            // completion backlog must not look like a dead worker to rescue.
+            // Only a bounded number of finished attempts may wait for
+            // finalization; while that queue is full the attempt also keeps
+            // its worker slot, so the dispatcher stops claiming.
+            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"))
+                    }
+                    _ = finalize_cancel.cancelled() => None,
+                },
+                None => None,
+            };
             if let Some(counter) = queue_in_flight.get(&job_queue) {
                 counter.fetch_sub(1, Ordering::SeqCst);
             }
             metrics.record_in_flight_change(&job_queue, -1);
 
-            let Some(result) = result else {
+            let (Some(result), Some(finalize_permit)) = (result, finalize_permit) else {
                 return;
             };
             let dlq_enabled = dlq_policy.enabled_for(&job_queue);
-            // Graceful shutdown waits for this before it stops the completion
-            // batcher: the attempt has already left `in_flight`, so nothing
-            // heartbeats it, and a finalize lost at exit leaves it to rescue.
+            // Graceful shutdown waits for this before stopping the completion
+            // batcher. The heartbeat guard keeps the attempt live until the
+            // durable write completes or this finalizer is abandoned.
             finalizers.spawn(async move {
                 // Every finalize path is one lease-guarded transaction, and a
                 // deadlock abort rolls the whole transaction back, so running
@@ -495,6 +531,8 @@ impl JobExecutor {
                         outcome => break outcome,
                     }
                 };
+                drop(heartbeat);
+                drop(finalize_permit);
 
                 match &outcome {
                     Ok(CompletionOutcome::Applied { terminal, .. }) => {
@@ -577,6 +615,19 @@ impl JobExecutor {
     }
 }
 
+/// Keeps a finished attempt in the heartbeat set until its finalize is over,
+/// and takes it out however the finalizer ends.
+struct HeartbeatUntilFinalized {
+    in_flight: InFlightMap,
+    key: (i64, i64),
+}
+
+impl Drop for HeartbeatUntilFinalized {
+    fn drop(&mut self) {
+        self.in_flight.remove(self.key);
+    }
+}
+
 /// How many times a finalize transaction is re-run after Postgres aborts it
 /// to break a lock cycle before the failure is surfaced.
 const COMPLETION_DEADLOCK_RETRIES: u32 = 3;

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;
@@ -76,8 +77,9 @@ impl HeartbeatService {
         fields(interval_ms = self.interval.as_millis() as u64)
     )]
     async fn heartbeat_once(&self) {
-        let mut all_keys: Vec<(i64, RunLease)> = self.in_flight.keys();
-        all_keys.sort_unstable();
+        let mut in_flight = self.in_flight.claim_refs();
+        in_flight.sort_unstable_by_key(|(key, _)| *key);
+        let all_keys: Vec<(i64, RunLease)> = in_flight.iter().map(|(key, _)| *key).collect();
 
         if all_keys.is_empty() {
             return;
@@ -104,7 +106,14 @@ impl HeartbeatService {
                     .await;
             }
             RuntimeStorage::QueueStorage(runtime) => {
-                self.heartbeat_once_queue_storage(runtime, &heartbeat_only, &pending_progress)
+                let (claimed, plain): (Vec<_>, Vec<_>) = in_flight
+                    .iter()
+                    .filter(|(key, _)| !progress_keys.contains(key))
+                    .partition(|(_, claim)| claim.is_some());
+                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)
                     .await;
             }
         }
@@ -131,7 +140,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()
@@ -178,7 +187,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(),
@@ -204,7 +213,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");
@@ -222,9 +237,26 @@ impl HeartbeatService {
     async fn heartbeat_once_queue_storage(
         &self,
         runtime: &QueueStorageRuntime,
+        claimed: &[ClaimHeartbeatRef],
         heartbeat_only: &[(i64, RunLease)],
         pending_progress: &[(i64, i64, u64, serde_json::Value)],
     ) {
+        for chunk in claimed.chunks(self.batch_size) {
+            match runtime
+                .store
+                .heartbeat_claims_batch(&self.pool, chunk)
+                .await
+            {
+                Ok(updated) => {
+                    self.metrics.heartbeat_batches.add(1, &[]);
+                    debug!(count = chunk.len(), updated, "Heartbeat claim batch sent");
+                }
+                Err(err) => {
+                    warn!(error = %err, "Failed to send heartbeat claim batch");
+                }
+            }
+        }
+
         for chunk in heartbeat_only.chunks(self.batch_size) {
             match runtime.store.heartbeat_batch(&self.pool, chunk).await {
                 Ok(updated) => {
@@ -258,7 +290,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");