use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use crate::mesh_clock::{Clock, MonoTime, RealClock};
#[derive(Debug, Clone)]
pub struct HeartbeatConfig {
pub probe_interval: Duration,
pub ping_timeout: Duration,
pub indirect_timeout: Duration,
pub indirect_witnesses: usize,
pub suspect_timeout: Duration,
}
impl Default for HeartbeatConfig {
fn default() -> Self {
Self {
probe_interval: Duration::from_secs(1),
ping_timeout: Duration::from_millis(500),
indirect_timeout: Duration::from_secs(1),
indirect_witnesses: 3,
suspect_timeout: Duration::from_secs(5),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PingRequest {
pub from_node: String,
pub sequence: u64,
pub indirect: bool,
pub on_behalf_of: Option<String>,
#[serde(default)]
pub mesh_time_us: u64,
#[serde(default)]
pub clock_source: ClockSource,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PingResponse {
pub from_node: String,
pub sequence: u64,
pub process_count: u32,
pub service_count: u32,
#[serde(default)]
pub mesh_time_us: u64,
#[serde(default)]
pub clock_source: ClockSource,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[derive(Default)]
pub enum ClockSource {
#[default]
Local = 0,
Mesh = 1,
Ntp = 2,
Tsf = 3,
Gps = 4,
}
#[derive(Debug, Clone)]
pub struct MeshClockSync {
pub offset_us: i64,
smoothed_offset: f64,
pub authority_id: Option<String>,
pub authority_source: ClockSource,
pub local_source: ClockSource,
pub sync_count: u64,
pub uncertainty_us: u64,
}
impl MeshClockSync {
pub fn new(local_source: ClockSource) -> Self {
Self {
offset_us: 0,
smoothed_offset: 0.0,
authority_id: None,
authority_source: ClockSource::Local,
local_source,
sync_count: 0,
uncertainty_us: 1_000_000, }
}
pub fn mesh_time_us(&self) -> u64 {
self.mesh_time_us_with(&RealClock)
}
pub fn mesh_time_us_with(&self, clock: &dyn Clock) -> u64 {
let local = clock.unix_time_us();
(local as i64 + self.offset_us) as u64
}
pub fn process_sync(
&mut self,
sender_id: &str,
sender_time_us: u64,
sender_source: ClockSource,
local_receive_time_us: u64,
) {
if sender_source < self.local_source && self.sync_count > 0 {
return;
}
if sender_source > self.authority_source
|| self.authority_id.is_none()
|| self.authority_id.as_deref() == Some(sender_id)
{
let raw_offset = sender_time_us as i64 - local_receive_time_us as i64;
if self.sync_count > 5 && (raw_offset - self.offset_us).unsigned_abs() > 100_000 {
return;
}
let alpha = if self.sync_count < 5 { 0.5 } else { 0.1 };
self.smoothed_offset = (1.0 - alpha) * self.smoothed_offset + alpha * raw_offset as f64;
self.offset_us = self.smoothed_offset as i64;
self.authority_id = Some(sender_id.to_string());
self.authority_source = sender_source;
self.sync_count += 1;
self.uncertainty_us = match self.sync_count {
0..=5 => 10_000, 6..=20 => 1_000, 21..=100 => 100, _ => 50, };
}
}
pub fn is_authority(&self, our_node_id: &str) -> bool {
self.authority_id.as_deref() == Some(our_node_id) || self.authority_id.is_none()
}
pub fn is_better_source(&self, other: ClockSource) -> bool {
self.local_source > other
}
}
pub fn system_time_us() -> u64 {
RealClock.unix_time_us()
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HeartbeatState {
Alive,
Suspect,
Dead,
}
#[derive(Debug, Clone)]
pub struct PeerHeartbeat {
pub state: HeartbeatState,
pub last_seen: MonoTime,
pub suspect_since: Option<MonoTime>,
pub missed_count: u32,
pub last_ping_seq: u64,
}
impl PeerHeartbeat {
pub fn new(clock: &dyn Clock) -> Self {
Self {
state: HeartbeatState::Alive,
last_seen: clock.now(),
suspect_since: None,
missed_count: 0,
last_ping_seq: 0,
}
}
pub fn record_alive(&mut self, clock: &dyn Clock) {
self.state = HeartbeatState::Alive;
self.last_seen = clock.now();
self.suspect_since = None;
self.missed_count = 0;
}
pub fn record_miss(&mut self, clock: &dyn Clock, config: &HeartbeatConfig) {
self.missed_count += 1;
match self.state {
HeartbeatState::Alive => {
self.state = HeartbeatState::Suspect;
self.suspect_since = Some(clock.now());
}
HeartbeatState::Suspect => {
if let Some(since) = self.suspect_since {
let elapsed = clock.now().elapsed_since(since);
if elapsed >= config.suspect_timeout {
self.state = HeartbeatState::Dead;
}
}
}
HeartbeatState::Dead => {} }
}
pub fn is_unreachable(&self, clock: &dyn Clock, config: &HeartbeatConfig) -> bool {
self.state == HeartbeatState::Dead
|| clock.now().elapsed_since(self.last_seen) > config.suspect_timeout * 3
}
}
impl Default for PeerHeartbeat {
fn default() -> Self {
Self::new(&RealClock)
}
}
pub struct HeartbeatTracker {
config: HeartbeatConfig,
peers: HashMap<String, PeerHeartbeat>,
next_sequence: u64,
clock: Arc<dyn Clock>,
}
impl HeartbeatTracker {
pub fn new(config: HeartbeatConfig) -> Self {
Self::with_clock(config, RealClock::shared())
}
pub fn with_clock(config: HeartbeatConfig, clock: Arc<dyn Clock>) -> Self {
Self {
config,
peers: HashMap::new(),
next_sequence: 1,
clock,
}
}
pub fn clock(&self) -> &Arc<dyn Clock> {
&self.clock
}
pub fn add_peer(&mut self, node_id: String) {
self.peers
.entry(node_id)
.or_insert_with(|| PeerHeartbeat::new(self.clock.as_ref()));
}
pub fn remove_peer(&mut self, node_id: &str) {
self.peers.remove(node_id);
}
pub fn record_alive(&mut self, node_id: &str) {
if let Some(peer) = self.peers.get_mut(node_id) {
peer.record_alive(self.clock.as_ref());
}
}
pub fn record_miss(&mut self, node_id: &str) {
if let Some(peer) = self.peers.get_mut(node_id) {
peer.record_miss(self.clock.as_ref(), &self.config);
}
}
pub fn next_sequence(&mut self) -> u64 {
let seq = self.next_sequence;
self.next_sequence += 1;
seq
}
pub fn suspect_peers(&self) -> Vec<&str> {
self.peers
.iter()
.filter(|(_, p)| p.state == HeartbeatState::Suspect)
.map(|(id, _)| id.as_str())
.collect()
}
pub fn dead_peers(&self) -> Vec<&str> {
self.peers
.iter()
.filter(|(_, p)| p.state == HeartbeatState::Dead)
.map(|(id, _)| id.as_str())
.collect()
}
pub fn peer_count(&self) -> usize {
self.peers.len()
}
pub fn peer_state(&self, node_id: &str) -> Option<HeartbeatState> {
self.peers.get(node_id).map(|p| p.state)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerMetrics {
pub node_id: String,
pub rtt_ms: f64,
pub messages_sent: u64,
pub messages_received: u64,
pub bytes_sent: u64,
pub bytes_received: u64,
pub errors: u64,
pub last_updated: u64,
}
impl PeerMetrics {
pub fn new(node_id: String) -> Self {
Self {
node_id,
rtt_ms: 0.0,
messages_sent: 0,
messages_received: 0,
bytes_sent: 0,
bytes_received: 0,
errors: 0,
last_updated: 0,
}
}
pub fn affinity_score(&self) -> f64 {
let error_rate = if self.messages_sent > 0 {
self.errors as f64 / self.messages_sent as f64
} else {
0.0
};
self.rtt_ms + error_rate * 1000.0
}
pub fn record_send(&mut self, bytes: u64) {
self.messages_sent += 1;
self.bytes_sent += bytes;
}
pub fn record_recv(&mut self, bytes: u64) {
self.messages_received += 1;
self.bytes_received += bytes;
}
pub fn record_error(&mut self) {
self.errors += 1;
}
pub fn update_rtt(&mut self, rtt_ms: f64) {
self.rtt_ms = self.rtt_ms * 0.8 + rtt_ms * 0.2;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::mesh_clock::MockClock;
#[test]
fn ping_request_serde_roundtrip() {
let req = PingRequest {
from_node: "node-a".to_string(),
sequence: 42,
indirect: true,
on_behalf_of: Some("node-b".to_string()),
mesh_time_us: system_time_us(),
clock_source: ClockSource::Ntp,
};
let json = serde_json::to_string(&req).unwrap();
let restored: PingRequest = serde_json::from_str(&json).unwrap();
assert_eq!(restored.from_node, "node-a");
assert_eq!(restored.sequence, 42);
assert!(restored.indirect);
assert_eq!(restored.on_behalf_of, Some("node-b".to_string()));
}
#[test]
fn ping_response_serde_roundtrip() {
let resp = PingResponse {
from_node: "node-b".to_string(),
sequence: 42,
process_count: 5,
service_count: 3,
mesh_time_us: system_time_us(),
clock_source: ClockSource::Local,
};
let json = serde_json::to_string(&resp).unwrap();
let restored: PingResponse = serde_json::from_str(&json).unwrap();
assert_eq!(restored.from_node, "node-b");
assert_eq!(restored.sequence, 42);
assert_eq!(restored.process_count, 5);
assert_eq!(restored.service_count, 3);
}
#[test]
fn heartbeat_config_defaults() {
let config = HeartbeatConfig::default();
assert_eq!(config.probe_interval, Duration::from_secs(1));
assert_eq!(config.ping_timeout, Duration::from_millis(500));
assert_eq!(config.indirect_timeout, Duration::from_secs(1));
assert_eq!(config.indirect_witnesses, 3);
assert_eq!(config.suspect_timeout, Duration::from_secs(5));
}
#[test]
fn peer_heartbeat_alive_to_suspect() {
let clock = MockClock::new();
let config = HeartbeatConfig::default();
let mut peer = PeerHeartbeat::new(&clock);
assert_eq!(peer.state, HeartbeatState::Alive);
peer.record_miss(&clock, &config);
assert_eq!(peer.state, HeartbeatState::Suspect);
assert_eq!(peer.missed_count, 1);
assert!(peer.suspect_since.is_some());
}
#[test]
fn peer_heartbeat_suspect_to_dead() {
let clock = MockClock::new();
let config = HeartbeatConfig {
suspect_timeout: Duration::from_secs(0),
..HeartbeatConfig::default()
};
let mut peer = PeerHeartbeat::new(&clock);
peer.record_miss(&clock, &config);
assert_eq!(peer.state, HeartbeatState::Suspect);
peer.record_miss(&clock, &config);
assert_eq!(peer.state, HeartbeatState::Dead);
}
#[test]
fn record_alive_resets_state() {
let clock = MockClock::new();
let config = HeartbeatConfig::default();
let mut peer = PeerHeartbeat::new(&clock);
peer.record_miss(&clock, &config);
assert_eq!(peer.state, HeartbeatState::Suspect);
peer.record_alive(&clock);
assert_eq!(peer.state, HeartbeatState::Alive);
assert_eq!(peer.missed_count, 0);
assert!(peer.suspect_since.is_none());
}
#[test]
fn mock_clock_deterministically_advances_heartbeat_ticks() {
let clock = MockClock::new();
let config = HeartbeatConfig {
suspect_timeout: Duration::from_secs(5),
..HeartbeatConfig::default()
};
let mut tracker = HeartbeatTracker::with_clock(config, Arc::new(clock.clone()));
tracker.add_peer("peer-a".into());
tracker.record_miss("peer-a");
assert_eq!(
tracker.peer_state("peer-a"),
Some(HeartbeatState::Suspect)
);
clock.advance(Duration::from_secs(4));
tracker.record_miss("peer-a");
assert_eq!(
tracker.peer_state("peer-a"),
Some(HeartbeatState::Suspect)
);
clock.advance(Duration::from_secs(2));
tracker.record_miss("peer-a");
assert_eq!(tracker.peer_state("peer-a"), Some(HeartbeatState::Dead));
assert!(tracker.dead_peers().contains(&"peer-a"));
}
#[test]
fn tracker_add_remove_peers() {
let mut tracker = HeartbeatTracker::new(HeartbeatConfig::default());
tracker.add_peer("node-a".to_string());
tracker.add_peer("node-b".to_string());
assert_eq!(tracker.peer_count(), 2);
tracker.remove_peer("node-a");
assert_eq!(tracker.peer_count(), 1);
assert!(tracker.peer_state("node-a").is_none());
assert_eq!(tracker.peer_state("node-b"), Some(HeartbeatState::Alive));
}
#[test]
fn tracker_suspect_and_dead_peers() {
let config = HeartbeatConfig {
suspect_timeout: Duration::from_secs(0),
..HeartbeatConfig::default()
};
let mut tracker = HeartbeatTracker::new(config);
tracker.add_peer("node-a".to_string());
tracker.add_peer("node-b".to_string());
tracker.add_peer("node-c".to_string());
tracker.record_miss("node-a");
tracker.record_miss("node-b");
tracker.record_miss("node-b");
let suspects = tracker.suspect_peers();
assert_eq!(suspects.len(), 1);
assert!(suspects.contains(&"node-a"));
let dead = tracker.dead_peers();
assert_eq!(dead.len(), 1);
assert!(dead.contains(&"node-b"));
}
#[test]
fn peer_metrics_affinity_score_zero_errors() {
let mut m = PeerMetrics::new("n1".into());
m.rtt_ms = 10.0;
m.messages_sent = 100;
m.errors = 0;
assert!((m.affinity_score() - 10.0).abs() < f64::EPSILON);
}
#[test]
fn peer_metrics_affinity_score_with_errors() {
let mut m = PeerMetrics::new("n1".into());
m.rtt_ms = 5.0;
m.messages_sent = 100;
m.errors = 10;
assert!((m.affinity_score() - 105.0).abs() < f64::EPSILON);
}
#[test]
fn peer_metrics_affinity_no_messages() {
let m = PeerMetrics::new("n1".into());
assert!((m.affinity_score() - 0.0).abs() < f64::EPSILON);
}
#[test]
fn peer_metrics_record_send_recv() {
let mut m = PeerMetrics::new("n1".into());
m.record_send(100);
m.record_send(200);
assert_eq!(m.messages_sent, 2);
assert_eq!(m.bytes_sent, 300);
m.record_recv(50);
assert_eq!(m.messages_received, 1);
assert_eq!(m.bytes_received, 50);
}
#[test]
fn peer_metrics_record_error() {
let mut m = PeerMetrics::new("n1".into());
m.record_error();
m.record_error();
assert_eq!(m.errors, 2);
}
#[test]
fn peer_metrics_update_rtt_ema() {
let mut m = PeerMetrics::new("n1".into());
m.rtt_ms = 100.0;
m.update_rtt(50.0);
assert!((m.rtt_ms - 90.0).abs() < f64::EPSILON);
}
#[test]
fn peer_metrics_serde_roundtrip() {
let mut m = PeerMetrics::new("node-x".into());
m.rtt_ms = 15.5;
m.messages_sent = 42;
let json = serde_json::to_string(&m).unwrap();
let restored: PeerMetrics = serde_json::from_str(&json).unwrap();
assert_eq!(restored.node_id, "node-x");
assert!((restored.rtt_ms - 15.5).abs() < f64::EPSILON);
assert_eq!(restored.messages_sent, 42);
}
#[test]
fn next_sequence_increments() {
let mut tracker = HeartbeatTracker::new(HeartbeatConfig::default());
assert_eq!(tracker.next_sequence(), 1);
assert_eq!(tracker.next_sequence(), 2);
assert_eq!(tracker.next_sequence(), 3);
}
#[test]
fn clock_source_ordering() {
assert!(ClockSource::Gps > ClockSource::Ntp);
assert!(ClockSource::Ntp > ClockSource::Mesh);
assert!(ClockSource::Mesh > ClockSource::Local);
assert!(ClockSource::Tsf > ClockSource::Ntp);
}
#[test]
fn mesh_clock_sync_initial_state() {
let sync = MeshClockSync::new(ClockSource::Local);
assert_eq!(sync.offset_us, 0);
assert_eq!(sync.sync_count, 0);
assert!(sync.authority_id.is_none());
assert_eq!(sync.uncertainty_us, 1_000_000);
}
#[test]
fn mesh_clock_sync_from_ntp_authority() {
let mut sync = MeshClockSync::new(ClockSource::Local);
let now = system_time_us();
let authority_time = now + 5_000;
sync.process_sync("authority-1", authority_time, ClockSource::Ntp, now);
assert_eq!(sync.authority_id.as_deref(), Some("authority-1"));
assert_eq!(sync.authority_source, ClockSource::Ntp);
assert_eq!(sync.sync_count, 1);
assert!((sync.offset_us - 2_500).unsigned_abs() < 1_000);
}
#[test]
fn mesh_clock_sync_ignores_worse_source() {
let mut sync = MeshClockSync::new(ClockSource::Ntp);
let now = system_time_us();
for i in 0..10 {
sync.process_sync("peer", now + i * 100, ClockSource::Local, now + i * 100);
}
let offset_before = sync.offset_us;
sync.process_sync("peer", now + 50_000, ClockSource::Local, now);
assert_eq!(sync.offset_us, offset_before);
}
#[test]
fn mesh_clock_sync_smooths_jitter() {
let mut sync = MeshClockSync::new(ClockSource::Local);
let now = system_time_us();
for i in 0..20u64 {
let jitter = (i % 3) as i64 * 50 - 50; let authority_time = (now as i64 + 1_000 + jitter) as u64;
sync.process_sync("auth", authority_time, ClockSource::Ntp, now);
}
assert!((sync.offset_us - 1_000).unsigned_abs() < 500);
assert!(sync.uncertainty_us < 10_000);
}
#[test]
fn system_time_us_returns_plausible_value() {
let t = system_time_us();
assert!(t > 1_577_836_800_000_000);
assert!(t < 2_524_608_000_000_000);
}
}