memlink_runtime/

health.rs

//! Module health checking and probes.
//!
//! Provides liveness and readiness probes for module health monitoring.

use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use tokio::sync::Mutex;

/// Health status.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HealthStatus {
    /// Module is healthy and ready.
    Healthy,
    /// Module is alive but not ready.
    Unready,
    /// Module is unhealthy.
    Unhealthy,
    /// Module health is unknown.
    Unknown,
}

/// Health check configuration.
#[derive(Debug, Clone)]
pub struct HealthConfig {
    /// Interval between health checks.
    pub check_interval: Duration,
    /// Timeout for health check calls.
    pub timeout: Duration,
    /// Number of failures before unhealthy.
    pub failure_threshold: u32,
    /// Number of successes before healthy.
    pub success_threshold: u32,
}

impl Default for HealthConfig {
    fn default() -> Self {
        Self {
            check_interval: Duration::from_secs(10),
            timeout: Duration::from_secs(5),
            failure_threshold: 3,
            success_threshold: 2,
        }
    }
}

/// Health check result.
#[derive(Debug, Clone)]
pub struct HealthCheckResult {
    /// Timestamp of the check.
    pub timestamp: Instant,
    /// Whether the check succeeded.
    pub success: bool,
    /// Response time in microseconds.
    pub latency_us: u64,
    /// Error message if failed.
    pub error: Option<String>,
}

/// Module health tracker.
#[derive(Debug)]
pub struct HealthTracker {
    /// Module name.
    module_name: String,
    /// Current health status.
    status: Mutex<HealthStatus>,
    /// Consecutive failures.
    failures: AtomicU32,
    /// Consecutive successes.
    successes: AtomicU32,
    /// Total health checks.
    total_checks: AtomicU64,
    /// Last successful check.
    last_success: Mutex<Option<Instant>>,
    /// Last failed check.
    last_failure: Mutex<Option<Instant>>,
    /// Configuration.
    config: HealthConfig,
}

impl HealthTracker {
    /// Creates a new health tracker.
    pub fn new(module_name: String, config: HealthConfig) -> Self {
        Self {
            module_name,
            status: Mutex::new(HealthStatus::Unknown),
            failures: AtomicU32::new(0),
            successes: AtomicU32::new(0),
            total_checks: AtomicU64::new(0),
            last_success: Mutex::new(None),
            last_failure: Mutex::new(None),
            config,
        }
    }

    /// Records a successful health check.
    pub async fn record_success(&self, _latency_us: u64) {
        self.total_checks.fetch_add(1, Ordering::Relaxed);
        self.successes.fetch_add(1, Ordering::Relaxed);
        self.failures.store(0, Ordering::Relaxed);
        *self.last_success.lock().await = Some(Instant::now());

        let mut status = self.status.lock().await;
        let current = *status;

        if current == HealthStatus::Unhealthy || current == HealthStatus::Unknown {
            let successes = self.successes.load(Ordering::Relaxed);
            if successes >= self.config.success_threshold {
                *status = HealthStatus::Healthy;
            } else if current == HealthStatus::Unhealthy {
                *status = HealthStatus::Unready;
            }
        }
    }

    /// Records a failed health check.
    pub async fn record_failure(&self, _error: Option<String>) {
        self.total_checks.fetch_add(1, Ordering::Relaxed);
        self.failures.fetch_add(1, Ordering::Relaxed);
        self.successes.store(0, Ordering::Relaxed);
        *self.last_failure.lock().await = Some(Instant::now());

        let mut status = self.status.lock().await;

        let failures = self.failures.load(Ordering::Relaxed);
        if failures >= self.config.failure_threshold {
            *status = HealthStatus::Unhealthy;
        }
    }

    /// Returns the current health status.
    pub async fn status(&self) -> HealthStatus {
        *self.status.lock().await
    }

    /// Returns whether the module is healthy.
    pub async fn is_healthy(&self) -> bool {
        *self.status.lock().await == HealthStatus::Healthy
    }

