1use 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#[derive(Clone, Debug, Eq, PartialEq)]
22pub struct InFlightActivity {
23 pub workflow_id: WorkflowId,
25 pub activity_id: ActivityId,
27 pub completion_token: CompletionToken,
29}
30
31#[derive(Clone, Debug, Eq, PartialEq)]
33pub struct TaskLiveness {
34 pub worker_id: WorkerId,
36 pub workflow_id: WorkflowId,
38 pub activity_id: ActivityId,
40 pub completion_token: CompletionToken,
42 pub heartbeat_window: Duration,
44 pub last_heartbeat_at: Instant,
46 pub last_progress: Option<Payload>,
48}
49
50#[derive(Clone, Debug, Eq, PartialEq)]
52pub struct HeartbeatUpdate {
53 pub liveness: TaskLiveness,
55}
56
57#[derive(Clone, Debug, Eq, PartialEq)]
59pub struct LostWorkerReport {
60 pub worker_id: WorkerId,
62 pub tasks: Vec<InFlightActivity>,
67 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 connections: HashMap<WorkerId, Instant>,
91 reachability: HashMap<WorkerId, Reachability>,
108}
109
110pub(crate) const DISPATCH_PROBATION_PINGS: u32 = 2;
132
133#[derive(Clone, Copy, Debug, Default)]
140struct Reachability {
141 consecutive_answers: u32,
142 proved_at: Option<Instant>,
143 ever_proved: bool,
155}
156
157impl Reachability {
158 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 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
191pub enum DispatchExclusion {
192 OpeningProbation {
197 answers: u32,
199 },
200 ReachabilityLost,
204}
205
206#[derive(Clone, Copy, Debug, PartialEq, Eq)]
208pub struct ExcludedWorker {
209 pub worker_id: WorkerId,
211 pub exclusion: DispatchExclusion,
213}
214
215#[derive(Clone, Debug)]
217pub struct HeartbeatTracker {
218 heartbeat_window: Duration,
219 inner: Arc<Mutex<HeartbeatState>>,
220 empty: Arc<Notify>,
221}
222
223impl HeartbeatTracker {
224 #[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 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 state
265 .reachability
266 .insert(worker_id, Reachability::default());
267 Ok(())
268 }
269
270 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 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 standing.ever_proved = true;
330 }
331 Ok(true)
332 }
333
334 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 Ok(true)
358 }
359
360 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 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 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 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 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 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 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 #[must_use]
533 pub const fn heartbeat_window(&self) -> Duration {
534 self.heartbeat_window
535 }
536
537 pub fn in_flight_count(&self) -> Result<usize, ServerError> {
543 Ok(self.state()?.tasks.len())
544 }
545
546 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 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 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 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 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 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 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 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 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 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 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#[must_use]
890pub fn sweep_interval(heartbeat_window: Duration) -> Duration {
891 const MINIMUM_PERIOD: Duration = Duration::from_millis(1);
894 const TARGET_FLOOR: Duration = Duration::from_secs(1);
897 let ceiling = heartbeat_window.max(MINIMUM_PERIOD);
898 (heartbeat_window / 4).clamp(TARGET_FLOOR.min(ceiling), ceiling)
900}
901
902pub struct HeartbeatSweeper<S> {
924 tracker: HeartbeatTracker,
925 registry: ConnectedWorkerRegistry,
926 sink: S,
927 drain: DrainState,
928 heartbeat_window: Duration,
929 interval: Duration,
930 queue_state: crate::worker::QueueServiceState,
935}
936
937impl<S> HeartbeatSweeper<S>
938where
939 S: ActivityCompletionSink + Send + Sync + 'static,
940{
941 #[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 #[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 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 if changed.is_err() || *shutdown.borrow() {
1002 break;
1003 }
1004 }
1005 }
1006 }
1007 info!("worker heartbeat sweeper stopped");
1008 }
1009
1010 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 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 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
1114fn 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 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 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 #[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 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 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 exclusion: DispatchExclusion::ReachabilityLost,
1219 }],
1220 "the worker must be named unreachable however alive its process looks"
1221 );
1222 Ok(())
1223 }
1224
1225 #[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 #[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 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 #[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 #[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 #[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 #[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 #[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 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 #[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 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 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 #[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 #[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(®istry, &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, ®istry, &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 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, ®istry, &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 #[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, ®istry, &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 let second = tracker.park_disconnected_worker(worker_id, ®istry, &sink)?;
1787 assert!(second.tasks.is_empty());
1788 let third = tracker.fail_disconnected_worker(worker_id, ®istry, &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 #[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 let notified = tracker.empty.notified();
1829 tokio::pin!(notified);
1830
1831 let reports = tracker.park_all_in_flight_workers(®istry, &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 #[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 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 #[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 #[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 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 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 assert_eq!(
2053 sweep_interval(Duration::from_secs(30)),
2054 Duration::from_millis(7_500)
2055 );
2056 assert_eq!(
2058 sweep_interval(Duration::from_secs(2)),
2059 Duration::from_secs(1)
2060 );
2061 assert_eq!(
2063 sweep_interval(Duration::from_secs(3_600)),
2064 Duration::from_secs(900)
2065 );
2066 assert_eq!(
2069 sweep_interval(Duration::from_millis(200)),
2070 Duration::from_millis(200)
2071 );
2072 assert_eq!(sweep_interval(Duration::ZERO), Duration::from_millis(1));
2075 }
2076}