Skip to main content

aion_server/worker/
heartbeat.rs

1//! Heartbeat window tracking and lost-worker failure surfacing.
2
3use chrono::{DateTime, Utc};
4use std::collections::{HashMap, HashSet};
5use std::sync::{Arc, Mutex, MutexGuard};
6use std::time::{Duration, Instant};
7use tokio::sync::{Notify, watch};
8use tracing::{error, info, warn};
9
10use aion_core::{ActivityId, Payload, WorkflowId};
11use aion_proto::{ProtoHeartbeat, WireError};
12
13use crate::error::ServerError;
14use crate::shutdown::DrainState;
15use crate::worker::dispatch::{
16    ActivityCompletion, ActivityCompletionOutcome, ActivityCompletionSink,
17};
18use crate::worker::envelope::CompletionToken;
19use crate::worker::registry::{ConnectedWorkerRegistry, WorkerId};
20
21/// In-flight activity assigned to a connected worker.
22#[derive(Clone, Debug, Eq, PartialEq)]
23pub struct InFlightActivity {
24    /// Owning workflow id.
25    pub workflow_id: WorkflowId,
26    /// Correlating activity id.
27    pub activity_id: ActivityId,
28    /// One-based delivery attempt this dispatch carries — the third axis of the
29    /// `(workflow, activity, attempt)` identity the transcript, the intervention
30    /// index, and history all key on.
31    ///
32    /// It is NOT part of the tracker's key: a completion is addressed by the
33    /// worker, workflow, and activity together, and widening the key would
34    /// break that addressing. It is carried so a reader asking about a SPECIFIC
35    /// attempt — the live describe join asking whose progress note this is —
36    /// can tell a tracked entry for the attempt it asked about from a lingering
37    /// entry for a superseded one, instead of attributing one attempt's note to
38    /// another.
39    pub attempt: u32,
40    /// Generation authorized to receive a result or synthesized loss.
41    pub completion_token: CompletionToken,
42}
43
44/// Observable liveness state for a single in-flight activity.
45#[derive(Clone, Debug, Eq, PartialEq)]
46pub struct TaskLiveness {
47    /// Worker currently responsible for the task.
48    pub worker_id: WorkerId,
49    /// Owning workflow id.
50    pub workflow_id: WorkflowId,
51    /// Correlating activity id.
52    pub activity_id: ActivityId,
53    /// One-based delivery attempt this tracked dispatch carries (see
54    /// [`InFlightActivity::attempt`]).
55    pub attempt: u32,
56    /// Generation authorized to receive a result or synthesized loss.
57    pub completion_token: CompletionToken,
58    /// Operator-configured heartbeat window used for expiry checks.
59    pub heartbeat_window: Duration,
60    /// Monotonic timestamp of assignment or the most recent heartbeat.
61    pub last_heartbeat_at: Instant,
62    /// Optional worker progress from the most recent heartbeat.
63    pub last_progress: Option<Payload>,
64    /// Wall-clock instant this process received the heartbeat that carried
65    /// [`Self::last_progress`]. `None` while no progress has been reported.
66    ///
67    /// A wall clock rather than the tracked [`Instant`] because this one is
68    /// REPORTED to operators, and a monotonic instant means nothing outside the
69    /// process that minted it. Server-side observability stamping only — the
70    /// determinism boundary governs workflow-visible time, and nothing here is
71    /// workflow-visible.
72    pub last_progress_at: Option<DateTime<Utc>>,
73}
74
75/// Result of accepting a heartbeat for an in-flight task.
76#[derive(Clone, Debug, Eq, PartialEq)]
77pub struct HeartbeatUpdate {
78    /// Updated liveness after recording the heartbeat.
79    pub liveness: TaskLiveness,
80}
81
82/// Tasks removed from tracking because a worker was declared lost.
83#[derive(Clone, Debug, Eq, PartialEq)]
84pub struct LostWorkerReport {
85    /// Lost worker removed from the connected-worker registry.
86    pub worker_id: WorkerId,
87    /// In-flight activities swept off the tracker: surfaced to the engine as
88    /// retryable failures on the `fail_*` paths, or parked for restart
89    /// recovery (nothing recorded, nothing delivered) on the graceful-drain
90    /// `park_*` paths (#207).
91    pub tasks: Vec<InFlightActivity>,
92    /// Task queue the lost worker was serving, captured from the registry
93    /// BEFORE deregistration (afterwards the handle is gone and the queue is
94    /// unknowable). `None` when the worker was already absent from the registry.
95    ///
96    /// This is what lets the deregistration log name the queue an operator has
97    /// to act on, rather than an opaque worker id.
98    pub task_queue: Option<String>,
99}
100
101#[derive(Clone, Debug, Eq, Hash, PartialEq)]
102struct TaskKey(WorkerId, WorkflowId, ActivityId);
103
104#[derive(Debug, Default)]
105struct HeartbeatState {
106    tasks: HashMap<TaskKey, TaskLiveness>,
107    /// Last frame observed on each live worker connection.
108    ///
109    /// This is the PROCESS-IS-ALIVE fact, and only that. It is advanced by
110    /// anything the worker sends — including the worker-side liveness pump,
111    /// which beats from a background task regardless of what its serve loop is
112    /// doing. A fresh entry here means "that process is running and its
113    /// worker-to-server direction works". It does NOT mean the server can
114    /// reach it.
115    connections: HashMap<WorkerId, Instant>,
116    /// Last time the server PROVED it can reach this worker's dispatch path,
117    /// i.e. the last answered liveness ping.
118    ///
119    /// This is the SERVER-CAN-REACH-THE-WORKER fact, and it is the only one
120    /// that is a dispatch precondition. It is advanced ONLY by an answered
121    /// ping, never by an inbound frame, because only the ping rides the same
122    /// server-to-worker leg a dispatch does.
123    ///
124    /// The two facts are separate because collapsing them hid a total outage:
125    /// on run `dfd2117c` the server could not push to a worker for fifteen
126    /// minutes while the worker's pump kept the single old lease perfectly
127    /// fresh, so the dead-man switch could not fire for the one failure it
128    /// exists to detect. `liminal_transport`'s own doc already forbids this —
129    /// the ping proves "the exact path a dispatch would take, not a parallel
130    /// one that could be healthy while the real one is not" — and the pump
131    /// feeding the same lease was exactly that parallel channel.
132    reachability: HashMap<WorkerId, Reachability>,
133}
134
135/// How many CONSECUTIVE answered pings re-admit a worker to dispatch.
136///
137/// A connection is a channel to prove reachability ON, never proof of it: the
138/// registration handshake's ack is SENT by the server, and a sent ack is not a
139/// received one — exactly the inference this whole lane exists to stop making.
140/// So eligibility is earned by measurement, on every connection including the
141/// first, and a redial re-seeds the measurement OPPORTUNITY rather than the
142/// verdict.
143///
144/// Why two and not one: one success re-admits a link that answered once by luck
145/// — a race, a buffer that happened to drain — so a link answering one probe in
146/// three would flap in and out of eligibility indefinitely, which is the defect
147/// this constant exists to remove rather than slow down. Two consecutive
148/// successes is the smallest number that distinguishes "answered" from
149/// "answering".
150///
151/// Why not three or more: the cost is paid on EVERY connect, in probe cadences.
152/// At the probe's cadence a fresh worker is undispatchable for `K` cadences
153/// while its first dispatches park, and that latency is charged to every honest
154/// worker to catch a dishonest one. Two buys the discrimination; three buys
155/// only delay.
156pub(crate) const DISPATCH_PROBATION_PINGS: u32 = 2;
157
158/// A worker's dispatch-path standing: how many consecutive pings it has
159/// answered, and when the most recent one landed.
160///
161/// `proved_at` is `None` until the probation is served, so a worker on
162/// probation is not merely stale — it has no proof at all, which is the honest
163/// description of a connection nothing has been measured on yet.
164#[derive(Clone, Copy, Debug, Default)]
165struct Reachability {
166    consecutive_answers: u32,
167    proved_at: Option<Instant>,
168    /// Whether this worker has EVER held dispatch eligibility on this
169    /// connection. Not a duplicate of the two fields above: they describe the
170    /// current standing, this describes the connection's history, and only the
171    /// history separates a worker still serving its opening probation from one
172    /// that earned eligibility and then lost it.
173    ///
174    /// Deliberately NOT cleared by [`HeartbeatTracker::record_dispatch_unreachable`]
175    /// — a failed ping ends the current proof, it does not un-happen the proof
176    /// that came before it. Cleared only by
177    /// [`HeartbeatTracker::register_connection`], because a new connection is a
178    /// new measurement and nothing earned on the old one carries across.
179    ever_proved: bool,
180}
181
182impl Reachability {
183    /// Whether the probation is served and the proof is still inside `window`.
184    fn is_proved(self, now: Instant, window: Duration) -> bool {
185        self.consecutive_answers >= DISPATCH_PROBATION_PINGS
186            && self.proved_at.is_some_and(|proved_at| {
187                now.checked_duration_since(proved_at)
188                    .is_none_or(|elapsed| elapsed <= window)
189            })
190    }
191
192    /// Why this standing does not currently permit dispatch.
193    ///
194    /// Only meaningful when [`Self::is_proved`] is false; the caller pairs them
195    /// so the classification and the membership test can never disagree about
196    /// which workers are excluded.
197    fn exclusion(self) -> DispatchExclusion {
198        if self.ever_proved {
199            DispatchExclusion::ReachabilityLost
200        } else {
201            DispatchExclusion::OpeningProbation {
202                answers: self.consecutive_answers,
203            }
204        }
205    }
206}
207
208/// Why a worker is currently excluded from dispatch selection.
209///
210/// These are DIFFERENT FACTS and an operator must be able to tell them apart —
211/// the same standard this module already holds the two ping failures to. One is
212/// the ordinary cost of connecting; the other is an incident. Reported as one
213/// value alongside the exclusion itself so nothing has to re-derive the reason
214/// from a second reading of the same state.
215#[derive(Clone, Copy, Debug, PartialEq, Eq)]
216pub enum DispatchExclusion {
217    /// The worker registered and has not yet answered
218    /// [`DISPATCH_PROBATION_PINGS`] consecutive pings, so eligibility has never
219    /// been earned on this connection. Expected on EVERY connect, including a
220    /// perfectly healthy one — this is the probation being served, not a fault.
221    OpeningProbation {
222        /// Consecutive answers banked so far, out of [`DISPATCH_PROBATION_PINGS`].
223        answers: u32,
224    },
225    /// The worker held dispatch eligibility on this connection and no longer
226    /// does: either a ping failed and restarted its probation, or the last
227    /// proof aged out of the heartbeat window. This one is an incident.
228    ReachabilityLost,
229}
230
231/// One worker excluded from dispatch, with the reason it is excluded.
232#[derive(Clone, Copy, Debug, PartialEq, Eq)]
233pub struct ExcludedWorker {
234    /// The worker selection must skip.
235    pub worker_id: WorkerId,
236    /// Why it is being skipped.
237    pub exclusion: DispatchExclusion,
238}
239
240/// Per-task liveness tracker for remote-worker streams.
241///
242/// It is also the server's ONLY store of worker progress notes, and that store
243/// is VOLATILE: notes live here and nowhere else, so a restart loses every note
244/// ever reported. [`Self::notes_held_since`] is the evidence a reader needs to
245/// tell "this worker has said nothing" from "this process was not running when
246/// the attempt began" — see `crate::worker::attempt_progress`.
247#[derive(Clone, Debug)]
248pub struct HeartbeatTracker {
249    heartbeat_window: Duration,
250    inner: Arc<Mutex<HeartbeatState>>,
251    empty: Arc<Notify>,
252    notes_held_since: DateTime<Utc>,
253}
254
255impl HeartbeatTracker {
256    /// Build a tracker using the operator-supplied heartbeat window.
257    #[must_use]
258    pub fn new(heartbeat_window: Duration) -> Self {
259        Self {
260            heartbeat_window,
261            inner: Arc::new(Mutex::new(HeartbeatState::default())),
262            empty: Arc::new(Notify::new()),
263            // The instant this volatile note store began holding notes. Read
264            // here rather than injected because it is exactly the tracker's own
265            // construction instant — there is no other value it could be, and a
266            // caller passing a different one would be reporting a fiction.
267            notes_held_since: Utc::now(),
268        }
269    }
270
271    /// When this tracker began holding progress notes.
272    ///
273    /// An attempt dispatched before this instant left its notes in a process
274    /// that no longer exists, so their absence here says nothing about what the
275    /// worker reported.
276    #[must_use]
277    pub const fn notes_held_since(&self) -> DateTime<Utc> {
278        self.notes_held_since
279    }
280
281    /// Start the connection-level lease for a newly registered worker, and open
282    /// its dispatch probation.
283    ///
284    /// The connection lease starts fresh — the worker's process is plainly
285    /// alive, it just registered. Dispatch reachability does NOT: a new
286    /// connection is a channel to prove reachability on, not proof of it.
287    ///
288    /// This deliberately reverses an earlier reading of mine, that the
289    /// registration handshake is "itself a completed server-to-worker round
290    /// trip". The ack is SENT by the server; nothing reports that it was
291    /// RECEIVED. Treating a send as a delivery is the same inference this lane
292    /// exists to stop making, and left unfixed it meant a worker the server
293    /// could never reach would re-seed itself on every redial and cycle in and
294    /// out of eligibility forever instead of settling out.
295    ///
296    /// See [`DISPATCH_PROBATION_PINGS`].
297    ///
298    /// # Errors
299    ///
300    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
301    pub fn register_connection(
302        &self,
303        worker_id: WorkerId,
304        now: Instant,
305    ) -> Result<(), ServerError> {
306        let mut state = self.state()?;
307        state.connections.insert(worker_id, now);
308        // A fresh, UNSERVED probation: zero answers, no proof. Inserted rather
309        // than left absent so the worker is carried by `unreachable_workers`
310        // and is therefore explicitly excluded, not merely unknown.
311        state
312            .reachability
313            .insert(worker_id, Reachability::default());
314        Ok(())
315    }
316
317    /// Advance a worker's connection lease after receiving any frame.
318    ///
319    /// Records ONLY that the worker's process is alive. It deliberately does
320    /// NOT advance dispatch reachability: an inbound frame — a heartbeat, a
321    /// pump beat, a completion — proves the worker-to-server direction and
322    /// says nothing about whether the server can push to it. Use
323    /// [`Self::record_dispatch_reachability`] for the fact that gates dispatch.
324    ///
325    /// Returns `false` if the worker has already been removed from lease tracking;
326    /// a frame racing deregistration must not resurrect it.
327    ///
328    /// # Errors
329    ///
330    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
331    pub fn record_connection_activity(
332        &self,
333        worker_id: WorkerId,
334        now: Instant,
335    ) -> Result<bool, ServerError> {
336        let mut state = self.state()?;
337        let Some(last_activity) = state.connections.get_mut(&worker_id) else {
338            return Ok(false);
339        };
340        *last_activity = now;
341        Ok(true)
342    }
343
344    /// Record proof that the server can reach this worker's dispatch path — an
345    /// ANSWERED liveness ping, and nothing else.
346    ///
347    /// Advances both facts, because an answered ping proves both: the worker
348    /// received a server push (reachability) and replied to it (alive). It also
349    /// serves one ping of the dispatch probation; eligibility returns once
350    /// [`DISPATCH_PROBATION_PINGS`] consecutive answers have landed.
351    ///
352    /// Returns `false` if the worker has already been removed from lease
353    /// tracking; a pong racing a reap must not resurrect it.
354    ///
355    /// # Errors
356    ///
357    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
358    pub fn record_dispatch_reachability(
359        &self,
360        worker_id: WorkerId,
361        now: Instant,
362    ) -> Result<bool, ServerError> {
363        let mut state = self.state()?;
364        let Some(last_activity) = state.connections.get_mut(&worker_id) else {
365            return Ok(false);
366        };
367        *last_activity = now;
368        let standing = state.reachability.entry(worker_id).or_default();
369        standing.consecutive_answers = standing.consecutive_answers.saturating_add(1);
370        standing.proved_at = Some(now);
371        if standing.consecutive_answers >= DISPATCH_PROBATION_PINGS {
372            // The probation is served. Recording it here — at the one place a
373            // probation can complete — is what lets a later exclusion say
374            // whether eligibility was ever held, without a second copy of the
375            // threshold anywhere else.
376            standing.ever_proved = true;
377        }
378        Ok(true)
379    }
380
381    /// Record that a liveness ping went UNANSWERED: the probation restarts.
382    ///
383    /// This is what makes the probation consecutive rather than cumulative. A
384    /// link answering one probe in three would otherwise accumulate its way to
385    /// eligibility and keep it, which is the flapping this design removes.
386    ///
387    /// Returns `false` if the worker has already been removed from lease
388    /// tracking.
389    ///
390    /// # Errors
391    ///
392    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
393    pub fn record_dispatch_unreachable(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
394        let mut state = self.state()?;
395        if !state.connections.contains_key(&worker_id) {
396            return Ok(false);
397        }
398        let standing = state.reachability.entry(worker_id).or_default();
399        standing.consecutive_answers = 0;
400        standing.proved_at = None;
401        // `ever_proved` deliberately survives: this connection DID earn
402        // eligibility once, and that is what makes the loss an incident rather
403        // than the ordinary cost of connecting.
404        Ok(true)
405    }
406
407    /// Whether the server has proved, within the heartbeat window, that it can
408    /// reach this worker's dispatch path — probation served AND the proof still
409    /// fresh.
410    ///
411    /// An untracked worker is not reachable: absence of proof is not proof.
412    ///
413    /// # Errors
414    ///
415    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
416    pub fn is_dispatch_reachable(
417        &self,
418        worker_id: WorkerId,
419        now: Instant,
420    ) -> Result<bool, ServerError> {
421        let state = self.state()?;
422        Ok(state
423            .reachability
424            .get(&worker_id)
425            .is_some_and(|standing| standing.is_proved(now, self.heartbeat_window)))
426    }
427
428    /// Every tracked worker the server has NOT been able to reach within the
429    /// heartbeat window, regardless of how alive its process looks, each paired
430    /// with WHY it is excluded.
431    ///
432    /// The reason travels with the membership rather than being recomputed by
433    /// the caller, so the set that gates dispatch and the reason an operator is
434    /// told can never describe different states.
435    ///
436    /// # Errors
437    ///
438    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
439    pub fn unreachable_workers(&self, now: Instant) -> Result<Vec<ExcludedWorker>, ServerError> {
440        let state = self.state()?;
441        let mut workers = state
442            .reachability
443            .iter()
444            .filter(|(_, standing)| !standing.is_proved(now, self.heartbeat_window))
445            .map(|(worker_id, standing)| ExcludedWorker {
446                worker_id: *worker_id,
447                exclusion: standing.exclusion(),
448            })
449            .collect::<Vec<_>>();
450        workers.sort_unstable_by_key(|excluded| excluded.worker_id);
451        Ok(workers)
452    }
453
454    /// End connection-lease tracking when a transport closes normally.
455    ///
456    /// # Errors
457    ///
458    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
459    pub fn unregister_connection(&self, worker_id: WorkerId) -> Result<(), ServerError> {
460        let mut state = self.state()?;
461        state.connections.remove(&worker_id);
462        state.reachability.remove(&worker_id);
463        Ok(())
464    }
465
466    /// Track a newly accepted in-flight activity for heartbeat expiry.
467    ///
468    /// # Errors
469    ///
470    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
471    pub fn track_task(
472        &self,
473        worker_id: WorkerId,
474        task: InFlightActivity,
475        now: Instant,
476    ) -> Result<(), ServerError> {
477        let key = TaskKey::new(
478            worker_id,
479            task.workflow_id.clone(),
480            task.activity_id.clone(),
481        );
482        let liveness = TaskLiveness {
483            worker_id,
484            workflow_id: task.workflow_id,
485            activity_id: task.activity_id,
486            attempt: task.attempt,
487            completion_token: task.completion_token,
488            heartbeat_window: self.heartbeat_window,
489            last_heartbeat_at: now,
490            last_progress: None,
491            last_progress_at: None,
492        };
493        let mut state = self.state()?;
494        state.tasks.insert(key, liveness);
495        state.connections.insert(worker_id, now);
496        Ok(())
497    }
498
499    /// Stop tracking a completed activity and wake drain waiters if this was the last task.
500    ///
501    /// Returns whether the task was still tracked when this ran: `true` means
502    /// THIS call retired the in-flight entry, `false` means another path (the
503    /// expiry sweep, a disconnect teardown, shutdown, or a completed dispatch)
504    /// already did. The liminal reply router uses that bool as its structural
505    /// gate for synthesizing a lost-worker failure — the exact mirror of the
506    /// gRPC sweep failing only still-tracked tasks.
507    ///
508    /// # Errors
509    ///
510    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
511    pub fn complete_task(
512        &self,
513        worker_id: WorkerId,
514        workflow_id: &WorkflowId,
515        activity_id: &ActivityId,
516    ) -> Result<bool, ServerError> {
517        let key = TaskKey::new(worker_id, workflow_id.clone(), activity_id.clone());
518        let (was_tracked, became_empty) = {
519            let mut state = self.state()?;
520            let was_tracked = state.tasks.remove(&key).is_some();
521            (was_tracked, state.tasks.is_empty())
522        };
523        if became_empty {
524            self.empty.notify_waiters();
525        }
526        Ok(was_tracked)
527    }
528
529    /// Whether the given in-flight task is still tracked (not yet completed,
530    /// swept, or drained). The liminal reply router polls this to bound its
531    /// wait: once the entry is gone the dispatch was resolved by another path,
532    /// so the router exits instead of parking on the connection forever.
533    ///
534    /// # Errors
535    ///
536    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
537    pub fn is_tracked(
538        &self,
539        worker_id: WorkerId,
540        workflow_id: &WorkflowId,
541        activity_id: &ActivityId,
542    ) -> Result<bool, ServerError> {
543        let key = TaskKey::new(worker_id, workflow_id.clone(), activity_id.clone());
544        Ok(self.state()?.tasks.contains_key(&key))
545    }
546
547    /// Refresh the liveness stamp of an in-flight task from a transport-level
548    /// liveness beat that carries no progress payload (the liminal worker's
549    /// automatic pump). Returns `true` when the task was tracked and refreshed,
550    /// `false` when it is not in flight — a benign outcome for a beat racing a
551    /// completion or covering an outbox dispatch the tracker never held.
552    ///
553    /// # Errors
554    ///
555    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
556    pub fn record_liveness(
557        &self,
558        worker_id: WorkerId,
559        workflow_id: &WorkflowId,
560        activity_id: &ActivityId,
561        now: Instant,
562    ) -> Result<bool, ServerError> {
563        let key = TaskKey::new(worker_id, workflow_id.clone(), activity_id.clone());
564        let mut state = self.state()?;
565        if !state.tasks.contains_key(&key) {
566            return Ok(false);
567        }
568        if let Some(last_activity) = state.connections.get_mut(&worker_id) {
569            *last_activity = now;
570        }
571        let Some(liveness) = state.tasks.get_mut(&key) else {
572            return Ok(false);
573        };
574        liveness.last_heartbeat_at = now;
575        Ok(true)
576    }
577
578    /// The operator-configured heartbeat window this tracker expires against.
579    /// The bridge stamps it onto each liminal dispatch so the worker's
580    /// automatic liveness pump beats at the matching quarter-window cadence.
581    #[must_use]
582    pub const fn heartbeat_window(&self) -> Duration {
583        self.heartbeat_window
584    }
585
586    /// Number of currently tracked in-flight activities.
587    ///
588    /// # Errors
589    ///
590    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
591    pub fn in_flight_count(&self) -> Result<usize, ServerError> {
592        Ok(self.state()?.tasks.len())
593    }
594
595    /// Record a worker heartbeat without completing the activity.
596    ///
597    /// Every heartbeat refreshes the task's liveness stamp. The progress
598    /// payload is only overwritten when the heartbeat CARRIES one: the worker
599    /// runtime's automatic liveness beats are payload-free and interleave
600    /// with explicit handler progress heartbeats, and a liveness beat must
601    /// never erase the handler's most recent progress report.
602    ///
603    /// # Errors
604    ///
605    /// Returns a stable wire error for malformed heartbeats or unknown in-flight tasks.
606    pub fn record_heartbeat(
607        &self,
608        worker_id: WorkerId,
609        heartbeat: ProtoHeartbeat,
610        now: Instant,
611    ) -> Result<HeartbeatUpdate, ServerError> {
612        let decoded = DecodedHeartbeat::try_from(heartbeat)?;
613        let key = TaskKey::new(worker_id, decoded.workflow_id, decoded.activity_id);
614        let mut state = self.state()?;
615        if !state.tasks.contains_key(&key) {
616            return Err(wire_error("heartbeat task is not in flight"));
617        }
618        if let Some(last_activity) = state.connections.get_mut(&worker_id) {
619            *last_activity = now;
620        }
621        let Some(liveness) = state.tasks.get_mut(&key) else {
622            return Err(wire_error("heartbeat task is not in flight"));
623        };
624        liveness.last_heartbeat_at = now;
625        if decoded.progress.is_some() {
626            liveness.last_progress = decoded.progress;
627            liveness.last_progress_at = Some(Utc::now());
628        }
629        Ok(HeartbeatUpdate {
630            liveness: liveness.clone(),
631        })
632    }
633
634    /// Return whether an in-flight task is still within its configured heartbeat window.
635    ///
636    /// # Errors
637    ///
638    /// Returns a stable wire error if the task is not tracked, or lock poison if state cannot be trusted.
639    pub fn is_live(
640        &self,
641        worker_id: WorkerId,
642        workflow_id: &WorkflowId,
643        activity_id: &ActivityId,
644        now: Instant,
645    ) -> Result<bool, ServerError> {
646        let key = TaskKey::new(worker_id, workflow_id.clone(), activity_id.clone());
647        let state = self.state()?;
648        let Some(liveness) = state.tasks.get(&key) else {
649            return Err(wire_error("heartbeat task is not in flight"));
650        };
651        Ok(!is_expired(liveness, now))
652    }
653
654    /// Return the workers that have at least one task beyond the configured heartbeat window.
655    ///
656    /// # Errors
657    ///
658    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
659    pub fn expired_workers(&self, now: Instant) -> Result<Vec<WorkerId>, ServerError> {
660        let state = self.state()?;
661        let mut seen = HashSet::new();
662        let mut workers = Vec::new();
663        for (worker_id, last_activity) in &state.connections {
664            if now
665                .checked_duration_since(*last_activity)
666                .is_some_and(|elapsed| elapsed > self.heartbeat_window)
667                && seen.insert(*worker_id)
668            {
669                workers.push(*worker_id);
670            }
671        }
672        for liveness in state.tasks.values() {
673            if is_expired(liveness, now) && seen.insert(liveness.worker_id) {
674                workers.push(liveness.worker_id);
675            }
676        }
677        workers.sort_unstable();
678        Ok(workers)
679    }
680
681    /// Mark all currently expired workers lost and fail their in-flight tasks through the engine sink.
682    ///
683    /// # Errors
684    ///
685    /// Returns registry, tracker, or sink errors without retrying or rescheduling activities.
686    pub fn fail_expired_workers(
687        &self,
688        registry: &ConnectedWorkerRegistry,
689        sink: &impl ActivityCompletionSink,
690        now: Instant,
691    ) -> Result<Vec<LostWorkerReport>, ServerError> {
692        let mut reports = Vec::new();
693        for worker_id in self.expired_workers(now)? {
694            let report = self.fail_lost_worker(worker_id, registry, sink)?;
695            reports.push(report);
696        }
697        Ok(reports)
698    }
699
700    /// Mark a disconnected worker lost and fail its in-flight tasks through the engine sink.
701    ///
702    /// # Errors
703    ///
704    /// Returns registry, tracker, or sink errors without retrying or rescheduling activities.
705    pub fn fail_disconnected_worker(
706        &self,
707        worker_id: WorkerId,
708        registry: &ConnectedWorkerRegistry,
709        sink: &impl ActivityCompletionSink,
710    ) -> Result<LostWorkerReport, ServerError> {
711        self.fail_lost_worker(worker_id, registry, sink)
712    }
713
714    /// Mark every currently in-flight worker lost and fail all remaining tasks through the sink.
715    ///
716    /// # Errors
717    ///
718    /// Returns registry, tracker, or sink errors without retrying or rescheduling activities.
719    pub fn fail_all_in_flight_workers(
720        &self,
721        registry: &ConnectedWorkerRegistry,
722        sink: &impl ActivityCompletionSink,
723    ) -> Result<Vec<LostWorkerReport>, ServerError> {
724        let worker_ids = {
725            let state = self.state()?;
726            let mut worker_ids = state
727                .tasks
728                .values()
729                .map(|liveness| liveness.worker_id)
730                .collect::<HashSet<_>>()
731                .into_iter()
732                .collect::<Vec<_>>();
733            worker_ids.sort_unstable();
734            worker_ids
735        };
736        let mut reports = Vec::new();
737        for worker_id in worker_ids {
738            let report = self.fail_lost_worker(worker_id, registry, sink)?;
739            if !report.tasks.is_empty() {
740                reports.push(report);
741            }
742        }
743        self.empty.notify_waiters();
744        Ok(reports)
745    }
746
747    /// Park a drain-disconnected worker's in-flight tasks for restart recovery
748    /// (#207): deregister the worker, remove its tracked tasks, and resolve
749    /// each pending waiter through [`ActivityCompletionSink::park_activity`].
750    ///
751    /// The graceful-drain counterpart of [`Self::fail_disconnected_worker`]:
752    /// same deregister-before-collect ordering (same closed dispatch/disconnect
753    /// race), but NO completion is synthesized — the durable log keeps its
754    /// dangling scheduled/started trail, byte-equivalent to a kill -9, and
755    /// restart recovery re-dispatches it. Deregistered with the honest
756    /// [`WorkerDeathReason::Disconnect`](aion_core::WorkerDeathReason::Disconnect):
757    /// the transport genuinely dropped (the worker obeyed the drain request).
758    ///
759    /// # Errors
760    ///
761    /// Returns registry, tracker, or sink errors without retrying or rescheduling activities.
762    pub fn park_disconnected_worker(
763        &self,
764        worker_id: WorkerId,
765        registry: &ConnectedWorkerRegistry,
766        sink: &impl ActivityCompletionSink,
767    ) -> Result<LostWorkerReport, ServerError> {
768        self.park_lost_worker(
769            worker_id,
770            registry,
771            sink,
772            aion_core::WorkerDeathReason::Disconnect,
773        )
774    }
775
776    /// Park EVERY currently in-flight worker's tasks for restart recovery
777    /// (#207) — the drain-timeout backstop's bulk counterpart of
778    /// [`Self::fail_all_in_flight_workers`].
779    ///
780    /// Deregistered with the honest
781    /// [`WorkerDeathReason::Timeout`](aion_core::WorkerDeathReason::Timeout):
782    /// the drain window genuinely expired on these workers. Wakes drain waiters
783    /// after the sweep so `wait_for_empty` observes the emptied tracker.
784    ///
785    /// # Errors
786    ///
787    /// Returns registry, tracker, or sink errors without retrying or rescheduling activities.
788    pub fn park_all_in_flight_workers(
789        &self,
790        registry: &ConnectedWorkerRegistry,
791        sink: &impl ActivityCompletionSink,
792    ) -> Result<Vec<LostWorkerReport>, ServerError> {
793        let worker_ids = {
794            let state = self.state()?;
795            let mut worker_ids = state
796                .tasks
797                .values()
798                .map(|liveness| liveness.worker_id)
799                .collect::<HashSet<_>>()
800                .into_iter()
801                .collect::<Vec<_>>();
802            worker_ids.sort_unstable();
803            worker_ids
804        };
805        let mut reports = Vec::new();
806        for worker_id in worker_ids {
807            let report = self.park_lost_worker(
808                worker_id,
809                registry,
810                sink,
811                aion_core::WorkerDeathReason::Timeout,
812            )?;
813            if !report.tasks.is_empty() {
814                reports.push(report);
815            }
816        }
817        self.empty.notify_waiters();
818        Ok(reports)
819    }
820
821    /// Shared park core (#207), structured exactly like [`Self::fail_lost_worker`]
822    /// — deregister BEFORE collecting tasks (see that method's race note) — but
823    /// resolving each waiter with the ephemeral parked sentinel instead of
824    /// synthesizing a lost-worker `ActivityFailed`. Idempotent for the same
825    /// reasons: `deregister_with_reason` no-ops on an already-removed worker and
826    /// each task is removed as it parks, so a second sweep (park or fail) sees
827    /// an empty report and resolves nothing.
828    fn park_lost_worker(
829        &self,
830        worker_id: WorkerId,
831        registry: &ConnectedWorkerRegistry,
832        sink: &impl ActivityCompletionSink,
833        reason: aion_core::WorkerDeathReason,
834    ) -> Result<LostWorkerReport, ServerError> {
835        let task_queue = task_queue_of(registry, worker_id);
836        registry.deregister_with_reason(worker_id, reason)?;
837        self.state()?.connections.remove(&worker_id);
838        let tasks = self.remove_worker_tasks(worker_id)?;
839        for task in &tasks {
840            sink.park_activity(&task.workflow_id, &task.activity_id)?;
841            info!(
842                worker_id = ?worker_id,
843                workflow_id = %task.workflow_id,
844                activity_id = %task.activity_id,
845                "activity parked for restart recovery"
846            );
847        }
848        Ok(LostWorkerReport {
849            worker_id,
850            tasks,
851            task_queue,
852        })
853    }
854
855    fn fail_lost_worker(
856        &self,
857        worker_id: WorkerId,
858        registry: &ConnectedWorkerRegistry,
859        sink: &impl ActivityCompletionSink,
860    ) -> Result<LostWorkerReport, ServerError> {
861        // Deregister BEFORE collecting tasks: the dispatch path tracks its
862        // task, sends, and then checks `registry.is_registered`. With this
863        // ordering, a dispatch that still sees the worker registered is
864        // guaranteed its tracked task is visible to any later sweep, so the
865        // unbounded completion wait always gets a lost-worker failure. (The
866        // reverse order leaves a window where a task tracked between the
867        // collection and the deregistration is never failed by anyone.)
868        // This is the liveness-timeout sweep: the proven reason is Timeout, the
869        // one finer-grained WS3 distinction this call site can honestly assert.
870        // The queue is read BEFORE the deregistration below, because afterwards
871        // the handle is gone and the log could no longer name it.
872        let task_queue = task_queue_of(registry, worker_id);
873        registry.deregister_with_reason(worker_id, aion_core::WorkerDeathReason::Timeout)?;
874        self.state()?.connections.remove(&worker_id);
875        let tasks = self.remove_worker_tasks(worker_id)?;
876        for task in &tasks {
877            sink.complete_activity(ActivityCompletion {
878                workflow_id: task.workflow_id.clone(),
879                activity_id: task.activity_id.clone(),
880                run_id: None,
881                completion_token: task.completion_token.clone(),
882                // A TRANSPORT-domain loss, not an activity failure: the
883                // activity never executed to a result. The sink classifies it
884                // (and applies the transport's own re-dispatch budget); this
885                // sweep only reports what it observed.
886                outcome: ActivityCompletionOutcome::WorkerLost { worker_id },
887            })?;
888        }
889        Ok(LostWorkerReport {
890            worker_id,
891            tasks,
892            task_queue,
893        })
894    }
895
896    fn remove_worker_tasks(
897        &self,
898        worker_id: WorkerId,
899    ) -> Result<Vec<InFlightActivity>, ServerError> {
900        let mut state = self.state()?;
901        let keys = state
902            .tasks
903            .keys()
904            .filter(|key| key.worker_id() == worker_id)
905            .cloned()
906            .collect::<Vec<_>>();
907        let mut tasks = Vec::with_capacity(keys.len());
908        for key in keys {
909            if let Some(liveness) = state.tasks.remove(&key) {
910                tasks.push(InFlightActivity {
911                    workflow_id: liveness.workflow_id,
912                    activity_id: liveness.activity_id,
913                    attempt: liveness.attempt,
914                    completion_token: liveness.completion_token,
915                });
916            }
917        }
918        Ok(tasks)
919    }
920
921    /// Every tracked in-flight entry for `(workflow, activity, attempt)`.
922    ///
923    /// Read from the SAME map the heartbeat path writes, so what a reader is
924    /// told and what the dispatch path knows cannot drift. Usually zero or one
925    /// entry; a within-attempt failover can briefly have a dying owner and its
926    /// adopter both tracked, which is why this returns them all and leaves the
927    /// choice to the caller (`crate::worker::attempt_progress`) rather than
928    /// silently picking one here.
929    ///
930    /// # Errors
931    ///
932    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
933    pub(in crate::worker) fn attempt_entries(
934        &self,
935        workflow_id: &WorkflowId,
936        activity_id: &ActivityId,
937        attempt: u32,
938    ) -> Result<Vec<TaskLiveness>, ServerError> {
939        Ok(self
940            .state()?
941            .tasks
942            .values()
943            .filter(|liveness| {
944                &liveness.workflow_id == workflow_id
945                    && &liveness.activity_id == activity_id
946                    && liveness.attempt == attempt
947            })
948            .cloned()
949            .collect())
950    }
951
952    /// Every activity of `workflow_id` a connected worker is holding RIGHT NOW
953    /// (#233).
954    ///
955    /// A tracked entry means a live worker owns that activity: entries are
956    /// added at dispatch ([`Self::track_task`]) and removed on completion,
957    /// disconnect, or heartbeat expiry. That is what makes this — rather than
958    /// the durable outbox row, which records no worker at all — the authority
959    /// on where a cancellation has to be sent. The worker id here is present
960    /// membership, not a recorded historical fact, so it cannot go stale the
961    /// way a worker id written down at claim time would.
962    ///
963    /// A within-attempt failover can briefly show a dying owner AND its adopter
964    /// for one activity. Both are returned: a cancel is addressed to a
965    /// `(workflow, activity)` key that either worker may or may not still hold,
966    /// and a worker that does not hold it ignores the message, so asking both
967    /// is correct and asking one would be a guess.
968    ///
969    /// # Errors
970    ///
971    /// Returns [`ServerError::LockPoisoned`] if tracker state cannot be trusted.
972    pub fn in_flight_for_workflow(
973        &self,
974        workflow_id: &WorkflowId,
975    ) -> Result<Vec<TaskLiveness>, ServerError> {
976        Ok(self
977            .state()?
978            .tasks
979            .values()
980            .filter(|liveness| &liveness.workflow_id == workflow_id)
981            .cloned()
982            .collect())
983    }
984
985    fn state(&self) -> Result<MutexGuard<'_, HeartbeatState>, ServerError> {
986        self.inner
987            .lock()
988            .map_err(|_| ServerError::lock_poisoned("worker heartbeat tracker"))
989    }
990}
991
992/// Sweep cadence derived from the operator's `worker.heartbeat_window`: a
993/// quarter of the window, clamped to `[1s, window]` (the default 30s window
994/// sweeps every 7.5s).
995///
996/// Deliberately derived rather than a separate config knob: the window is the
997/// operational contract ("a silent worker is dead after this long"), and the
998/// sweep cadence is an implementation detail of enforcing it — a quarter-window
999/// cadence bounds detection latency at `window + window/4` while keeping the
1000/// sweep cheap. A window shorter than one second (test configurations) sweeps
1001/// once per window rather than sub-second-spinning, and a zero window is
1002/// floored at one millisecond because `tokio::time::interval` rejects a zero
1003/// period.
1004#[must_use]
1005pub fn sweep_interval(heartbeat_window: Duration) -> Duration {
1006    /// `tokio::time::interval` panics on a zero period, so even a
1007    /// (misconfigured) zero window gets a positive cadence.
1008    const MINIMUM_PERIOD: Duration = Duration::from_millis(1);
1009    /// Target lower bound: sweeping more often than once a second buys no
1010    /// meaningful detection latency against real heartbeat windows.
1011    const TARGET_FLOOR: Duration = Duration::from_secs(1);
1012    let ceiling = heartbeat_window.max(MINIMUM_PERIOD);
1013    // The floor never exceeds the ceiling, so `clamp` cannot panic.
1014    (heartbeat_window / 4).clamp(TARGET_FLOOR.min(ceiling), ceiling)
1015}
1016
1017/// Production driver of [`HeartbeatTracker::fail_expired_workers`] (#176).
1018///
1019/// The tracker records connection and per-task liveness, while the stream-teardown
1020/// sweep fails a worker whose stream ENDS. A worker whose stream stays open while
1021/// its process wedges is caught by the connection lease even when it is idle.
1022/// This interval task expires every silent connection or task, deregistering it
1023/// with the provable
1024/// [`WorkerDeathReason::Timeout`](aion_core::WorkerDeathReason::Timeout) and
1025/// surfacing its tasks as TRANSPORT losses through the shared completion sink
1026/// — the `lost:` class the engine re-dispatches attempt-neutrally, never the
1027/// action's retry vocabulary. It shares the server's shutdown watch, so it drains with
1028/// the transports (mirroring
1029/// [`OutboxDispatcher::run`](crate::worker::OutboxDispatcher::run)).
1030///
1031/// Double-fail safety: this sweep and the stream-teardown path
1032/// ([`HeartbeatTracker::fail_disconnected_worker`]) can both observe the same
1033/// dead worker. Both funnel into the same idempotent core —
1034/// `deregister_with_reason` is a no-op for an already-removed worker (no
1035/// duplicate WS3 delta, no metrics double-count) and the tracker removes each
1036/// task as it fails it — so whichever path runs second sees an empty report and
1037/// never double-completes an activity.
1038pub struct HeartbeatSweeper<S> {
1039    tracker: HeartbeatTracker,
1040    registry: ConnectedWorkerRegistry,
1041    sink: S,
1042    drain: DrainState,
1043    heartbeat_window: Duration,
1044    interval: Duration,
1045    /// Live unserved-queue state, read only to state the CONSEQUENCE of a
1046    /// deregistration in the same log line as its cause: how many dispatches are
1047    /// already parked on the queue the reaped worker was serving. Default-empty
1048    /// in wirings that have no queue service, where the count reads zero.
1049    queue_state: crate::worker::QueueServiceState,
1050}
1051
1052impl<S> HeartbeatSweeper<S>
1053where
1054    S: ActivityCompletionSink + Send + Sync + 'static,
1055{
1056    /// Build a sweeper over the server's shared liveness tracker, worker
1057    /// registry, completion sink, and drain gate. The cadence is derived from
1058    /// `heartbeat_window` by [`sweep_interval`].
1059    #[must_use]
1060    pub fn new(
1061        tracker: HeartbeatTracker,
1062        registry: ConnectedWorkerRegistry,
1063        sink: S,
1064        drain: DrainState,
1065        heartbeat_window: Duration,
1066    ) -> Self {
1067        let interval = sweep_interval(heartbeat_window);
1068        Self {
1069            tracker,
1070            registry,
1071            sink,
1072            drain,
1073            heartbeat_window,
1074            interval,
1075            queue_state: crate::worker::QueueServiceState::default(),
1076        }
1077    }
1078
1079    /// Share the live unserved-queue state so a deregistration log can state how
1080    /// many dispatches are already parked on the queue the reaped worker served.
1081    ///
1082    /// Without it the count reads zero — honest for a wiring with no queue
1083    /// service, and never a reason to withhold the deregistration itself.
1084    #[must_use]
1085    pub fn with_queue_state(mut self, queue_state: crate::worker::QueueServiceState) -> Self {
1086        self.queue_state = queue_state;
1087        self
1088    }
1089
1090    /// Run the expiry sweep until `shutdown` flips to `true`.
1091    ///
1092    /// A tracker/registry error during a sweep is logged and retried next tick
1093    /// rather than tearing the task down — a transient failure must not
1094    /// silently stop dead-worker detection. Shutdown is observed both while
1095    /// waiting for the next tick and re-checked before each sweep, exactly like
1096    /// the outbox dispatcher's run loop.
1097    pub async fn run(self, mut shutdown: watch::Receiver<bool>) {
1098        info!(
1099            sweep_interval_ms = self.interval.as_millis(),
1100            heartbeat_window_ms = self.heartbeat_window.as_millis(),
1101            "worker heartbeat sweeper started"
1102        );
1103        let mut ticks = tokio::time::interval(self.interval);
1104        ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
1105        loop {
1106            tokio::select! {
1107                _ = ticks.tick() => {
1108                    if *shutdown.borrow() {
1109                        break;
1110                    }
1111                    self.sweep_once(Instant::now());
1112                }
1113                changed = shutdown.changed() => {
1114                    // A receive error means every sender dropped; treat that as
1115                    // a shutdown request rather than spinning.
1116                    if changed.is_err() || *shutdown.borrow() {
1117                        break;
1118                    }
1119                }
1120            }
1121        }
1122        info!("worker heartbeat sweeper stopped");
1123    }
1124
1125    /// Fail every currently-expired worker once, logging each lost-worker
1126    /// report at warn (mirroring the stream-teardown sweep's logging).
1127    fn sweep_once(&self, now: Instant) {
1128        let reports = match self
1129            .tracker
1130            .fail_expired_workers(&self.registry, &self.sink, now)
1131        {
1132            Ok(reports) => reports,
1133            Err(sweep_error) => {
1134                error!(
1135                    error = %sweep_error,
1136                    "heartbeat expiry sweep failed; retrying next tick"
1137                );
1138                return;
1139            }
1140        };
1141        for report in &reports {
1142            let task_queue = report.task_queue.as_deref().unwrap_or("<unregistered>");
1143            // The consequence, stated with the cause: a queue whose worker just
1144            // died and which already holds parked dispatches is an alertable
1145            // condition, and the operator should not have to join two log lines
1146            // to see it. A read failure reports `None` rather than suppressing
1147            // the deregistration line.
1148            let parked = report
1149                .task_queue
1150                .as_deref()
1151                .map(|queue| self.queue_state.parked_on_queue(queue))
1152                .transpose()
1153                .unwrap_or_else(|error| {
1154                    error!(%error, "could not read parked-dispatch count for a reaped worker");
1155                    None
1156                })
1157                .unwrap_or(0);
1158            if report.tasks.is_empty() {
1159                warn!(
1160                    worker_id = ?report.worker_id,
1161                    task_queue,
1162                    parked_dispatches = parked,
1163                    heartbeat_window_ms = self.heartbeat_window.as_millis(),
1164                    "idle worker connection lease expired; worker deregistered"
1165                );
1166            } else {
1167                warn!(
1168                    worker_id = ?report.worker_id,
1169                    task_queue,
1170                    parked_dispatches = parked,
1171                    failed_tasks = report.tasks.len(),
1172                    heartbeat_window_ms = self.heartbeat_window.as_millis(),
1173                    "worker heartbeat window expired with in-flight activities; \
1174                     deregistered and surfaced as transport losses, to be \
1175                     re-dispatched attempt-neutrally"
1176                );
1177            }
1178        }
1179        if !reports.is_empty() {
1180            // In-flight accounting may have just reached zero; wake any drain
1181            // waiter so shutdown does not sit out its full timeout (mirrors
1182            // the stream-teardown sweep).
1183            self.drain.notify_activity_drained();
1184        }
1185    }
1186}
1187
1188impl TaskKey {
1189    fn new(worker_id: WorkerId, workflow_id: WorkflowId, activity_id: ActivityId) -> Self {
1190        Self(worker_id, workflow_id, activity_id)
1191    }
1192
1193    const fn worker_id(&self) -> WorkerId {
1194        self.0
1195    }
1196}
1197
1198struct DecodedHeartbeat {
1199    workflow_id: WorkflowId,
1200    activity_id: ActivityId,
1201    progress: Option<Payload>,
1202}
1203
1204impl TryFrom<ProtoHeartbeat> for DecodedHeartbeat {
1205    type Error = ServerError;
1206
1207    fn try_from(value: ProtoHeartbeat) -> Result<Self, Self::Error> {
1208        let workflow_id = value
1209            .workflow_id
1210            .ok_or_else(|| wire_error("heartbeat workflow id is missing"))
1211            .and_then(|id| WorkflowId::try_from(id).map_err(ServerError::from))?;
1212        let activity_id = value
1213            .activity_id
1214            .ok_or_else(|| wire_error("heartbeat activity id is missing"))
1215            .map(ActivityId::from)?;
1216        let progress = value
1217            .progress
1218            .map(Payload::try_from)
1219            .transpose()
1220            .map_err(ServerError::from)?;
1221        Ok(Self {
1222            workflow_id,
1223            activity_id,
1224            progress,
1225        })
1226    }
1227}
1228
1229/// The task queue a still-registered worker serves, for the deregistration log.
1230///
1231/// A registry read failure (poisoned lock) yields `None` rather than aborting
1232/// the sweep: losing the queue NAME must never stop a dead worker being reaped.
1233fn task_queue_of(registry: &ConnectedWorkerRegistry, worker_id: WorkerId) -> Option<String> {
1234    registry
1235        .worker_by_id(worker_id)
1236        .ok()
1237        .flatten()
1238        .map(|handle| handle.task_queue().to_owned())
1239}
1240
1241fn is_expired(liveness: &TaskLiveness, now: Instant) -> bool {
1242    now.checked_duration_since(liveness.last_heartbeat_at)
1243        .is_some_and(|elapsed| elapsed > liveness.heartbeat_window)
1244}
1245
1246fn wire_error(message: &'static str) -> ServerError {
1247    ServerError::Wire {
1248        wire: WireError::backend(message),
1249    }
1250}
1251
1252#[cfg(test)]
1253mod reachability_tests {
1254    use std::time::{Duration, Instant};
1255
1256    use super::{
1257        DISPATCH_PROBATION_PINGS, DispatchExclusion, ExcludedWorker, HeartbeatTracker, ServerError,
1258        WorkerId,
1259    };
1260
1261    const WINDOW: Duration = Duration::from_secs(30);
1262
1263    /// Every test returns `Result` and uses `?` rather than unwrapping: a lock
1264    /// fault inside the tracker is a real failure mode of the code under test,
1265    /// and it should surface as a failed test carrying the typed error, not as a
1266    /// panic message written by the test.
1267    type TestResult = Result<(), ServerError>;
1268
1269    fn tracker_with_worker(now: Instant) -> Result<(HeartbeatTracker, WorkerId), ServerError> {
1270        let tracker = HeartbeatTracker::new(WINDOW);
1271        let worker = WorkerId::from_value(1);
1272        tracker.register_connection(worker, now)?;
1273        Ok((tracker, worker))
1274    }
1275
1276    /// Answer the full probation, so the worker is genuinely eligible. Tests
1277    /// about staleness, pump beats, or withdrawal must start from a worker that
1278    /// HAS eligibility — otherwise they pass on a worker that never had any and
1279    /// prove nothing about the behaviour they name.
1280    fn serve_probation(
1281        tracker: &HeartbeatTracker,
1282        worker: WorkerId,
1283        at: Instant,
1284    ) -> Result<(), ServerError> {
1285        for _ in 0..DISPATCH_PROBATION_PINGS {
1286            assert!(
1287                tracker.record_dispatch_reachability(worker, at)?,
1288                "the worker must still be tracked while it serves its probation"
1289            );
1290        }
1291        Ok(())
1292    }
1293
1294    /// THE REGRESSION. This is the defect that made run `dfd2117c` invisible:
1295    /// the worker's liveness pump beat from a background task, refreshed the one
1296    /// shared lease, and the dead-man switch could not fire while the server had
1297    /// been unable to push to that worker for fifteen minutes.
1298    ///
1299    /// An inbound frame must prove the worker is ALIVE and must NOT prove the
1300    /// server can REACH it.
1301    #[test]
1302    fn an_inbound_frame_cannot_prove_dispatch_reachability() -> TestResult {
1303        let start = Instant::now();
1304        let (tracker, worker) = tracker_with_worker(start)?;
1305        // The worker STARTS eligible, earned honestly. Without this the test
1306        // would pass on a worker that never had eligibility to lose, which says
1307        // nothing about whether a pump beat can preserve it.
1308        serve_probation(&tracker, worker, start)?;
1309        assert!(
1310            tracker.is_dispatch_reachable(worker, start)?,
1311            "precondition: the worker is eligible before the connection goes one-way"
1312        );
1313
1314        // Well past the window, with the pump beating throughout — exactly what
1315        // a busy worker on a poisoned connection looks like.
1316        let much_later = start + WINDOW * 4;
1317        assert!(
1318            tracker.record_connection_activity(worker, much_later)?,
1319            "the worker is still tracked"
1320        );
1321
1322        assert!(
1323            !tracker.is_dispatch_reachable(worker, much_later)?,
1324            "a pump beat must NOT make a worker the server cannot push to look reachable"
1325        );
1326        assert_eq!(
1327            tracker.unreachable_workers(much_later)?,
1328            vec![ExcludedWorker {
1329                worker_id: worker,
1330                // It HELD eligibility (served above) and lost it to a stale
1331                // proof. Classifying this as an opening probation would tell an
1332                // operator a poisoned connection is an ordinary worker start.
1333                exclusion: DispatchExclusion::ReachabilityLost,
1334            }],
1335            "the worker must be named unreachable however alive its process looks"
1336        );
1337        Ok(())
1338    }
1339
1340    /// The control for the test above: without it, a tracker that reported
1341    /// EVERYTHING unreachable would satisfy that assertion and prove nothing.
1342    #[test]
1343    fn an_answered_ping_does_prove_dispatch_reachability() -> TestResult {
1344        let start = Instant::now();
1345        let (tracker, worker) = tracker_with_worker(start)?;
1346
1347        let much_later = start + WINDOW * 4;
1348        serve_probation(&tracker, worker, much_later)?;
1349
1350        assert!(
1351            tracker.is_dispatch_reachable(worker, much_later)?,
1352            "answered pings are the one thing that proves the push leg works"
1353        );
1354        assert!(
1355            tracker.unreachable_workers(much_later)?.is_empty(),
1356            "a worker answering pings is never unreachable"
1357        );
1358        Ok(())
1359    }
1360
1361    /// Registration opens a PROBATION and grants nothing. The handshake ack is
1362    /// SENT by this server; nothing reports that it was RECEIVED, so a
1363    /// connection is a channel, not proof that the channel carries. Eligibility
1364    /// is earned by answered pings only.
1365    #[test]
1366    fn registration_opens_a_probation_and_does_not_grant_eligibility() -> TestResult {
1367        let start = Instant::now();
1368        let (tracker, worker) = tracker_with_worker(start)?;
1369
1370        assert!(
1371            !tracker.is_dispatch_reachable(worker, start)?,
1372            "a brand-new connection has proved nothing about the push leg"
1373        );
1374        assert_eq!(
1375            tracker.unreachable_workers(start)?,
1376            vec![ExcludedWorker {
1377                worker_id: worker,
1378                // And it is carried as a PROBATION, not as a reachability
1379                // failure. This is the distinction that stopped an ordinary
1380                // worker start from being announced to the operator as an
1381                // unreachable dispatch path.
1382                exclusion: DispatchExclusion::OpeningProbation { answers: 0 },
1383            }],
1384            "a worker serving its probation is carried in the census as unreachable"
1385        );
1386        Ok(())
1387    }
1388
1389    /// The case that produced a FALSE ALARM on every healthy worker start.
1390    ///
1391    /// One answer banked out of two: the server has demonstrably reached this
1392    /// worker — moments ago — and is merely waiting for the second consecutive
1393    /// answer. Reporting that as a reachability failure told Tom's operator log
1394    /// his worker's dispatch path was dead when the opposite had just been
1395    /// measured. The exclusion is real; the REASON is an opening probation.
1396    #[test]
1397    fn a_part_served_probation_is_a_probation_and_not_a_reachability_failure() -> TestResult {
1398        let start = Instant::now();
1399        let (tracker, worker) = tracker_with_worker(start)?;
1400        const {
1401            assert!(
1402                DISPATCH_PROBATION_PINGS > 1,
1403                "this test is only meaningful while the probation takes more than one answer"
1404            );
1405        }
1406        assert!(tracker.record_dispatch_reachability(worker, start)?);
1407
1408        assert_eq!(
1409            tracker.unreachable_workers(start)?,
1410            vec![ExcludedWorker {
1411                worker_id: worker,
1412                exclusion: DispatchExclusion::OpeningProbation { answers: 1 },
1413            }],
1414            "a worker that has answered part of its opening probation is still excluded, but it \
1415             must not be described as one the server cannot reach — it answered"
1416        );
1417        Ok(())
1418    }
1419
1420    /// The other side of the same discrimination, and the control for the test
1421    /// above: once eligibility has actually been HELD, losing it is an incident
1422    /// and must classify differently. Without this, a classifier that answered
1423    /// `OpeningProbation` unconditionally would satisfy the test above.
1424    #[test]
1425    fn losing_held_eligibility_is_reported_as_a_loss_not_as_a_fresh_probation() -> TestResult {
1426        let start = Instant::now();
1427        let (tracker, worker) = tracker_with_worker(start)?;
1428        serve_probation(&tracker, worker, start)?;
1429        assert!(
1430            tracker.is_dispatch_reachable(worker, start)?,
1431            "precondition: eligibility was genuinely held before it was lost"
1432        );
1433
1434        assert!(
1435            tracker.record_dispatch_unreachable(worker)?,
1436            "the worker is still tracked when its ping fails"
1437        );
1438
1439        assert_eq!(
1440            tracker.unreachable_workers(start)?,
1441            vec![ExcludedWorker {
1442                worker_id: worker,
1443                exclusion: DispatchExclusion::ReachabilityLost,
1444            }],
1445            "a failed ping on a worker that HAD eligibility is an incident, and must not be \
1446             filed as the ordinary probation every fresh connection serves"
1447        );
1448        Ok(())
1449    }
1450
1451    /// A REDIAL is a new measurement. The previous connection's proof must not
1452    /// make the new connection's ordinary probation look like an incident —
1453    /// otherwise every reconnect of a healthy worker would raise the alarm that
1454    /// is supposed to mean something has gone wrong.
1455    #[test]
1456    fn a_reconnect_starts_a_fresh_probation_not_a_lost_eligibility() -> TestResult {
1457        let start = Instant::now();
1458        let (tracker, worker) = tracker_with_worker(start)?;
1459        serve_probation(&tracker, worker, start)?;
1460
1461        tracker.unregister_connection(worker)?;
1462        tracker.register_connection(worker, start)?;
1463
1464        assert_eq!(
1465            tracker.unreachable_workers(start)?,
1466            vec![ExcludedWorker {
1467                worker_id: worker,
1468                exclusion: DispatchExclusion::OpeningProbation { answers: 0 },
1469            }],
1470            "nothing earned on the old connection carries across to the new one"
1471        );
1472        Ok(())
1473    }
1474
1475    /// The probation must be SERVED IN FULL. One answered ping can be luck — a
1476    /// link that answers one probe in three would otherwise accrue eligibility
1477    /// and then flap. This pins the boundary from below: K-1 answers is not
1478    /// enough, and the very next one is.
1479    #[test]
1480    fn one_ping_short_of_the_probation_earns_nothing() -> TestResult {
1481        let start = Instant::now();
1482        let (tracker, worker) = tracker_with_worker(start)?;
1483
1484        for _ in 0..DISPATCH_PROBATION_PINGS - 1 {
1485            assert!(tracker.record_dispatch_reachability(worker, start)?);
1486            assert!(
1487                !tracker.is_dispatch_reachable(worker, start)?,
1488                "eligibility must not be granted before the probation is served in full"
1489            );
1490        }
1491
1492        assert!(tracker.record_dispatch_reachability(worker, start)?);
1493        assert!(
1494            tracker.is_dispatch_reachable(worker, start)?,
1495            "the ping that completes the probation must grant eligibility — otherwise this test \
1496             would pass on a tracker that never grants it at all"
1497        );
1498        Ok(())
1499    }
1500
1501    /// A failed probe RESETS the run. Eligibility is withdrawn immediately, not
1502    /// when the window later expires: an unanswered probe is direct evidence
1503    /// about the push leg, and direct negative evidence must weigh at least as
1504    /// much as silence.
1505    #[test]
1506    fn a_failed_probe_withdraws_eligibility_at_once_and_restarts_the_probation() -> TestResult {
1507        let start = Instant::now();
1508        let (tracker, worker) = tracker_with_worker(start)?;
1509        serve_probation(&tracker, worker, start)?;
1510        assert!(
1511            tracker.is_dispatch_reachable(worker, start)?,
1512            "precondition"
1513        );
1514
1515        assert!(
1516            tracker.record_dispatch_unreachable(worker)?,
1517            "the worker is still tracked"
1518        );
1519        assert!(
1520            !tracker.is_dispatch_reachable(worker, start)?,
1521            "a failed probe withdraws eligibility on the spot, inside the window"
1522        );
1523
1524        // And the run restarts from zero rather than resuming: one answer does
1525        // not restore what a full probation earned.
1526        assert!(tracker.record_dispatch_reachability(worker, start)?);
1527        assert!(
1528            !tracker.is_dispatch_reachable(worker, start)?,
1529            "a single answer after a failure must not restore eligibility"
1530        );
1531        Ok(())
1532    }
1533
1534    /// 🔴 THE FLAPPING PIN. A link that answers every other probe must NEVER
1535    /// become eligible. Cumulative counting would let it accrue, and eligibility
1536    /// would switch on and off under a running fleet — the intermittent evidence
1537    /// that costs hours to attribute. Consecutiveness is what forbids it.
1538    #[test]
1539    fn a_link_that_answers_every_other_probe_never_becomes_eligible() -> TestResult {
1540        let start = Instant::now();
1541        let (tracker, worker) = tracker_with_worker(start)?;
1542
1543        // Far more probes than the probation demands, alternating.
1544        for probe in 0..DISPATCH_PROBATION_PINGS * 10 {
1545            let now = start + Duration::from_millis(u64::from(probe));
1546            if probe % 2 == 0 {
1547                assert!(tracker.record_dispatch_reachability(worker, now)?);
1548            } else {
1549                assert!(tracker.record_dispatch_unreachable(worker)?);
1550            }
1551            assert!(
1552                !tracker.is_dispatch_reachable(worker, now)?,
1553                "a flapping link must never hold dispatch eligibility, at any probe (probe {probe})"
1554            );
1555        }
1556
1557        // The control: the same worker, answering consecutively, DOES become
1558        // eligible — so this test cannot pass on a tracker that grants nothing.
1559        let now = start + Duration::from_secs(1);
1560        serve_probation(&tracker, worker, now)?;
1561        assert!(
1562            tracker.is_dispatch_reachable(worker, now)?,
1563            "consecutive answers must still earn eligibility"
1564        );
1565        Ok(())
1566    }
1567
1568    /// Reachability must EXPIRE on its own clock. If it were only ever advanced
1569    /// and never allowed to go stale, eligibility could never be withdrawn.
1570    #[test]
1571    fn reachability_goes_stale_once_the_window_passes() -> TestResult {
1572        let start = Instant::now();
1573        let (tracker, worker) = tracker_with_worker(start)?;
1574        serve_probation(&tracker, worker, start)?;
1575
1576        assert!(
1577            tracker.is_dispatch_reachable(worker, start + WINDOW)?,
1578            "still inside the window"
1579        );
1580        assert!(
1581            !tracker.is_dispatch_reachable(worker, start + WINDOW + Duration::from_millis(1))?,
1582            "one millisecond past the window is stale"
1583        );
1584        Ok(())
1585    }
1586
1587    /// An untracked worker is not reachable: absence of proof is not proof. A
1588    /// pong racing a reap must not resurrect it either.
1589    #[test]
1590    fn an_unregistered_worker_is_never_reachable_and_cannot_be_resurrected() -> TestResult {
1591        let start = Instant::now();
1592        let (tracker, worker) = tracker_with_worker(start)?;
1593        tracker.unregister_connection(worker)?;
1594
1595        assert!(
1596            !tracker.is_dispatch_reachable(worker, start)?,
1597            "a deregistered worker is not reachable"
1598        );
1599        assert!(
1600            !tracker.record_dispatch_reachability(worker, start)?,
1601            "a late pong must not resurrect a deregistered worker"
1602        );
1603        assert!(
1604            !tracker.record_dispatch_unreachable(worker)?,
1605            "a late probe FAILURE must not resurrect a deregistered worker either — the reset \
1606             path allocates an entry, so it has to refuse an untracked worker as firmly as the \
1607             success path does"
1608        );
1609        assert!(
1610            tracker.unreachable_workers(start)?.is_empty(),
1611            "an untracked worker is not carried in the census either"
1612        );
1613        Ok(())
1614    }
1615}
1616
1617#[cfg(test)]
1618mod tests {
1619    use std::sync::Mutex;
1620
1621    use aion_core::ContentType;
1622    use aion_proto::{ProtoActivityId, ProtoPayload, ProtoWorkflowId};
1623    use serde_json::json;
1624    use uuid::Uuid;
1625
1626    use crate::worker::registry::WorkerRegistration;
1627
1628    use super::*;
1629
1630    #[derive(Default)]
1631    struct RecordingSink {
1632        completions: Mutex<Vec<ActivityCompletion>>,
1633        parks: Mutex<Vec<(WorkflowId, ActivityId)>>,
1634    }
1635
1636    impl ActivityCompletionSink for RecordingSink {
1637        fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
1638            self.completions
1639                .lock()
1640                .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1641                .push(completion);
1642            Ok(())
1643        }
1644
1645        fn park_activity(
1646            &self,
1647            workflow_id: &WorkflowId,
1648            activity_id: &ActivityId,
1649        ) -> Result<(), ServerError> {
1650            self.parks
1651                .lock()
1652                .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1653                .push((workflow_id.clone(), activity_id.clone()));
1654            Ok(())
1655        }
1656    }
1657
1658    fn workflow_id() -> WorkflowId {
1659        WorkflowId::new(Uuid::nil())
1660    }
1661
1662    fn activity_id(position: u64) -> ActivityId {
1663        ActivityId::from_sequence_position(position)
1664    }
1665
1666    fn payload(value: &serde_json::Value) -> Result<Payload, Box<dyn std::error::Error>> {
1667        Ok(Payload::from_json(value)?)
1668    }
1669
1670    fn heartbeat(
1671        workflow_id: WorkflowId,
1672        activity_id: ActivityId,
1673        progress: Option<Payload>,
1674    ) -> ProtoHeartbeat {
1675        ProtoHeartbeat {
1676            workflow_id: Some(ProtoWorkflowId::from(workflow_id)),
1677            activity_id: Some(ProtoActivityId::from(activity_id)),
1678            progress: progress.map(ProtoPayload::from),
1679        }
1680    }
1681
1682    fn registry_with_worker()
1683    -> Result<(ConnectedWorkerRegistry, WorkerRegistration, WorkerId), ServerError> {
1684        let registry = ConnectedWorkerRegistry::default();
1685        let (tx, _rx) = tokio::sync::mpsc::channel(1);
1686        let activity_types = [String::from("charge-card")];
1687        let registration = registry.register("tenant-a", activity_types.iter(), tx)?;
1688        let worker_id = registration
1689            .worker_id()
1690            .ok_or_else(|| ServerError::lock_poisoned("test worker registration"))?;
1691        Ok((registry, registration, worker_id))
1692    }
1693
1694    #[test]
1695    fn heartbeat_refresh_keeps_task_live_across_window() -> Result<(), Box<dyn std::error::Error>> {
1696        let window = Duration::from_secs(5);
1697        let tracker = HeartbeatTracker::new(window);
1698        let worker_id = WorkerIdForTest::registered()?;
1699        let workflow_id = workflow_id();
1700        let activity_id = activity_id(10);
1701        let start = Instant::now();
1702
1703        tracker.track_task(
1704            worker_id,
1705            InFlightActivity {
1706                workflow_id: workflow_id.clone(),
1707                activity_id: activity_id.clone(),
1708                attempt: 1,
1709                completion_token: crate::worker::CompletionToken::for_test(),
1710            },
1711            start,
1712        )?;
1713        assert!(tracker.is_live(worker_id, &workflow_id, &activity_id, start + window)?);
1714
1715        let progress = payload(&json!({"percent": 50}))?;
1716        let update = tracker.record_heartbeat(
1717            worker_id,
1718            heartbeat(
1719                workflow_id.clone(),
1720                activity_id.clone(),
1721                Some(progress.clone()),
1722            ),
1723            start + window,
1724        )?;
1725
1726        assert_eq!(update.liveness.last_progress, Some(progress));
1727        assert!(tracker.is_live(
1728            worker_id,
1729            &workflow_id,
1730            &activity_id,
1731            start + window + window
1732        )?);
1733        assert!(tracker.expired_workers(start + window + window)?.is_empty());
1734        Ok(())
1735    }
1736
1737    #[test]
1738    fn missed_heartbeat_deregisters_worker_and_fails_in_flight_once()
1739    -> Result<(), Box<dyn std::error::Error>> {
1740        let (registry, _registration, worker_id) = registry_with_worker()?;
1741        let sink = RecordingSink::default();
1742        let tracker = HeartbeatTracker::new(Duration::from_secs(5));
1743        let workflow_id = workflow_id();
1744        let activity_id = activity_id(11);
1745        let start = Instant::now();
1746
1747        tracker.track_task(
1748            worker_id,
1749            InFlightActivity {
1750                workflow_id: workflow_id.clone(),
1751                activity_id: activity_id.clone(),
1752                attempt: 1,
1753                completion_token: crate::worker::CompletionToken::for_test(),
1754            },
1755            start,
1756        )?;
1757
1758        let reports =
1759            tracker.fail_expired_workers(&registry, &sink, start + Duration::from_secs(6))?;
1760        assert_eq!(reports.len(), 1);
1761        assert_eq!(reports[0].worker_id, worker_id);
1762        assert_eq!(reports[0].tasks.len(), 1);
1763        assert!(
1764            registry
1765                .workers_for("tenant-a", "default", "charge-card", None)?
1766                .is_empty()
1767        );
1768
1769        let second = tracker.fail_disconnected_worker(worker_id, &registry, &sink)?;
1770        assert!(second.tasks.is_empty());
1771        let completions = sink
1772            .completions
1773            .lock()
1774            .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
1775        assert_eq!(completions.len(), 1);
1776        assert_eq!(completions[0].workflow_id, workflow_id);
1777        assert_eq!(completions[0].activity_id, activity_id);
1778        // The sweep reports a TRANSPORT-domain loss, not an activity failure:
1779        // the activity never executed to a result, so the sink (not this sweep)
1780        // classifies it and applies the transport's own re-dispatch budget.
1781        // Before this distinction existed the sweep synthesized a `Retryable`
1782        // `ActivityError` that the engine then delivered as a TERMINAL failure
1783        // whenever the activity carried no authored retry policy.
1784        match &completions[0].outcome {
1785            ActivityCompletionOutcome::WorkerLost { worker_id: lost } => {
1786                assert_eq!(*lost, worker_id);
1787            }
1788            other => {
1789                return Err(format!("expected a lost-worker outcome, got {other:?}").into());
1790            }
1791        }
1792        Ok(())
1793    }
1794
1795    #[test]
1796    fn disconnected_worker_fails_each_in_flight_task_once() -> Result<(), Box<dyn std::error::Error>>
1797    {
1798        let (registry, _registration, worker_id) = registry_with_worker()?;
1799        let sink = RecordingSink::default();
1800        let tracker = HeartbeatTracker::new(Duration::from_secs(5));
1801        let workflow_id = workflow_id();
1802        let start = Instant::now();
1803
1804        tracker.track_task(
1805            worker_id,
1806            InFlightActivity {
1807                workflow_id: workflow_id.clone(),
1808                activity_id: activity_id(21),
1809                attempt: 1,
1810                completion_token: crate::worker::CompletionToken::for_test(),
1811            },
1812            start,
1813        )?;
1814        tracker.track_task(
1815            worker_id,
1816            InFlightActivity {
1817                workflow_id,
1818                activity_id: activity_id(22),
1819                attempt: 1,
1820                completion_token: crate::worker::CompletionToken::for_test(),
1821            },
1822            start,
1823        )?;
1824
1825        let report = tracker.fail_disconnected_worker(worker_id, &registry, &sink)?;
1826        assert_eq!(report.tasks.len(), 2);
1827        assert!(
1828            registry
1829                .workers_for("tenant-a", "default", "charge-card", None)?
1830                .is_empty()
1831        );
1832
1833        let completions = sink
1834            .completions
1835            .lock()
1836            .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
1837        assert_eq!(completions.len(), 2);
1838        assert!(completions.iter().all(|completion| matches!(
1839            &completion.outcome,
1840            ActivityCompletionOutcome::WorkerLost { .. }
1841        )));
1842        Ok(())
1843    }
1844
1845    /// #207: parking a drain-disconnected worker removes its tasks, deregisters
1846    /// it, and PARKS each task through the sink — zero completions synthesized,
1847    /// so the durable log stays byte-equivalent to a kill -9. A second park (or
1848    /// a later fail sweep) finds nothing: the idempotent-deregister discipline
1849    /// the fail path already proves holds for parks too.
1850    #[test]
1851    fn park_disconnected_worker_parks_tasks_without_synthesizing_completions()
1852    -> Result<(), Box<dyn std::error::Error>> {
1853        let (registry, _registration, worker_id) = registry_with_worker()?;
1854        let sink = RecordingSink::default();
1855        let tracker = HeartbeatTracker::new(Duration::from_secs(5));
1856        let workflow_id = workflow_id();
1857        let start = Instant::now();
1858        tracker.track_task(
1859            worker_id,
1860            InFlightActivity {
1861                workflow_id: workflow_id.clone(),
1862                activity_id: activity_id(60),
1863                attempt: 1,
1864                completion_token: crate::worker::CompletionToken::for_test(),
1865            },
1866            start,
1867        )?;
1868        tracker.track_task(
1869            worker_id,
1870            InFlightActivity {
1871                workflow_id: workflow_id.clone(),
1872                activity_id: activity_id(61),
1873                attempt: 1,
1874                completion_token: crate::worker::CompletionToken::for_test(),
1875            },
1876            start,
1877        )?;
1878
1879        let report = tracker.park_disconnected_worker(worker_id, &registry, &sink)?;
1880        assert_eq!(report.tasks.len(), 2);
1881        assert_eq!(
1882            tracker.in_flight_count()?,
1883            0,
1884            "parking must remove every tracked task so drain accounting reaches zero"
1885        );
1886        assert!(
1887            registry
1888                .workers_for("tenant-a", "default", "charge-card", None)?
1889                .is_empty(),
1890            "the parked worker must be deregistered from routing"
1891        );
1892        let parks = sink
1893            .parks
1894            .lock()
1895            .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
1896        assert_eq!(parks.len(), 2, "each task must be parked exactly once");
1897        drop(parks);
1898        assert!(
1899            sink.completions
1900                .lock()
1901                .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1902                .is_empty(),
1903            "parking must never synthesize an activity completion"
1904        );
1905
1906        // Double-park and park-after-fail are no-ops: the idempotent core.
1907        let second = tracker.park_disconnected_worker(worker_id, &registry, &sink)?;
1908        assert!(second.tasks.is_empty());
1909        let third = tracker.fail_disconnected_worker(worker_id, &registry, &sink)?;
1910        assert!(third.tasks.is_empty());
1911        assert_eq!(
1912            sink.parks
1913                .lock()
1914                .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1915                .len(),
1916            2,
1917            "re-sweeping a parked worker must park nothing further"
1918        );
1919        assert!(
1920            sink.completions
1921                .lock()
1922                .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1923                .is_empty(),
1924            "a fail sweep after the park must fail nothing"
1925        );
1926        Ok(())
1927    }
1928
1929    /// #207 drain-timeout backstop: the bulk park removes every worker's tasks,
1930    /// parks each through the sink, and wakes drain waiters — never
1931    /// synthesizing a completion.
1932    #[tokio::test]
1933    async fn park_all_in_flight_workers_parks_everything_and_wakes_drain_waiters()
1934    -> Result<(), Box<dyn std::error::Error>> {
1935        let (registry, _registration, worker_id) = registry_with_worker()?;
1936        let sink = RecordingSink::default();
1937        let tracker = HeartbeatTracker::new(Duration::from_secs(5));
1938        let workflow_id = workflow_id();
1939        tracker.track_task(
1940            worker_id,
1941            InFlightActivity {
1942                workflow_id: workflow_id.clone(),
1943                activity_id: activity_id(70),
1944                attempt: 1,
1945                completion_token: crate::worker::CompletionToken::for_test(),
1946            },
1947            Instant::now(),
1948        )?;
1949        // Arm a waiter on the tracker's empty notify BEFORE the bulk park.
1950        let notified = tracker.empty.notified();
1951        tokio::pin!(notified);
1952
1953        let reports = tracker.park_all_in_flight_workers(&registry, &sink)?;
1954        assert_eq!(reports.len(), 1);
1955        assert_eq!(reports[0].worker_id, worker_id);
1956        assert_eq!(reports[0].tasks.len(), 1);
1957        assert_eq!(tracker.in_flight_count()?, 0);
1958        assert_eq!(
1959            sink.parks
1960                .lock()
1961                .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1962                .len(),
1963            1
1964        );
1965        assert!(
1966            sink.completions
1967                .lock()
1968                .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1969                .is_empty(),
1970            "the bulk park must never synthesize a completion"
1971        );
1972        tokio::time::timeout(Duration::from_millis(200), notified)
1973            .await
1974            .map_err(|_| "the bulk park must wake drain waiters")?;
1975        Ok(())
1976    }
1977
1978    /// The worker runtime's AUTOMATIC liveness beats carry no payload and
1979    /// interleave with explicit handler progress heartbeats: a payload-free
1980    /// beat must refresh the liveness stamp WITHOUT erasing the handler's
1981    /// most recent progress report.
1982    #[test]
1983    fn payload_free_heartbeat_refreshes_liveness_without_clearing_progress()
1984    -> Result<(), Box<dyn std::error::Error>> {
1985        let window = Duration::from_secs(5);
1986        let tracker = HeartbeatTracker::new(window);
1987        let worker_id = WorkerIdForTest::registered()?;
1988        let workflow_id = workflow_id();
1989        let activity_id = activity_id(12);
1990        let start = Instant::now();
1991
1992        tracker.track_task(
1993            worker_id,
1994            InFlightActivity {
1995                workflow_id: workflow_id.clone(),
1996                activity_id: activity_id.clone(),
1997                attempt: 1,
1998                completion_token: crate::worker::CompletionToken::for_test(),
1999            },
2000            start,
2001        )?;
2002        let progress = payload(&json!({"percent": 80}))?;
2003        tracker.record_heartbeat(
2004            worker_id,
2005            heartbeat(
2006                workflow_id.clone(),
2007                activity_id.clone(),
2008                Some(progress.clone()),
2009            ),
2010            start + Duration::from_secs(1),
2011        )?;
2012
2013        // An automatic liveness beat: no payload, later timestamp.
2014        let update = tracker.record_heartbeat(
2015            worker_id,
2016            heartbeat(workflow_id.clone(), activity_id.clone(), None),
2017            start + Duration::from_secs(4),
2018        )?;
2019
2020        assert_eq!(
2021            update.liveness.last_progress,
2022            Some(progress),
2023            "a payload-free liveness beat must not erase handler progress"
2024        );
2025        assert!(
2026            tracker.is_live(
2027                worker_id,
2028                &workflow_id,
2029                &activity_id,
2030                start + Duration::from_secs(8)
2031            )?,
2032            "the payload-free beat must still refresh the liveness stamp"
2033        );
2034        Ok(())
2035    }
2036
2037    #[test]
2038    fn malformed_heartbeat_missing_ids_is_wire_error() -> Result<(), Box<dyn std::error::Error>> {
2039        let worker_id = WorkerIdForTest::registered()?;
2040        let tracker = HeartbeatTracker::new(Duration::from_secs(5));
2041        let missing = ProtoHeartbeat {
2042            workflow_id: None,
2043            activity_id: Some(ProtoActivityId::from(activity_id(30))),
2044            progress: None,
2045        };
2046
2047        let result = tracker.record_heartbeat(worker_id, missing, Instant::now());
2048        assert!(matches!(result, Err(ServerError::Wire { .. })));
2049        Ok(())
2050    }
2051
2052    #[test]
2053    fn heartbeat_progress_is_not_reported_as_activity_result()
2054    -> Result<(), Box<dyn std::error::Error>> {
2055        let sink = RecordingSink::default();
2056        let worker_id = WorkerIdForTest::registered()?;
2057        let tracker = HeartbeatTracker::new(Duration::from_secs(5));
2058        let workflow_id = workflow_id();
2059        let activity_id = activity_id(40);
2060        let now = Instant::now();
2061
2062        tracker.track_task(
2063            worker_id,
2064            InFlightActivity {
2065                workflow_id: workflow_id.clone(),
2066                activity_id: activity_id.clone(),
2067                attempt: 1,
2068                completion_token: crate::worker::CompletionToken::for_test(),
2069            },
2070            now,
2071        )?;
2072        tracker.record_heartbeat(
2073            worker_id,
2074            heartbeat(
2075                workflow_id,
2076                activity_id,
2077                Some(Payload::new(
2078                    ContentType::Json,
2079                    b"{\"progress\":1}".to_vec(),
2080                )),
2081            ),
2082            now,
2083        )?;
2084
2085        let completions = sink
2086            .completions
2087            .lock()
2088            .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
2089        assert!(completions.is_empty());
2090        Ok(())
2091    }
2092
2093    struct WorkerIdForTest;
2094
2095    impl WorkerIdForTest {
2096        fn registered() -> Result<WorkerId, ServerError> {
2097            let (_registry, _registration, worker_id) = registry_with_worker()?;
2098            Ok(worker_id)
2099        }
2100    }
2101
2102    /// `complete_task` reports whether THIS call retired the entry — the
2103    /// structural gate the liminal reply router uses to synthesize a
2104    /// lost-worker failure only for a dispatch nobody else resolved.
2105    #[test]
2106    fn complete_task_reports_whether_the_entry_was_tracked()
2107    -> Result<(), Box<dyn std::error::Error>> {
2108        let tracker = HeartbeatTracker::new(Duration::from_secs(5));
2109        let worker_id = WorkerIdForTest::registered()?;
2110        let workflow_id = workflow_id();
2111        let id = activity_id(50);
2112        tracker.track_task(
2113            worker_id,
2114            InFlightActivity {
2115                workflow_id: workflow_id.clone(),
2116                activity_id: id.clone(),
2117                attempt: 1,
2118                completion_token: crate::worker::CompletionToken::for_test(),
2119            },
2120            Instant::now(),
2121        )?;
2122
2123        assert!(tracker.is_tracked(worker_id, &workflow_id, &id)?);
2124        assert!(
2125            tracker.complete_task(worker_id, &workflow_id, &id)?,
2126            "the first completion retires the tracked entry"
2127        );
2128        assert!(!tracker.is_tracked(worker_id, &workflow_id, &id)?);
2129        assert!(
2130            !tracker.complete_task(worker_id, &workflow_id, &id)?,
2131            "a second completion finds nothing to retire"
2132        );
2133        Ok(())
2134    }
2135
2136    /// A liveness beat (the liminal worker's automatic pump) refreshes the
2137    /// task's expiry stamp — keeping a genuinely-running over-window activity
2138    /// out of the sweep — and reports an untracked task benignly.
2139    #[test]
2140    fn record_liveness_refreshes_stamp_and_ignores_untracked_tasks()
2141    -> Result<(), Box<dyn std::error::Error>> {
2142        let window = Duration::from_secs(5);
2143        let tracker = HeartbeatTracker::new(window);
2144        let worker_id = WorkerIdForTest::registered()?;
2145        let workflow_id = workflow_id();
2146        let id = activity_id(51);
2147        let start = Instant::now();
2148        tracker.track_task(
2149            worker_id,
2150            InFlightActivity {
2151                workflow_id: workflow_id.clone(),
2152                activity_id: id.clone(),
2153                attempt: 1,
2154                completion_token: crate::worker::CompletionToken::for_test(),
2155            },
2156            start,
2157        )?;
2158
2159        // Beaten at the window edge, the task survives past the original expiry.
2160        assert!(tracker.record_liveness(worker_id, &workflow_id, &id, start + window)?);
2161        assert!(tracker.is_live(worker_id, &workflow_id, &id, start + window + window)?);
2162        assert!(tracker.expired_workers(start + window + window)?.is_empty());
2163
2164        // An untracked beat (an outbox dispatch, or a beat racing completion)
2165        // is a benign false, never an error.
2166        assert!(!tracker.record_liveness(
2167            worker_id,
2168            &workflow_id,
2169            &activity_id(52),
2170            start + window
2171        )?);
2172        Ok(())
2173    }
2174
2175    #[test]
2176    fn sweep_interval_is_quarter_window_clamped_to_one_second_and_window() {
2177        // The default 30s window sweeps every 7.5s (quarter-window).
2178        assert_eq!(
2179            sweep_interval(Duration::from_secs(30)),
2180            Duration::from_millis(7_500)
2181        );
2182        // A short window's quarter (500ms) is floored at 1s.
2183        assert_eq!(
2184            sweep_interval(Duration::from_secs(2)),
2185            Duration::from_secs(1)
2186        );
2187        // A very long window's quarter stays within the [1s, window] band.
2188        assert_eq!(
2189            sweep_interval(Duration::from_secs(3_600)),
2190            Duration::from_secs(900)
2191        );
2192        // A sub-second (test) window sweeps once per window, never spinning
2193        // sub-window nor waiting longer than the window itself.
2194        assert_eq!(
2195            sweep_interval(Duration::from_millis(200)),
2196            Duration::from_millis(200)
2197        );
2198        // A zero window is floored at the minimum positive period rather than
2199        // producing the zero interval `tokio::time::interval` rejects.
2200        assert_eq!(sweep_interval(Duration::ZERO), Duration::from_millis(1));
2201    }
2202}