    /// Returns whether the module is ready.
    pub async fn is_ready(&self) -> bool {
        let status = *self.status.lock().await;
        status == HealthStatus::Healthy || status == HealthStatus::Unready
    }

    /// Returns statistics.
    pub fn stats(&self) -> HealthStats {
        HealthStats {
            module_name: self.module_name.clone(),
            total_checks: self.total_checks.load(Ordering::Relaxed),
            consecutive_failures: self.failures.load(Ordering::Relaxed),
            consecutive_successes: self.successes.load(Ordering::Relaxed),
        }
    }
}

/// Health statistics.
#[derive(Debug, Clone)]
pub struct HealthStats {
    pub module_name: String,
    pub total_checks: u64,
    pub consecutive_failures: u32,
    pub consecutive_successes: u32,
}

/// Health check registry for all modules.
#[derive(Debug)]
pub struct HealthRegistry {
    /// Health trackers by module name.
    trackers: DashMap<String, Arc<HealthTracker>>,
    /// Default configuration.
    default_config: HealthConfig,
}

impl HealthRegistry {
    /// Creates a new health registry.
    pub fn new() -> Self {
        Self {
            trackers: DashMap::new(),
            default_config: HealthConfig::default(),
        }
    }

    /// Gets or creates a health tracker.
    pub fn get_or_create(&self, module_name: &str) -> Arc<HealthTracker> {
        self.trackers
            .entry(module_name.to_string())
            .or_insert_with(|| {
                Arc::new(HealthTracker::new(
                    module_name.to_string(),
                    self.default_config.clone(),
                ))
            })
            .clone()
    }

    /// Returns all unhealthy modules.
    pub async fn unhealthy_modules(&self) -> Vec<String> {
        let mut result = Vec::new();
        for entry in self.trackers.iter() {
            if !entry.value().is_healthy().await {
                result.push(entry.key().clone());
            }
        }
        result
    }

    /// Returns overall health summary.
    pub async fn summary(&self) -> HealthSummary {
        let mut healthy = 0;
        let mut unready = 0;
        let mut unhealthy = 0;
        let mut unknown = 0;

        for entry in self.trackers.iter() {
            match entry.value().status().await {
                HealthStatus::Healthy => healthy += 1,
                HealthStatus::Unready => unready += 1,
                HealthStatus::Unhealthy => unhealthy += 1,
                HealthStatus::Unknown => unknown += 1,
            }
        }

        HealthSummary {
            total: self.trackers.len(),
            healthy,
            unready,
            unhealthy,
            unknown,
        }
    }
}

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

/// Overall health summary.
#[derive(Debug, Clone)]
pub struct HealthSummary {
    pub total: usize,
    pub healthy: usize,
    pub unready: usize,
    pub unhealthy: usize,
    pub unknown: usize,
}

use dashmap::DashMap;

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

    #[tokio::test]
    async fn test_health_tracker_transitions() {
        let tracker = HealthTracker::new("test".to_string(), HealthConfig::default());

        // Initially unknown
        assert_eq!(tracker.status().await, HealthStatus::Unknown);

        // Record successes
        tracker.record_success(100).await;
        tracker.record_success(100).await;

        assert_eq!(tracker.status().await, HealthStatus::Healthy);
        assert!(tracker.is_healthy().await);
        assert!(tracker.is_ready().await);
    }

    #[tokio::test]
    async fn test_health_tracker_failure() {
        let config = HealthConfig {
            failure_threshold: 2,
            ..HealthConfig::default()
        };
        let tracker = HealthTracker::new("test".to_string(), config);

        // Record failures
        tracker.record_failure(Some("error".to_string())).await;
        tracker.record_failure(Some("error".to_string())).await;
        tracker.record_failure(Some("error".to_string())).await;

        assert_eq!(tracker.status().await, HealthStatus::Unhealthy);
        assert!(!tracker.is_healthy().await);
        assert!(!tracker.is_ready().await);
    }
}