1use dashmap::DashMap;
10use parking_lot::{Mutex, RwLock};
11use std::collections::VecDeque;
12use std::net::SocketAddr;
13use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
14use std::sync::Arc;
15use std::time::{Duration, Instant};
16
17#[derive(Debug, Clone)]
19pub struct FailureDetectorConfig {
20 pub timeout: Duration,
22 pub miss_threshold: u32,
24 pub suspicion_threshold: u32,
26 pub cleanup_interval: Duration,
28}
29
30impl Default for FailureDetectorConfig {
31 fn default() -> Self {
32 Self {
33 timeout: Duration::from_secs(5),
34 miss_threshold: 3,
35 suspicion_threshold: 2,
36 cleanup_interval: Duration::from_secs(30),
37 }
38 }
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum NodeStatus {
44 Healthy,
46 Suspected,
48 Failed,
50 Unknown,
52}
53
54#[derive(Debug)]
56struct NodeState {
57 last_heartbeat: Instant,
59 missed_count: u32,
61 status: NodeStatus,
63 #[allow(dead_code)]
65 addr: SocketAddr,
66 total_heartbeats: u64,
68 #[allow(dead_code)]
70 first_seen: Instant,
71}
72
73impl NodeState {
74 fn new(addr: SocketAddr) -> Self {
75 let now = Instant::now();
76 Self {
77 last_heartbeat: now,
78 missed_count: 0,
79 status: NodeStatus::Healthy,
80 addr,
81 total_heartbeats: 1,
82 first_seen: now,
83 }
84 }
85
86 fn on_heartbeat(&mut self) {
87 self.last_heartbeat = Instant::now();
88 self.missed_count = 0;
89 self.status = NodeStatus::Healthy;
90 self.total_heartbeats += 1;
91 }
92
93 fn check(
94 &mut self,
95 now: Instant,
96 timeout: Duration,
97 suspicion_threshold: u32,
98 miss_threshold: u32,
99 ) {
100 let elapsed = now.saturating_duration_since(self.last_heartbeat);
103
104 if elapsed > timeout {
105 let timeout_nanos = timeout.as_nanos().max(1);
110 self.missed_count = (elapsed.as_nanos() / timeout_nanos) as u32;
111
112 if self.missed_count >= miss_threshold {
113 self.status = NodeStatus::Failed;
114 } else if self.missed_count >= suspicion_threshold {
115 self.status = NodeStatus::Suspected;
116 }
117 }
118 }
119}
120
121#[derive(Debug, Clone, Default)]
123pub struct FailureStats {
124 pub nodes_tracked: usize,
126 pub nodes_healthy: usize,
128 pub nodes_suspected: usize,
130 pub nodes_failed: usize,
132 pub total_failures: u64,
134 pub total_recoveries: u64,
136}
137
138pub struct FailureDetector {
142 config: FailureDetectorConfig,
144 nodes: DashMap<u64, NodeState>,
146 on_failure: Option<Arc<dyn Fn(u64) + Send + Sync>>,
148 on_recovery: Option<Arc<dyn Fn(u64) + Send + Sync>>,
150 total_failures: AtomicU64,
152 total_recoveries: AtomicU64,
154 num_nodes: AtomicUsize,
159 last_cleanup: Mutex<Instant>,
161}
162
163impl FailureDetector {
164 pub fn new() -> Self {
166 Self::with_config(FailureDetectorConfig::default())
167 }
168
169 pub fn with_config(config: FailureDetectorConfig) -> Self {
171 Self {
172 config,
173 nodes: DashMap::new(),
174 on_failure: None,
175 on_recovery: None,
176 total_failures: AtomicU64::new(0),
177 total_recoveries: AtomicU64::new(0),
178 num_nodes: AtomicUsize::new(0),
179 last_cleanup: Mutex::new(Instant::now()),
180 }
181 }
182
183 pub fn on_failure<F>(mut self, f: F) -> Self
185 where
186 F: Fn(u64) + Send + Sync + 'static,
187 {
188 self.on_failure = Some(Arc::new(f));
189 self
190 }
191
192 pub fn on_recovery<F>(mut self, f: F) -> Self
194 where
195 F: Fn(u64) + Send + Sync + 'static,
196 {
197 self.on_recovery = Some(Arc::new(f));
198 self
199 }
200
201 pub fn heartbeat(&self, node_id: u64, addr: SocketAddr) {
213 let mut should_notify_recovery = false;
214 let mut node_inserted = false;
215 self.nodes
216 .entry(node_id)
217 .and_modify(|state| {
218 let was_failed = state.status == NodeStatus::Failed;
219 state.on_heartbeat();
220
221 if was_failed {
222 self.total_recoveries.fetch_add(1, Ordering::Relaxed);
223 should_notify_recovery = true;
224 }
225 })
226 .or_insert_with(|| {
227 node_inserted = true;
228 NodeState::new(addr)
229 });
230 if node_inserted {
231 self.num_nodes.fetch_add(1, Ordering::Relaxed);
232 }
233
234 if should_notify_recovery {
235 if let Some(ref cb) = self.on_recovery {
236 cb(node_id);
237 }
238 }
239 }
240
241 pub fn check_all(&self) -> Vec<u64> {
252 let mut newly_failed = Vec::new();
253
254 let now = Instant::now();
256 for mut entry in self.nodes.iter_mut() {
257 let prev_status = entry.status;
258 entry.check(
259 now,
260 self.config.timeout,
261 self.config.suspicion_threshold,
262 self.config.miss_threshold,
263 );
264
265 if entry.status == NodeStatus::Failed && prev_status != NodeStatus::Failed {
266 newly_failed.push(*entry.key());
267 self.total_failures.fetch_add(1, Ordering::Relaxed);
268 }
269 }
270
271 if let Some(ref cb) = self.on_failure {
272 for id in &newly_failed {
273 cb(*id);
274 }
275 }
276
277 newly_failed
278 }
279
280 pub fn status(&self, node_id: u64) -> NodeStatus {
282 self.nodes
283 .get(&node_id)
284 .map(|s| s.status)
285 .unwrap_or(NodeStatus::Unknown)
286 }
287
288 pub fn failed_nodes(&self) -> Vec<u64> {
290 self.nodes
291 .iter()
292 .filter(|r| r.status == NodeStatus::Failed)
293 .map(|r| *r.key())
294 .collect()
295 }
296
297 pub fn suspected_nodes(&self) -> Vec<u64> {
299 self.nodes
300 .iter()
301 .filter(|r| r.status == NodeStatus::Suspected)
302 .map(|r| *r.key())
303 .collect()
304 }
305
306 pub fn healthy_nodes(&self) -> Vec<u64> {
308 self.nodes
309 .iter()
310 .filter(|r| r.status == NodeStatus::Healthy)
311 .map(|r| *r.key())
312 .collect()
313 }
314
315 pub fn remove(&self, node_id: u64) {
317 if self.nodes.remove(&node_id).is_some() {
318 self.num_nodes.fetch_sub(1, Ordering::Relaxed);
319 }
320 }
321
322 pub fn cleanup(&self) -> usize {
324 let mut last = self.last_cleanup.lock();
331 if last.elapsed() < self.config.cleanup_interval {
332 return 0;
333 }
334 *last = Instant::now();
335 drop(last);
336
337 let stale_threshold = self.config.timeout * 10; let mut removed = 0;
339
340 self.nodes.retain(|_, state| {
341 if state.status == NodeStatus::Failed
342 && state.last_heartbeat.elapsed() > stale_threshold
343 {
344 removed += 1;
345 false
346 } else {
347 true
348 }
349 });
350
351 self.num_nodes.fetch_sub(removed, Ordering::Relaxed);
352 removed
353 }
354
355 pub fn stats(&self) -> FailureStats {
364 let mut healthy = 0;
365 let mut suspected = 0;
366 let mut failed = 0;
367
368 for entry in self.nodes.iter() {
369 match entry.status {
370 NodeStatus::Healthy => healthy += 1,
371 NodeStatus::Suspected => suspected += 1,
372 NodeStatus::Failed => failed += 1,
373 NodeStatus::Unknown => {}
374 }
375 }
376
377 FailureStats {
378 nodes_tracked: self.num_nodes.load(Ordering::Relaxed),
379 nodes_healthy: healthy,
380 nodes_suspected: suspected,
381 nodes_failed: failed,
382 total_failures: self.total_failures.load(Ordering::Relaxed),
383 total_recoveries: self.total_recoveries.load(Ordering::Relaxed),
384 }
385 }
386
387 pub fn node_count(&self) -> usize {
389 self.num_nodes.load(Ordering::Relaxed)
390 }
391}
392
393impl Default for FailureDetector {
394 fn default() -> Self {
395 Self::new()
396 }
397}
398
399pub struct LossSimulator {
403 loss_rate: f32,
405 in_burst: AtomicBool,
407 burst_prob: f32,
409 burst_length: u32,
411 burst_remaining: AtomicU64,
413 rng_state: AtomicU64,
415 total_packets: AtomicU64,
417 total_dropped: AtomicU64,
419}
420
421impl LossSimulator {
422 pub fn new(loss_rate: f32) -> Self {
424 Self {
425 loss_rate: loss_rate.clamp(0.0, 1.0),
426 in_burst: AtomicBool::new(false),
427 burst_prob: 0.0,
428 burst_length: 0,
429 burst_remaining: AtomicU64::new(0),
430 rng_state: AtomicU64::new(
431 std::time::SystemTime::now()
432 .duration_since(std::time::UNIX_EPOCH)
433 .unwrap_or_default()
434 .as_nanos() as u64,
435 ),
436 total_packets: AtomicU64::new(0),
437 total_dropped: AtomicU64::new(0),
438 }
439 }
440
441 pub fn with_bursts(mut self, burst_prob: f32, burst_length: u32) -> Self {
443 self.burst_prob = burst_prob.clamp(0.0, 1.0);
444 self.burst_length = burst_length;
445 self
446 }
447
448 pub fn should_drop(&self) -> bool {
450 self.total_packets.fetch_add(1, Ordering::Relaxed);
451
452 loop {
455 let remaining = self.burst_remaining.load(Ordering::Relaxed);
456 if remaining == 0 {
457 break;
458 }
459 match self.burst_remaining.compare_exchange_weak(
460 remaining,
461 remaining - 1,
462 Ordering::Relaxed,
463 Ordering::Relaxed,
464 ) {
465 Ok(_) => {
466 self.total_dropped.fetch_add(1, Ordering::Relaxed);
467 return true;
468 }
469 Err(_) => continue, }
471 }
472
473 let r = self.next_random();
475
476 if self.burst_prob > 0.0 && r < self.burst_prob {
478 self.burst_remaining.store(
481 self.burst_length.saturating_sub(1) as u64,
482 Ordering::Relaxed,
483 );
484 self.in_burst.store(true, Ordering::Relaxed);
485 self.total_dropped.fetch_add(1, Ordering::Relaxed);
486 return true;
487 }
488
489 if r < self.loss_rate {
491 self.total_dropped.fetch_add(1, Ordering::Relaxed);
492 return true;
493 }
494
495 false
496 }
497
498 pub fn effective_loss_rate(&self) -> f32 {
500 let total = self.total_packets.load(Ordering::Relaxed);
501 let dropped = self.total_dropped.load(Ordering::Relaxed);
502 if total == 0 {
503 return 0.0;
504 }
505 dropped as f32 / total as f32
506 }
507
508 pub fn reset(&self) {
510 self.total_packets.store(0, Ordering::Relaxed);
511 self.total_dropped.store(0, Ordering::Relaxed);
512 self.burst_remaining.store(0, Ordering::Relaxed);
513 self.in_burst.store(false, Ordering::Relaxed);
514 }
515
516 pub fn stats(&self) -> (u64, u64) {
518 (
519 self.total_packets.load(Ordering::Relaxed),
520 self.total_dropped.load(Ordering::Relaxed),
521 )
522 }
523
524 #[expect(
529 clippy::unwrap_used,
530 reason = "closure always returns Some, so fetch_update never returns Err"
531 )]
532 fn next_random(&self) -> f32 {
533 let prev = self
534 .rng_state
535 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |s| {
536 Some(s.wrapping_mul(6364136223846793005).wrapping_add(1))
537 })
538 .unwrap();
539 let new_state = prev.wrapping_mul(6364136223846793005).wrapping_add(1);
540 (new_state >> 33) as f32 / (1u64 << 31) as f32
541 }
542}
543
544#[derive(Debug, Clone, Copy, PartialEq, Eq)]
546pub enum CircuitState {
547 Closed,
549 Open,
551 HalfOpen,
553}
554
555pub struct CircuitBreaker {
557 state: RwLock<CircuitState>,
559 failure_count: AtomicU64,
561 success_count: AtomicU64,
563 failure_threshold: u64,
565 success_threshold: u64,
567 reset_timeout: Duration,
569 last_state_change: Mutex<Instant>,
571 total_trips: AtomicU64,
573}
574
575impl CircuitBreaker {
576 pub fn new(failure_threshold: u64, success_threshold: u64, reset_timeout: Duration) -> Self {
578 Self {
579 state: RwLock::new(CircuitState::Closed),
580 failure_count: AtomicU64::new(0),
581 success_count: AtomicU64::new(0),
582 failure_threshold,
583 success_threshold,
584 reset_timeout,
585 last_state_change: Mutex::new(Instant::now()),
586 total_trips: AtomicU64::new(0),
587 }
588 }
589
590 pub fn allow(&self) -> bool {
592 {
595 let state = *self.state.read();
596 match state {
597 CircuitState::Closed | CircuitState::HalfOpen => return true,
598 CircuitState::Open => {} }
600 }
601 let mut state = self.state.write();
611 match *state {
612 CircuitState::Closed | CircuitState::HalfOpen => true,
613 CircuitState::Open => {
614 let elapsed = self.last_state_change.lock().elapsed();
615 if elapsed >= self.reset_timeout {
616 Self::transition_locked(
617 &mut state,
618 CircuitState::HalfOpen,
619 &self.failure_count,
620 &self.success_count,
621 &self.last_state_change,
622 &self.total_trips,
623 );
624 true
625 } else {
626 false
627 }
628 }
629 }
630 }
631
632 pub fn record_success(&self) {
634 let mut state = self.state.write();
638 match *state {
639 CircuitState::Closed => {
640 self.failure_count.store(0, Ordering::Relaxed);
642 }
643 CircuitState::HalfOpen => {
644 let count = self.success_count.fetch_add(1, Ordering::Relaxed) + 1;
645 if count >= self.success_threshold {
646 Self::transition_locked(
647 &mut state,
648 CircuitState::Closed,
649 &self.failure_count,
650 &self.success_count,
651 &self.last_state_change,
652 &self.total_trips,
653 );
654 }
655 }
656 CircuitState::Open => {}
657 }
658 }
659
660 pub fn record_failure(&self) {
662 let mut state = self.state.write();
666 match *state {
667 CircuitState::Closed => {
668 let count = self.failure_count.fetch_add(1, Ordering::Relaxed) + 1;
669 if count >= self.failure_threshold {
670 Self::transition_locked(
671 &mut state,
672 CircuitState::Open,
673 &self.failure_count,
674 &self.success_count,
675 &self.last_state_change,
676 &self.total_trips,
677 );
678 }
679 }
680 CircuitState::HalfOpen => {
681 Self::transition_locked(
683 &mut state,
684 CircuitState::Open,
685 &self.failure_count,
686 &self.success_count,
687 &self.last_state_change,
688 &self.total_trips,
689 );
690 }
691 CircuitState::Open => {}
692 }
693 }
694
695 pub fn state(&self) -> CircuitState {
697 *self.state.read()
698 }
699
700 pub fn total_trips(&self) -> u64 {
702 self.total_trips.load(Ordering::Relaxed)
703 }
704
705 pub fn reset(&self) {
707 self.transition_to(CircuitState::Closed);
708 self.failure_count.store(0, Ordering::Relaxed);
709 self.success_count.store(0, Ordering::Relaxed);
710 }
711
712 fn transition_to(&self, new_state: CircuitState) {
713 let mut state = self.state.write();
714 Self::transition_locked(
715 &mut state,
716 new_state,
717 &self.failure_count,
718 &self.success_count,
719 &self.last_state_change,
720 &self.total_trips,
721 );
722 }
723
724 fn transition_locked(
727 state: &mut CircuitState,
728 new_state: CircuitState,
729 failure_count: &AtomicU64,
730 success_count: &AtomicU64,
731 last_state_change: &Mutex<Instant>,
732 total_trips: &AtomicU64,
733 ) {
734 let old_state = *state;
735 if old_state != new_state {
736 *state = new_state;
737 *last_state_change.lock() = Instant::now();
738
739 failure_count.store(0, Ordering::Relaxed);
741 success_count.store(0, Ordering::Relaxed);
742
743 if new_state == CircuitState::Open {
745 total_trips.fetch_add(1, Ordering::Relaxed);
746 }
747 }
748 }
749}
750
751#[derive(Debug, Clone)]
753pub enum RecoveryAction {
754 Reroute {
756 via: Vec<u64>,
758 },
759 Retry {
761 delay_ms: u64,
763 },
764 Drop {
766 reason: String,
768 },
769 Queue,
771}
772
773#[derive(Debug, Clone, Default)]
775pub struct RecoveryStats {
776 pub reroutes: u64,
778 pub retries: u64,
780 pub dropped: u64,
782 pub queued: u64,
784 pub avg_recovery_ms: u64,
786}
787
788pub struct RecoveryManager {
790 failed_nodes: DashMap<u64, FailedNodeState>,
792 recovery_queue: Mutex<VecDeque<(u64, Instant)>>,
794 reroutes: AtomicU64,
796 retries: AtomicU64,
797 dropped: AtomicU64,
798 queued: AtomicU64,
799 total_recovery_time_ms: AtomicU64,
800 recovery_count: AtomicU64,
801}
802
803#[derive(Debug)]
804struct FailedNodeState {
805 failed_at: Instant,
807 retry_count: u32,
809 alternates: Vec<u64>,
811}
812
813impl RecoveryManager {
814 pub fn new() -> Self {
816 Self {
817 failed_nodes: DashMap::new(),
818 recovery_queue: Mutex::new(VecDeque::new()),
819 reroutes: AtomicU64::new(0),
820 retries: AtomicU64::new(0),
821 dropped: AtomicU64::new(0),
822 queued: AtomicU64::new(0),
823 total_recovery_time_ms: AtomicU64::new(0),
824 recovery_count: AtomicU64::new(0),
825 }
826 }
827
828 pub fn on_failure(&self, node_id: u64, alternates: Vec<u64>) -> RecoveryAction {
830 self.failed_nodes
838 .entry(node_id)
839 .and_modify(|s| {
840 if !alternates.is_empty() {
841 s.alternates = alternates.clone();
842 }
843 })
844 .or_insert_with(|| FailedNodeState {
845 failed_at: Instant::now(),
846 retry_count: 0,
847 alternates: alternates.clone(),
848 });
849
850 if !alternates.is_empty() {
851 self.reroutes.fetch_add(1, Ordering::Relaxed);
852 RecoveryAction::Reroute { via: alternates }
853 } else {
854 self.queued.fetch_add(1, Ordering::Relaxed);
855 self.recovery_queue
856 .lock()
857 .push_back((node_id, Instant::now()));
858 RecoveryAction::Queue
859 }
860 }
861
862 pub fn on_recovery(&self, node_id: u64) {
864 if let Some((_, state)) = self.failed_nodes.remove(&node_id) {
865 let recovery_time = state.failed_at.elapsed().as_millis() as u64;
866 self.total_recovery_time_ms
867 .fetch_add(recovery_time, Ordering::Relaxed);
868 self.recovery_count.fetch_add(1, Ordering::Relaxed);
869 }
870 }
871
872 pub fn get_action(&self, node_id: u64, max_retries: u32) -> RecoveryAction {
874 if let Some(mut state) = self.failed_nodes.get_mut(&node_id) {
875 if !state.alternates.is_empty() {
876 return RecoveryAction::Reroute {
877 via: state.alternates.clone(),
878 };
879 }
880
881 if state.retry_count < max_retries {
882 state.retry_count += 1;
883 self.retries.fetch_add(1, Ordering::Relaxed);
884 let delay = 100 * (1 << state.retry_count.min(6)); return RecoveryAction::Retry { delay_ms: delay };
886 }
887
888 self.dropped.fetch_add(1, Ordering::Relaxed);
889 RecoveryAction::Drop {
890 reason: "max retries exceeded".into(),
891 }
892 } else {
893 RecoveryAction::Retry { delay_ms: 100 }
904 }
905 }
906
907 pub fn is_failed(&self, node_id: u64) -> bool {
909 self.failed_nodes.contains_key(&node_id)
910 }
911
912 pub fn stats(&self) -> RecoveryStats {
914 let count = self.recovery_count.load(Ordering::Relaxed);
915 let total_time = self.total_recovery_time_ms.load(Ordering::Relaxed);
916 let avg = total_time.checked_div(count).unwrap_or(0);
917
918 RecoveryStats {
919 reroutes: self.reroutes.load(Ordering::Relaxed),
920 retries: self.retries.load(Ordering::Relaxed),
921 dropped: self.dropped.load(Ordering::Relaxed),
922 queued: self.queued.load(Ordering::Relaxed),
923 avg_recovery_ms: avg,
924 }
925 }
926
927 pub fn failed_count(&self) -> usize {
929 self.failed_nodes.len()
930 }
931}
932
933impl Default for RecoveryManager {
934 fn default() -> Self {
935 Self::new()
936 }
937}
938
939#[cfg(test)]
940mod tests {
941 use super::*;
942
943 #[test]
944 fn test_failure_detector_basic() {
945 let detector = FailureDetector::with_config(FailureDetectorConfig {
946 timeout: Duration::from_millis(100),
947 miss_threshold: 2,
948 suspicion_threshold: 1,
949 cleanup_interval: Duration::from_secs(60),
950 });
951
952 let addr: SocketAddr = "127.0.0.1:9000".parse().unwrap();
953 detector.heartbeat(0x1234, addr);
954
955 assert_eq!(detector.status(0x1234), NodeStatus::Healthy);
956 assert_eq!(detector.node_count(), 1);
957 }
958
959 #[test]
963 fn node_count_tracks_heartbeats_and_removal() {
964 let detector = FailureDetector::new();
965 let addr: SocketAddr = "127.0.0.1:9000".parse().unwrap();
966
967 detector.heartbeat(1, addr);
968 detector.heartbeat(2, addr);
969 detector.heartbeat(3, addr);
970 assert_eq!(detector.node_count(), 3);
971 assert_eq!(detector.stats().nodes_tracked, 3);
972
973 detector.heartbeat(1, addr);
975 assert_eq!(detector.node_count(), 3, "re-heartbeat must not grow count");
976
977 detector.remove(2);
978 assert_eq!(detector.node_count(), 2);
979 assert_eq!(detector.stats().nodes_tracked, 2);
980
981 detector.remove(999);
983 assert_eq!(detector.node_count(), 2);
984 }
985
986 #[test]
987 fn test_failure_detector_failure() {
988 let detector = FailureDetector::with_config(FailureDetectorConfig {
999 timeout: Duration::from_millis(100),
1000 miss_threshold: 2,
1001 suspicion_threshold: 1,
1002 cleanup_interval: Duration::from_secs(60),
1003 });
1004
1005 let addr: SocketAddr = "127.0.0.1:9000".parse().unwrap();
1006 detector.heartbeat(0x1234, addr);
1007
1008 std::thread::sleep(Duration::from_millis(150));
1010
1011 detector.check_all();
1013 assert_eq!(detector.status(0x1234), NodeStatus::Suspected);
1014
1015 std::thread::sleep(Duration::from_millis(150));
1017
1018 let failed = detector.check_all();
1020 assert_eq!(failed.len(), 1);
1021 assert_eq!(failed[0], 0x1234);
1022 assert_eq!(detector.status(0x1234), NodeStatus::Failed);
1023 }
1024
1025 #[test]
1026 fn test_failure_detector_recovery() {
1027 let detector = FailureDetector::with_config(FailureDetectorConfig {
1028 timeout: Duration::from_millis(10),
1029 miss_threshold: 1,
1030 suspicion_threshold: 1,
1031 cleanup_interval: Duration::from_secs(60),
1032 });
1033
1034 let addr: SocketAddr = "127.0.0.1:9000".parse().unwrap();
1035 detector.heartbeat(0x1234, addr);
1036
1037 std::thread::sleep(Duration::from_millis(15));
1038 detector.check_all();
1039 assert_eq!(detector.status(0x1234), NodeStatus::Failed);
1040
1041 detector.heartbeat(0x1234, addr);
1043 assert_eq!(detector.status(0x1234), NodeStatus::Healthy);
1044
1045 let stats = detector.stats();
1046 assert_eq!(stats.total_failures, 1);
1047 assert_eq!(stats.total_recoveries, 1);
1048 }
1049
1050 #[test]
1051 fn test_failure_detector_elapsed_based_missed_count() {
1052 let detector = FailureDetector::with_config(FailureDetectorConfig {
1057 timeout: Duration::from_millis(10),
1058 miss_threshold: 3,
1059 suspicion_threshold: 2,
1060 cleanup_interval: Duration::from_secs(60),
1061 });
1062
1063 let addr: SocketAddr = "127.0.0.1:9000".parse().unwrap();
1064 detector.heartbeat(0x1234, addr);
1065
1066 std::thread::sleep(Duration::from_millis(35));
1068
1069 let failed = detector.check_all();
1073 assert_eq!(
1074 detector.status(0x1234),
1075 NodeStatus::Failed,
1076 "node should be Failed after 3+ timeout intervals, even with one check call"
1077 );
1078 assert_eq!(failed.len(), 1);
1079 }
1080
1081 #[test]
1082 fn test_loss_simulator() {
1083 let sim = LossSimulator::new(0.5);
1084
1085 let mut dropped = 0;
1086 for _ in 0..1000 {
1087 if sim.should_drop() {
1088 dropped += 1;
1089 }
1090 }
1091
1092 assert!(dropped > 300 && dropped < 700);
1094 }
1095
1096 #[test]
1097 fn test_loss_simulator_burst() {
1098 let sim = LossSimulator::new(0.0).with_bursts(0.1, 5);
1099
1100 let mut total_bursts = 0;
1101 let mut in_burst = false;
1102 for _ in 0..1000 {
1103 if sim.should_drop() {
1104 if !in_burst {
1105 in_burst = true;
1106 total_bursts += 1;
1107 }
1108 } else {
1109 in_burst = false;
1110 }
1111 }
1112
1113 assert!(total_bursts > 0);
1115 }
1116
1117 #[test]
1118 fn test_burst_drops_exactly_burst_length_packets() {
1119 let burst_len = 5u32;
1125 let sim = LossSimulator::new(0.0).with_bursts(1.0, burst_len);
1126
1127 assert!(sim.should_drop());
1129 let remaining = sim.burst_remaining.load(Ordering::Relaxed);
1131 assert_eq!(
1132 remaining,
1133 (burst_len - 1) as u64,
1134 "after trigger, burst_remaining should be burst_length - 1, \
1135 not burst_length (which would cause burst_length + 1 total drops)"
1136 );
1137
1138 for _ in 0..remaining {
1140 assert!(sim.should_drop());
1141 }
1142
1143 assert_eq!(sim.burst_remaining.load(Ordering::Relaxed), 0);
1145 assert_eq!(sim.total_dropped.load(Ordering::Relaxed), burst_len as u64);
1146 }
1147
1148 #[test]
1149 fn test_circuit_breaker() {
1150 let cb = CircuitBreaker::new(3, 2, Duration::from_millis(50));
1151
1152 assert_eq!(cb.state(), CircuitState::Closed);
1153 assert!(cb.allow());
1154
1155 cb.record_failure();
1157 cb.record_failure();
1158 cb.record_failure();
1159
1160 assert_eq!(cb.state(), CircuitState::Open);
1161 assert!(!cb.allow());
1162
1163 std::thread::sleep(Duration::from_millis(60));
1165
1166 assert!(cb.allow());
1168 assert_eq!(cb.state(), CircuitState::HalfOpen);
1169
1170 cb.record_success();
1172 cb.record_success();
1173 assert_eq!(cb.state(), CircuitState::Closed);
1174 }
1175
1176 #[test]
1177 fn test_regression_loss_simulator_burst_no_underflow() {
1178 use std::sync::Arc;
1182
1183 let sim = Arc::new(LossSimulator::new(0.0).with_bursts(0.3, 10));
1184 let threads: Vec<_> = (0..8)
1185 .map(|_| {
1186 let sim = Arc::clone(&sim);
1187 std::thread::spawn(move || {
1188 for _ in 0..5_000 {
1189 sim.should_drop();
1190 }
1191 })
1192 })
1193 .collect();
1194
1195 for t in threads {
1196 t.join().unwrap();
1197 }
1198
1199 let (total, dropped) = sim.stats();
1200 assert!(
1203 dropped <= total,
1204 "dropped ({dropped}) must not exceed total ({total}) — \
1205 would indicate burst_remaining underflow"
1206 );
1207 assert_eq!(total, 8 * 5_000);
1209 }
1210
1211 #[test]
1212 fn test_regression_circuit_breaker_concurrent_transitions() {
1213 use std::sync::Arc;
1218
1219 let cb = Arc::new(CircuitBreaker::new(3, 2, Duration::from_millis(10)));
1220
1221 let threads: Vec<_> = (0..8)
1222 .map(|i| {
1223 let cb = Arc::clone(&cb);
1224 std::thread::spawn(move || {
1225 for _ in 0..2_000 {
1226 if i % 2 == 0 {
1227 cb.record_failure();
1228 } else {
1229 cb.record_success();
1230 }
1231 }
1232 })
1233 })
1234 .collect();
1235
1236 for t in threads {
1237 t.join().unwrap();
1238 }
1239
1240 let state = cb.state();
1242 assert!(
1243 state == CircuitState::Closed
1244 || state == CircuitState::Open
1245 || state == CircuitState::HalfOpen,
1246 "circuit breaker state is invalid after concurrent access"
1247 );
1248 let trips = cb.total_trips();
1250 assert!(
1252 trips <= 8_000,
1253 "total_trips ({trips}) is unreasonably high, suggests corruption"
1254 );
1255 }
1256
1257 #[test]
1258 fn test_regression_allow_does_not_undo_reset() {
1259 use std::sync::atomic::{AtomicU8, Ordering};
1280 use std::sync::Arc;
1281 use std::thread;
1282
1283 const TRIALS: u32 = 5_000;
1284
1285 let cb = Arc::new(CircuitBreaker::new(1, 1, Duration::from_nanos(1)));
1286 let signal = Arc::new(AtomicU8::new(0)); let cb_observer = cb.clone();
1289 let signal_observer = signal.clone();
1290 let observer = thread::spawn(move || loop {
1291 match signal_observer.load(Ordering::Acquire) {
1292 0 => std::hint::spin_loop(),
1293 1 => {
1294 cb_observer.allow();
1295 signal_observer.store(0, Ordering::Release);
1296 }
1297 _ => return,
1298 }
1299 });
1300
1301 let mut bug_count = 0u32;
1302 for _ in 0..TRIALS {
1303 cb.record_failure();
1305 assert_eq!(cb.state(), CircuitState::Open);
1306
1307 signal.store(1, Ordering::Release);
1309 cb.reset();
1310 while signal.load(Ordering::Acquire) != 0 {
1311 std::hint::spin_loop();
1312 }
1313
1314 if cb.state() != CircuitState::Closed {
1315 bug_count += 1;
1316 cb.reset();
1319 }
1320 }
1321
1322 signal.store(2, Ordering::Release);
1323 observer.join().unwrap();
1324
1325 assert_eq!(
1326 bug_count, 0,
1327 "{bug_count} of {TRIALS} trials ended in non-Closed state — \
1328 allow() transitioned to HalfOpen on top of a fresh reset()"
1329 );
1330 }
1331
1332 #[test]
1333 fn test_recovery_manager() {
1334 let mgr = RecoveryManager::new();
1335
1336 let action = mgr.on_failure(0x1234, vec![0x5678, 0x9ABC]);
1338 match action {
1339 RecoveryAction::Reroute { via } => {
1340 assert_eq!(via, vec![0x5678, 0x9ABC]);
1341 }
1342 _ => panic!("expected reroute"),
1343 }
1344
1345 let action = mgr.on_failure(0x2222, vec![]);
1347 match action {
1348 RecoveryAction::Queue => {}
1349 _ => panic!("expected queue"),
1350 }
1351
1352 assert!(mgr.is_failed(0x1234));
1353 assert!(mgr.is_failed(0x2222));
1354
1355 mgr.on_recovery(0x1234);
1357 assert!(!mgr.is_failed(0x1234));
1358
1359 let stats = mgr.stats();
1360 assert_eq!(stats.reroutes, 1);
1361 assert_eq!(stats.queued, 1);
1362 }
1363
1364 #[test]
1370 fn on_failure_preserves_retry_count_on_repeat() {
1371 let mgr = RecoveryManager::new();
1372 let node = 0x42u64;
1373 let max_retries = 3u32;
1374
1375 let action = mgr.on_failure(node, vec![]);
1378 assert!(matches!(action, RecoveryAction::Queue));
1379
1380 for expected_count in 1..=max_retries {
1382 match mgr.get_action(node, max_retries) {
1383 RecoveryAction::Retry { .. } => {}
1384 other => panic!(
1385 "expected Retry on attempt {} (count would become {}), got {:?}",
1386 expected_count, expected_count, other
1387 ),
1388 }
1389 }
1390
1391 let _ = mgr.on_failure(node, vec![]);
1396
1397 match mgr.get_action(node, max_retries) {
1400 RecoveryAction::Drop { .. } => {}
1401 other => panic!(
1402 "expected Drop after exhausting retries across a flap; got {:?} \
1403 (pre-fix on_failure reset retry_count to 0 on repeat)",
1404 other
1405 ),
1406 }
1407 }
1408
1409 #[test]
1414 fn on_failure_repeat_updates_alternates_without_resetting_count() {
1415 let mgr = RecoveryManager::new();
1416 let node = 0x99u64;
1417 let max_retries = 2u32;
1418
1419 let _ = mgr.on_failure(node, vec![]);
1421 let _ = mgr.get_action(node, max_retries);
1423
1424 let action = mgr.on_failure(node, vec![0xDEAD]);
1428 match action {
1429 RecoveryAction::Reroute { via } => assert_eq!(via, vec![0xDEAD]),
1430 other => panic!("expected Reroute, got {:?}", other),
1431 }
1432
1433 if let Some(mut s) = mgr.failed_nodes.get_mut(&node) {
1438 s.alternates.clear();
1439 }
1440 let _ = mgr.get_action(node, max_retries); match mgr.get_action(node, max_retries) {
1442 RecoveryAction::Drop { .. } => {}
1443 other => panic!("expected Drop after exhausting retries; got {:?}", other),
1444 }
1445 }
1446
1447 #[test]
1464 fn callbacks_run_after_shard_lock_release() {
1465 use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};
1466 use std::sync::Arc;
1467
1468 let detector = Arc::new(FailureDetector::with_config(FailureDetectorConfig {
1469 timeout: Duration::from_millis(10),
1470 miss_threshold: 1,
1471 suspicion_threshold: 1,
1472 cleanup_interval: Duration::from_secs(60),
1473 }));
1474
1475 let detector_for_cb = Arc::clone(&detector);
1476 let observed = Arc::new(AtomicBool::new(false));
1477 let observed_clone = Arc::clone(&observed);
1478
1479 let detector_arc = Arc::new(
1485 FailureDetector::with_config(FailureDetectorConfig {
1488 timeout: Duration::from_millis(10),
1489 miss_threshold: 1,
1490 suspicion_threshold: 1,
1491 cleanup_interval: Duration::from_secs(60),
1492 })
1493 .on_recovery(move |id| {
1494 let _ = detector_for_cb.status(id);
1497 observed_clone.store(true, AtomicOrdering::SeqCst);
1498 }),
1499 );
1500
1501 let addr: SocketAddr = "127.0.0.1:9000".parse().unwrap();
1502 detector_arc.heartbeat(0x4242, addr);
1504 std::thread::sleep(Duration::from_millis(25));
1505 let _ = detector_arc.check_all();
1506 assert_eq!(detector_arc.status(0x4242), NodeStatus::Failed);
1507
1508 detector_arc.heartbeat(0x4242, addr);
1510
1511 assert!(
1512 observed.load(AtomicOrdering::SeqCst),
1513 "recovery callback must have run (and re-entered status()) — \
1514 a deadlock here would manifest as the test hanging (#14)"
1515 );
1516 let _ = detector;
1517 }
1518
1519 #[test]
1528 fn get_action_on_unfailed_node_does_not_busy_loop() {
1529 let mgr = RecoveryManager::new();
1530 let untracked = 0xDEAD_BEEFu64;
1531
1532 assert!(
1534 !mgr.is_failed(untracked),
1535 "precondition: node must not be tracked as failed"
1536 );
1537
1538 let action = mgr.get_action(untracked, 3);
1539 match action {
1540 RecoveryAction::Retry { delay_ms } => {
1541 assert!(
1542 delay_ms > 0,
1543 "regression: get_action on an unfailed node returned \
1544 Retry {{ delay_ms: 0 }} — a delay-respecting caller \
1545 would busy-loop on this and saturate a CPU"
1546 );
1547 assert_eq!(
1548 delay_ms, 100,
1549 "first-step backoff should match the failed-node \
1550 path's retry-1 delay (100ms) so callers pace \
1551 consistently across both branches"
1552 );
1553 }
1554 other => panic!("unfailed-node branch must return Retry, got {:?}", other),
1555 }
1556 }
1557
1558 #[test]
1568 fn circuit_breaker_half_open_failure_trips_back_to_open() {
1569 let cb = CircuitBreaker::new(2, 2, Duration::from_nanos(1));
1579
1580 cb.record_failure();
1582 cb.record_failure();
1583 assert_eq!(cb.state(), CircuitState::Open);
1584
1585 assert!(
1587 cb.allow(),
1588 "expected allow() to admit a probe after reset_timeout"
1589 );
1590 assert_eq!(cb.state(), CircuitState::HalfOpen);
1591
1592 cb.record_failure();
1594 assert_eq!(
1595 cb.state(),
1596 CircuitState::Open,
1597 "HalfOpen + failure must snap back to Open; \
1598 a regression here keeps probing a broken backend",
1599 );
1600 }
1601
1602 #[test]
1611 fn failure_detector_config_default_values() {
1612 let cfg = FailureDetectorConfig::default();
1613 assert_eq!(cfg.timeout, Duration::from_secs(5));
1614 assert_eq!(cfg.miss_threshold, 3);
1615 assert_eq!(cfg.suspicion_threshold, 2);
1616 assert_eq!(cfg.cleanup_interval, Duration::from_secs(30));
1617 }
1618
1619 #[test]
1625 fn suspected_and_healthy_nodes_filter_by_status() {
1626 let detector = FailureDetector::new();
1627 let addr: SocketAddr = "127.0.0.1:9000".parse().unwrap();
1632 detector.heartbeat(1, addr);
1633 detector.heartbeat(2, addr);
1634 detector.heartbeat(3, addr);
1635 detector.nodes.get_mut(&2).unwrap().status = NodeStatus::Suspected;
1637 detector.nodes.get_mut(&3).unwrap().status = NodeStatus::Failed;
1638
1639 let mut healthy = detector.healthy_nodes();
1640 healthy.sort_unstable();
1641 assert_eq!(healthy, vec![1]);
1642
1643 let suspected = detector.suspected_nodes();
1644 assert_eq!(suspected, vec![2]);
1645 }
1646
1647 #[test]
1652 fn cleanup_returns_zero_within_cleanup_interval() {
1653 let detector = FailureDetector::with_config(FailureDetectorConfig {
1654 cleanup_interval: Duration::from_secs(60),
1655 ..Default::default()
1656 });
1657 *detector.last_cleanup.lock() = Instant::now();
1660 assert_eq!(
1661 detector.cleanup(),
1662 0,
1663 "cleanup called inside the rate-limit window must return 0 without scanning",
1664 );
1665 }
1666
1667 #[test]
1673 fn loss_simulator_effective_loss_rate_handles_div_by_zero_and_ratio() {
1674 let sim = LossSimulator::new(1.0); assert_eq!(sim.effective_loss_rate(), 0.0, "no packets → 0, not NaN");
1676
1677 for _ in 0..4 {
1680 let _ = sim.should_drop();
1681 }
1682 let rate = sim.effective_loss_rate();
1683 assert!(
1684 (rate - 1.0).abs() < 1e-6,
1685 "expected loss_rate ≈ 1.0 after 4 drops; got {rate}",
1686 );
1687
1688 sim.reset();
1692 assert_eq!(sim.total_packets.load(Ordering::Relaxed), 0);
1693 assert_eq!(sim.total_dropped.load(Ordering::Relaxed), 0);
1694 assert_eq!(sim.burst_remaining.load(Ordering::Relaxed), 0);
1695 assert!(!sim.in_burst.load(Ordering::Relaxed));
1696 assert_eq!(sim.effective_loss_rate(), 0.0);
1697 }
1698}