1use std::collections::HashMap;
27
28#[derive(Clone, Debug, PartialEq)]
34pub enum ConnectionEvent {
35 KeepaliveSuccess { rtt_ms: f64 },
37 KeepaliveFailed,
39 DataSent { bytes: u64 },
41 DataReceived { bytes: u64 },
43 Reset { reason: String },
45 Reconnected,
47}
48
49#[derive(Clone, Debug, PartialEq)]
55pub enum ConnectionHealthState {
56 Healthy,
58 Degraded { consecutive_failures: u32 },
61 Dead { reason: String },
63 Reconnecting,
66}
67
68const MAX_EVENTS: usize = 50;
74
75const RTT_ALPHA: f64 = 0.2;
77
78#[derive(Clone, Debug)]
80pub struct ConnectionRecord {
81 pub conn_id: u64,
83 pub peer_id: String,
85 pub state: ConnectionHealthState,
87 pub events: Vec<ConnectionEvent>,
89 pub bytes_sent: u64,
91 pub bytes_received: u64,
93 pub consecutive_failures: u32,
95 pub avg_rtt_ms: f64,
98}
99
100impl ConnectionRecord {
101 pub fn is_alive(&self) -> bool {
103 !matches!(self.state, ConnectionHealthState::Dead { .. })
104 }
105
106 fn push_event(&mut self, event: ConnectionEvent) {
108 if self.events.len() >= MAX_EVENTS {
109 self.events.remove(0);
110 }
111 self.events.push(event);
112 }
113
114 fn update_rtt(&mut self, rtt_ms: f64) {
119 if self.avg_rtt_ms == 0.0 {
120 self.avg_rtt_ms = rtt_ms;
121 } else {
122 self.avg_rtt_ms = RTT_ALPHA * rtt_ms + (1.0 - RTT_ALPHA) * self.avg_rtt_ms;
123 }
124 }
125}
126
127#[derive(Clone, Debug)]
133pub struct HealthCheckerConfig {
134 pub failure_threshold: u32,
137 pub dead_threshold: u32,
140 pub rtt_spike_factor: f64,
145}
146
147impl Default for HealthCheckerConfig {
148 fn default() -> Self {
149 Self {
150 failure_threshold: 3,
151 dead_threshold: 10,
152 rtt_spike_factor: 2.0,
153 }
154 }
155}
156
157pub struct ConnectionHealthChecker {
170 pub connections: HashMap<u64, ConnectionRecord>,
172 pub config: HealthCheckerConfig,
174 pub next_conn_id: u64,
176}
177
178impl ConnectionHealthChecker {
179 pub fn new(config: HealthCheckerConfig) -> Self {
181 Self {
182 connections: HashMap::new(),
183 config,
184 next_conn_id: 1,
185 }
186 }
187
188 pub fn open_connection(&mut self, peer_id: String) -> u64 {
193 let conn_id = self.next_conn_id;
194 self.next_conn_id += 1;
195
196 let record = ConnectionRecord {
197 conn_id,
198 peer_id,
199 state: ConnectionHealthState::Healthy,
200 events: Vec::new(),
201 bytes_sent: 0,
202 bytes_received: 0,
203 consecutive_failures: 0,
204 avg_rtt_ms: 0.0,
205 };
206 self.connections.insert(conn_id, record);
207 conn_id
208 }
209
210 pub fn record_event(&mut self, conn_id: u64, event: ConnectionEvent) {
215 let Some(record) = self.connections.get_mut(&conn_id) else {
216 return;
217 };
218
219 record.push_event(event.clone());
221
222 match event {
223 ConnectionEvent::KeepaliveFailed => {
224 record.consecutive_failures += 1;
225
226 if record.consecutive_failures >= self.config.dead_threshold {
227 record.state = ConnectionHealthState::Dead {
228 reason: format!(
229 "exceeded dead threshold ({} consecutive failures)",
230 record.consecutive_failures
231 ),
232 };
233 } else if record.consecutive_failures >= self.config.failure_threshold {
234 record.state = ConnectionHealthState::Degraded {
235 consecutive_failures: record.consecutive_failures,
236 };
237 }
238 }
239
240 ConnectionEvent::KeepaliveSuccess { rtt_ms } => {
241 record.consecutive_failures = 0;
242 record.update_rtt(rtt_ms);
243 record.state = ConnectionHealthState::Healthy;
244 }
245
246 ConnectionEvent::Reset { reason } => {
247 record.state = ConnectionHealthState::Dead { reason };
248 }
249
250 ConnectionEvent::Reconnected => {
251 record.consecutive_failures = 0;
252 record.state = ConnectionHealthState::Reconnecting;
256 record.state = ConnectionHealthState::Healthy;
257 }
258
259 ConnectionEvent::DataSent { bytes } => {
260 record.bytes_sent = record.bytes_sent.saturating_add(bytes);
261 }
262
263 ConnectionEvent::DataReceived { bytes } => {
264 record.bytes_received = record.bytes_received.saturating_add(bytes);
265 }
266 }
267 }
268
269 pub fn close_connection(&mut self, conn_id: u64) -> bool {
274 self.connections.remove(&conn_id).is_some()
275 }
276
277 pub fn alive_connections(&self) -> Vec<&ConnectionRecord> {
280 self.connections.values().filter(|r| r.is_alive()).collect()
281 }
282
283 pub fn dead_connections(&self) -> Vec<&ConnectionRecord> {
286 self.connections
287 .values()
288 .filter(|r| !r.is_alive())
289 .collect()
290 }
291
292 pub fn stats(&self) -> (usize, usize, usize) {
294 let total = self.connections.len();
295 let dead = self.connections.values().filter(|r| !r.is_alive()).count();
296 let alive = total - dead;
297 (total, alive, dead)
298 }
299}
300
301#[cfg(test)]
306mod tests {
307 use super::*;
308
309 fn default_checker() -> ConnectionHealthChecker {
310 ConnectionHealthChecker::new(HealthCheckerConfig::default())
311 }
312
313 #[test]
315 fn test_open_connection_creates_record() {
316 let mut checker = default_checker();
317 let id = checker.open_connection("peer-1".to_string());
318 let record = checker.connections.get(&id).expect("record must exist");
319 assert_eq!(record.peer_id, "peer-1");
320 assert_eq!(record.conn_id, id);
321 }
322
323 #[test]
325 fn test_open_connection_starts_healthy() {
326 let mut checker = default_checker();
327 let id = checker.open_connection("peer-2".to_string());
328 let record = checker.connections.get(&id).expect("record must exist");
329 assert!(matches!(record.state, ConnectionHealthState::Healthy));
330 }
331
332 #[test]
334 fn test_open_connection_initial_counters_zero() {
335 let mut checker = default_checker();
336 let id = checker.open_connection("peer-3".to_string());
337 let record = checker.connections.get(&id).expect("record must exist");
338 assert_eq!(record.bytes_sent, 0);
339 assert_eq!(record.bytes_received, 0);
340 assert_eq!(record.consecutive_failures, 0);
341 assert_eq!(record.avg_rtt_ms, 0.0);
342 }
343
344 #[test]
346 fn test_open_connection_unique_ids() {
347 let mut checker = default_checker();
348 let id1 = checker.open_connection("peer-a".to_string());
349 let id2 = checker.open_connection("peer-b".to_string());
350 let id3 = checker.open_connection("peer-c".to_string());
351 assert_ne!(id1, id2);
352 assert_ne!(id2, id3);
353 assert_ne!(id1, id3);
354 }
355
356 #[test]
358 fn test_keepalive_failed_increments_failures() {
359 let mut checker = default_checker();
360 let id = checker.open_connection("peer".to_string());
361 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
362 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
363 let record = checker.connections.get(&id).expect("record must exist");
364 assert_eq!(record.consecutive_failures, 2);
365 }
366
367 #[test]
369 fn test_degraded_after_failure_threshold() {
370 let mut checker = ConnectionHealthChecker::new(HealthCheckerConfig {
371 failure_threshold: 3,
372 dead_threshold: 10,
373 rtt_spike_factor: 2.0,
374 });
375 let id = checker.open_connection("peer".to_string());
376 for _ in 0..3 {
377 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
378 }
379 let record = checker.connections.get(&id).expect("record must exist");
380 assert!(
381 matches!(
382 record.state,
383 ConnectionHealthState::Degraded {
384 consecutive_failures: 3
385 }
386 ),
387 "expected Degraded, got {:?}",
388 record.state
389 );
390 }
391
392 #[test]
394 fn test_dead_after_dead_threshold() {
395 let mut checker = ConnectionHealthChecker::new(HealthCheckerConfig {
396 failure_threshold: 3,
397 dead_threshold: 5,
398 rtt_spike_factor: 2.0,
399 });
400 let id = checker.open_connection("peer".to_string());
401 for _ in 0..5 {
402 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
403 }
404 let record = checker.connections.get(&id).expect("record must exist");
405 assert!(
406 matches!(record.state, ConnectionHealthState::Dead { .. }),
407 "expected Dead, got {:?}",
408 record.state
409 );
410 assert!(!record.is_alive());
411 }
412
413 #[test]
415 fn test_keepalive_success_resets_failures() {
416 let mut checker = default_checker();
417 let id = checker.open_connection("peer".to_string());
418 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
419 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
420 checker.record_event(id, ConnectionEvent::KeepaliveSuccess { rtt_ms: 20.0 });
421 let record = checker.connections.get(&id).expect("record must exist");
422 assert_eq!(record.consecutive_failures, 0);
423 assert!(matches!(record.state, ConnectionHealthState::Healthy));
424 }
425
426 #[test]
428 fn test_avg_rtt_ewma_first_sample() {
429 let mut checker = default_checker();
430 let id = checker.open_connection("peer".to_string());
431 checker.record_event(id, ConnectionEvent::KeepaliveSuccess { rtt_ms: 50.0 });
432 let record = checker.connections.get(&id).expect("record must exist");
433 assert!((record.avg_rtt_ms - 50.0).abs() < 1e-9);
435 }
436
437 #[test]
439 fn test_avg_rtt_ewma_second_sample() {
440 let mut checker = default_checker();
441 let id = checker.open_connection("peer".to_string());
442 checker.record_event(id, ConnectionEvent::KeepaliveSuccess { rtt_ms: 100.0 });
443 checker.record_event(id, ConnectionEvent::KeepaliveSuccess { rtt_ms: 200.0 });
444 let record = checker.connections.get(&id).expect("record must exist");
445 let expected = 0.2 * 200.0 + 0.8 * 100.0;
447 assert!(
448 (record.avg_rtt_ms - expected).abs() < 1e-6,
449 "avg_rtt_ms={} expected={}",
450 record.avg_rtt_ms,
451 expected
452 );
453 }
454
455 #[test]
457 fn test_avg_rtt_ewma_convergence() {
458 let mut checker = default_checker();
459 let id = checker.open_connection("peer".to_string());
460 for _ in 0..100 {
462 checker.record_event(id, ConnectionEvent::KeepaliveSuccess { rtt_ms: 30.0 });
463 }
464 let record = checker.connections.get(&id).expect("record must exist");
465 assert!(
466 (record.avg_rtt_ms - 30.0).abs() < 0.01,
467 "avg_rtt_ms={} should converge to 30.0",
468 record.avg_rtt_ms
469 );
470 }
471
472 #[test]
474 fn test_reset_sets_dead() {
475 let mut checker = default_checker();
476 let id = checker.open_connection("peer".to_string());
477 checker.record_event(
478 id,
479 ConnectionEvent::Reset {
480 reason: "TCP RST received".to_string(),
481 },
482 );
483 let record = checker.connections.get(&id).expect("record must exist");
484 match &record.state {
485 ConnectionHealthState::Dead { reason } => {
486 assert_eq!(reason, "TCP RST received");
487 }
488 other => panic!("expected Dead, got {other:?}"),
489 }
490 assert!(!record.is_alive());
491 }
492
493 #[test]
495 fn test_reconnected_restores_healthy() {
496 let mut checker = ConnectionHealthChecker::new(HealthCheckerConfig {
497 failure_threshold: 2,
498 dead_threshold: 5,
499 rtt_spike_factor: 2.0,
500 });
501 let id = checker.open_connection("peer".to_string());
502 for _ in 0..5 {
504 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
505 }
506 checker.record_event(id, ConnectionEvent::Reconnected);
508 let record = checker.connections.get(&id).expect("record must exist");
509 assert!(matches!(record.state, ConnectionHealthState::Healthy));
510 assert_eq!(record.consecutive_failures, 0);
511 assert!(record.is_alive());
512 }
513
514 #[test]
516 fn test_data_sent_updates_bytes() {
517 let mut checker = default_checker();
518 let id = checker.open_connection("peer".to_string());
519 checker.record_event(id, ConnectionEvent::DataSent { bytes: 512 });
520 checker.record_event(id, ConnectionEvent::DataSent { bytes: 1024 });
521 let record = checker.connections.get(&id).expect("record must exist");
522 assert_eq!(record.bytes_sent, 1536);
523 assert_eq!(record.bytes_received, 0);
524 }
525
526 #[test]
528 fn test_data_received_updates_bytes() {
529 let mut checker = default_checker();
530 let id = checker.open_connection("peer".to_string());
531 checker.record_event(id, ConnectionEvent::DataReceived { bytes: 4096 });
532 let record = checker.connections.get(&id).expect("record must exist");
533 assert_eq!(record.bytes_received, 4096);
534 assert_eq!(record.bytes_sent, 0);
535 }
536
537 #[test]
539 fn test_close_connection_removes_record() {
540 let mut checker = default_checker();
541 let id = checker.open_connection("peer".to_string());
542 assert!(checker.close_connection(id));
543 assert!(!checker.connections.contains_key(&id));
544 }
545
546 #[test]
548 fn test_close_connection_unknown_id() {
549 let mut checker = default_checker();
550 assert!(!checker.close_connection(9999));
551 }
552
553 #[test]
555 fn test_alive_dead_counts() {
556 let mut checker = ConnectionHealthChecker::new(HealthCheckerConfig {
557 failure_threshold: 2,
558 dead_threshold: 3,
559 rtt_spike_factor: 2.0,
560 });
561 let id_alive = checker.open_connection("peer-alive".to_string());
562 let id_dead = checker.open_connection("peer-dead".to_string());
563
564 for _ in 0..3 {
566 checker.record_event(id_dead, ConnectionEvent::KeepaliveFailed);
567 }
568 checker.record_event(id_alive, ConnectionEvent::KeepaliveSuccess { rtt_ms: 5.0 });
570
571 let alive = checker.alive_connections();
572 let dead = checker.dead_connections();
573 assert_eq!(alive.len(), 1);
574 assert_eq!(dead.len(), 1);
575 assert_eq!(alive[0].conn_id, id_alive);
576 assert_eq!(dead[0].conn_id, id_dead);
577 }
578
579 #[test]
581 fn test_stats_tuple() {
582 let mut checker = ConnectionHealthChecker::new(HealthCheckerConfig {
583 failure_threshold: 2,
584 dead_threshold: 3,
585 rtt_spike_factor: 2.0,
586 });
587 let id1 = checker.open_connection("p1".to_string());
588 let id2 = checker.open_connection("p2".to_string());
589 let _id3 = checker.open_connection("p3".to_string());
590
591 for _ in 0..3 {
593 checker.record_event(id1, ConnectionEvent::KeepaliveFailed);
594 checker.record_event(id2, ConnectionEvent::KeepaliveFailed);
595 }
596
597 let (total, alive, dead) = checker.stats();
598 assert_eq!(total, 3);
599 assert_eq!(alive, 1);
600 assert_eq!(dead, 2);
601 }
602
603 #[test]
605 fn test_event_cap_at_50() {
606 let mut checker = default_checker();
607 let id = checker.open_connection("peer".to_string());
608 for _ in 0..70 {
610 checker.record_event(id, ConnectionEvent::DataSent { bytes: 1 });
611 }
612 let record = checker.connections.get(&id).expect("record must exist");
613 assert_eq!(record.events.len(), MAX_EVENTS);
614 }
615
616 #[test]
618 fn test_event_ring_evicts_oldest() {
619 let mut checker = default_checker();
620 let id = checker.open_connection("peer".to_string());
621 for _ in 0..50 {
623 checker.record_event(id, ConnectionEvent::DataSent { bytes: 1 });
624 }
625 checker.record_event(id, ConnectionEvent::DataReceived { bytes: 999 });
627 let record = checker.connections.get(&id).expect("record must exist");
628 assert_eq!(record.events.len(), MAX_EVENTS);
629 assert_eq!(
631 record.events.last(),
632 Some(&ConnectionEvent::DataReceived { bytes: 999 })
633 );
634 }
635
636 #[test]
638 fn test_record_event_unknown_conn_id_ignored() {
639 let mut checker = default_checker();
640 checker.record_event(42, ConnectionEvent::KeepaliveFailed);
642 assert!(checker.connections.is_empty());
643 }
644
645 #[test]
647 fn test_is_alive_degraded_is_still_alive() {
648 let mut checker = ConnectionHealthChecker::new(HealthCheckerConfig {
649 failure_threshold: 2,
650 dead_threshold: 10,
651 rtt_spike_factor: 2.0,
652 });
653 let id = checker.open_connection("peer".to_string());
654 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
656 checker.record_event(id, ConnectionEvent::KeepaliveFailed);
657 let record = checker.connections.get(&id).expect("record must exist");
658 assert!(matches!(
659 record.state,
660 ConnectionHealthState::Degraded { .. }
661 ));
662 assert!(record.is_alive());
663 }
664
665 #[test]
667 fn test_multiple_connections_independent() {
668 let mut checker = ConnectionHealthChecker::new(HealthCheckerConfig {
669 failure_threshold: 2,
670 dead_threshold: 3,
671 rtt_spike_factor: 2.0,
672 });
673 let id_a = checker.open_connection("a".to_string());
674 let id_b = checker.open_connection("b".to_string());
675
676 checker.record_event(id_a, ConnectionEvent::KeepaliveSuccess { rtt_ms: 10.0 });
677 for _ in 0..3 {
678 checker.record_event(id_b, ConnectionEvent::KeepaliveFailed);
679 }
680
681 let rec_a = checker.connections.get(&id_a).expect("a must exist");
682 let rec_b = checker.connections.get(&id_b).expect("b must exist");
683
684 assert!(matches!(rec_a.state, ConnectionHealthState::Healthy));
685 assert!(matches!(rec_b.state, ConnectionHealthState::Dead { .. }));
686 }
687
688 #[test]
690 fn test_reset_not_in_alive_connections() {
691 let mut checker = default_checker();
692 let id = checker.open_connection("peer".to_string());
693 checker.record_event(
694 id,
695 ConnectionEvent::Reset {
696 reason: "timeout".to_string(),
697 },
698 );
699 assert_eq!(checker.alive_connections().len(), 0);
700 assert_eq!(checker.dead_connections().len(), 1);
701 }
702}