fix(worker): drain finalization on graceful shutdown (0.6)

CHANGELOG.md

@@ -6,6 +6,23 @@ Notable changes between releases. Detailed migration notes for storage transitio
 
 ### Fixed
 
+- **A graceful shutdown no longer leaves a finished attempt for stale-heartbeat
+  rescue to run again.** Once a handler returned, its finalize (completion,
+  retry, cancel or failure, and any ADR-029 follow-up inserts) ran in a task
+  that `Client::shutdown` did not wait for. When the last handler returned
+  just as the drain ended, shutdown stopped the completion batcher before
+  that finalize reached it, and the direct fallback raced the process exit.
+  The attempt had already left the heartbeat set, so it went stale and was
+  rescued and executed again `heartbeat_staleness` later; a deploy could
+  re-run a few jobs on every rolling restart. Shutdown now waits for
+  pending finalization, within the same drain timeout, before it stops the
+  completion batcher, heartbeat and maintenance services. A finalize that
+  outlives the timeout is still left to rescue, as before.
+- **Point reads show the live queue and priority after a lane move.** Retained
+  ready rows from before a priority or queue change could win a read's tie
+  after the ring wrapped, returning the old fields even though the batch
+  operation succeeded. Point reads now exclude tombstoned source rows.
+
 - **Callback completion, `get_job`, lifecycle events and the admin reads
   find running jobs that priority aging promoted on queue storage** (backport
   of #600). Claims and leases record a job's aged effective priority while its

awa-model/src/queue_storage.rs

@@ -11695,8 +11695,17 @@ impl QueueStorage {
                 unique_key,
                 unique_states,
                 COALESCE(payload, '{{}}'::jsonb) AS payload
-            FROM {schema}.ready_entries
-            WHERE job_id = $1
+            FROM {schema}.ready_entries AS ready
+            WHERE ready.job_id = $1
+              AND NOT EXISTS (
+                  SELECT 1 FROM {schema}.ready_tombstones AS tomb
+                  WHERE tomb.ready_slot = ready.ready_slot
+                    AND tomb.ready_generation = ready.ready_generation
+                    AND tomb.queue = ready.queue
+                    AND tomb.priority = ready.priority
+                    AND tomb.enqueue_shard = ready.enqueue_shard
+                    AND tomb.lane_seq = ready.lane_seq
+              )
             ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC
             "#,
         ))

awa-worker/Cargo.toml

@@ -21,7 +21,7 @@ sqlx.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 tokio.workspace = true
-tokio-util.workspace = true
+tokio-util = { workspace = true, features = ["rt"] }
 async-trait.workspace = true
 thiserror.workspace = true
 chrono.workspace = true

awa-worker/src/client.rs

@@ -27,6 +27,7 @@ use std::time::Duration;
 use tokio::sync::{Mutex, RwLock};
 use tokio::task::JoinSet;
 use tokio_util::sync::CancellationToken;
+use tokio_util::task::TaskTracker;
 use tracing::{info, warn};
 use uuid::Uuid;
 
