Skip to main content

aion_server/worker/
heartbeat.rs

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