firecloud_net/

health.rs

//! Peer health monitoring system
//!
//! Tracks the health and availability of peers in the network using:
//! - Periodic ping measurements (RTT)
//! - Success/failure tracking
//! - Health score calculation
//! - Automatic peer status updates

use libp2p::PeerId;
use std::collections::HashMap;
use std::time::{Duration, Instant};
use tracing::{debug, warn};

/// Health status of a peer
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PeerHealth {
    /// Peer is healthy and responsive
    Healthy,
    /// Peer is degraded (high latency or some failures)
    Degraded,
    /// Peer is unhealthy (many failures or timeouts)
    Unhealthy,
    /// Peer status is unknown (newly discovered)
    Unknown,
}

impl PeerHealth {
    /// Convert health to a numeric score (0-100)
    pub fn score(&self) -> u8 {
        match self {
            PeerHealth::Healthy => 100,
            PeerHealth::Degraded => 50,
            PeerHealth::Unhealthy => 10,
            PeerHealth::Unknown => 0,
        }
    }
}

/// Statistics for a single peer
#[derive(Debug, Clone)]
pub struct PeerStats {
    /// Unique peer identifier
    pub peer_id: PeerId,
    /// Last time we successfully communicated with this peer
    pub last_seen: Instant,
    /// Number of successful ping responses
    pub success_count: u32,
    /// Number of failed pings or timeouts
    pub fail_count: u32,
    /// Average round-trip time in milliseconds
    pub avg_rtt_ms: Option<u64>,
    /// Last measured RTT
    pub last_rtt_ms: Option<u64>,
    /// Current health status
    pub health: PeerHealth,
    /// Whether this peer is on the local network
    pub is_local: bool,
    /// Total number of bytes sent to this peer
    pub bytes_sent: u64,
    /// Total number of bytes received from this peer
    pub bytes_received: u64,
}

impl PeerStats {
    /// Create new peer stats for a discovered peer
    pub fn new(peer_id: PeerId, is_local: bool) -> Self {
        Self {
            peer_id,
            last_seen: Instant::now(),
            success_count: 0,
            fail_count: 0,
            avg_rtt_ms: None,
            last_rtt_ms: None,
            health: PeerHealth::Unknown,
            is_local,
            bytes_sent: 0,
            bytes_received: 0,
        }
    }

    /// Record a successful ping response
    pub fn record_success(&mut self, rtt: Duration) {
        self.last_seen = Instant::now();
        self.success_count += 1;
        
        let rtt_ms = rtt.as_millis() as u64;
        self.last_rtt_ms = Some(rtt_ms);
        
        // Update rolling average RTT
        if let Some(avg) = self.avg_rtt_ms {
            // Exponential moving average: new_avg = 0.7 * old_avg + 0.3 * new_value
            self.avg_rtt_ms = Some((avg * 7 + rtt_ms * 3) / 10);
        } else {
            self.avg_rtt_ms = Some(rtt_ms);
        }
        
        self.update_health();
    }

    /// Record a failed ping or timeout
    pub fn record_failure(&mut self) {
        self.fail_count += 1;
        self.update_health();
    }

    /// Update health status based on current statistics
    fn update_health(&mut self) {
        let total = self.success_count + self.fail_count;
        if total == 0 {
            self.health = PeerHealth::Unknown;
            return;
        }

        let fail_rate = self.fail_count as f64 / total as f64;
        let avg_rtt = self.avg_rtt_ms.unwrap_or(0);

        // Calculate health score (0-100)
        // Penalize: high failure rate and high RTT
        let health_score = 100.0 - (fail_rate * 50.0) - (avg_rtt as f64 / 10.0).min(50.0);

        self.health = if health_score >= 70.0 {
            PeerHealth::Healthy
        } else if health_score >= 40.0 {
            PeerHealth::Degraded
        } else {
            PeerHealth::Unhealthy
        };

        debug!(
            "Peer {:?} health updated: {:?} (score: {:.1}, fail_rate: {:.2}, avg_rtt: {}ms)",
            self.peer_id, self.health, health_score, fail_rate, avg_rtt
        );
    }

