use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
use serde::Serialize;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
const DEFAULT_HEALTH_CHECK_TIMEOUT: Duration = Duration::from_secs(5);
#[derive(Debug, Serialize)]
pub struct HealthResponse {
pub status: &'static str,
pub service: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub details: Option<String>,
}
#[derive(Debug, Clone)]
pub struct HealthResult {
pub healthy: bool,
pub details: Option<String>,
}
impl HealthResult {
#[must_use]
pub fn healthy() -> Self {
Self {
healthy: true,
details: None,
}
}
pub fn unhealthy(details: impl Into<String>) -> Self {
Self {
healthy: false,
details: Some(details.into()),
}
}
}
pub type AsyncHealthCheckFn =
Box<dyn Fn() -> Pin<Box<dyn Future<Output = HealthResult> + Send>> + Send + Sync>;
pub struct HealthCheck {
pub name: String,
sync_check: Option<Box<dyn Fn() -> bool + Send + Sync>>,
async_check: Option<AsyncHealthCheckFn>,
timeout: Duration,
}
impl std::fmt::Debug for HealthCheck {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HealthCheck")
.field("name", &self.name)
.field("sync_check", &self.sync_check.as_ref().map(|_| "<fn>"))
.field(
"async_check",
&self.async_check.as_ref().map(|_| "<async fn>"),
)
.field("timeout", &self.timeout)
.finish()
}
}
impl HealthCheck {
pub fn new(
name: impl Into<String>,
check_fn: impl Fn() -> bool + Send + Sync + 'static,
) -> Self {
Self {
name: name.into(),
sync_check: Some(Box::new(check_fn)),
async_check: None,
timeout: DEFAULT_HEALTH_CHECK_TIMEOUT,
}
}
pub fn with_async<F, Fut>(name: impl Into<String>, check_fn: F) -> Self
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = HealthResult> + Send + 'static,
{
Self {
name: name.into(),
sync_check: None,
async_check: Some(Box::new(move || Box::pin(check_fn()))),
timeout: DEFAULT_HEALTH_CHECK_TIMEOUT,
}
}
pub fn always_healthy(name: impl Into<String>) -> Self {
Self::new(name, || true)
}
#[must_use]
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
#[must_use]
pub fn is_healthy(&self) -> bool {
if let Some(ref check) = self.sync_check {
check()
} else {
true }
}
pub async fn check_async(&self) -> HealthResult {
if let Some(ref async_check) = self.async_check {
let check_future = async_check();
match tokio::time::timeout(self.timeout, check_future).await {
Ok(result) => result,
Err(_) => HealthResult::unhealthy(format!(
"Health check timed out after {:?}",
self.timeout
)),
}
} else if let Some(ref sync_check) = self.sync_check {
if sync_check() {
HealthResult::healthy()
} else {
HealthResult::unhealthy("Sync health check failed")
}
} else {
HealthResult::healthy()
}
}
}
pub async fn health_handler(State(health): State<Arc<HealthCheck>>) -> impl IntoResponse {
let result = health.check_async().await;
let response = HealthResponse {
status: if result.healthy {
"healthy"
} else {
"unhealthy"
},
service: health.name.clone(),
details: result.details,
};
let status = if result.healthy {
StatusCode::OK
} else {
StatusCode::SERVICE_UNAVAILABLE
};
(status, Json(response))
}
pub async fn simple_health_handler() -> impl IntoResponse {
Json(serde_json::json!({
"status": "healthy"
}))
}