use std::sync::{Arc, Mutex, OnceLock};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
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",
}
}
}
type HealthCheck = Arc<dyn Fn() -> HealthStatus + Send + Sync>;
struct HealthEntry {
name: String,
check: HealthCheck,
}
pub struct HealthRegistry {
components: Mutex<Vec<HealthEntry>>,
}
static REGISTRY: OnceLock<HealthRegistry> = OnceLock::new();
impl HealthRegistry {
fn new() -> Self {
Self {
components: Mutex::new(Vec::new()),
}
}
fn global() -> &'static Self {
REGISTRY.get_or_init(Self::new)
}
pub fn register(
name: impl Into<String>,
check: impl Fn() -> HealthStatus + Send + Sync + 'static,
) {
let registry = Self::global();
if let Ok(mut components) = registry.components.lock() {
components.push(HealthEntry {
name: name.into(),
check: Arc::new(check),
});
}
}
#[must_use]
pub fn is_healthy() -> bool {
let registry = Self::global();
let Ok(components) = registry.components.lock() else {
return false;
};
components
.iter()
.all(|c| (c.check)() == HealthStatus::Healthy)
}
#[must_use]
pub fn is_ready() -> bool {
let registry = Self::global();
let Ok(components) = registry.components.lock() else {
return false;
};
components
.iter()
.all(|c| (c.check)() != HealthStatus::Unhealthy)
}
#[must_use]
pub fn components() -> Vec<(String, HealthStatus)> {
let registry = Self::global();
let Ok(components) = registry.components.lock() else {
return Vec::new();
};
components
.iter()
.map(|c| (c.name.clone(), (c.check)()))
.collect()
}
#[cfg(feature = "serde_json")]
#[must_use]
pub fn to_json() -> serde_json::Value {
let components = Self::components();
let overall = if Self::is_healthy() {
"healthy"
} else if Self::is_ready() {
"degraded"
} else {
"unhealthy"
};
serde_json::json!({
"status": overall,
"components": components.iter().map(|(name, status)| {
serde_json::json!({
"name": name,
"status": status.as_str(),
})
}).collect::<Vec<_>>()
})
}
#[cfg(test)]
pub(crate) fn reset() {
let registry = Self::global();
if let Ok(mut components) = registry.components.lock() {
components.clear();
}
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicU8, Ordering};
use super::*;
static TEST_LOCK: Mutex<()> = Mutex::new(());
macro_rules! serial_test {
() => {
let _guard = TEST_LOCK.lock().unwrap();
HealthRegistry::reset();
};
}
#[test]
fn empty_registry_is_healthy() {
serial_test!();
assert!(HealthRegistry::is_healthy());
assert!(HealthRegistry::is_ready());
assert!(HealthRegistry::components().is_empty());
}
#[test]
fn register_and_check_healthy() {
serial_test!();
HealthRegistry::register("transport", || HealthStatus::Healthy);
HealthRegistry::register("database", || HealthStatus::Healthy);
assert!(HealthRegistry::is_healthy());
assert!(HealthRegistry::is_ready());
let components = HealthRegistry::components();
assert_eq!(components.len(), 2);
assert_eq!(components[0].0, "transport");
assert_eq!(components[0].1, HealthStatus::Healthy);
assert_eq!(components[1].0, "database");
assert_eq!(components[1].1, HealthStatus::Healthy);
}
#[test]
fn unhealthy_component_fails_check() {
serial_test!();
HealthRegistry::register("transport", || HealthStatus::Healthy);
HealthRegistry::register("database", || HealthStatus::Unhealthy);
assert!(!HealthRegistry::is_healthy());
assert!(!HealthRegistry::is_ready());
}
#[test]
fn degraded_is_ready_but_not_healthy() {
serial_test!();
HealthRegistry::register("transport", || HealthStatus::Healthy);
HealthRegistry::register("circuit_breaker", || HealthStatus::Degraded);
assert!(!HealthRegistry::is_healthy());
assert!(HealthRegistry::is_ready());
}
#[test]
fn dynamic_health_check_reflects_state_changes() {
serial_test!();
let state = Arc::new(AtomicU8::new(0)); let state_clone = state.clone();
HealthRegistry::register("dynamic", move || {
match state_clone.load(Ordering::Relaxed) {
0 => HealthStatus::Healthy,
1 => HealthStatus::Degraded,
_ => HealthStatus::Unhealthy,
}
});
assert!(HealthRegistry::is_healthy());
assert!(HealthRegistry::is_ready());
state.store(1, Ordering::Relaxed);
assert!(!HealthRegistry::is_healthy());
assert!(HealthRegistry::is_ready());
state.store(2, Ordering::Relaxed);
assert!(!HealthRegistry::is_healthy());
assert!(!HealthRegistry::is_ready());
state.store(0, Ordering::Relaxed);
assert!(HealthRegistry::is_healthy());
assert!(HealthRegistry::is_ready());
}
#[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]
#[cfg(feature = "serde_json")]
fn to_json_includes_all_components() {
serial_test!();
HealthRegistry::register("kafka", || HealthStatus::Healthy);
HealthRegistry::register("clickhouse", || HealthStatus::Degraded);
let json = HealthRegistry::to_json();
assert_eq!(json["status"], "degraded");
let components = json["components"].as_array().unwrap();
assert_eq!(components.len(), 2);
assert_eq!(components[0]["name"], "kafka");
assert_eq!(components[0]["status"], "healthy");
assert_eq!(components[1]["name"], "clickhouse");
assert_eq!(components[1]["status"], "degraded");
}
#[test]
#[cfg(feature = "serde_json")]
fn to_json_empty_registry() {
serial_test!();
let json = HealthRegistry::to_json();
assert_eq!(json["status"], "healthy");
assert!(json["components"].as_array().unwrap().is_empty());
}
#[test]
#[cfg(feature = "serde_json")]
fn to_json_unhealthy_status() {
serial_test!();
HealthRegistry::register("broken", || HealthStatus::Unhealthy);
let json = HealthRegistry::to_json();
assert_eq!(json["status"], "unhealthy");
}
}