cloudlysis/aggregate/src/server/health.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::RwLock;
use std::time::Instant;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum HealthStatus {
    Healthy,
    Degraded { issues: Vec<String> },
    Unhealthy { reasons: Vec<String> },
}

impl HealthStatus {
    pub fn is_healthy(&self) -> bool {
        matches!(self, Self::Healthy)
    }

    pub fn is_degraded(&self) -> bool {
        matches!(self, Self::Degraded { .. })
    }

    pub fn is_unhealthy(&self) -> bool {
        matches!(self, Self::Unhealthy { .. })
    }
}

#[derive(Debug, Clone)]
pub struct ComponentHealth {
    pub name: String,
    pub status: HealthStatus,
    pub last_check: Instant,
    pub details: HashMap<String, String>,
}

impl ComponentHealth {
    pub fn healthy(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            status: HealthStatus::Healthy,
            last_check: Instant::now(),
            details: HashMap::new(),
        }
    }

    pub fn degraded(name: impl Into<String>, issues: Vec<String>) -> Self {
        Self {
            name: name.into(),
            status: HealthStatus::Degraded { issues },
            last_check: Instant::now(),
            details: HashMap::new(),
        }
    }

    pub fn unhealthy(name: impl Into<String>, reasons: Vec<String>) -> Self {
        Self {
            name: name.into(),
            status: HealthStatus::Unhealthy { reasons },
            last_check: Instant::now(),
            details: HashMap::new(),
        }
    }

    pub fn with_detail(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.details.insert(key.into(), value.into());
        self
    }
}

pub struct HealthChecker {
    storage_healthy: AtomicBool,
    stream_healthy: AtomicBool,
    components: RwLock<HashMap<String, ComponentHealth>>,
}

impl HealthChecker {
    pub fn new() -> Self {
        Self {
            storage_healthy: AtomicBool::new(true),
            stream_healthy: AtomicBool::new(true),
            components: RwLock::new(HashMap::new()),
        }
    }

    pub fn storage_healthy(&self) -> bool {
        self.storage_healthy.load(Ordering::Relaxed)
    }

    pub fn stream_healthy(&self) -> bool {
        self.stream_healthy.load(Ordering::Relaxed)
    }

    pub fn set_storage_healthy(&self, healthy: bool) {
        self.storage_healthy.store(healthy, Ordering::Relaxed);
        self.update_component(
            "storage",
            healthy,
            if healthy { "connected" } else { "disconnected" },
        );
    }

    pub fn set_stream_healthy(&self, healthy: bool) {
        self.stream_healthy.store(healthy, Ordering::Relaxed);
        self.update_component(
            "stream",
            healthy,
            if healthy { "connected" } else { "disconnected" },
        );
    }

    fn update_component(&self, name: &str, healthy: bool, status: &str) {
        let mut components = self.components.write().unwrap();
        let health = if healthy {
            ComponentHealth::healthy(name).with_detail("status", status)
        } else {
            ComponentHealth::unhealthy(name, vec![format!("status: {}", status)])
        };
        components.insert(name.to_string(), health);
    }

    pub fn check(&self) -> HealthStatus {
        let storage = self.storage_healthy.load(Ordering::Relaxed);
        let stream = self.stream_healthy.load(Ordering::Relaxed);

        match (storage, stream) {
            (true, true) => HealthStatus::Healthy,
            (true, false) | (false, true) => {
                let mut issues = Vec::new();
                if !storage {
                    issues.push("storage disconnected".to_string());
                }
                if !stream {
                    issues.push("stream disconnected".to_string());
                }
                HealthStatus::Degraded { issues }
            }
            (false, false) => HealthStatus::Unhealthy {
                reasons: vec![
                    "storage disconnected".to_string(),
                    "stream disconnected".to_string(),
                ],
            },
        }
    }

    pub fn is_ready(&self) -> bool {
        let status = self.check();
        status.is_healthy() || status.is_degraded()
    }

    pub fn is_live(&self) -> bool {
        true
    }

    pub fn components(&self) -> HashMap<String, ComponentHealth> {
        self.components.read().unwrap().clone()
    }
}

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

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

    #[test]
    fn health_status_checks() {
        let healthy = HealthStatus::Healthy;
        assert!(healthy.is_healthy());
        assert!(!healthy.is_degraded());
        assert!(!healthy.is_unhealthy());

        let degraded = HealthStatus::Degraded {
            issues: vec!["test".to_string()],
        };
        assert!(!degraded.is_healthy());
        assert!(degraded.is_degraded());
        assert!(!degraded.is_unhealthy());

        let unhealthy = HealthStatus::Unhealthy {
            reasons: vec!["test".to_string()],
        };
        assert!(!unhealthy.is_healthy());
        assert!(!unhealthy.is_degraded());
        assert!(unhealthy.is_unhealthy());
    }

    #[test]
    fn component_health_builders() {
        let healthy = ComponentHealth::healthy("storage");
        assert_eq!(healthy.name, "storage");
        assert!(healthy.status.is_healthy());

        let degraded = ComponentHealth::degraded("stream", vec!["slow".to_string()]);
        assert!(degraded.status.is_degraded());

        let unhealthy = ComponentHealth::unhealthy("db", vec!["down".to_string()]);
        assert!(unhealthy.status.is_unhealthy());
    }

    #[test]
    fn health_checker_starts_healthy() {
        let checker = HealthChecker::new();
        assert!(checker.check().is_healthy());
    }

    #[test]
    fn health_checker_storage_failure() {
        let checker = HealthChecker::new();
        checker.set_storage_healthy(false);

        let status = checker.check();
        assert!(status.is_degraded());
    }

    #[test]
    fn health_checker_all_failures() {
        let checker = HealthChecker::new();
        checker.set_storage_healthy(false);
        checker.set_stream_healthy(false);

        let status = checker.check();
        assert!(status.is_unhealthy());
    }

    #[test]
    fn health_checker_is_ready() {
        let checker = HealthChecker::new();
        assert!(checker.is_ready());

        checker.set_storage_healthy(false);
        assert!(checker.is_ready());
    }

    #[test]
    fn health_checker_is_live() {
        let checker = HealthChecker::new();
        assert!(checker.is_live());

        checker.set_storage_healthy(false);
        checker.set_stream_healthy(false);
        assert!(checker.is_live());
    }

    #[test]
    fn health_checker_tracks_components() {
        let checker = HealthChecker::new();
        checker.set_storage_healthy(true);
        checker.set_stream_healthy(true);

        let components = checker.components();
        assert!(components.contains_key("storage"));
        assert!(components.contains_key("stream"));
    }
}