Skip to main content

rustapi_core/
health.rs

1//! Health check system for monitoring application health
2//!
3//! This module provides a flexible health check system for monitoring
4//! the health and readiness of your application and its dependencies.
5//!
6//! # Example
7//!
8//! ```rust,no_run
9//! use rustapi_core::health::{HealthCheck, HealthCheckBuilder, HealthStatus};
10//!
11//! #[tokio::main]
12//! async fn main() {
13//!     let health = HealthCheckBuilder::new(true)
14//!         .add_check("database", || async {
15//!             // Check database connection
16//!             HealthStatus::healthy()
17//!         })
18//!         .add_check("redis", || async {
19//!             // Check Redis connection
20//!             HealthStatus::healthy()
21//!         })
22//!         .build();
23//!
24//!     // Use health.execute().await to get results
25//! }
26//! ```
27
28use crate::response::{Body, IntoResponse, Response};
29use http::{header, StatusCode};
30use rustapi_openapi::{MediaType, Operation, ResponseModifier, ResponseSpec, SchemaRef};
31use serde::{Deserialize, Serialize};
32use serde_json::json;
33use std::collections::HashMap;
34use std::future::Future;
35use std::pin::Pin;
36use std::sync::Arc;
37
38/// Health status of a component
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
40#[serde(rename_all = "lowercase")]
41pub enum HealthStatus {
42    /// Component is healthy
43    #[serde(rename = "healthy")]
44    Healthy,
45    /// Component is unhealthy
46    #[serde(rename = "unhealthy")]
47    Unhealthy { reason: String },
48    /// Component is degraded but functional
49    #[serde(rename = "degraded")]
50    Degraded { reason: String },
51}
52
53impl HealthStatus {
54    /// Create a healthy status
55    pub fn healthy() -> Self {
56        Self::Healthy
57    }
58
59    /// Create an unhealthy status with a reason
60    pub fn unhealthy(reason: impl Into<String>) -> Self {
61        Self::Unhealthy {
62            reason: reason.into(),
63        }
64    }
65
66    /// Create a degraded status with a reason
67    pub fn degraded(reason: impl Into<String>) -> Self {
68        Self::Degraded {
69            reason: reason.into(),
70        }
71    }
72
73    /// Check if the status is healthy
74    pub fn is_healthy(&self) -> bool {
75        matches!(self, Self::Healthy)
76    }
77
78    /// Check if the status is unhealthy
79    pub fn is_unhealthy(&self) -> bool {
80        matches!(self, Self::Unhealthy { .. })
81    }
82
83    /// Check if the status is degraded
84    pub fn is_degraded(&self) -> bool {
85        matches!(self, Self::Degraded { .. })
86    }
87}
88
89/// Overall health check result
90#[derive(Debug, Serialize, Deserialize)]
91pub struct HealthCheckResult {
92    /// Overall status
93    pub status: HealthStatus,
94    /// Individual component checks
95    pub checks: HashMap<String, HealthStatus>,
96    /// Application version (if provided)
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub version: Option<String>,
99    /// Timestamp of check (ISO 8601)
100    pub timestamp: String,
101}
102
103/// Configuration for built-in health endpoints.
104///
105/// By default RustAPI exposes three endpoints when enabled:
106/// - `/health` - aggregated dependency health
107/// - `/ready` - readiness probe for orchestrators/load balancers
108/// - `/live` - lightweight liveness probe
109#[derive(Debug, Clone, PartialEq, Eq)]
110pub struct HealthEndpointConfig {
111    /// Path for the aggregated health endpoint.
112    pub health_path: String,
113    /// Path for the readiness endpoint.
114    pub readiness_path: String,
115    /// Path for the liveness endpoint.
116    pub liveness_path: String,
117}
118
119impl HealthEndpointConfig {
120    /// Create a new configuration with default paths.
121    pub fn new() -> Self {
122        Self::default()
123    }
124
125    /// Override the health endpoint path.
126    pub fn health_path(mut self, path: impl Into<String>) -> Self {
127        self.health_path = path.into();
128        self
129    }
130
131    /// Override the readiness endpoint path.
132    pub fn readiness_path(mut self, path: impl Into<String>) -> Self {
133        self.readiness_path = path.into();
134        self
135    }
136
137    /// Override the liveness endpoint path.
138    pub fn liveness_path(mut self, path: impl Into<String>) -> Self {
139        self.liveness_path = path.into();
140        self
141    }
142}
143
144impl Default for HealthEndpointConfig {
145    fn default() -> Self {
146        Self {
147            health_path: "/health".to_string(),
148            readiness_path: "/ready".to_string(),
149            liveness_path: "/live".to_string(),
150        }
151    }
152}
153
154/// JSON health response used by RustAPI's built-in health endpoints.
155#[derive(Debug, Clone)]
156pub struct HealthResponse {
157    status: StatusCode,
158    body: serde_json::Value,
159}
160
161impl HealthResponse {
162    /// Create a new health response from an HTTP status and JSON body.
163    pub fn new(status: StatusCode, body: serde_json::Value) -> Self {
164        Self { status, body }
165    }
166
167    /// Create a health response from a health check result.
168    pub fn from_result(result: HealthCheckResult) -> Self {
169        let status = if result.status.is_unhealthy() {
170            StatusCode::SERVICE_UNAVAILABLE
171        } else {
172            StatusCode::OK
173        };
174
175        let body = serde_json::to_value(result).unwrap_or_else(|_| {
176            json!({
177                "status": { "unhealthy": { "reason": "failed to serialize health result" } }
178            })
179        });
180
181        Self { status, body }
182    }
183}
184
185impl IntoResponse for HealthResponse {
186    fn into_response(self) -> Response {
187        match serde_json::to_vec(&self.body) {
188            Ok(body) => http::Response::builder()
189                .status(self.status)
190                .header(header::CONTENT_TYPE, "application/json")
191                .body(Body::from(body))
192                .unwrap(),
193            Err(err) => crate::error::ApiError::internal(format!(
194                "Failed to serialize health response: {}",
195                err
196            ))
197            .into_response(),
198        }
199    }
200}
201
202impl ResponseModifier for HealthResponse {
203    fn update_response(op: &mut Operation) {
204        let mut content = std::collections::BTreeMap::new();
205        content.insert(
206            "application/json".to_string(),
207            MediaType {
208                schema: Some(SchemaRef::Inline(json!({
209                    "type": "object",
210                    "additionalProperties": true
211                }))),
212                example: Some(json!({
213                    "status": "healthy",
214                    "checks": {
215                        "self": "healthy"
216                    },
217                    "timestamp": "1741411200.000000000Z"
218                })),
219            },
220        );
221
222        op.responses.insert(
223            "200".to_string(),
224            ResponseSpec {
225                description: "Service is healthy or ready".to_string(),
226                content: content.clone(),
227                headers: Default::default(),
228            },
229        );
230
231        op.responses.insert(
232            "503".to_string(),
233            ResponseSpec {
234                description: "Service or one of its dependencies is unhealthy".to_string(),
235                content,
236                headers: Default::default(),
237            },
238        );
239    }
240}
241
242/// Type alias for async health check functions
243pub type HealthCheckFn =
244    Arc<dyn Fn() -> Pin<Box<dyn Future<Output = HealthStatus> + Send>> + Send + Sync>;
245
246/// Health check configuration
247#[derive(Clone)]
248pub struct HealthCheck {
249    checks: HashMap<String, HealthCheckFn>,
250    version: Option<String>,
251}
252
253impl HealthCheck {
254    /// Execute all health checks
255    pub async fn execute(&self) -> HealthCheckResult {
256        let mut results = HashMap::new();
257        let mut overall_status = HealthStatus::Healthy;
258
259        for (name, check) in &self.checks {
260            let status = check().await;
261
262            // Determine overall status
263            match &status {
264                HealthStatus::Unhealthy { .. } => {
265                    overall_status = HealthStatus::unhealthy("one or more checks failed");
266                }
267                HealthStatus::Degraded { .. } if overall_status.is_healthy() => {
268                    overall_status = HealthStatus::degraded("one or more checks degraded");
269                }
270                HealthStatus::Degraded { .. } => {}
271                _ => {}
272            }
273
274            results.insert(name.clone(), status);
275        }
276
277        // Use UTC timestamp formatted as ISO 8601
278        let timestamp = std::time::SystemTime::now()
279            .duration_since(std::time::UNIX_EPOCH)
280            .map(|d| {
281                let secs = d.as_secs();
282                let nanos = d.subsec_nanos();
283                format!("{}.{:09}Z", secs, nanos)
284            })
285            .unwrap_or_else(|_| "unknown".to_string());
286
287        HealthCheckResult {
288            status: overall_status,
289            checks: results,
290            version: self.version.clone(),
291            timestamp,
292        }
293    }
294}
295
296/// Execute an aggregated health check and return an HTTP-friendly response.
297pub async fn health_response(health: HealthCheck) -> HealthResponse {
298    HealthResponse::from_result(health.execute().await)
299}
300
301/// Execute a readiness probe based on the configured health checks.
302///
303/// Readiness currently shares the same dependency checks as the aggregated
304/// health endpoint; unhealthy dependencies return `503 Service Unavailable`.
305pub async fn readiness_response(health: HealthCheck) -> HealthResponse {
306    HealthResponse::from_result(health.execute().await)
307}
308
309/// Return a lightweight liveness probe response.
310pub async fn liveness_response() -> HealthResponse {
311    let result = HealthCheckBuilder::default().build().execute().await;
312    HealthResponse::from_result(result)
313}
314
315/// Builder for health check configuration
316pub struct HealthCheckBuilder {
317    checks: HashMap<String, HealthCheckFn>,
318    version: Option<String>,
319}
320
321impl HealthCheckBuilder {
322    /// Create a new health check builder
323    ///
324    /// # Arguments
325    ///
326    /// * `include_default` - Whether to include a default "self" check that always returns healthy
327    pub fn new(include_default: bool) -> Self {
328        let mut checks = HashMap::new();
329
330        if include_default {
331            let check: HealthCheckFn = Arc::new(|| Box::pin(async { HealthStatus::healthy() }));
332            checks.insert("self".to_string(), check);
333        }
334
335        Self {
336            checks,
337            version: None,
338        }
339    }
340
341    /// Add a health check
342    ///
343    /// # Example
344    ///
345    /// ```rust
346    /// use rustapi_core::health::{HealthCheckBuilder, HealthStatus};
347    ///
348    /// let health = HealthCheckBuilder::new(false)
349    ///     .add_check("database", || async {
350    ///         // Simulate database check
351    ///         HealthStatus::healthy()
352    ///     })
353    ///     .build();
354    /// ```
355    pub fn add_check<F, Fut>(mut self, name: impl Into<String>, check: F) -> Self
356    where
357        F: Fn() -> Fut + Send + Sync + 'static,
358        Fut: Future<Output = HealthStatus> + Send + 'static,
359    {
360        let check_fn = Arc::new(move || {
361            Box::pin(check()) as Pin<Box<dyn Future<Output = HealthStatus> + Send>>
362        });
363        self.checks.insert(name.into(), check_fn);
364        self
365    }
366
367    /// Set the application version
368    pub fn version(mut self, version: impl Into<String>) -> Self {
369        self.version = Some(version.into());
370        self
371    }
372
373    /// Build the health check
374    pub fn build(self) -> HealthCheck {
375        HealthCheck {
376            checks: self.checks,
377            version: self.version,
378        }
379    }
380}
381
382impl Default for HealthCheckBuilder {
383    fn default() -> Self {
384        Self::new(true)
385    }
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391
392    #[tokio::test]
393    async fn health_check_all_healthy() {
394        let health = HealthCheckBuilder::new(false)
395            .add_check("db", || async { HealthStatus::healthy() })
396            .add_check("cache", || async { HealthStatus::healthy() })
397            .version("1.0.0")
398            .build();
399
400        let result = health.execute().await;
401
402        assert!(result.status.is_healthy());
403        assert_eq!(result.checks.len(), 2);
404        assert_eq!(result.version, Some("1.0.0".to_string()));
405    }
406
407    #[tokio::test]
408    async fn health_check_one_unhealthy() {
409        let health = HealthCheckBuilder::new(false)
410            .add_check("db", || async { HealthStatus::healthy() })
411            .add_check("cache", || async {
412                HealthStatus::unhealthy("connection failed")
413            })
414            .build();
415
416        let result = health.execute().await;
417
418        assert!(result.status.is_unhealthy());
419        assert_eq!(result.checks.len(), 2);
420    }
421
422    #[tokio::test]
423    async fn health_check_one_degraded() {
424        let health = HealthCheckBuilder::new(false)
425            .add_check("db", || async { HealthStatus::healthy() })
426            .add_check("cache", || async { HealthStatus::degraded("high latency") })
427            .build();
428
429        let result = health.execute().await;
430
431        assert!(result.status.is_degraded());
432        assert_eq!(result.checks.len(), 2);
433    }
434
435    #[tokio::test]
436    async fn health_check_with_default() {
437        let health = HealthCheckBuilder::new(true).build();
438
439        let result = health.execute().await;
440
441        assert!(result.status.is_healthy());
442        assert_eq!(result.checks.len(), 1);
443        assert!(result.checks.contains_key("self"));
444    }
445
446    #[tokio::test]
447    async fn readiness_response_returns_service_unavailable_for_unhealthy_checks() {
448        let health = HealthCheckBuilder::new(false)
449            .add_check("db", || async { HealthStatus::unhealthy("db offline") })
450            .build();
451
452        let response = readiness_response(health).await.into_response();
453
454        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
455    }
456}