    /// Check if peer should be considered offline
    pub fn is_offline(&self, timeout: Duration) -> bool {
        self.last_seen.elapsed() > timeout
    }

    /// Record data transfer statistics
    pub fn record_transfer(&mut self, sent: u64, received: u64) {
        self.bytes_sent += sent;
        self.bytes_received += received;
        self.last_seen = Instant::now();
    }

    /// Calculate the overall health score (0-100)
    pub fn health_score(&self) -> u8 {
        let total = self.success_count + self.fail_count;
        if total == 0 {
            return 0;
        }

        let fail_rate = self.fail_count as f64 / total as f64;
        let avg_rtt = self.avg_rtt_ms.unwrap_or(1000);
        
        let score = 100.0 - (fail_rate * 50.0) - (avg_rtt as f64 / 10.0).min(50.0);
        score.max(0.0).min(100.0) as u8
    }
}

/// Health monitor for tracking all peers in the network
pub struct HealthMonitor {
    /// Map of peer_id to their statistics
    peers: HashMap<PeerId, PeerStats>,
    /// How long before considering a peer offline
    offline_timeout: Duration,
    /// Minimum health score to consider a peer usable
    min_health_score: u8,
}

impl HealthMonitor {
    /// Create a new health monitor
    pub fn new() -> Self {
        Self {
            peers: HashMap::new(),
            offline_timeout: Duration::from_secs(120), // 2 minutes
            min_health_score: 30, // Minimum score of 30/100
        }
    }

    /// Add or update a peer
    pub fn add_peer(&mut self, peer_id: PeerId, is_local: bool) {
        self.peers
            .entry(peer_id)
            .or_insert_with(|| PeerStats::new(peer_id, is_local));
    }

    /// Remove a peer from tracking
    pub fn remove_peer(&mut self, peer_id: &PeerId) {
        self.peers.remove(peer_id);
    }

    /// Record a successful ping to a peer
    pub fn record_ping_success(&mut self, peer_id: &PeerId, rtt: Duration) {
        if let Some(stats) = self.peers.get_mut(peer_id) {
            stats.record_success(rtt);
        }
    }

    /// Record a failed ping to a peer
    pub fn record_ping_failure(&mut self, peer_id: &PeerId) {
        if let Some(stats) = self.peers.get_mut(peer_id) {
            stats.record_failure();
        }
    }

    /// Record data transfer with a peer
    pub fn record_transfer(&mut self, peer_id: &PeerId, sent: u64, received: u64) {
        if let Some(stats) = self.peers.get_mut(peer_id) {
            stats.record_transfer(sent, received);
        }
    }

    /// Get statistics for a specific peer
    pub fn get_stats(&self, peer_id: &PeerId) -> Option<&PeerStats> {
        self.peers.get(peer_id)
    }

    /// Get all healthy peers, sorted by health score
    pub fn get_healthy_peers(&self) -> Vec<PeerId> {
        let mut peers: Vec<_> = self
            .peers
            .values()
            .filter(|stats| {
                !stats.is_offline(self.offline_timeout)
                    && stats.health_score() >= self.min_health_score
            })
            .map(|stats| (stats.peer_id, stats.health_score()))
            .collect();

        // Sort by health score (descending)
        peers.sort_by(|a, b| b.1.cmp(&a.1));

        peers.into_iter().map(|(peer_id, _)| peer_id).collect()
    }

    /// Get peers on local network, sorted by health
    pub fn get_local_peers(&self) -> Vec<PeerId> {
        let mut peers: Vec<_> = self
            .peers
            .values()
            .filter(|stats| {
                stats.is_local
                    && !stats.is_offline(self.offline_timeout)
                    && stats.health_score() >= self.min_health_score
            })
            .map(|stats| (stats.peer_id, stats.health_score()))
            .collect();

        peers.sort_by(|a, b| b.1.cmp(&a.1));
        peers.into_iter().map(|(peer_id, _)| peer_id).collect()
    }

