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 }, Unhealthy { reasons: Vec }, } 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, } impl ComponentHealth { pub fn healthy(name: impl Into) -> Self { Self { name: name.into(), status: HealthStatus::Healthy, last_check: Instant::now(), details: HashMap::new(), } } pub fn degraded(name: impl Into, issues: Vec) -> Self { Self { name: name.into(), status: HealthStatus::Degraded { issues }, last_check: Instant::now(), details: HashMap::new(), } } pub fn unhealthy(name: impl Into, reasons: Vec) -> 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, value: impl Into) -> Self { self.details.insert(key.into(), value.into()); self } } pub struct HealthChecker { storage_healthy: AtomicBool, stream_healthy: AtomicBool, components: RwLock>, } 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 { 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")); } }