@@ -1085,6 +1086,7 @@ impl ClientBuilder {
             dispatcher_handles: RwLock::new(Vec::new()),
             service_handles: RwLock::new(Vec::new()),
             job_set: Arc::new(Mutex::new(JoinSet::new())),
+            finalizers: TaskTracker::new(),
             in_flight: Arc::new(InFlightRegistry::default()),
             queue_in_flight,
             dispatcher_alive,
@@ -1185,6 +1187,8 @@ pub struct Client {
     service_handles: RwLock<Vec<tokio::task::JoinHandle<()>>>,
     /// JoinSet tracking in-flight job tasks for graceful drain.
     job_set: Arc<Mutex<JoinSet<()>>>,
+    /// Finalization of attempts whose handler has returned, drained on shutdown.
+    finalizers: TaskTracker,
     in_flight: InFlightMap,
     queue_in_flight: Arc<HashMap<String, Arc<AtomicU32>>>,
     dispatcher_alive: Arc<HashMap<String, Arc<AtomicBool>>>,
@@ -1566,6 +1570,7 @@ impl Client {
             completion_handle,
             effective_storage.clone(),
             self.dlq_policy.clone(),
+            self.finalizers.clone(),
         ));
 
         // Admin cancellation listener: fires the in-flight cancel flag
@@ -1804,8 +1809,9 @@ impl Client {
     /// 1. Stop dispatchers (no new jobs claimed)
     /// 2. Signal in-flight jobs to cancel
     /// 3. Wait for dispatchers to exit
-    /// 4. Drain in-flight jobs (heartbeat + maintenance still alive!)
-    /// 5. Stop heartbeat + maintenance
+    /// 4. Drain in-flight jobs and their finalization (heartbeat,
+    ///    maintenance and the completion batcher still alive!)
+    /// 5. Stop heartbeat, maintenance and the completion batcher
     pub async fn shutdown(&self, timeout: Duration) {
         info!("Initiating graceful shutdown");
 
@@ -1863,6 +1869,14 @@ impl Client {
             );
         }
 
+        self.finalizers.close();
+        if tokio::time::timeout_at(drain_deadline, self.finalizers.wait())
+            .await
+            .is_err()
+        {
+            warn!("Shutdown drain timeout exceeded waiting for durable finalization");
+        }
+
         // Phase 5: Stop background services (heartbeat + maintenance)
         self.service_cancel.cancel();
         let service_handles: Vec<_> = {

awa-worker/src/dispatcher.rs

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

awa-worker/src/executor.rs

@@ -11,6 +11,7 @@ use std::collections::HashMap;
 use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
+use tokio_util::task::TaskTracker;
 use tracing::{error, info, info_span, warn, Instrument};
 
 /// Result of executing a job handler.
@@ -267,6 +268,7 @@ pub struct JobExecutor {
     completion_batcher: CompletionBatcherHandle,
     storage: RuntimeStorage,
     dlq_policy: DlqPolicy,
+    finalizers: TaskTracker,
 }
 
 impl JobExecutor {
@@ -288,6 +290,7 @@ impl JobExecutor {
         completion_batcher: CompletionBatcherHandle,
         storage: RuntimeStorage,
         dlq_policy: DlqPolicy,
+        finalizers: TaskTracker,
     ) -> Self {
         Self {
             pool,
@@ -301,6 +304,7 @@ impl JobExecutor {
             completion_batcher,
             storage,
             dlq_policy,
+            finalizers,
         }
     }
 
@@ -327,6 +331,7 @@ impl JobExecutor {
         let completion_batcher = self.completion_batcher.clone();
         let storage = self.storage.clone();
         let dlq_policy = self.dlq_policy.clone();
+        let finalizers = self.finalizers.clone();
         let job_id = job.id;
         let job_run_lease = job.run_lease;
         let job_kind = job.kind.clone();
@@ -444,7 +449,10 @@ impl JobExecutor {
                 return;
             };
             let dlq_enabled = dlq_policy.enabled_for(&job_queue);
-            tokio::spawn(async move {
+            // 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.
+            finalizers.spawn(async move {
                 // Every finalize path is one lease-guarded transaction, and a
                 // deadlock abort rolls the whole transaction back, so running
                 // it again is safe. Without the retry the job would sit

awa/tests/ci_timing.rs

@@ -0,0 +1,186 @@
+//! Shared CI contention scaling for the nightly chaos and benchmark suites.
+//!
+//! The nightly suites run on shared GitHub runners whose CPU allocation
+//! varies run to run. Three assertion shapes are sensitive to that and
+//! flaked repeatedly through 2026-07 (#399, #434):
+//!
+//!   1. **Wall-clock waits.** A test waits N seconds for a state the
+//!      runtime reaches in milliseconds when it has a core to itself.
+//!      Scale with [`scaled_timeout`].
+//!
+//!   2. **Aggressive rescue cadences.** A chaos client sets
+//!      `heartbeat_staleness` a few multiples above `heartbeat_interval`
+//!      so its own backdating triggers rescue promptly. Under contention
+//!      a *live* worker's heartbeat can miss that window, so the runtime
+//!      correctly rescues a healthy attempt — and the test sees a genuine
+//!      duplicate completion, which reads as a correctness failure. Scale
+//!      with [`scaled_staleness`]: the interval stays fast so rescue is
+//!      still exercised, but the staleness window gains the same margin
+//!      as the waits around it.
+//!
+//!   3. **Minimum-progress floors.** A gate asserts "at least N of these
+//!      happened" to catch a stalled subsystem. The floor must sit far
+//!      below the *observed* operating point, not just below the nominal
+//!      one. Scale with [`contention_floor`].
+//!
+//! Scaling only ever loosens a bound, and only on CI. A local run keeps
+//! the strict values, so a real regression still fails fast on a
+//! developer machine. Override with `AWA_CHAOS_TIMEOUT_MULTIPLIER`.
+#![allow(dead_code)]
+
+use std::time::Duration;
+
+/// Upper bound on the override. `Duration::mul_f64` panics on a non-finite
+/// or overflowing result, so an unbounded override would turn a typo into a
+/// failed nightly rather than a loose bound.
+const MAX_MULTIPLIER: f64 = 100.0;
+
+/// How much slack to give contention-sensitive bounds. `1.0` locally,
+/// `3.0` on CI, or the `AWA_CHAOS_TIMEOUT_MULTIPLIER` override (clamped to
+/// `1.0 ..= MAX_MULTIPLIER` so the override can only ever loosen, and can
+/// never produce a duration `mul_f64` refuses).
+pub fn chaos_timeout_multiplier() -> f64 {
+    let override_var = std::env::var("AWA_CHAOS_TIMEOUT_MULTIPLIER").ok();
+    multiplier_from(override_var.as_deref(), std::env::var_os("CI").is_some())
+}
+
+/// The multiplier decision, with the environment passed in so it can be
+/// tested without mutating process-global state.
+fn multiplier_from(override_var: Option<&str>, is_ci: bool) -> f64 {
+    // `"inf"` and `"nan"` parse successfully as f64, and both reach
+    // `Duration::mul_f64`, which panics rather than saturating. Treat a
+    // non-finite override as no override at all.
+    if let Some(parsed) = override_var
+        .and_then(|raw| raw.parse::<f64>().ok())
+        .filter(|parsed| parsed.is_finite())
+    {
+        return parsed.clamp(1.0, MAX_MULTIPLIER);
+    }
+
+    if is_ci {
+        3.0
+    } else {
+        1.0
+    }
+}
+
+/// Grow a wait deadline by the contention multiplier.
+pub fn scaled_timeout(timeout: Duration) -> Duration {
+    timeout.mul_f64(chaos_timeout_multiplier())
+}
+
+/// Grow a heartbeat-staleness window by the contention multiplier.
+///
+/// Distinct from [`scaled_timeout`] only in intent: this one is passed to
+/// `ClientBuilder::heartbeat_staleness`, where the cost of being too tight
+/// is a spurious rescue of a live attempt rather than a timeout. Paired
+/// heartbeat/rescue *intervals* are deliberately left unscaled so the
+/// rescue path still runs at chaos cadence.
+pub fn scaled_staleness(staleness: Duration) -> Duration {
+    staleness.mul_f64(chaos_timeout_multiplier())
+}
+
+/// Shrink a minimum-progress floor by the contention multiplier.
+///
+/// Use for "this subsystem must have advanced at least N times" gates.
+/// The returned floor is at least 1: the regression these gates exist to
+/// catch is a fully stalled subsystem, so zero progress must still fail.
+pub fn contention_floor(nominal: i64) -> i64 {
+    let scaled = (nominal as f64 / chaos_timeout_multiplier()).floor() as i64;
+    scaled.max(1)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn local_runs_keep_strict_bounds() {
+        assert_eq!(multiplier_from(None, false), 1.0);
+    }
+
+    #[test]
+    fn ci_runs_get_margin() {
+        assert_eq!(multiplier_from(None, true), 3.0);
+    }
+
+    #[test]
+    fn override_can_only_loosen() {
+        assert_eq!(multiplier_from(Some("0.1"), true), 1.0);
+        assert_eq!(multiplier_from(Some("-5"), true), 1.0);
+        assert_eq!(multiplier_from(Some("6"), false), 6.0);
+    }
+
+    #[test]
+    fn unparseable_override_falls_back_to_the_environment() {
+        assert_eq!(multiplier_from(Some("banana"), true), 3.0);
+        assert_eq!(multiplier_from(Some(""), false), 1.0);
+    }
+
+    #[test]
+    fn non_finite_override_falls_back_to_the_environment() {
+        // These all parse as f64 and would otherwise reach `mul_f64`, which
+        // panics on a non-finite result instead of saturating.
+        for raw in ["inf", "-inf", "infinity", "NaN", "nan"] {
+            assert_eq!(multiplier_from(Some(raw), true), 3.0, "raw={raw}");
+            assert_eq!(multiplier_from(Some(raw), false), 1.0, "raw={raw}");
+        }
+    }
+
+    #[test]
+    fn oversized_override_is_clamped() {
+        assert_eq!(multiplier_from(Some("1e308"), true), MAX_MULTIPLIER);
+        assert_eq!(multiplier_from(Some("10000"), true), MAX_MULTIPLIER);
+    }
+
+    #[test]
+    fn every_multiplier_produces_a_usable_duration() {
+        // `Duration::mul_f64` panics rather than saturating, so the clamp is
+        // what keeps a bad override from failing the nightly it was set to
+        // rescue.
+        for raw in [
+            "inf", "-inf", "NaN", "1e308", "10000", "0", "-1", "banana", "",
+        ] {
+            for is_ci in [true, false] {
+                let m = multiplier_from(Some(raw), is_ci);
+                assert!(m.is_finite() && (1.0..=MAX_MULTIPLIER).contains(&m));
+                // Would panic if the multiplier were unbounded.
+                let _ = Duration::from_secs(60).mul_f64(m);
+            }
+        }
+    }
+
+    #[test]
+    fn scaling_a_timeout_grows_it() {
+        // 250ms staleness against a 50ms heartbeat interval is a 5x margin
+        // locally; on CI it becomes 15x, which is what stops a contended
+        // runner's live worker from being rescued as if it had died.
+        assert_eq!(
+            Duration::from_millis(250).mul_f64(multiplier_from(None, true)),
+            Duration::from_millis(750)
+        );
+    }
+
+    #[test]
+    fn progress_floor_keeps_margin_under_the_observed_rate() {
+        // The receipt gate's nominal floor is duration_secs / 4 = 45 for a
+        // 180s run. #399 saw a healthy run produce 41, so the CI floor has
+        // to sit well under that.
+        let ci_floor = (45.0_f64 / multiplier_from(None, true)).floor() as i64;
+        assert_eq!(ci_floor, 15);
+        assert!(
+            ci_floor < 41,
+            "floor must clear the observed operating point"
+        );
+    }
+
+    #[test]
+    fn progress_floor_never_reaches_zero() {
+        // Whatever the multiplier, a fully pinned ring (zero rotations) has
+        // to keep failing — that is the regression the gate exists for.
+        assert_eq!(contention_floor(45), contention_floor(45).max(1));
+        let absurd = (45.0_f64 / multiplier_from(Some("10000"), true)).floor() as i64;
+        assert_eq!(absurd.max(1), 1);
+        assert_eq!(contention_floor(0), 1);
+    }
+}