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 { .. } => {
268                    if overall_status.is_healthy() {
269                        overall_status = HealthStatus::degraded("one or more checks degraded");
270                    }
271                }
272                _ => {}
273            }
274
275            results.insert(name.clone(), status);
276        }
277
278        // Use UTC timestamp formatted as ISO 8601
279        let timestamp = std::time::SystemTime::now()
280            .duration_since(std::time::UNIX_EPOCH)
281            .map(|d| {
282                let secs = d.as_secs();
283                let nanos = d.subsec_nanos();
284                format!("{}.{:09}Z", secs, nanos)
285            })
286            .unwrap_or_else(|_| "unknown".to_string());
287
288        HealthCheckResult {
289            status: overall_status,
290            checks: results,
291            version: self.version.clone(),
292            timestamp,
293        }
294    }
295}
296
297/// Execute an aggregated health check and return an HTTP-friendly response.
298pub async fn health_response(health: HealthCheck) -> HealthResponse {
299    HealthResponse::from_result(health.execute().await)
300}
301
302/// Execute a readiness probe based on the configured health checks.
303///
304/// Readiness currently shares the same dependency checks as the aggregated
305/// health endpoint; unhealthy dependencies return `503 Service Unavailable`.
306pub async fn readiness_response(health: HealthCheck) -> HealthResponse {
307    HealthResponse::from_result(health.execute().await)
308}
309
310/// Return a lightweight liveness probe response.
311pub async fn liveness_response() -> HealthResponse {
312    let result = HealthCheckBuilder::default().build().execute().await;
313    HealthResponse::from_result(result)
314}
315
316/// Builder for health check configuration
317pub struct HealthCheckBuilder {
318    checks: HashMap<String, HealthCheckFn>,
319    version: Option<String>,
320}
321
322impl HealthCheckBuilder {
323    /// Create a new health check builder
324    ///
325    /// # Arguments
326    ///
327    /// * `include_default` - Whether to include a default "self" check that always returns healthy
328    pub fn new(include_default: bool) -> Self {
329        let mut checks = HashMap::new();
330
331        if include_default {
332            let check: HealthCheckFn = Arc::new(|| Box::pin(async { HealthStatus::healthy() }));
333            checks.insert("self".to_string(), check);
334        }
335
336        Self {
337            checks,
338            version: None,
339        }
340    }
341
342    /// Add a health check
343    ///
344    /// # Example
345    ///
346    /// ```rust
347    /// use rustapi_core::health::{HealthCheckBuilder, HealthStatus};
348    ///
349    /// let health = HealthCheckBuilder::new(false)
350    ///     .add_check("database", || async {
351    ///         // Simulate database check
352    ///         HealthStatus::healthy()
353    ///     })
354    ///     .build();
355    /// ```
356    pub fn add_check<F, Fut>(mut self, name: impl Into<String>, check: F) -> Self
357    where
358        F: Fn() -> Fut + Send + Sync + 'static,
359        Fut: Future<Output = HealthStatus> + Send + 'static,
360    {
361        let check_fn = Arc::new(move || {
362            Box::pin(check()) as Pin<Box<dyn Future<Output = HealthStatus> + Send>>
363        });
364        self.checks.insert(name.into(), check_fn);
365        self
366    }
367
368    /// Set the application version
369    pub fn version(mut self, version: impl Into<String>) -> Self {
370        self.version = Some(version.into());
371        self
372    }
373
374    /// Build the health check
375    pub fn build(self) -> HealthCheck {
376        HealthCheck {
377            checks: self.checks,
378            version: self.version,
379        }
380    }
381}
382
383impl Default for HealthCheckBuilder {
384    fn default() -> Self {
385        Self::new(true)
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392
393    #[tokio::test]
394    async fn health_check_all_healthy() {
395        let health = HealthCheckBuilder::new(false)
396            .add_check("db", || async { HealthStatus::healthy() })
397            .add_check("cache", || async { HealthStatus::healthy() })
398            .version("1.0.0")
399            .build();
400
401        let result = health.execute().await;
402
403        assert!(result.status.is_healthy());
404        assert_eq!(result.checks.len(), 2);
405        assert_eq!(result.version, Some("1.0.0".to_string()));
406    }
407
408    #[tokio::test]
409    async fn health_check_one_unhealthy() {
410        let health = HealthCheckBuilder::new(false)
411            .add_check("db", || async { HealthStatus::healthy() })
412            .add_check("cache", || async {
413                HealthStatus::unhealthy("connection failed")
414            })
415            .build();
416
417        let result = health.execute().await;
418
419        assert!(result.status.is_unhealthy());
420        assert_eq!(result.checks.len(), 2);
421    }
422
423    #[tokio::test]
424    async fn health_check_one_degraded() {
425        let health = HealthCheckBuilder::new(false)
426            .add_check("db", || async { HealthStatus::healthy() })
427            .add_check("cache", || async { HealthStatus::degraded("high latency") })
428            .build();
429
430        let result = health.execute().await;
431
432        assert!(result.status.is_degraded());
433        assert_eq!(result.checks.len(), 2);
434    }
435
436    #[tokio::test]
437    async fn health_check_with_default() {
438        let health = HealthCheckBuilder::new(true).build();
439
440        let result = health.execute().await;
441
442        assert!(result.status.is_healthy());
443        assert_eq!(result.checks.len(), 1);
444        assert!(result.checks.contains_key("self"));
445    }
446
447    #[tokio::test]
448    async fn readiness_response_returns_service_unavailable_for_unhealthy_checks() {
449        let health = HealthCheckBuilder::new(false)
450            .add_check("db", || async { HealthStatus::unhealthy("db offline") })
451            .build();
452
453        let response = readiness_response(health).await.into_response();
454
455        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
456    }
457}