use std::collections::HashMap;
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum HealthTier {
Healthy = 0,
Degraded = 1,
Unhealthy = 2,
Dead = 3,
}
#[derive(Clone, Debug)]
pub struct HeartbeatRecord {
pub peer_id: String,
pub last_seen_tick: u64,
pub consecutive_misses: u32,
pub total_heartbeats: u64,
pub total_misses: u64,
}
impl HeartbeatRecord {
fn new(peer_id: &str) -> Self {
Self {
peer_id: peer_id.to_owned(),
last_seen_tick: 0,
consecutive_misses: 0,
total_heartbeats: 0,
total_misses: 0,
}
}
pub fn miss_rate(&self) -> f64 {
let total = self.total_heartbeats + self.total_misses;
if total == 0 {
0.0
} else {
self.total_misses as f64 / total as f64
}
}
pub fn health_score(&self) -> f64 {
1.0 - self.miss_rate()
}
}
#[derive(Clone, Debug)]
pub struct PeerHealthCheckerConfig {
pub heartbeat_interval_ticks: u64,
pub degraded_miss_threshold: u32,
pub unhealthy_miss_threshold: u32,
pub dead_miss_threshold: u32,
}
impl Default for PeerHealthCheckerConfig {
fn default() -> Self {
Self {
heartbeat_interval_ticks: 30,
degraded_miss_threshold: 2,
unhealthy_miss_threshold: 5,
dead_miss_threshold: 10,
}
}
}
#[derive(Clone, Debug)]
pub struct HealthCheckerStats {
pub total_peers: usize,
pub by_tier: HashMap<HealthTier, usize>,
pub total_heartbeats_received: u64,
pub total_misses_recorded: u64,
}
pub struct PeerHealthChecker {
pub records: HashMap<String, HeartbeatRecord>,
pub config: PeerHealthCheckerConfig,
}
impl PeerHealthChecker {
pub fn new(config: PeerHealthCheckerConfig) -> Self {
Self {
records: HashMap::new(),
config,
}
}
pub fn record_heartbeat(&mut self, peer_id: &str, current_tick: u64) {
let record = self
.records
.entry(peer_id.to_owned())
.or_insert_with(|| HeartbeatRecord::new(peer_id));
record.last_seen_tick = current_tick;
record.consecutive_misses = 0;
record.total_heartbeats += 1;
}
pub fn record_miss(&mut self, peer_id: &str) {
let record = self
.records
.entry(peer_id.to_owned())
.or_insert_with(|| HeartbeatRecord::new(peer_id));
record.consecutive_misses += 1;
record.total_misses += 1;
}
pub fn tier(&self, peer_id: &str) -> HealthTier {
match self.records.get(peer_id) {
None => HealthTier::Dead,
Some(rec) => {
if rec.consecutive_misses >= self.config.dead_miss_threshold {
HealthTier::Dead
} else if rec.consecutive_misses >= self.config.unhealthy_miss_threshold {
HealthTier::Unhealthy
} else if rec.consecutive_misses >= self.config.degraded_miss_threshold {
HealthTier::Degraded
} else {
HealthTier::Healthy
}
}
}
}
pub fn tick_check(&mut self, current_tick: u64) -> Vec<String> {
let interval = self.config.heartbeat_interval_ticks;
let stale: Vec<String> = self
.records
.values()
.filter(|rec| current_tick.saturating_sub(rec.last_seen_tick) > interval)
.map(|rec| rec.peer_id.clone())
.collect();
for peer_id in &stale {
if let Some(rec) = self.records.get_mut(peer_id) {
rec.consecutive_misses += 1;
rec.total_misses += 1;
}
}
let mut sorted = stale;
sorted.sort_unstable();
sorted
}
pub fn remove_peer(&mut self, peer_id: &str) -> bool {
self.records.remove(peer_id).is_some()
}
pub fn health_score(&self, peer_id: &str) -> f64 {
self.records
.get(peer_id)
.map(|rec| rec.health_score())
.unwrap_or(0.0)
}
pub fn healthy_peers(&self) -> Vec<&str> {
let mut peers: Vec<&str> = self
.records
.keys()
.filter(|id| self.tier(id) == HealthTier::Healthy)
.map(|s| s.as_str())
.collect();
peers.sort_unstable();
peers
}
pub fn stats(&self) -> HealthCheckerStats {
let mut by_tier: HashMap<HealthTier, usize> = HashMap::new();
let mut total_heartbeats_received: u64 = 0;
let mut total_misses_recorded: u64 = 0;
for rec in self.records.values() {
*by_tier.entry(self.tier(&rec.peer_id)).or_insert(0) += 1;
total_heartbeats_received += rec.total_heartbeats;
total_misses_recorded += rec.total_misses;
}
HealthCheckerStats {
total_peers: self.records.len(),
by_tier,
total_heartbeats_received,
total_misses_recorded,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn default_checker() -> PeerHealthChecker {
PeerHealthChecker::new(PeerHealthCheckerConfig::default())
}
#[test]
fn miss_rate_zero_when_no_data() {
let rec = HeartbeatRecord::new("p");
assert_eq!(rec.miss_rate(), 0.0);
}
#[test]
fn miss_rate_calculation() {
let mut rec = HeartbeatRecord::new("p");
rec.total_heartbeats = 3;
rec.total_misses = 1;
let diff = (rec.miss_rate() - 0.25_f64).abs();
assert!(
diff < 1e-10,
"miss_rate should be 0.25, got {}",
rec.miss_rate()
);
}
#[test]
fn health_score_equals_one_minus_miss_rate() {
let mut rec = HeartbeatRecord::new("p");
rec.total_heartbeats = 3;
rec.total_misses = 1;
let expected = 1.0 - rec.miss_rate();
let diff = (rec.health_score() - expected).abs();
assert!(diff < 1e-10);
}
#[test]
fn health_score_is_one_when_no_misses() {
let mut rec = HeartbeatRecord::new("p");
rec.total_heartbeats = 10;
let diff = (rec.health_score() - 1.0_f64).abs();
assert!(diff < 1e-10);
}
#[test]
fn health_score_is_one_when_no_data() {
let rec = HeartbeatRecord::new("p");
let diff = (rec.health_score() - 1.0_f64).abs();
assert!(
diff < 1e-10,
"empty record should have score 1.0, got {}",
rec.health_score()
);
}
#[test]
fn record_heartbeat_creates_record() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 50);
assert!(checker.records.contains_key("peer-1"));
}
#[test]
fn record_heartbeat_updates_last_seen_tick() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 42);
assert_eq!(checker.records["peer-1"].last_seen_tick, 42);
}
#[test]
fn record_heartbeat_resets_consecutive_misses() {
let mut checker = default_checker();
checker.record_miss("peer-1");
checker.record_miss("peer-1");
checker.record_heartbeat("peer-1", 100);
assert_eq!(checker.records["peer-1"].consecutive_misses, 0);
}
#[test]
fn record_heartbeat_increments_total_heartbeats() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 10);
checker.record_heartbeat("peer-1", 20);
assert_eq!(checker.records["peer-1"].total_heartbeats, 2);
}
#[test]
fn record_miss_creates_record() {
let mut checker = default_checker();
checker.record_miss("peer-2");
assert!(checker.records.contains_key("peer-2"));
}
#[test]
fn record_miss_increments_consecutive_misses() {
let mut checker = default_checker();
checker.record_miss("peer-2");
checker.record_miss("peer-2");
assert_eq!(checker.records["peer-2"].consecutive_misses, 2);
}
#[test]
fn record_miss_increments_total_misses() {
let mut checker = default_checker();
checker.record_miss("peer-2");
checker.record_miss("peer-2");
assert_eq!(checker.records["peer-2"].total_misses, 2);
}
#[test]
fn tier_dead_for_missing_peer() {
let checker = default_checker();
assert_eq!(checker.tier("unknown"), HealthTier::Dead);
}
#[test]
fn tier_healthy_on_zero_misses() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 10);
assert_eq!(checker.tier("peer-1"), HealthTier::Healthy);
}
#[test]
fn tier_degraded_at_threshold() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 10);
checker.record_miss("peer-1");
checker.record_miss("peer-1");
assert_eq!(checker.tier("peer-1"), HealthTier::Degraded);
}
#[test]
fn tier_unhealthy_at_threshold() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 10);
for _ in 0..5 {
checker.record_miss("peer-1");
}
assert_eq!(checker.tier("peer-1"), HealthTier::Unhealthy);
}
#[test]
fn tier_dead_at_threshold() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 10);
for _ in 0..10 {
checker.record_miss("peer-1");
}
assert_eq!(checker.tier("peer-1"), HealthTier::Dead);
}
#[test]
fn tier_healthy_after_heartbeat_resets_misses() {
let mut checker = default_checker();
for _ in 0..8 {
checker.record_miss("peer-1");
}
checker.record_heartbeat("peer-1", 500);
assert_eq!(checker.tier("peer-1"), HealthTier::Healthy);
}
#[test]
fn tick_check_detects_stale_peers() {
let mut checker = default_checker();
checker.record_heartbeat("peer-a", 0);
let missed = checker.tick_check(31);
assert!(missed.contains(&"peer-a".to_owned()));
}
#[test]
fn tick_check_does_not_flag_recent_peers() {
let mut checker = default_checker();
checker.record_heartbeat("peer-b", 20);
let missed = checker.tick_check(30);
assert!(!missed.contains(&"peer-b".to_owned()));
}
#[test]
fn tick_check_returns_sorted_alphabetically() {
let mut checker = default_checker();
checker.record_heartbeat("peer-z", 0);
checker.record_heartbeat("peer-a", 0);
checker.record_heartbeat("peer-m", 0);
let missed = checker.tick_check(100);
let mut expected = missed.clone();
expected.sort();
assert_eq!(missed, expected);
}
#[test]
fn tick_check_increments_consecutive_misses_for_stale_peer() {
let mut checker = default_checker();
checker.record_heartbeat("peer-c", 0);
checker.tick_check(100);
assert!(checker.records["peer-c"].consecutive_misses > 0);
}
#[test]
fn tick_check_returns_empty_when_no_peers_are_stale() {
let mut checker = default_checker();
checker.record_heartbeat("peer-d", 100);
let missed = checker.tick_check(110);
assert!(missed.is_empty());
}
#[test]
fn remove_peer_returns_true_when_present() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 10);
assert!(checker.remove_peer("peer-1"));
}
#[test]
fn remove_peer_returns_false_when_absent() {
let mut checker = default_checker();
assert!(!checker.remove_peer("ghost"));
}
#[test]
fn remove_peer_actually_removes_record() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 10);
checker.remove_peer("peer-1");
assert!(!checker.records.contains_key("peer-1"));
}
#[test]
fn health_score_returns_zero_for_unknown_peer() {
let checker = default_checker();
let diff = checker.health_score("unknown").abs();
assert!(diff < 1e-10);
}
#[test]
fn health_score_returns_record_score() {
let mut checker = default_checker();
checker.record_heartbeat("peer-1", 10);
checker.record_heartbeat("peer-1", 20);
checker.record_miss("peer-1");
let score = checker.health_score("peer-1");
let expected = checker.records["peer-1"].health_score();
let diff = (score - expected).abs();
assert!(diff < 1e-10);
}
#[test]
fn healthy_peers_sorted_alphabetically() {
let mut checker = default_checker();
checker.record_heartbeat("peer-z", 10);
checker.record_heartbeat("peer-a", 10);
checker.record_heartbeat("peer-m", 10);
let peers = checker.healthy_peers();
let mut expected = peers.clone();
expected.sort_unstable();
assert_eq!(peers, expected);
}
#[test]
fn healthy_peers_excludes_degraded_peers() {
let mut checker = default_checker();
checker.record_heartbeat("peer-good", 10);
checker.record_heartbeat("peer-bad", 10);
checker.record_miss("peer-bad");
checker.record_miss("peer-bad");
let peers = checker.healthy_peers();
assert!(peers.contains(&"peer-good"));
assert!(!peers.contains(&"peer-bad"));
}
#[test]
fn stats_total_peers_count() {
let mut checker = default_checker();
checker.record_heartbeat("p1", 1);
checker.record_heartbeat("p2", 1);
let s = checker.stats();
assert_eq!(s.total_peers, 2);
}
#[test]
fn stats_by_tier_counts_healthy_and_degraded() {
let mut checker = default_checker();
checker.record_heartbeat("healthy", 10);
checker.record_heartbeat("degraded", 10);
checker.record_miss("degraded");
checker.record_miss("degraded");
let s = checker.stats();
assert_eq!(s.by_tier.get(&HealthTier::Healthy).copied().unwrap_or(0), 1);
assert_eq!(
s.by_tier.get(&HealthTier::Degraded).copied().unwrap_or(0),
1
);
}
#[test]
fn stats_total_heartbeats_received() {
let mut checker = default_checker();
checker.record_heartbeat("p1", 1);
checker.record_heartbeat("p1", 2);
checker.record_heartbeat("p2", 1);
let s = checker.stats();
assert_eq!(s.total_heartbeats_received, 3);
}
#[test]
fn stats_total_misses_recorded() {
let mut checker = default_checker();
checker.record_heartbeat("p1", 1);
checker.record_miss("p1");
checker.record_miss("p1");
checker.record_miss("p2");
let s = checker.stats();
assert_eq!(s.total_misses_recorded, 3);
}
#[test]
fn stats_by_tier_dead_count() {
let mut checker = default_checker();
checker.record_heartbeat("p1", 10);
for _ in 0..10 {
checker.record_miss("p1");
}
let s = checker.stats();
assert_eq!(s.by_tier.get(&HealthTier::Dead).copied().unwrap_or(0), 1);
}
}