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 { .. } => {
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 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
297pub async fn health_response(health: HealthCheck) -> HealthResponse {
299 HealthResponse::from_result(health.execute().await)
300}
301
302pub async fn readiness_response(health: HealthCheck) -> HealthResponse {
307 HealthResponse::from_result(health.execute().await)
308}
309
310pub async fn liveness_response() -> HealthResponse {
312 let result = HealthCheckBuilder::default().build().execute().await;
313 HealthResponse::from_result(result)
314}
315
316pub struct HealthCheckBuilder {
318 checks: HashMap<String, HealthCheckFn>,
319 version: Option<String>,
320}
321
322impl HealthCheckBuilder {
323 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 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 pub fn version(mut self, version: impl Into<String>) -> Self {
370 self.version = Some(version.into());
371 self
372 }
373
374 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}