    /// Get best peer for data transfer (prioritizes local + healthy)
    pub fn get_best_peer(&self) -> Option<PeerId> {
        // First try local peers
        if let Some(peer) = self.get_local_peers().first() {
            return Some(*peer);
        }

        // Fall back to any healthy peer
        self.get_healthy_peers().first().copied()
    }

    /// Get all peers with their statistics
    pub fn get_all_stats(&self) -> Vec<&PeerStats> {
        self.peers.values().collect()
    }

    /// Remove offline peers
    pub fn cleanup_offline_peers(&mut self) {
        let offline_peers: Vec<_> = self
            .peers
            .iter()
            .filter(|(_, stats)| stats.is_offline(self.offline_timeout))
            .map(|(peer_id, _)| *peer_id)
            .collect();

        for peer_id in offline_peers {
            warn!("Removing offline peer: {:?}", peer_id);
            self.peers.remove(&peer_id);
        }
    }

    /// Get network health summary
    pub fn get_network_summary(&self) -> NetworkHealthSummary {
        let total_peers = self.peers.len();
        let healthy = self
            .peers
            .values()
            .filter(|s| s.health == PeerHealth::Healthy)
            .count();
        let degraded = self
            .peers
            .values()
            .filter(|s| s.health == PeerHealth::Degraded)
            .count();
        let unhealthy = self
            .peers
            .values()
            .filter(|s| s.health == PeerHealth::Unhealthy)
            .count();
        let local_peers = self.peers.values().filter(|s| s.is_local).count();

        NetworkHealthSummary {
            total_peers,
            healthy_peers: healthy,
            degraded_peers: degraded,
            unhealthy_peers: unhealthy,
            local_peers,
            remote_peers: total_peers - local_peers,
        }
    }
}

impl Default for HealthMonitor {
    fn default() -> Self {
        Self::new()
    }
}

/// Summary of network health
#[derive(Debug, Clone)]
pub struct NetworkHealthSummary {
    pub total_peers: usize,
    pub healthy_peers: usize,
    pub degraded_peers: usize,
    pub unhealthy_peers: usize,
    pub local_peers: usize,
    pub remote_peers: usize,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_peer_stats_success() {
        let peer_id = PeerId::random();
        let mut stats = PeerStats::new(peer_id, true);

        // Record several successful pings
        stats.record_success(Duration::from_millis(50));
        assert_eq!(stats.success_count, 1);
        assert_eq!(stats.last_rtt_ms, Some(50));
        
        stats.record_success(Duration::from_millis(60));
        assert_eq!(stats.success_count, 2);
        
        // Should be healthy with low RTT
        assert_eq!(stats.health, PeerHealth::Healthy);
    }

    #[test]
    fn test_peer_stats_failure() {
        let peer_id = PeerId::random();
        let mut stats = PeerStats::new(peer_id, false);

        // Record mostly failures
        stats.record_failure();
        stats.record_failure();
        stats.record_success(Duration::from_millis(100));
        
        // Should be unhealthy due to high failure rate
        assert!(matches!(
            stats.health,
            PeerHealth::Unhealthy | PeerHealth::Degraded
        ));
    }

    #[test]
    fn test_health_monitor() {
        let mut monitor = HealthMonitor::new();
        let peer1 = PeerId::random();
        let peer2 = PeerId::random();

        monitor.add_peer(peer1, true);
        monitor.add_peer(peer2, false);

        // Make peer1 healthy
        monitor.record_ping_success(&peer1, Duration::from_millis(30));
        monitor.record_ping_success(&peer1, Duration::from_millis(40));

        // Make peer2 degraded
        monitor.record_ping_success(&peer2, Duration::from_millis(200));
        monitor.record_ping_failure(&peer2);

        let healthy = monitor.get_healthy_peers();
        assert!(healthy.contains(&peer1));
        
        let local = monitor.get_local_peers();
        assert_eq!(local.len(), 1);
        assert_eq!(local[0], peer1);
    }
}