github_mcp/core/
health_check_manager.rs

1// GitHub v3 REST API MCP server — generated by mcpify. Do not hand-edit.
2
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::Arc;
6use std::time::Duration;
7
8use tokio::sync::Mutex;
9
10use super::component_registry::{ComponentRegistry, ComponentStatus};
11
12pub type HealthCheckFuture = Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send>>;
13pub type HealthCheckFn = Box<dyn Fn() -> HealthCheckFuture + Send + Sync>;
14
15const DEFAULT_INTERVAL: Duration = Duration::from_secs(30);
16const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
17
18struct RegisteredCheck {
19    name: String,
20    check: HealthCheckFn,
21}
22
23/// Runs registered health checks periodically, updating a shared
24/// `ComponentRegistry` so degraded-mode startup and a health-check endpoint
25/// (plus the `test-connection` CLI command) can report accurately —
26/// REQ-2.3.4.
27pub struct HealthCheckManager {
28    registry: Arc<Mutex<ComponentRegistry>>,
29    checks: Vec<RegisteredCheck>,
30    interval: Duration,
31    timeout: Duration,
32}
33
34impl HealthCheckManager {
35    pub fn new(registry: Arc<Mutex<ComponentRegistry>>) -> Self {
36        Self {
37            registry,
38            checks: Vec::new(),
39            interval: DEFAULT_INTERVAL,
40            timeout: DEFAULT_TIMEOUT,
41        }
42    }
43
44    pub async fn register(&mut self, name: &str, critical: bool, check: HealthCheckFn) {
45        self.registry.lock().await.register(name, critical);
46        self.checks.push(RegisteredCheck {
47            name: name.to_string(),
48            check,
49        });
50    }
51
52    pub async fn run_once(&self) {
53        for RegisteredCheck { name, check } in &self.checks {
54            let outcome = tokio::time::timeout(self.timeout, check()).await;
55            let mut registry = self.registry.lock().await;
56            match outcome {
57                Ok(Ok(())) => registry.report(name, ComponentStatus::Healthy, None),
58                Ok(Err(err)) => {
59                    tracing::warn!(component = name, error = %err, "health check failed");
60                    registry.report(name, ComponentStatus::Unhealthy, Some(err.to_string()));
61                }
62                Err(_) => {
63                    tracing::warn!(component = name, "health check timed out");
64                    registry.report(
65                        name,
66                        ComponentStatus::Unhealthy,
67                        Some("timed out".to_string()),
68                    );
69                }
70            }
71        }
72    }
73
74    /// Spawns the periodic check loop; the returned handle can be aborted
75    /// during shutdown.
76    pub fn start(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
77        tokio::spawn(async move {
78            let mut ticker = tokio::time::interval(self.interval);
79            loop {
80                ticker.tick().await;
81                self.run_once().await;
82            }
83        })
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[tokio::test]
92    async fn a_passing_check_reports_healthy() {
93        let registry = Arc::new(Mutex::new(ComponentRegistry::new()));
94        let mut manager = HealthCheckManager::new(registry.clone());
95        manager
96            .register("db", true, Box::new(|| Box::pin(async { Ok(()) })))
97            .await;
98
99        manager.run_once().await;
100
101        assert_eq!(
102            registry.lock().await.overall_status(),
103            ComponentStatus::Healthy
104        );
105    }
106
107    #[tokio::test]
108    async fn a_failing_check_reports_unhealthy() {
109        let registry = Arc::new(Mutex::new(ComponentRegistry::new()));
110        let mut manager = HealthCheckManager::new(registry.clone());
111        manager
112            .register(
113                "db",
114                true,
115                Box::new(|| Box::pin(async { Err(anyhow::anyhow!("connection refused")) })),
116            )
117            .await;
118
119        manager.run_once().await;
120
121        assert_eq!(
122            registry.lock().await.overall_status(),
123            ComponentStatus::Unhealthy
124        );
125    }
126}
github_mcp/core/health_check_manager.rs

github_mcp/core/
health_check_manager.rs
