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        let checks = match self.checks.lock() {
67            Ok(c) => c.clone(),
68            Err(_) => return false,
69        };
70        checks.iter().all(|(_, f)| f().status != "fail")
71    }
72
73    /// 返回所有健康检查项的当前快照。
74    pub fn snapshot(&self) -> Vec<HealthCheckEntry> {
75        let checks = match self.checks.lock() {
76            Ok(c) => c.clone(),
77            Err(_) => return Vec::new(),
78        };
79        checks
80            .iter()
81            .map(|(name, f)| {
82                let s = f();
83                HealthCheckEntry {
84                    name: name.clone(),
85                    status: s.status.to_string(),
86                    detail: s.detail.clone(),
87                }
88            })
89            .collect()
90    }
91
92    /// 计算整体状态:fail 优先于 warn 优先于 pass;空列表视为 pass。
93    /// Mutex poisoned 时保守返回 "fail"。
94    pub fn overall_status(&self) -> &'static str {
95        let checks = match self.checks.lock() {
96            Ok(c) => c.clone(),
97            Err(_) => return "fail",
98        };
99        if checks.is_empty() {
100            "pass"
101        } else {
102            let any_fail = checks.iter().any(|(_, f)| f().status == "fail");
103            let any_warn = checks.iter().any(|(_, f)| f().status == "warn");
104            if any_fail {
105                "fail"
106            } else if any_warn {
107                "warn"
108            } else {
109                "pass"
110            }
111        }
112    }
113}
114
115/// Build an RFC 8407 `application/health+json` body from the registry.
116///
117/// When the registry has no checks, returns `{"status":"pass"}`.
118/// Otherwise returns `{"status":<overall>,"checks":[...]}`.
119pub fn build_health_response(registry: &HealthCheckRegistry) -> Vec<u8> {
120    let entries = registry.snapshot();
121    let overall = registry.overall_status();
122    let body = if entries.is_empty() {
123        serde_json::json!({ "status": overall })
124    } else {
125        serde_json::json!({
126            "status": overall,
127            "checks": entries
128        })
129    };
130    serde_json::to_vec(&body).unwrap_or_default()
131}
132
133/// HTTP status for a health overall status: `fail` → 503, otherwise 200.
134pub fn health_http_status(overall: &str) -> u16 {
135    if overall == "fail" {
136        503
137    } else {
138        200
139    }
140}
141
142/// Dynamic health endpoint — evaluates registered probes on each request.
143pub struct HealthCheckEndpoint {
144    registry: Arc<HealthCheckRegistry>,
145}
146
147impl HealthCheckEndpoint {
148    pub fn new(registry: Arc<HealthCheckRegistry>) -> Self {
149        Self { registry }
150    }
151}
152
153#[async_trait::async_trait]
154impl IEndpoint for HealthCheckEndpoint {
155    async fn handle(&self, ctx: &mut dyn IHttpContext) -> Result<()> {
156        let overall = self.registry.overall_status();
157        let body = build_health_response(&self.registry);
158        ctx.response_mut()
159            .set_status(health_http_status(overall));
160        ctx.response_mut()
161            .set_header("content-type", "application/health+json");
162        ctx.response_mut().write_bytes(body).await
163    }
164}