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 /// Poisons the tracker's state lock, for tests that pin the fail-open
992 /// contract of paths that must never let a poisoned tracker withhold a
993 /// completion. Unwinds across a held guard inside `catch_unwind`, which is
994 /// exactly what real poison is.
995 #[cfg(test)]
996 pub(crate) fn poison_for_tests(&self) {
997 let inner = Arc::clone(&self.inner);
998 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
999 let _guard = inner.lock();
1000 // `resume_unwind` raises the unwind directly (no panic hook, no
1001 // panic machinery in a non-test path): the guard drops while the
1002 // thread is unwinding, which is precisely how a mutex is poisoned.
1003 std::panic::resume_unwind(Box::new(
1004 "poisoning the heartbeat tracker for a fail-open test",
1005 ));
1006 }));
1007 }
1008}
1009
1010/// Sweep cadence derived from the operator's `worker.heartbeat_window`: a
1011/// quarter of the window, clamped to `[1s, window]` (the default 30s window
1012/// sweeps every 7.5s).
1013///
1014/// Deliberately derived rather than a separate config knob: the window is the
1015/// operational contract ("a silent worker is dead after this long"), and the
1016/// sweep cadence is an implementation detail of enforcing it — a quarter-window
1017/// cadence bounds detection latency at `window + window/4` while keeping the
1018/// sweep cheap. A window shorter than one second (test configurations) sweeps
1019/// once per window rather than sub-second-spinning, and a zero window is
1020/// floored at one millisecond because `tokio::time::interval` rejects a zero
1021/// period.
1022#[must_use]
1023pub fn sweep_interval(heartbeat_window: Duration) -> Duration {
1024 /// `tokio::time::interval` panics on a zero period, so even a
1025 /// (misconfigured) zero window gets a positive cadence.
1026 const MINIMUM_PERIOD: Duration = Duration::from_millis(1);
1027 /// Target lower bound: sweeping more often than once a second buys no
1028 /// meaningful detection latency against real heartbeat windows.
1029 const TARGET_FLOOR: Duration = Duration::from_secs(1);
1030 let ceiling = heartbeat_window.max(MINIMUM_PERIOD);
1031 // The floor never exceeds the ceiling, so `clamp` cannot panic.
1032 (heartbeat_window / 4).clamp(TARGET_FLOOR.min(ceiling), ceiling)
1033}
1034
1035/// Production driver of [`HeartbeatTracker::fail_expired_workers`] (#176).
1036///
1037/// The tracker records connection and per-task liveness, while the stream-teardown
1038/// sweep fails a worker whose stream ENDS. A worker whose stream stays open while
1039/// its process wedges is caught by the connection lease even when it is idle.
1040/// This interval task expires every silent connection or task, deregistering it
1041/// with the provable
1042/// [`WorkerDeathReason::Timeout`](aion_core::WorkerDeathReason::Timeout) and
1043/// surfacing its tasks as TRANSPORT losses through the shared completion sink
1044/// — the `lost:` class the engine re-dispatches attempt-neutrally, never the
1045/// action's retry vocabulary. It shares the server's shutdown watch, so it drains with
1046/// the transports (mirroring
1047/// [`OutboxDispatcher::run`](crate::worker::OutboxDispatcher::run)).
1048///
1049/// Double-fail safety: this sweep and the stream-teardown path
1050/// ([`HeartbeatTracker::fail_disconnected_worker`]) can both observe the same
1051/// dead worker. Both funnel into the same idempotent core —
1052/// `deregister_with_reason` is a no-op for an already-removed worker (no
1053/// duplicate WS3 delta, no metrics double-count) and the tracker removes each
1054/// task as it fails it — so whichever path runs second sees an empty report and
1055/// never double-completes an activity.
1056pub struct HeartbeatSweeper<S> {
1057 tracker: HeartbeatTracker,
1058 registry: ConnectedWorkerRegistry,
1059 sink: S,
1060 drain: DrainState,
1061 heartbeat_window: Duration,
1062 interval: Duration,
1063 /// Live unserved-queue state, read only to state the CONSEQUENCE of a
1064 /// deregistration in the same log line as its cause: how many dispatches are
1065 /// already parked on the queue the reaped worker was serving. Default-empty
1066 /// in wirings that have no queue service, where the count reads zero.
1067 queue_state: crate::worker::QueueServiceState,
1068}
1069
1070impl<S> HeartbeatSweeper<S>
1071where
1072 S: ActivityCompletionSink + Send + Sync + 'static,
1073{
1074 /// Build a sweeper over the server's shared liveness tracker, worker
1075 /// registry, completion sink, and drain gate. The cadence is derived from
1076 /// `heartbeat_window` by [`sweep_interval`].
1077 #[must_use]
1078 pub fn new(
1079 tracker: HeartbeatTracker,
1080 registry: ConnectedWorkerRegistry,
1081 sink: S,
1082 drain: DrainState,
1083 heartbeat_window: Duration,
1084 ) -> Self {
1085 let interval = sweep_interval(heartbeat_window);
1086 Self {
1087 tracker,
1088 registry,
1089 sink,
1090 drain,
1091 heartbeat_window,
1092 interval,
1093 queue_state: crate::worker::QueueServiceState::default(),
1094 }
1095 }
1096
1097 /// Share the live unserved-queue state so a deregistration log can state how
1098 /// many dispatches are already parked on the queue the reaped worker served.
1099 ///
1100 /// Without it the count reads zero — honest for a wiring with no queue
1101 /// service, and never a reason to withhold the deregistration itself.
1102 #[must_use]
1103 pub fn with_queue_state(mut self, queue_state: crate::worker::QueueServiceState) -> Self {
1104 self.queue_state = queue_state;
1105 self
1106 }
1107
1108 /// Run the expiry sweep until `shutdown` flips to `true`.
1109 ///
1110 /// A tracker/registry error during a sweep is logged and retried next tick
1111 /// rather than tearing the task down — a transient failure must not
1112 /// silently stop dead-worker detection. Shutdown is observed both while
1113 /// waiting for the next tick and re-checked before each sweep, exactly like
1114 /// the outbox dispatcher's run loop.
1115 pub async fn run(self, mut shutdown: watch::Receiver<bool>) {
1116 info!(
1117 sweep_interval_ms = self.interval.as_millis(),
1118 heartbeat_window_ms = self.heartbeat_window.as_millis(),
1119 "worker heartbeat sweeper started"
1120 );
1121 let mut ticks = tokio::time::interval(self.interval);
1122 ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
1123 loop {
1124 tokio::select! {
1125 _ = ticks.tick() => {
1126 if *shutdown.borrow() {
1127 break;
1128 }
1129 self.sweep_once(Instant::now());
1130 }
1131 changed = shutdown.changed() => {
1132 // A receive error means every sender dropped; treat that as
1133 // a shutdown request rather than spinning.
1134 if changed.is_err() || *shutdown.borrow() {
1135 break;
1136 }
1137 }
1138 }
1139 }
1140 info!("worker heartbeat sweeper stopped");
1141 }
1142
1143 /// Fail every currently-expired worker once, logging each lost-worker
1144 /// report at warn (mirroring the stream-teardown sweep's logging).
1145 fn sweep_once(&self, now: Instant) {
1146 let reports = match self
1147 .tracker
1148 .fail_expired_workers(&self.registry, &self.sink, now)
1149 {
1150 Ok(reports) => reports,
1151 Err(sweep_error) => {
1152 error!(
1153 error = %sweep_error,
1154 "heartbeat expiry sweep failed; retrying next tick"
1155 );
1156 return;
1157 }
1158 };
1159 for report in &reports {
1160 let task_queue = report.task_queue.as_deref().unwrap_or("<unregistered>");
1161 // The consequence, stated with the cause: a queue whose worker just
1162 // died and which already holds parked dispatches is an alertable
1163 // condition, and the operator should not have to join two log lines
1164 // to see it. A read failure reports `None` rather than suppressing
1165 // the deregistration line.
1166 let parked = report
1167 .task_queue
1168 .as_deref()
1169 .map(|queue| self.queue_state.parked_on_queue(queue))
1170 .transpose()
1171 .unwrap_or_else(|error| {
1172 error!(%error, "could not read parked-dispatch count for a reaped worker");
1173 None
1174 })
1175 .unwrap_or(0);
1176 if report.tasks.is_empty() {
1177 warn!(
1178 worker_id = ?report.worker_id,
1179 task_queue,
1180 parked_dispatches = parked,
1181 heartbeat_window_ms = self.heartbeat_window.as_millis(),
1182 "idle worker connection lease expired; worker deregistered"
1183 );
1184 } else {
1185 warn!(
1186 worker_id = ?report.worker_id,
1187 task_queue,
1188 parked_dispatches = parked,
1189 failed_tasks = report.tasks.len(),
1190 heartbeat_window_ms = self.heartbeat_window.as_millis(),
1191 "worker heartbeat window expired with in-flight activities; \
1192 deregistered and surfaced as transport losses, to be \
1193 re-dispatched attempt-neutrally"
1194 );
1195 }
1196 }
1197 if !reports.is_empty() {
1198 // In-flight accounting may have just reached zero; wake any drain
1199 // waiter so shutdown does not sit out its full timeout (mirrors
1200 // the stream-teardown sweep).
1201 self.drain.notify_activity_drained();
1202 }
1203 }
1204}
1205
1206impl TaskKey {
1207 fn new(worker_id: WorkerId, workflow_id: WorkflowId, activity_id: ActivityId) -> Self {
1208 Self(worker_id, workflow_id, activity_id)
1209 }
1210
1211 const fn worker_id(&self) -> WorkerId {
1212 self.0
1213 }
1214}
1215
1216struct DecodedHeartbeat {
1217 workflow_id: WorkflowId,
1218 activity_id: ActivityId,
1219 progress: Option<Payload>,
1220}
1221
1222impl TryFrom<ProtoHeartbeat> for DecodedHeartbeat {
1223 type Error = ServerError;
1224
1225 fn try_from(value: ProtoHeartbeat) -> Result<Self, Self::Error> {
1226 let workflow_id = value
1227 .workflow_id
1228 .ok_or_else(|| wire_error("heartbeat workflow id is missing"))
1229 .and_then(|id| WorkflowId::try_from(id).map_err(ServerError::from))?;
1230 let activity_id = value
1231 .activity_id
1232 .ok_or_else(|| wire_error("heartbeat activity id is missing"))
1233 .map(ActivityId::from)?;
1234 let progress = value
1235 .progress
1236 .map(Payload::try_from)
1237 .transpose()
1238 .map_err(ServerError::from)?;
1239 Ok(Self {
1240 workflow_id,
1241 activity_id,
1242 progress,
1243 })
1244 }
1245}
1246
1247/// The task queue a still-registered worker serves, for the deregistration log.
1248///
1249/// A registry read failure (poisoned lock) yields `None` rather than aborting
1250/// the sweep: losing the queue NAME must never stop a dead worker being reaped.
1251fn task_queue_of(registry: &ConnectedWorkerRegistry, worker_id: WorkerId) -> Option<String> {
1252 registry
1253 .worker_by_id(worker_id)
1254 .ok()
1255 .flatten()
1256 .map(|handle| handle.task_queue().to_owned())
1257}
1258
1259fn is_expired(liveness: &TaskLiveness, now: Instant) -> bool {
1260 now.checked_duration_since(liveness.last_heartbeat_at)
1261 .is_some_and(|elapsed| elapsed > liveness.heartbeat_window)
1262}
1263
1264fn wire_error(message: &'static str) -> ServerError {
1265 ServerError::Wire {
1266 wire: WireError::backend(message),
1267 }
1268}
1269
1270#[cfg(test)]
1271mod reachability_tests {
1272 use std::time::{Duration, Instant};
1273
1274 use super::{
1275 DISPATCH_PROBATION_PINGS, DispatchExclusion, ExcludedWorker, HeartbeatTracker, ServerError,
1276 WorkerId,
1277 };
1278
1279 const WINDOW: Duration = Duration::from_secs(30);
1280
1281 /// Every test returns `Result` and uses `?` rather than unwrapping: a lock
1282 /// fault inside the tracker is a real failure mode of the code under test,
1283 /// and it should surface as a failed test carrying the typed error, not as a
1284 /// panic message written by the test.
1285 type TestResult = Result<(), ServerError>;
1286
1287 fn tracker_with_worker(now: Instant) -> Result<(HeartbeatTracker, WorkerId), ServerError> {
1288 let tracker = HeartbeatTracker::new(WINDOW);
1289 let worker = WorkerId::from_value(1);
1290 tracker.register_connection(worker, now)?;
1291 Ok((tracker, worker))
1292 }
1293
1294 /// Answer the full probation, so the worker is genuinely eligible. Tests
1295 /// about staleness, pump beats, or withdrawal must start from a worker that
1296 /// HAS eligibility — otherwise they pass on a worker that never had any and
1297 /// prove nothing about the behaviour they name.
1298 fn serve_probation(
1299 tracker: &HeartbeatTracker,
1300 worker: WorkerId,
1301 at: Instant,
1302 ) -> Result<(), ServerError> {
1303 for _ in 0..DISPATCH_PROBATION_PINGS {
1304 assert!(
1305 tracker.record_dispatch_reachability(worker, at)?,
1306 "the worker must still be tracked while it serves its probation"
1307 );
1308 }
1309 Ok(())
1310 }
1311
1312 /// THE REGRESSION. This is the defect that made run `dfd2117c` invisible:
1313 /// the worker's liveness pump beat from a background task, refreshed the one
1314 /// shared lease, and the dead-man switch could not fire while the server had
1315 /// been unable to push to that worker for fifteen minutes.
1316 ///
1317 /// An inbound frame must prove the worker is ALIVE and must NOT prove the
1318 /// server can REACH it.
1319 #[test]
1320 fn an_inbound_frame_cannot_prove_dispatch_reachability() -> TestResult {
1321 let start = Instant::now();
1322 let (tracker, worker) = tracker_with_worker(start)?;
1323 // The worker STARTS eligible, earned honestly. Without this the test
1324 // would pass on a worker that never had eligibility to lose, which says
1325 // nothing about whether a pump beat can preserve it.
1326 serve_probation(&tracker, worker, start)?;
1327 assert!(
1328 tracker.is_dispatch_reachable(worker, start)?,
1329 "precondition: the worker is eligible before the connection goes one-way"
1330 );
1331
1332 // Well past the window, with the pump beating throughout — exactly what
1333 // a busy worker on a poisoned connection looks like.
1334 let much_later = start + WINDOW * 4;
1335 assert!(
1336 tracker.record_connection_activity(worker, much_later)?,
1337 "the worker is still tracked"
1338 );
1339
1340 assert!(
1341 !tracker.is_dispatch_reachable(worker, much_later)?,
1342 "a pump beat must NOT make a worker the server cannot push to look reachable"
1343 );
1344 assert_eq!(
1345 tracker.unreachable_workers(much_later)?,
1346 vec![ExcludedWorker {
1347 worker_id: worker,
1348 // It HELD eligibility (served above) and lost it to a stale
1349 // proof. Classifying this as an opening probation would tell an
1350 // operator a poisoned connection is an ordinary worker start.
1351 exclusion: DispatchExclusion::ReachabilityLost,
1352 }],
1353 "the worker must be named unreachable however alive its process looks"
1354 );
1355 Ok(())
1356 }
1357
1358 /// The control for the test above: without it, a tracker that reported
1359 /// EVERYTHING unreachable would satisfy that assertion and prove nothing.
1360 #[test]
1361 fn an_answered_ping_does_prove_dispatch_reachability() -> TestResult {
1362 let start = Instant::now();
1363 let (tracker, worker) = tracker_with_worker(start)?;
1364
1365 let much_later = start + WINDOW * 4;
1366 serve_probation(&tracker, worker, much_later)?;
1367
1368 assert!(
1369 tracker.is_dispatch_reachable(worker, much_later)?,
1370 "answered pings are the one thing that proves the push leg works"
1371 );
1372 assert!(
1373 tracker.unreachable_workers(much_later)?.is_empty(),
1374 "a worker answering pings is never unreachable"
1375 );
1376 Ok(())
1377 }
1378
1379 /// Registration opens a PROBATION and grants nothing. The handshake ack is
1380 /// SENT by this server; nothing reports that it was RECEIVED, so a
1381 /// connection is a channel, not proof that the channel carries. Eligibility
1382 /// is earned by answered pings only.
1383 #[test]
1384 fn registration_opens_a_probation_and_does_not_grant_eligibility() -> TestResult {
1385 let start = Instant::now();
1386 let (tracker, worker) = tracker_with_worker(start)?;
1387
1388 assert!(
1389 !tracker.is_dispatch_reachable(worker, start)?,
1390 "a brand-new connection has proved nothing about the push leg"
1391 );
1392 assert_eq!(
1393 tracker.unreachable_workers(start)?,
1394 vec![ExcludedWorker {
1395 worker_id: worker,
1396 // And it is carried as a PROBATION, not as a reachability
1397 // failure. This is the distinction that stopped an ordinary
1398 // worker start from being announced to the operator as an
1399 // unreachable dispatch path.
1400 exclusion: DispatchExclusion::OpeningProbation { answers: 0 },
1401 }],
1402 "a worker serving its probation is carried in the census as unreachable"
1403 );
1404 Ok(())
1405 }
1406
1407 /// The case that produced a FALSE ALARM on every healthy worker start.
1408 ///
1409 /// One answer banked out of two: the server has demonstrably reached this
1410 /// worker — moments ago — and is merely waiting for the second consecutive
1411 /// answer. Reporting that as a reachability failure told Tom's operator log
1412 /// his worker's dispatch path was dead when the opposite had just been
1413 /// measured. The exclusion is real; the REASON is an opening probation.
1414 #[test]
1415 fn a_part_served_probation_is_a_probation_and_not_a_reachability_failure() -> TestResult {
1416 let start = Instant::now();
1417 let (tracker, worker) = tracker_with_worker(start)?;
1418 const {
1419 assert!(
1420 DISPATCH_PROBATION_PINGS > 1,
1421 "this test is only meaningful while the probation takes more than one answer"
1422 );
1423 }
1424 assert!(tracker.record_dispatch_reachability(worker, start)?);
1425
1426 assert_eq!(
1427 tracker.unreachable_workers(start)?,
1428 vec![ExcludedWorker {
1429 worker_id: worker,
1430 exclusion: DispatchExclusion::OpeningProbation { answers: 1 },
1431 }],
1432 "a worker that has answered part of its opening probation is still excluded, but it \
1433 must not be described as one the server cannot reach — it answered"
1434 );
1435 Ok(())
1436 }
1437
1438 /// The other side of the same discrimination, and the control for the test
1439 /// above: once eligibility has actually been HELD, losing it is an incident
1440 /// and must classify differently. Without this, a classifier that answered
1441 /// `OpeningProbation` unconditionally would satisfy the test above.
1442 #[test]
1443 fn losing_held_eligibility_is_reported_as_a_loss_not_as_a_fresh_probation() -> TestResult {
1444 let start = Instant::now();
1445 let (tracker, worker) = tracker_with_worker(start)?;
1446 serve_probation(&tracker, worker, start)?;
1447 assert!(
1448 tracker.is_dispatch_reachable(worker, start)?,
1449 "precondition: eligibility was genuinely held before it was lost"
1450 );
1451
1452 assert!(
1453 tracker.record_dispatch_unreachable(worker)?,
1454 "the worker is still tracked when its ping fails"
1455 );
1456
1457 assert_eq!(
1458 tracker.unreachable_workers(start)?,
1459 vec![ExcludedWorker {
1460 worker_id: worker,
1461 exclusion: DispatchExclusion::ReachabilityLost,
1462 }],
1463 "a failed ping on a worker that HAD eligibility is an incident, and must not be \
1464 filed as the ordinary probation every fresh connection serves"
1465 );
1466 Ok(())
1467 }
1468
1469 /// A REDIAL is a new measurement. The previous connection's proof must not
1470 /// make the new connection's ordinary probation look like an incident —
1471 /// otherwise every reconnect of a healthy worker would raise the alarm that
1472 /// is supposed to mean something has gone wrong.
1473 #[test]
1474 fn a_reconnect_starts_a_fresh_probation_not_a_lost_eligibility() -> TestResult {
1475 let start = Instant::now();
1476 let (tracker, worker) = tracker_with_worker(start)?;
1477 serve_probation(&tracker, worker, start)?;
1478
1479 tracker.unregister_connection(worker)?;
1480 tracker.register_connection(worker, start)?;
1481
1482 assert_eq!(
1483 tracker.unreachable_workers(start)?,
1484 vec![ExcludedWorker {
1485 worker_id: worker,
1486 exclusion: DispatchExclusion::OpeningProbation { answers: 0 },
1487 }],
1488 "nothing earned on the old connection carries across to the new one"
1489 );
1490 Ok(())
1491 }
1492
1493 /// The probation must be SERVED IN FULL. One answered ping can be luck — a
1494 /// link that answers one probe in three would otherwise accrue eligibility
1495 /// and then flap. This pins the boundary from below: K-1 answers is not
1496 /// enough, and the very next one is.
1497 #[test]
1498 fn one_ping_short_of_the_probation_earns_nothing() -> TestResult {
1499 let start = Instant::now();
1500 let (tracker, worker) = tracker_with_worker(start)?;
1501
1502 for _ in 0..DISPATCH_PROBATION_PINGS - 1 {
1503 assert!(tracker.record_dispatch_reachability(worker, start)?);
1504 assert!(
1505 !tracker.is_dispatch_reachable(worker, start)?,
1506 "eligibility must not be granted before the probation is served in full"
1507 );
1508 }
1509
1510 assert!(tracker.record_dispatch_reachability(worker, start)?);
1511 assert!(
1512 tracker.is_dispatch_reachable(worker, start)?,
1513 "the ping that completes the probation must grant eligibility — otherwise this test \
1514 would pass on a tracker that never grants it at all"
1515 );
1516 Ok(())
1517 }
1518
1519 /// A failed probe RESETS the run. Eligibility is withdrawn immediately, not
1520 /// when the window later expires: an unanswered probe is direct evidence
1521 /// about the push leg, and direct negative evidence must weigh at least as
1522 /// much as silence.
1523 #[test]
1524 fn a_failed_probe_withdraws_eligibility_at_once_and_restarts_the_probation() -> TestResult {
1525 let start = Instant::now();
1526 let (tracker, worker) = tracker_with_worker(start)?;
1527 serve_probation(&tracker, worker, start)?;
1528 assert!(
1529 tracker.is_dispatch_reachable(worker, start)?,
1530 "precondition"
1531 );
1532
1533 assert!(
1534 tracker.record_dispatch_unreachable(worker)?,
1535 "the worker is still tracked"
1536 );
1537 assert!(
1538 !tracker.is_dispatch_reachable(worker, start)?,
1539 "a failed probe withdraws eligibility on the spot, inside the window"
1540 );
1541
1542 // And the run restarts from zero rather than resuming: one answer does
1543 // not restore what a full probation earned.
1544 assert!(tracker.record_dispatch_reachability(worker, start)?);
1545 assert!(
1546 !tracker.is_dispatch_reachable(worker, start)?,
1547 "a single answer after a failure must not restore eligibility"
1548 );
1549 Ok(())
1550 }
1551
1552 /// 🔴 THE FLAPPING PIN. A link that answers every other probe must NEVER
1553 /// become eligible. Cumulative counting would let it accrue, and eligibility
1554 /// would switch on and off under a running fleet — the intermittent evidence
1555 /// that costs hours to attribute. Consecutiveness is what forbids it.
1556 #[test]
1557 fn a_link_that_answers_every_other_probe_never_becomes_eligible() -> TestResult {
1558 let start = Instant::now();
1559 let (tracker, worker) = tracker_with_worker(start)?;
1560
1561 // Far more probes than the probation demands, alternating.
1562 for probe in 0..DISPATCH_PROBATION_PINGS * 10 {
1563 let now = start + Duration::from_millis(u64::from(probe));
1564 if probe % 2 == 0 {
1565 assert!(tracker.record_dispatch_reachability(worker, now)?);
1566 } else {
1567 assert!(tracker.record_dispatch_unreachable(worker)?);
1568 }
1569 assert!(
1570 !tracker.is_dispatch_reachable(worker, now)?,
1571 "a flapping link must never hold dispatch eligibility, at any probe (probe {probe})"
1572 );
1573 }
1574
1575 // The control: the same worker, answering consecutively, DOES become
1576 // eligible — so this test cannot pass on a tracker that grants nothing.
1577 let now = start + Duration::from_secs(1);
1578 serve_probation(&tracker, worker, now)?;
1579 assert!(
1580 tracker.is_dispatch_reachable(worker, now)?,
1581 "consecutive answers must still earn eligibility"
1582 );
1583 Ok(())
1584 }
1585
1586 /// Reachability must EXPIRE on its own clock. If it were only ever advanced
1587 /// and never allowed to go stale, eligibility could never be withdrawn.
1588 #[test]
1589 fn reachability_goes_stale_once_the_window_passes() -> TestResult {
1590 let start = Instant::now();
1591 let (tracker, worker) = tracker_with_worker(start)?;
1592 serve_probation(&tracker, worker, start)?;
1593
1594 assert!(
1595 tracker.is_dispatch_reachable(worker, start + WINDOW)?,
1596 "still inside the window"
1597 );
1598 assert!(
1599 !tracker.is_dispatch_reachable(worker, start + WINDOW + Duration::from_millis(1))?,
1600 "one millisecond past the window is stale"
1601 );
1602 Ok(())
1603 }
1604
1605 /// An untracked worker is not reachable: absence of proof is not proof. A
1606 /// pong racing a reap must not resurrect it either.
1607 #[test]
1608 fn an_unregistered_worker_is_never_reachable_and_cannot_be_resurrected() -> TestResult {
1609 let start = Instant::now();
1610 let (tracker, worker) = tracker_with_worker(start)?;
1611 tracker.unregister_connection(worker)?;
1612
1613 assert!(
1614 !tracker.is_dispatch_reachable(worker, start)?,
1615 "a deregistered worker is not reachable"
1616 );
1617 assert!(
1618 !tracker.record_dispatch_reachability(worker, start)?,
1619 "a late pong must not resurrect a deregistered worker"
1620 );
1621 assert!(
1622 !tracker.record_dispatch_unreachable(worker)?,
1623 "a late probe FAILURE must not resurrect a deregistered worker either — the reset \
1624 path allocates an entry, so it has to refuse an untracked worker as firmly as the \
1625 success path does"
1626 );
1627 assert!(
1628 tracker.unreachable_workers(start)?.is_empty(),
1629 "an untracked worker is not carried in the census either"
1630 );
1631 Ok(())
1632 }
1633}
1634
1635#[cfg(test)]
1636mod tests {
1637 use std::sync::Mutex;
1638
1639 use aion_core::ContentType;
1640 use aion_proto::{ProtoActivityId, ProtoPayload, ProtoWorkflowId};
1641 use serde_json::json;
1642 use uuid::Uuid;
1643
1644 use crate::worker::registry::WorkerRegistration;
1645
1646 use super::*;
1647
1648 #[derive(Default)]
1649 struct RecordingSink {
1650 completions: Mutex<Vec<ActivityCompletion>>,
1651 parks: Mutex<Vec<(WorkflowId, ActivityId)>>,
1652 }
1653
1654 impl ActivityCompletionSink for RecordingSink {
1655 fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
1656 self.completions
1657 .lock()
1658 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1659 .push(completion);
1660 Ok(())
1661 }
1662
1663 fn park_activity(
1664 &self,
1665 workflow_id: &WorkflowId,
1666 activity_id: &ActivityId,
1667 ) -> Result<(), ServerError> {
1668 self.parks
1669 .lock()
1670 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1671 .push((workflow_id.clone(), activity_id.clone()));
1672 Ok(())
1673 }
1674 }
1675
1676 fn workflow_id() -> WorkflowId {
1677 WorkflowId::new(Uuid::nil())
1678 }
1679
1680 fn activity_id(position: u64) -> ActivityId {
1681 ActivityId::from_sequence_position(position)
1682 }
1683
1684 fn payload(value: &serde_json::Value) -> Result<Payload, Box<dyn std::error::Error>> {
1685 Ok(Payload::from_json(value)?)
1686 }
1687
1688 fn heartbeat(
1689 workflow_id: WorkflowId,
1690 activity_id: ActivityId,
1691 progress: Option<Payload>,
1692 ) -> ProtoHeartbeat {
1693 ProtoHeartbeat {
1694 workflow_id: Some(ProtoWorkflowId::from(workflow_id)),
1695 activity_id: Some(ProtoActivityId::from(activity_id)),
1696 progress: progress.map(ProtoPayload::from),
1697 }
1698 }
1699
1700 fn registry_with_worker()
1701 -> Result<(ConnectedWorkerRegistry, WorkerRegistration, WorkerId), ServerError> {
1702 let registry = ConnectedWorkerRegistry::default();
1703 let (tx, _rx) = tokio::sync::mpsc::channel(1);
1704 let activity_types = [String::from("charge-card")];
1705 let registration = registry.register("tenant-a", activity_types.iter(), tx)?;
1706 let worker_id = registration
1707 .worker_id()
1708 .ok_or_else(|| ServerError::lock_poisoned("test worker registration"))?;
1709 Ok((registry, registration, worker_id))
1710 }
1711
1712 #[test]
1713 fn heartbeat_refresh_keeps_task_live_across_window() -> Result<(), Box<dyn std::error::Error>> {
1714 let window = Duration::from_secs(5);
1715 let tracker = HeartbeatTracker::new(window);
1716 let worker_id = WorkerIdForTest::registered()?;
1717 let workflow_id = workflow_id();
1718 let activity_id = activity_id(10);
1719 let start = Instant::now();
1720
1721 tracker.track_task(
1722 worker_id,
1723 InFlightActivity {
1724 workflow_id: workflow_id.clone(),
1725 activity_id: activity_id.clone(),
1726 attempt: 1,
1727 completion_token: crate::worker::CompletionToken::for_test(),
1728 },
1729 start,
1730 )?;
1731 assert!(tracker.is_live(worker_id, &workflow_id, &activity_id, start + window)?);
1732
1733 let progress = payload(&json!({"percent": 50}))?;
1734 let update = tracker.record_heartbeat(
1735 worker_id,
1736 heartbeat(
1737 workflow_id.clone(),
1738 activity_id.clone(),
1739 Some(progress.clone()),
1740 ),
1741 start + window,
1742 )?;
1743
1744 assert_eq!(update.liveness.last_progress, Some(progress));
1745 assert!(tracker.is_live(
1746 worker_id,
1747 &workflow_id,
1748 &activity_id,
1749 start + window + window
1750 )?);
1751 assert!(tracker.expired_workers(start + window + window)?.is_empty());
1752 Ok(())
1753 }
1754
1755 #[test]
1756 fn missed_heartbeat_deregisters_worker_and_fails_in_flight_once()
1757 -> Result<(), Box<dyn std::error::Error>> {
1758 let (registry, _registration, worker_id) = registry_with_worker()?;
1759 let sink = RecordingSink::default();
1760 let tracker = HeartbeatTracker::new(Duration::from_secs(5));
1761 let workflow_id = workflow_id();
1762 let activity_id = activity_id(11);
1763 let start = Instant::now();
1764
1765 tracker.track_task(
1766 worker_id,
1767 InFlightActivity {
1768 workflow_id: workflow_id.clone(),
1769 activity_id: activity_id.clone(),
1770 attempt: 1,
1771 completion_token: crate::worker::CompletionToken::for_test(),
1772 },
1773 start,
1774 )?;
1775
1776 let reports =
1777 tracker.fail_expired_workers(®istry, &sink, start + Duration::from_secs(6))?;
1778 assert_eq!(reports.len(), 1);
1779 assert_eq!(reports[0].worker_id, worker_id);
1780 assert_eq!(reports[0].tasks.len(), 1);
1781 assert!(
1782 registry
1783 .workers_for("tenant-a", "default", "charge-card", None)?
1784 .is_empty()
1785 );
1786
1787 let second = tracker.fail_disconnected_worker(worker_id, ®istry, &sink)?;
1788 assert!(second.tasks.is_empty());
1789 let completions = sink
1790 .completions
1791 .lock()
1792 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
1793 assert_eq!(completions.len(), 1);
1794 assert_eq!(completions[0].workflow_id, workflow_id);
1795 assert_eq!(completions[0].activity_id, activity_id);
1796 // The sweep reports a TRANSPORT-domain loss, not an activity failure:
1797 // the activity never executed to a result, so the sink (not this sweep)
1798 // classifies it and applies the transport's own re-dispatch budget.
1799 // Before this distinction existed the sweep synthesized a `Retryable`
1800 // `ActivityError` that the engine then delivered as a TERMINAL failure
1801 // whenever the activity carried no authored retry policy.
1802 match &completions[0].outcome {
1803 ActivityCompletionOutcome::WorkerLost { worker_id: lost } => {
1804 assert_eq!(*lost, worker_id);
1805 }
1806 other => {
1807 return Err(format!("expected a lost-worker outcome, got {other:?}").into());
1808 }
1809 }
1810 Ok(())
1811 }
1812
1813 #[test]
1814 fn disconnected_worker_fails_each_in_flight_task_once() -> Result<(), Box<dyn std::error::Error>>
1815 {
1816 let (registry, _registration, worker_id) = registry_with_worker()?;
1817 let sink = RecordingSink::default();
1818 let tracker = HeartbeatTracker::new(Duration::from_secs(5));
1819 let workflow_id = workflow_id();
1820 let start = Instant::now();
1821
1822 tracker.track_task(
1823 worker_id,
1824 InFlightActivity {
1825 workflow_id: workflow_id.clone(),
1826 activity_id: activity_id(21),
1827 attempt: 1,
1828 completion_token: crate::worker::CompletionToken::for_test(),
1829 },
1830 start,
1831 )?;
1832 tracker.track_task(
1833 worker_id,
1834 InFlightActivity {
1835 workflow_id,
1836 activity_id: activity_id(22),
1837 attempt: 1,
1838 completion_token: crate::worker::CompletionToken::for_test(),
1839 },
1840 start,
1841 )?;
1842
1843 let report = tracker.fail_disconnected_worker(worker_id, ®istry, &sink)?;
1844 assert_eq!(report.tasks.len(), 2);
1845 assert!(
1846 registry
1847 .workers_for("tenant-a", "default", "charge-card", None)?
1848 .is_empty()
1849 );
1850
1851 let completions = sink
1852 .completions
1853 .lock()
1854 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
1855 assert_eq!(completions.len(), 2);
1856 assert!(completions.iter().all(|completion| matches!(
1857 &completion.outcome,
1858 ActivityCompletionOutcome::WorkerLost { .. }
1859 )));
1860 Ok(())
1861 }
1862
1863 /// #207: parking a drain-disconnected worker removes its tasks, deregisters
1864 /// it, and PARKS each task through the sink — zero completions synthesized,
1865 /// so the durable log stays byte-equivalent to a kill -9. A second park (or
1866 /// a later fail sweep) finds nothing: the idempotent-deregister discipline
1867 /// the fail path already proves holds for parks too.
1868 #[test]
1869 fn park_disconnected_worker_parks_tasks_without_synthesizing_completions()
1870 -> Result<(), Box<dyn std::error::Error>> {
1871 let (registry, _registration, worker_id) = registry_with_worker()?;
1872 let sink = RecordingSink::default();
1873 let tracker = HeartbeatTracker::new(Duration::from_secs(5));
1874 let workflow_id = workflow_id();
1875 let start = Instant::now();
1876 tracker.track_task(
1877 worker_id,
1878 InFlightActivity {
1879 workflow_id: workflow_id.clone(),
1880 activity_id: activity_id(60),
1881 attempt: 1,
1882 completion_token: crate::worker::CompletionToken::for_test(),
1883 },
1884 start,
1885 )?;
1886 tracker.track_task(
1887 worker_id,
1888 InFlightActivity {
1889 workflow_id: workflow_id.clone(),
1890 activity_id: activity_id(61),
1891 attempt: 1,
1892 completion_token: crate::worker::CompletionToken::for_test(),
1893 },
1894 start,
1895 )?;
1896
1897 let report = tracker.park_disconnected_worker(worker_id, ®istry, &sink)?;
1898 assert_eq!(report.tasks.len(), 2);
1899 assert_eq!(
1900 tracker.in_flight_count()?,
1901 0,
1902 "parking must remove every tracked task so drain accounting reaches zero"
1903 );
1904 assert!(
1905 registry
1906 .workers_for("tenant-a", "default", "charge-card", None)?
1907 .is_empty(),
1908 "the parked worker must be deregistered from routing"
1909 );
1910 let parks = sink
1911 .parks
1912 .lock()
1913 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
1914 assert_eq!(parks.len(), 2, "each task must be parked exactly once");
1915 drop(parks);
1916 assert!(
1917 sink.completions
1918 .lock()
1919 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1920 .is_empty(),
1921 "parking must never synthesize an activity completion"
1922 );
1923
1924 // Double-park and park-after-fail are no-ops: the idempotent core.
1925 let second = tracker.park_disconnected_worker(worker_id, ®istry, &sink)?;
1926 assert!(second.tasks.is_empty());
1927 let third = tracker.fail_disconnected_worker(worker_id, ®istry, &sink)?;
1928 assert!(third.tasks.is_empty());
1929 assert_eq!(
1930 sink.parks
1931 .lock()
1932 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1933 .len(),
1934 2,
1935 "re-sweeping a parked worker must park nothing further"
1936 );
1937 assert!(
1938 sink.completions
1939 .lock()
1940 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1941 .is_empty(),
1942 "a fail sweep after the park must fail nothing"
1943 );
1944 Ok(())
1945 }
1946
1947 /// #207 drain-timeout backstop: the bulk park removes every worker's tasks,
1948 /// parks each through the sink, and wakes drain waiters — never
1949 /// synthesizing a completion.
1950 #[tokio::test]
1951 async fn park_all_in_flight_workers_parks_everything_and_wakes_drain_waiters()
1952 -> Result<(), Box<dyn std::error::Error>> {
1953 let (registry, _registration, worker_id) = registry_with_worker()?;
1954 let sink = RecordingSink::default();
1955 let tracker = HeartbeatTracker::new(Duration::from_secs(5));
1956 let workflow_id = workflow_id();
1957 tracker.track_task(
1958 worker_id,
1959 InFlightActivity {
1960 workflow_id: workflow_id.clone(),
1961 activity_id: activity_id(70),
1962 attempt: 1,
1963 completion_token: crate::worker::CompletionToken::for_test(),
1964 },
1965 Instant::now(),
1966 )?;
1967 // Arm a waiter on the tracker's empty notify BEFORE the bulk park.
1968 let notified = tracker.empty.notified();
1969 tokio::pin!(notified);
1970
1971 let reports = tracker.park_all_in_flight_workers(®istry, &sink)?;
1972 assert_eq!(reports.len(), 1);
1973 assert_eq!(reports[0].worker_id, worker_id);
1974 assert_eq!(reports[0].tasks.len(), 1);
1975 assert_eq!(tracker.in_flight_count()?, 0);
1976 assert_eq!(
1977 sink.parks
1978 .lock()
1979 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1980 .len(),
1981 1
1982 );
1983 assert!(
1984 sink.completions
1985 .lock()
1986 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1987 .is_empty(),
1988 "the bulk park must never synthesize a completion"
1989 );
1990 tokio::time::timeout(Duration::from_millis(200), notified)
1991 .await
1992 .map_err(|_| "the bulk park must wake drain waiters")?;
1993 Ok(())
1994 }
1995
1996 /// The worker runtime's AUTOMATIC liveness beats carry no payload and
1997 /// interleave with explicit handler progress heartbeats: a payload-free
1998 /// beat must refresh the liveness stamp WITHOUT erasing the handler's
1999 /// most recent progress report.
2000 #[test]
2001 fn payload_free_heartbeat_refreshes_liveness_without_clearing_progress()
2002 -> Result<(), Box<dyn std::error::Error>> {
2003 let window = Duration::from_secs(5);
2004 let tracker = HeartbeatTracker::new(window);
2005 let worker_id = WorkerIdForTest::registered()?;
2006 let workflow_id = workflow_id();
2007 let activity_id = activity_id(12);
2008 let start = Instant::now();
2009
2010 tracker.track_task(
2011 worker_id,
2012 InFlightActivity {
2013 workflow_id: workflow_id.clone(),
2014 activity_id: activity_id.clone(),
2015 attempt: 1,
2016 completion_token: crate::worker::CompletionToken::for_test(),
2017 },
2018 start,
2019 )?;
2020 let progress = payload(&json!({"percent": 80}))?;
2021 tracker.record_heartbeat(
2022 worker_id,
2023 heartbeat(
2024 workflow_id.clone(),
2025 activity_id.clone(),
2026 Some(progress.clone()),
2027 ),
2028 start + Duration::from_secs(1),
2029 )?;
2030
2031 // An automatic liveness beat: no payload, later timestamp.
2032 let update = tracker.record_heartbeat(
2033 worker_id,
2034 heartbeat(workflow_id.clone(), activity_id.clone(), None),
2035 start + Duration::from_secs(4),
2036 )?;
2037
2038 assert_eq!(
2039 update.liveness.last_progress,
2040 Some(progress),
2041 "a payload-free liveness beat must not erase handler progress"
2042 );
2043 assert!(
2044 tracker.is_live(
2045 worker_id,
2046 &workflow_id,
2047 &activity_id,
2048 start + Duration::from_secs(8)
2049 )?,
2050 "the payload-free beat must still refresh the liveness stamp"
2051 );
2052 Ok(())
2053 }
2054
2055 #[test]
2056 fn malformed_heartbeat_missing_ids_is_wire_error() -> Result<(), Box<dyn std::error::Error>> {
2057 let worker_id = WorkerIdForTest::registered()?;
2058 let tracker = HeartbeatTracker::new(Duration::from_secs(5));
2059 let missing = ProtoHeartbeat {
2060 workflow_id: None,
2061 activity_id: Some(ProtoActivityId::from(activity_id(30))),
2062 progress: None,
2063 };
2064
2065 let result = tracker.record_heartbeat(worker_id, missing, Instant::now());
2066 assert!(matches!(result, Err(ServerError::Wire { .. })));
2067 Ok(())
2068 }
2069
2070 #[test]
2071 fn heartbeat_progress_is_not_reported_as_activity_result()
2072 -> Result<(), Box<dyn std::error::Error>> {
2073 let sink = RecordingSink::default();
2074 let worker_id = WorkerIdForTest::registered()?;
2075 let tracker = HeartbeatTracker::new(Duration::from_secs(5));
2076 let workflow_id = workflow_id();
2077 let activity_id = activity_id(40);
2078 let now = Instant::now();
2079
2080 tracker.track_task(
2081 worker_id,
2082 InFlightActivity {
2083 workflow_id: workflow_id.clone(),
2084 activity_id: activity_id.clone(),
2085 attempt: 1,
2086 completion_token: crate::worker::CompletionToken::for_test(),
2087 },
2088 now,
2089 )?;
2090 tracker.record_heartbeat(
2091 worker_id,
2092 heartbeat(
2093 workflow_id,
2094 activity_id,
2095 Some(Payload::new(
2096 ContentType::Json,
2097 b"{\"progress\":1}".to_vec(),
2098 )),
2099 ),
2100 now,
2101 )?;
2102
2103 let completions = sink
2104 .completions
2105 .lock()
2106 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
2107 assert!(completions.is_empty());
2108 Ok(())
2109 }
2110
2111 struct WorkerIdForTest;
2112
2113 impl WorkerIdForTest {
2114 fn registered() -> Result<WorkerId, ServerError> {
2115 let (_registry, _registration, worker_id) = registry_with_worker()?;
2116 Ok(worker_id)
2117 }
2118 }
2119
2120 /// `complete_task` reports whether THIS call retired the entry — the
2121 /// structural gate the liminal reply router uses to synthesize a
2122 /// lost-worker failure only for a dispatch nobody else resolved.
2123 #[test]
2124 fn complete_task_reports_whether_the_entry_was_tracked()
2125 -> Result<(), Box<dyn std::error::Error>> {
2126 let tracker = HeartbeatTracker::new(Duration::from_secs(5));
2127 let worker_id = WorkerIdForTest::registered()?;
2128 let workflow_id = workflow_id();
2129 let id = activity_id(50);
2130 tracker.track_task(
2131 worker_id,
2132 InFlightActivity {
2133 workflow_id: workflow_id.clone(),
2134 activity_id: id.clone(),
2135 attempt: 1,
2136 completion_token: crate::worker::CompletionToken::for_test(),
2137 },
2138 Instant::now(),
2139 )?;
2140
2141 assert!(tracker.is_tracked(worker_id, &workflow_id, &id)?);
2142 assert!(
2143 tracker.complete_task(worker_id, &workflow_id, &id)?,
2144 "the first completion retires the tracked entry"
2145 );
2146 assert!(!tracker.is_tracked(worker_id, &workflow_id, &id)?);
2147 assert!(
2148 !tracker.complete_task(worker_id, &workflow_id, &id)?,
2149 "a second completion finds nothing to retire"
2150 );
2151 Ok(())
2152 }
2153
2154 /// A liveness beat (the liminal worker's automatic pump) refreshes the
2155 /// task's expiry stamp — keeping a genuinely-running over-window activity
2156 /// out of the sweep — and reports an untracked task benignly.
2157 #[test]
2158 fn record_liveness_refreshes_stamp_and_ignores_untracked_tasks()
2159 -> Result<(), Box<dyn std::error::Error>> {
2160 let window = Duration::from_secs(5);
2161 let tracker = HeartbeatTracker::new(window);
2162 let worker_id = WorkerIdForTest::registered()?;
2163 let workflow_id = workflow_id();
2164 let id = activity_id(51);
2165 let start = Instant::now();
2166 tracker.track_task(
2167 worker_id,
2168 InFlightActivity {
2169 workflow_id: workflow_id.clone(),
2170 activity_id: id.clone(),
2171 attempt: 1,
2172 completion_token: crate::worker::CompletionToken::for_test(),
2173 },
2174 start,
2175 )?;
2176
2177 // Beaten at the window edge, the task survives past the original expiry.
2178 assert!(tracker.record_liveness(worker_id, &workflow_id, &id, start + window)?);
2179 assert!(tracker.is_live(worker_id, &workflow_id, &id, start + window + window)?);
2180 assert!(tracker.expired_workers(start + window + window)?.is_empty());
2181
2182 // An untracked beat (an outbox dispatch, or a beat racing completion)
2183 // is a benign false, never an error.
2184 assert!(!tracker.record_liveness(
2185 worker_id,
2186 &workflow_id,
2187 &activity_id(52),
2188 start + window
2189 )?);
2190 Ok(())
2191 }
2192
2193 #[test]
2194 fn sweep_interval_is_quarter_window_clamped_to_one_second_and_window() {
2195 // The default 30s window sweeps every 7.5s (quarter-window).
2196 assert_eq!(
2197 sweep_interval(Duration::from_secs(30)),
2198 Duration::from_millis(7_500)
2199 );
2200 // A short window's quarter (500ms) is floored at 1s.
2201 assert_eq!(
2202 sweep_interval(Duration::from_secs(2)),
2203 Duration::from_secs(1)
2204 );
2205 // A very long window's quarter stays within the [1s, window] band.
2206 assert_eq!(
2207 sweep_interval(Duration::from_secs(3_600)),
2208 Duration::from_secs(900)
2209 );
2210 // A sub-second (test) window sweeps once per window, never spinning
2211 // sub-window nor waiting longer than the window itself.
2212 assert_eq!(
2213 sweep_interval(Duration::from_millis(200)),
2214 Duration::from_millis(200)
2215 );
2216 // A zero window is floored at the minimum positive period rather than
2217 // producing the zero interval `tokio::time::interval` rejects.
2218 assert_eq!(sweep_interval(Duration::ZERO), Duration::from_millis(1));
2219 }
2220}