use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Instant;
use crate::context::RequestContext;
use crate::request::Request;
use crate::response::{Response, ResponseBody, StatusCode};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum HealthStatus {
Healthy,
Degraded,
Unhealthy,
}
impl HealthStatus {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::Healthy => "healthy",
Self::Degraded => "degraded",
Self::Unhealthy => "unhealthy",
}
}
#[must_use]
pub fn status_code(self) -> StatusCode {
match self {
Self::Healthy | Self::Degraded => StatusCode::OK,
Self::Unhealthy => StatusCode::SERVICE_UNAVAILABLE,
}
}
}
impl std::fmt::Display for HealthStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone)]
pub struct HealthCheckResult {
pub name: String,
pub passed: bool,
pub critical: bool,
pub latency_ms: u64,
pub error: Option<String>,
}
#[derive(Debug, Clone)]
pub struct HealthReport {
pub status: HealthStatus,
pub checks: Vec<HealthCheckResult>,
pub total_latency_ms: u64,
}
impl HealthReport {
#[must_use]
pub fn to_json(&self) -> Vec<u8> {
let mut json = String::with_capacity(256);
json.push_str("{\"status\":\"");
json.push_str(self.status.as_str());
json.push_str("\",\"total_latency_ms\":");
json.push_str(&self.total_latency_ms.to_string());
json.push_str(",\"checks\":[");
for (i, check) in self.checks.iter().enumerate() {
if i > 0 {
json.push(',');
}
json.push_str("{\"name\":\"");
json.push_str(&check.name);
json.push_str("\",\"passed\":");
json.push_str(if check.passed { "true" } else { "false" });
json.push_str(",\"critical\":");
json.push_str(if check.critical { "true" } else { "false" });
json.push_str(",\"latency_ms\":");
json.push_str(&check.latency_ms.to_string());
if let Some(ref err) = check.error {
json.push_str(",\"error\":\"");
for ch in err.chars() {
match ch {
'"' => json.push_str("\\\""),
'\\' => json.push_str("\\\\"),
'\n' => json.push_str("\\n"),
'\r' => json.push_str("\\r"),
'\t' => json.push_str("\\t"),
c => json.push(c),
}
}
json.push('"');
}
json.push('}');
}
json.push_str("]}");
json.into_bytes()
}
}
type CheckFn =
Arc<dyn Fn() -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>;
struct HealthCheckEntry {
name: String,
critical: bool,
check_fn: CheckFn,
}
pub struct HealthCheckRegistry {
checks: Vec<HealthCheckEntry>,
}
impl HealthCheckRegistry {
#[must_use]
pub fn new() -> Self {
Self { checks: Vec::new() }
}
pub fn add<F, Fut>(&mut self, name: impl Into<String>, critical: bool, check_fn: F)
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<(), String>> + Send + 'static,
{
let check_fn = Arc::new(move || {
let fut = check_fn();
Box::pin(fut) as Pin<Box<dyn Future<Output = Result<(), String>> + Send>>
}) as CheckFn;
self.checks.push(HealthCheckEntry {
name: name.into(),
critical,
check_fn,
});
}
pub async fn check_all(&self) -> HealthReport {
let total_start = Instant::now();
let mut results = Vec::with_capacity(self.checks.len());
let mut has_critical_failure = false;
let mut has_non_critical_failure = false;
for entry in &self.checks {
let start = Instant::now();
let outcome = (entry.check_fn)().await;
let latency_ms = start.elapsed().as_millis() as u64;
let (passed, error) = match outcome {
Ok(()) => (true, None),
Err(msg) => {
if entry.critical {
has_critical_failure = true;
} else {
has_non_critical_failure = true;
}
(false, Some(msg))
}
};
results.push(HealthCheckResult {
name: entry.name.clone(),
passed,
critical: entry.critical,
latency_ms,
error,
});
}
let status = if has_critical_failure {
HealthStatus::Unhealthy
} else if has_non_critical_failure {
HealthStatus::Degraded
} else {
HealthStatus::Healthy
};
HealthReport {
status,
checks: results,
total_latency_ms: total_start.elapsed().as_millis() as u64,
}
}
#[must_use]
pub fn len(&self) -> usize {
self.checks.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.checks.is_empty()
}
}
impl Default for HealthCheckRegistry {
fn default() -> Self {
Self::new()
}
}
pub fn basic_health_handler()
-> impl Fn(&RequestContext, &mut Request) -> std::future::Ready<Response> + Send + Sync + 'static {
|_ctx: &RequestContext, _req: &mut Request| {
let body = br#"{"status":"healthy"}"#.to_vec();
std::future::ready(
Response::with_status(StatusCode::OK)
.header("Content-Type", b"application/json".to_vec())
.body(ResponseBody::Bytes(body)),
)
}
}
pub fn detailed_health_handler(
registry: Arc<HealthCheckRegistry>,
) -> impl Fn(&RequestContext, &mut Request) -> Pin<Box<dyn Future<Output = Response> + Send>>
+ Send
+ Sync
+ 'static {
move |_ctx: &RequestContext, _req: &mut Request| {
let registry = Arc::clone(®istry);
Box::pin(async move {
let report = registry.check_all().await;
let status_code = report.status.status_code();
let body = report.to_json();
Response::with_status(status_code)
.header("Content-Type", b"application/json".to_vec())
.header("Cache-Control", b"no-cache, no-store".to_vec())
.body(ResponseBody::Bytes(body))
})
}
}
pub fn liveness_handler()
-> impl Fn(&RequestContext, &mut Request) -> std::future::Ready<Response> + Send + Sync + 'static {
basic_health_handler()
}
pub fn readiness_handler(
registry: Arc<HealthCheckRegistry>,
) -> impl Fn(&RequestContext, &mut Request) -> Pin<Box<dyn Future<Output = Response> + Send>>
+ Send
+ Sync
+ 'static {
detailed_health_handler(registry)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::request::Method;
fn run_handler<F, Fut>(handler: F) -> Response
where
F: Fn(&RequestContext, &mut Request) -> Fut,
Fut: Future<Output = Response>,
{
let ctx = RequestContext::new(asupersync::Cx::for_testing(), 1);
let mut req = Request::new(Method::Get, "/health");
futures_executor::block_on(handler(&ctx, &mut req))
}
#[test]
fn health_status_as_str() {
assert_eq!(HealthStatus::Healthy.as_str(), "healthy");
assert_eq!(HealthStatus::Degraded.as_str(), "degraded");
assert_eq!(HealthStatus::Unhealthy.as_str(), "unhealthy");
}
#[test]
fn health_status_display() {
assert_eq!(format!("{}", HealthStatus::Healthy), "healthy");
assert_eq!(format!("{}", HealthStatus::Unhealthy), "unhealthy");
}
#[test]
fn health_status_code() {
assert_eq!(HealthStatus::Healthy.status_code(), StatusCode::OK);
assert_eq!(HealthStatus::Degraded.status_code(), StatusCode::OK);
assert_eq!(
HealthStatus::Unhealthy.status_code(),
StatusCode::SERVICE_UNAVAILABLE
);
}
#[test]
fn empty_registry_is_healthy() {
let registry = HealthCheckRegistry::new();
assert!(registry.is_empty());
let report = futures_executor::block_on(registry.check_all());
assert_eq!(report.status, HealthStatus::Healthy);
assert!(report.checks.is_empty());
}
#[test]
fn all_checks_pass_is_healthy() {
let mut registry = HealthCheckRegistry::new();
registry.add("database", true, || async { Ok(()) });
registry.add("cache", false, || async { Ok(()) });
let report = futures_executor::block_on(registry.check_all());
assert_eq!(report.status, HealthStatus::Healthy);
assert_eq!(report.checks.len(), 2);
assert!(report.checks[0].passed);
assert!(report.checks[1].passed);
}
#[test]
fn non_critical_failure_is_degraded() {
let mut registry = HealthCheckRegistry::new();
registry.add("database", true, || async { Ok(()) });
registry.add("cache", false, || async {
Err("Cache connection refused".to_string())
});
let report = futures_executor::block_on(registry.check_all());
assert_eq!(report.status, HealthStatus::Degraded);
assert!(report.checks[0].passed);
assert!(!report.checks[1].passed);
assert_eq!(
report.checks[1].error.as_deref(),
Some("Cache connection refused")
);
}
#[test]
fn critical_failure_is_unhealthy() {
let mut registry = HealthCheckRegistry::new();
registry.add("database", true, || async {
Err("Connection timeout".to_string())
});
registry.add("cache", false, || async { Ok(()) });
let report = futures_executor::block_on(registry.check_all());
assert_eq!(report.status, HealthStatus::Unhealthy);
assert!(!report.checks[0].passed);
assert!(report.checks[0].critical);
assert!(report.checks[1].passed);
}
#[test]
fn latency_is_measured() {
let mut registry = HealthCheckRegistry::new();
registry.add("fast", true, || async { Ok(()) });
let report = futures_executor::block_on(registry.check_all());
assert!(report.checks[0].latency_ms < 1000);
assert!(report.total_latency_ms < 1000);
}
#[test]
fn report_to_json_healthy() {
let report = HealthReport {
status: HealthStatus::Healthy,
checks: vec![HealthCheckResult {
name: "database".to_string(),
passed: true,
critical: true,
latency_ms: 5,
error: None,
}],
total_latency_ms: 5,
};
let json = String::from_utf8(report.to_json()).unwrap();
assert!(json.contains("\"status\":\"healthy\""));
assert!(json.contains("\"name\":\"database\""));
assert!(json.contains("\"passed\":true"));
assert!(json.contains("\"critical\":true"));
assert!(json.contains("\"latency_ms\":5"));
}
#[test]
fn report_to_json_unhealthy_with_error() {
let report = HealthReport {
status: HealthStatus::Unhealthy,
checks: vec![HealthCheckResult {
name: "database".to_string(),
passed: false,
critical: true,
latency_ms: 3000,
error: Some("Connection refused".to_string()),
}],
total_latency_ms: 3000,
};
let json = String::from_utf8(report.to_json()).unwrap();
assert!(json.contains("\"status\":\"unhealthy\""));
assert!(json.contains("\"passed\":false"));
assert!(json.contains("\"error\":\"Connection refused\""));
}
#[test]
fn report_to_json_escapes_special_chars() {
let report = HealthReport {
status: HealthStatus::Unhealthy,
checks: vec![HealthCheckResult {
name: "test".to_string(),
passed: false,
critical: true,
latency_ms: 0,
error: Some("Error with \"quotes\" and \\backslash".to_string()),
}],
total_latency_ms: 0,
};
let json = String::from_utf8(report.to_json()).unwrap();
assert!(json.contains(r#"\"quotes\""#));
assert!(json.contains(r"\\backslash"));
}
#[test]
fn basic_health_handler_returns_200() {
let handler = basic_health_handler();
let resp = run_handler(handler);
assert_eq!(resp.status(), StatusCode::OK);
if let ResponseBody::Bytes(body) = resp.body_ref() {
let body_str = std::str::from_utf8(body).unwrap();
assert!(body_str.contains("\"status\":\"healthy\""));
} else {
panic!("expected Bytes body");
}
}
#[test]
fn detailed_health_handler_healthy() {
let mut registry = HealthCheckRegistry::new();
registry.add("db", true, || async { Ok(()) });
let handler = detailed_health_handler(Arc::new(registry));
let ctx = RequestContext::new(asupersync::Cx::for_testing(), 1);
let mut req = Request::new(Method::Get, "/health");
let resp = futures_executor::block_on(handler(&ctx, &mut req));
assert_eq!(resp.status(), StatusCode::OK);
if let ResponseBody::Bytes(body) = resp.body_ref() {
let body_str = std::str::from_utf8(body).unwrap();
assert!(body_str.contains("\"status\":\"healthy\""));
assert!(body_str.contains("\"name\":\"db\""));
} else {
panic!("expected Bytes body");
}
}
#[test]
fn detailed_health_handler_unhealthy_returns_503() {
let mut registry = HealthCheckRegistry::new();
registry.add("db", true, || async { Err("down".to_string()) });
let handler = detailed_health_handler(Arc::new(registry));
let ctx = RequestContext::new(asupersync::Cx::for_testing(), 1);
let mut req = Request::new(Method::Get, "/health");
let resp = futures_executor::block_on(handler(&ctx, &mut req));
assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
if let ResponseBody::Bytes(body) = resp.body_ref() {
let body_str = std::str::from_utf8(body).unwrap();
assert!(body_str.contains("\"status\":\"unhealthy\""));
} else {
panic!("expected Bytes body");
}
}
#[test]
fn detailed_health_handler_has_cache_headers() {
let registry = HealthCheckRegistry::new();
let handler = detailed_health_handler(Arc::new(registry));
let ctx = RequestContext::new(asupersync::Cx::for_testing(), 1);
let mut req = Request::new(Method::Get, "/health");
let resp = futures_executor::block_on(handler(&ctx, &mut req));
let has_no_cache = resp
.headers()
.iter()
.any(|(n, v)| n.eq_ignore_ascii_case("cache-control") && v == b"no-cache, no-store");
assert!(
has_no_cache,
"should have Cache-Control: no-cache, no-store"
);
}
#[test]
fn registry_len_and_is_empty() {
let mut registry = HealthCheckRegistry::new();
assert!(registry.is_empty());
assert_eq!(registry.len(), 0);
registry.add("db", true, || async { Ok(()) });
assert!(!registry.is_empty());
assert_eq!(registry.len(), 1);
}
#[test]
fn liveness_handler_returns_200() {
let handler = liveness_handler();
let resp = run_handler(handler);
assert_eq!(resp.status(), StatusCode::OK);
}
#[test]
fn readiness_handler_healthy_returns_200() {
let mut registry = HealthCheckRegistry::new();
registry.add("db", true, || async { Ok(()) });
let handler = readiness_handler(Arc::new(registry));
let ctx = RequestContext::new(asupersync::Cx::for_testing(), 1);
let mut req = Request::new(Method::Get, "/ready");
let resp = futures_executor::block_on(handler(&ctx, &mut req));
assert_eq!(resp.status(), StatusCode::OK);
}
#[test]
fn readiness_handler_unhealthy_returns_503() {
let mut registry = HealthCheckRegistry::new();
registry.add("db", true, || async { Err("down".to_string()) });
let handler = readiness_handler(Arc::new(registry));
let ctx = RequestContext::new(asupersync::Cx::for_testing(), 1);
let mut req = Request::new(Method::Get, "/ready");
let resp = futures_executor::block_on(handler(&ctx, &mut req));
assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
}
#[test]
fn multiple_checks_mixed_results() {
let mut registry = HealthCheckRegistry::new();
registry.add("database", true, || async { Ok(()) });
registry.add("cache", false, || async {
Err("Cache timeout".to_string())
});
registry.add("search", false, || async { Ok(()) });
let report = futures_executor::block_on(registry.check_all());
assert_eq!(report.status, HealthStatus::Degraded);
assert_eq!(report.checks.len(), 3);
assert!(report.checks[0].passed);
assert!(report.checks[0].critical);
assert!(!report.checks[1].passed);
assert!(!report.checks[1].critical);
assert!(report.checks[1].error.is_some());
assert!(report.checks[2].passed);
}
}