Skip to main content

mecha10_runtime/
health.rs

1//! Health checking infrastructure
2
3use mecha10_core::rwlock_collections::RwLockMapExt;
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::time::Duration;
7use tokio::sync::RwLock;
8
9#[derive(Clone, Debug, PartialEq)]
10pub enum HealthStatus {
11    Healthy,
12
13    Degraded { reason: String },
14
15    Unhealthy { reason: String },
16}
17
18impl HealthStatus {
19    /// Check if status is healthy
20    pub fn is_healthy(&self) -> bool {
21        matches!(self, HealthStatus::Healthy)
22    }
23
24    /// Check if status is degraded
25    pub fn is_degraded(&self) -> bool {
26        matches!(self, HealthStatus::Degraded { .. })
27    }
28
29    /// Check if status is unhealthy
30    pub fn is_unhealthy(&self) -> bool {
31        matches!(self, HealthStatus::Unhealthy { .. })
32    }
33
34    /// Get reason if not healthy
35    pub fn reason(&self) -> Option<&str> {
36        match self {
37            HealthStatus::Healthy => None,
38            HealthStatus::Degraded { reason } => Some(reason),
39            HealthStatus::Unhealthy { reason } => Some(reason),
40        }
41    }
42}
43
44pub struct HealthChecker {
45    checks: Arc<RwLock<HashMap<String, HealthStatus>>>,
46    interval: Duration,
47}
48
49impl HealthChecker {
50    /// Create a new health checker with specified interval
51    pub fn new(interval: Duration) -> Self {
52        Self {
53            checks: Arc::new(RwLock::new(HashMap::new())),
54            interval,
55        }
56    }
57
58    /// Register a node for health checking
59    pub async fn register(&self, name: String, initial_status: HealthStatus) {
60        let mut checks = self.checks.write().await;
61        checks.insert(name, initial_status);
62    }
63
64    /// Update health status for a node
65    pub async fn update(&self, name: &str, status: HealthStatus) {
66        let mut checks = self.checks.write().await;
67        checks.insert(name.to_string(), status);
68    }
69
70    /// Get health status for a specific node
71    pub async fn check_one(&self, name: &str) -> Option<HealthStatus> {
72        self.checks.cloned_get(&name.to_string()).await
73    }
74
75    /// Get health status for all nodes
76    pub async fn check_all(&self) -> HashMap<String, HealthStatus> {
77        let checks = self.checks.read().await;
78        checks.clone()
79    }
80
81    /// Get the check interval
82    pub fn interval(&self) -> Duration {
83        self.interval
84    }
85
86    /// Remove a node from health checking
87    pub async fn unregister(&self, name: &str) {
88        let mut checks = self.checks.write().await;
89        checks.remove(name);
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[tokio::test]
98    async fn check_one_returns_none_before_registration() {
99        let checker = HealthChecker::new(Duration::from_secs(1));
100
101        assert_eq!(checker.check_one("camera").await, None);
102    }
103
104    #[tokio::test]
105    async fn check_one_returns_the_registered_status() {
106        let checker = HealthChecker::new(Duration::from_secs(1));
107
108        checker.register("camera".to_string(), HealthStatus::Healthy).await;
109
110        assert_eq!(checker.check_one("camera").await, Some(HealthStatus::Healthy));
111    }
112
113    #[tokio::test]
114    async fn update_overwrites_the_status_seen_by_check_one_and_check_all() {
115        let checker = HealthChecker::new(Duration::from_secs(1));
116        checker.register("camera".to_string(), HealthStatus::Healthy).await;
117
118        checker
119            .update(
120                "camera",
121                HealthStatus::Degraded {
122                    reason: "low battery".to_string(),
123                },
124            )
125            .await;
126
127        assert_eq!(
128            checker.check_one("camera").await,
129            Some(HealthStatus::Degraded {
130                reason: "low battery".to_string()
131            })
132        );
133        assert_eq!(checker.check_all().await.len(), 1);
134    }
135
136    #[tokio::test]
137    async fn unregister_removes_the_node_from_check_one_and_check_all() {
138        let checker = HealthChecker::new(Duration::from_secs(1));
139        checker.register("camera".to_string(), HealthStatus::Healthy).await;
140
141        checker.unregister("camera").await;
142
143        assert_eq!(checker.check_one("camera").await, None);
144        assert!(checker.check_all().await.is_empty());
145    }
146}