fix(worker): bound finalization and speed up metadata completions (0.6)

CHANGELOG.md

@@ -30,6 +30,23 @@ 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.
+
+- **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-model/src/queue_storage.rs

@@ -6722,14 +6722,22 @@ impl QueueStorage {
                             .expect("receipt-backed slow completion requires receipt_id")
                     })
                     .collect();
+                let receipt_claim_batch_ids: Vec<Option<i64>> = receipt_claimed
+                    .iter()
+                    .map(|entry| entry.claim.claim_batch_id)
+                    .collect();
+                let receipt_claim_batch_indices: Vec<Option<i32>> = receipt_claimed
+                    .iter()
+                    .map(|entry| entry.claim.claim_batch_index)
+                    .collect();
                 let closure_rel = format!("{schema}.lease_claim_closures");
                 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
                 let closed_evidence =
                     receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
                 let updated: Vec<(i64, i64)> = sqlx::query_as(&format!(
                     r#"
-                    WITH completed(claim_slot, job_id, run_lease, receipt_id) AS (
-                        SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[])
+                    WITH completed(claim_slot, job_id, run_lease, receipt_id, claim_batch_id, claim_batch_index) AS (
+                        SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[], $5::bigint[], $6::int[])
                     ),
                     locked_row_claims AS (
                         SELECT claims.claim_slot, claims.job_id, claims.run_lease
@@ -6756,18 +6764,23 @@ impl QueueStorage {
                             items.job_id,
                             items.run_lease,
                             items.receipt_id
-                        FROM {schema}.lease_claim_batches AS claim_batches
-                        CROSS JOIN LATERAL unnest(
-                            claim_batches.job_ids,
-                            claim_batches.run_leases,
-                            claim_batches.receipt_ids
-                        ) AS items(job_id, run_lease, receipt_id)
-                        JOIN completed
-                          ON completed.claim_slot = claim_batches.claim_slot
-                         AND completed.job_id = items.job_id
-                         AND completed.run_lease = items.run_lease
-                         AND completed.receipt_id = items.receipt_id
-                        WHERE NOT EXISTS (
+                        FROM completed
+                        JOIN {schema}.lease_claim_batches AS claim_batches
+                          ON claim_batches.claim_slot = completed.claim_slot
+                         AND claim_batches.batch_id = completed.claim_batch_id
+                        CROSS JOIN LATERAL (
+                            SELECT
+                                claim_batches.job_ids[completed.claim_batch_index] AS job_id,
+                                claim_batches.run_leases[completed.claim_batch_index] AS run_lease,
+                                claim_batches.receipt_ids[completed.claim_batch_index] AS receipt_id
+                        ) AS items
+                        WHERE completed.claim_batch_id IS NOT NULL
+                          AND completed.claim_batch_index IS NOT NULL
+                          AND completed.claim_batch_index BETWEEN 1 AND claim_batches.claimed_count
+                          AND completed.job_id = items.job_id
+                          AND completed.run_lease = items.run_lease
+                          AND completed.receipt_id = items.receipt_id
+                          AND NOT EXISTS (
                               SELECT 1
                               FROM {schema}.lease_claim_closures AS closures
                               WHERE closures.claim_slot = claim_batches.claim_slot
@@ -6804,7 +6817,6 @@ impl QueueStorage {
                               WHERE dlq.job_id = items.job_id
                                 AND dlq.run_lease = items.run_lease
                           )
-                        FOR UPDATE OF claim_batches
                     ),
                     locked_claims AS (
                         SELECT claim_slot, job_id, run_lease FROM locked_row_claims
@@ -6896,6 +6908,8 @@ impl QueueStorage {
                 .bind(&receipt_job_ids)
                 .bind(&receipt_run_leases)
                 .bind(&receipt_receipt_ids)
+                .bind(&receipt_claim_batch_ids)
+                .bind(&receipt_claim_batch_indices)
                 .fetch_all(tx.as_mut())
                 .await
                 .map_err(map_sqlx_error)?;

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();
@@ -435,17 +448,28 @@ 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.
+            // Leaving `in_flight` ends the attempt's heartbeat, so only a
+            // bounded number of finished attempts may wait for finalization.
+            // While that queue is full the attempt keeps its worker slot and
+            // its heartbeat, and the dispatcher stops claiming: a completion
+            // backlog cannot outgrow `heartbeat_staleness` and get finished
+            // attempts rescued and run again.
+            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,
+            };
             in_flight.remove((job_id, job_run_lease));
             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);
@@ -495,6 +519,7 @@ impl JobExecutor {
                         outcome => break outcome,
                     }
                 };
+                drop(finalize_permit);
 
                 match &outcome {
                     Ok(CompletionOutcome::Applied { terminal, .. }) => {