rusty_common/websocket/

heartbeat.rs

//! Heartbeat monitoring for WebSocket connections
//!
//! This module provides a heartbeat mechanism for monitoring the health of WebSocket connections
//! and triggering reconnection when necessary.

use log::{debug, warn};
use parking_lot::RwLock;
use quanta::{Clock, Instant};
use std::sync::Arc;
use std::time::Duration;

use super::client::ConnectionState;
use super::stats::ConnectionStats;

/// Heartbeat status
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HeartbeatStatus {
    /// Connection is healthy
    Healthy,
    /// Connection is unhealthy (missed some heartbeats)
    Unhealthy,
    /// Connection is dead (missed too many heartbeats)
    Dead,
}

/// Heartbeat monitor for WebSocket connections
#[derive(Debug, Clone)]
pub struct HeartbeatMonitor {
    /// Last time a message was received
    last_message_time: Arc<RwLock<Instant>>,
    /// Last time a heartbeat check was performed
    last_check_time: Arc<RwLock<Instant>>,
    /// Heartbeat interval in milliseconds
    heartbeat_interval_milliseconds: u64,
    /// Heartbeat timeout in milliseconds
    heartbeat_timeout_milliseconds: u64,
    /// Maximum number of consecutive missed heartbeats before considering the connection dead
    max_missed_heartbeats: u32,
    /// Number of consecutive missed heartbeats
    missed_heartbeats: Arc<RwLock<u32>>,
    /// Connection status
    connection_status: Arc<RwLock<ConnectionState>>,
    /// Connection statistics
    stats: Arc<RwLock<ConnectionStats>>,
    /// High-precision clock
    clock: Clock,
}

impl HeartbeatMonitor {
    /// Create a new heartbeat monitor
    #[must_use]
    pub fn new(
        heartbeat_interval_milliseconds: u64,
        heartbeat_timeout_milliseconds: u64,
        max_missed_heartbeats: u32,
        connection_status: Arc<RwLock<ConnectionState>>,
        stats: Arc<RwLock<ConnectionStats>>,
    ) -> Self {
        let clock = Clock::new();
        let now = clock.now();
        Self {
            last_message_time: Arc::new(RwLock::new(now)),
            last_check_time: Arc::new(RwLock::new(now)),
            heartbeat_interval_milliseconds,
            heartbeat_timeout_milliseconds,
            max_missed_heartbeats,
            missed_heartbeats: Arc::new(RwLock::new(0)),
            connection_status,
            stats,
            clock,
        }
    }

    /// Update the last message time
    pub fn update_last_message_time(&self) {
        *self.last_message_time.write() = self.clock.now();
        *self.missed_heartbeats.write() = 0;
    }

    /// Check the heartbeat status
    pub fn check_heartbeat(&self) -> HeartbeatStatus {
        // If heartbeat is disabled, always return Healthy
        if self.heartbeat_interval_milliseconds == 0 {
            return HeartbeatStatus::Healthy;
        }

        let now = self.clock.now();
        let last_message_time = *self.last_message_time.read();
        let last_check_time = *self.last_check_time.read();
        let elapsed_since_last_message = now.duration_since(last_message_time);
        let elapsed_since_last_check = now.duration_since(last_check_time);
        let interval = Duration::from_millis(self.heartbeat_interval_milliseconds);
        let timeout = Duration::from_millis(self.heartbeat_timeout_milliseconds);

        // Only check heartbeat if enough time has passed since the last check
        if elapsed_since_last_check < interval {
            // If we haven't reached the heartbeat interval yet, return the current status
            let missed = *self.missed_heartbeats.read();
            return if missed >= self.max_missed_heartbeats {
                HeartbeatStatus::Dead
            } else if missed > 0 {
                HeartbeatStatus::Unhealthy
            } else {
                HeartbeatStatus::Healthy
            };
        }

        // Update last check time
        *self.last_check_time.write() = now;

        // Check if we've exceeded the heartbeat timeout
        if elapsed_since_last_message > timeout {
            // Increment missed heartbeats
            let mut missed_heartbeats = self.missed_heartbeats.write();
            *missed_heartbeats += 1;

            // Log the missed heartbeat
            debug!(
                "Missed heartbeat: {} of {} (timeout: {}ms, elapsed: {}ms)",
                *missed_heartbeats,
                self.max_missed_heartbeats,
                self.heartbeat_timeout_milliseconds,
                elapsed_since_last_message.as_millis()
            );

            // Check if we've missed too many heartbeats
            if *missed_heartbeats >= self.max_missed_heartbeats {
                warn!(
                    "Connection dead: missed {} heartbeats (max: {})",
                    *missed_heartbeats, self.max_missed_heartbeats
                );
                *self.connection_status.write() = ConnectionState::Error;
                self.stats.write().errors += 1;
                return HeartbeatStatus::Dead;
            }

            // Connection is unhealthy but not dead yet
            return HeartbeatStatus::Unhealthy;
        }

        // Connection is healthy
        *self.missed_heartbeats.write() = 0;
        HeartbeatStatus::Healthy
    }

    /// Reset the heartbeat monitor
    pub fn reset(&self) {
        let now = self.clock.now();
        *self.last_message_time.write() = now;
        *self.last_check_time.write() = now;
        *self.missed_heartbeats.write() = 0;
    }

    /// Get the number of consecutive missed heartbeats
    pub fn missed_heartbeats(&self) -> u32 {
        *self.missed_heartbeats.read()
    }

    /// Get the time elapsed since the last message
    pub fn time_since_last_message(&self) -> Duration {
        self.clock
            .now()
            .duration_since(*self.last_message_time.read())
    }
}

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

    #[test]
    fn test_heartbeat_monitor() {
        let connection_status = Arc::new(RwLock::new(ConnectionState::Connected));
        let stats = Arc::new(RwLock::new(ConnectionStats::default()));

        // Create a heartbeat monitor with a short timeout for testing
        let monitor = HeartbeatMonitor::new(
            100, // 100ms heartbeat interval
            200, // 200ms heartbeat timeout
            3,   // 3 missed heartbeats before connection is dead
            connection_status.clone(),
            stats.clone(),
        );

        // Initially, the connection should be healthy
        assert_eq!(monitor.check_heartbeat(), HeartbeatStatus::Healthy);

        // Update the last message time
        monitor.update_last_message_time();
        assert_eq!(monitor.check_heartbeat(), HeartbeatStatus::Healthy);

        // Wait for the heartbeat timeout to expire
        thread::sleep(Duration::from_millis(300));

        // After the timeout, the connection should be unhealthy
        assert_eq!(monitor.check_heartbeat(), HeartbeatStatus::Unhealthy);
        assert_eq!(monitor.missed_heartbeats(), 1);

        // Wait for more timeouts
        thread::sleep(Duration::from_millis(300));
        assert_eq!(monitor.check_heartbeat(), HeartbeatStatus::Unhealthy);
        assert_eq!(monitor.missed_heartbeats(), 2);

        thread::sleep(Duration::from_millis(300));

        // After 3 missed heartbeats, the connection should be dead
        assert_eq!(monitor.check_heartbeat(), HeartbeatStatus::Dead);
        assert_eq!(monitor.missed_heartbeats(), 3);
        assert_eq!(*connection_status.read(), ConnectionState::Error);

        // Reset the monitor
        monitor.reset();
        assert_eq!(monitor.missed_heartbeats(), 0);
        assert_eq!(monitor.check_heartbeat(), HeartbeatStatus::Healthy);
    }
}