Skip to main content

rust_webx_host/
health.rs

1//! Enhanced health check endpoints for Kubernetes.
2//!
3//! Provides HealthCheckRegistry and typed HealthStatus.
4//! Endpoints follow RFC 8407 (`application/health+json` content type).
5
6use rust_webx_core::error::Result;
7use rust_webx_core::http::IHttpContext;
8use rust_webx_core::routing::IEndpoint;
9use std::sync::{Arc, Mutex};
10
11pub type HealthCheckFn = Arc<dyn Fn() -> HealthStatus + Send + Sync>;
12
13#[derive(Debug, Clone, serde::Serialize)]
14pub struct HealthStatus {
15    pub status: &'static str,
16    #[serde(skip_serializing_if = "Option::is_none")]
17    pub detail: Option<String>,
18}
19
20impl HealthStatus {
21    pub fn pass() -> Self {
22        Self {
23            status: "pass",
24            detail: None,
25        }
26    }
27    pub fn warn(d: impl Into<String>) -> Self {
28        Self {
29            status: "warn",
30            detail: Some(d.into()),
31        }
32    }
33    pub fn fail(d: impl Into<String>) -> Self {
34        Self {
35            status: "fail",
36            detail: Some(d.into()),
37        }
38    }
39}
40
41/// A single health check entry in a registry snapshot.
42#[derive(Debug, Clone, serde::Serialize)]
43pub struct HealthCheckEntry {
44    pub name: String,
45    pub status: String,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub detail: Option<String>,
48}
49
50pub struct HealthCheckRegistry {
51    checks: Mutex<Vec<(String, HealthCheckFn)>>,
52}
53
54impl HealthCheckRegistry {
55    pub fn new() -> Arc<Self> {
56        Arc::new(Self {
57            checks: Mutex::new(Vec::new()),
58        })
59    }
60    pub fn register(self: &Arc<Self>, name: impl Into<String>, check: HealthCheckFn) {
61        if let Ok(mut checks) = self.checks.lock() {
62            checks.push((name.into(), check));
63        }
64    }
65    pub fn all_healthy(&self) -> bool {
66        self.checks
67            .lock()
68            .map(|c| c.iter().all(|(_, f)| f().status != "fail"))
69            .unwrap_or(false)
70    }
71
72    /// 返回所有健康检查项的当前快照。
73    pub fn snapshot(&self) -> Vec<HealthCheckEntry> {
74        self.checks
75            .lock()
76            .map(|c| {
77                c.iter()
78                    .map(|(name, f)| {
79                        let s = f();
80                        HealthCheckEntry {
81                            name: name.clone(),
82                            status: s.status.to_string(),
83                            detail: s.detail.clone(),
84                        }
85                    })
86                    .collect()
87            })
88            .unwrap_or_default()
89    }
90
91    /// 计算整体状态:fail 优先于 warn 优先于 pass;空列表视为 pass。
92    /// Mutex poisoned 时保守返回 "fail"。
93    pub fn overall_status(&self) -> &'static str {
94        self.checks
95            .lock()
96            .map(|c| {
97                let entries: Vec<_> = c.iter().map(|(_, f)| f()).collect();
98                if entries.is_empty() {
99                    "pass"
100                } else if entries.iter().any(|e| e.status == "fail") {
101                    "fail"
102                } else if entries.iter().any(|e| e.status == "warn") {
103                    "warn"
104                } else {
105                    "pass"
106                }
107            })
108            .unwrap_or("fail")
109    }
110}
111
112/// Build an RFC 8407 `application/health+json` body from the registry.
113///
114/// When the registry has no checks, returns `{"status":"pass"}`.
115/// Otherwise returns `{"status":<overall>,"checks":[...]}`.
116pub fn build_health_response(registry: &HealthCheckRegistry) -> Vec<u8> {
117    let entries = registry.snapshot();
118    let overall = registry.overall_status();
119    let body = if entries.is_empty() {
120        serde_json::json!({ "status": overall })
121    } else {
122        serde_json::json!({
123            "status": overall,
124            "checks": entries
125        })
126    };
127    serde_json::to_vec(&body).unwrap_or_default()
128}
129
130/// HTTP status for a health overall status: `fail` → 503, otherwise 200.
131pub fn health_http_status(overall: &str) -> u16 {
132    if overall == "fail" {
133        503
134    } else {
135        200
136    }
137}
138
139/// Dynamic health endpoint — evaluates registered probes on each request.
140pub struct HealthCheckEndpoint {
141    registry: Arc<HealthCheckRegistry>,
142}
143
144impl HealthCheckEndpoint {
145    pub fn new(registry: Arc<HealthCheckRegistry>) -> Self {
146        Self { registry }
147    }
148}
149
150#[async_trait::async_trait]
151impl IEndpoint for HealthCheckEndpoint {
152    async fn handle(&self, ctx: &mut dyn IHttpContext) -> Result<()> {
153        let overall = self.registry.overall_status();
154        let body = build_health_response(&self.registry);
155        ctx.response_mut()
156            .set_status(health_http_status(overall));
157        ctx.response_mut()
158            .set_header("content-type", "application/health+json");
159        ctx.response_mut().write_bytes(body).await
160    }
161}