1use 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
40#[serde(rename_all = "lowercase")]
41pub enum HealthStatus {
42 #[serde(rename = "healthy")]
44 Healthy,
45 #[serde(rename = "unhealthy")]
47 Unhealthy { reason: String },
48 #[serde(rename = "degraded")]
50 Degraded { reason: String },
51}
52
53impl HealthStatus {
54 pub fn healthy() -> Self {
56 Self::Healthy
57 }
58
59 pub fn unhealthy(reason: impl Into<String>) -> Self {
61 Self::Unhealthy {
62 reason: reason.into(),
63 }
64 }
65
66 pub fn degraded(reason: impl Into<String>) -> Self {
68 Self::Degraded {
69 reason: reason.into(),
70 }
71 }
72
73 pub fn is_healthy(&self) -> bool {
75 matches!(self, Self::Healthy)
76 }
77
78 pub fn is_unhealthy(&self) -> bool {
80 matches!(self, Self::Unhealthy { .. })
81 }
82
83 pub fn is_degraded(&self) -> bool {
85 matches!(self, Self::Degraded { .. })
86 }
87}
88
89#[derive(Debug, Serialize, Deserialize)]
91pub struct HealthCheckResult {
92 pub status: HealthStatus,
94 pub checks: HashMap<String, HealthStatus>,
96 #[serde(skip_serializing_if = "Option::is_none")]
98 pub version: Option<String>,
99 pub timestamp: String,
101}
102
103#[derive(Debug, Clone, PartialEq, Eq)]
110pub struct HealthEndpointConfig {
111 pub health_path: String,
113 pub readiness_path: String,
115 pub liveness_path: String,
117}
118
119impl HealthEndpointConfig {
120 pub fn new() -> Self {
122 Self::default()
123 }
124
125 pub fn health_path(mut self, path: impl Into<String>) -> Self {
127 self.health_path = path.into();
128 self
129 }
130
131 pub fn readiness_path(mut self, path: impl Into<String>) -> Self {
133 self.readiness_path = path.into();
134 self
135 }
136
137 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#[derive(Debug, Clone)]
156pub struct HealthResponse {
157 status: StatusCode,
158 body: serde_json::Value,
159}
160
161impl HealthResponse {
162 pub fn new(status: StatusCode, body: serde_json::Value) -> Self {
164 Self { status, body }
165 }
166
167 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
242pub type HealthCheckFn =
244 Arc<dyn Fn() -> Pin<Box<dyn Future<Output = HealthStatus> + Send>> + Send + Sync>;
245
246#[derive(Clone)]
248pub struct HealthCheck {
249 checks: HashMap<String, HealthCheckFn>,
250 version: Option<String>,
251}
252
253impl HealthCheck {
254 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 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 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
296pub async fn health_response(health: HealthCheck) -> HealthResponse {
298 HealthResponse::from_result(health.execute().await)
299}
300
301pub async fn readiness_response(health: HealthCheck) -> HealthResponse {
306 HealthResponse::from_result(health.execute().await)
307}
308
309pub async fn liveness_response() -> HealthResponse {
311 let result = HealthCheckBuilder::default().build().execute().await;
312 HealthResponse::from_result(result)
313}
314
315pub struct HealthCheckBuilder {
317 checks: HashMap<String, HealthCheckFn>,
318 version: Option<String>,
319}
320
321impl HealthCheckBuilder {
322 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 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 pub fn version(mut self, version: impl Into<String>) -> Self {
369 self.version = Some(version.into());
370 self
371 }
372
373 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}