use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum HealthStatus {
Healthy,
Degraded,
Unhealthy,
}
impl HealthStatus {
pub fn is_healthy(&self) -> bool {
matches!(self, HealthStatus::Healthy)
}
pub fn is_degraded(&self) -> bool {
matches!(self, HealthStatus::Degraded)
}
pub fn is_unhealthy(&self) -> bool {
matches!(self, HealthStatus::Unhealthy)
}
pub fn is_ready(&self) -> bool {
!self.is_unhealthy()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthCheckResult {
pub status: HealthStatus,
pub component: String,
pub message: String,
pub checked_at: String,
pub duration_ms: u64,
pub metadata: HashMap<String, String>,
}
impl HealthCheckResult {
pub fn healthy(component: String, message: String, duration: Duration) -> Self {
Self {
status: HealthStatus::Healthy,
component,
message,
checked_at: chrono::Utc::now().to_rfc3339(),
duration_ms: duration.as_millis() as u64,
metadata: HashMap::new(),
}
}
pub fn degraded(component: String, message: String, duration: Duration) -> Self {
Self {
status: HealthStatus::Degraded,
component,
message,
checked_at: chrono::Utc::now().to_rfc3339(),
duration_ms: duration.as_millis() as u64,
metadata: HashMap::new(),
}
}
pub fn unhealthy(component: String, message: String, duration: Duration) -> Self {
Self {
status: HealthStatus::Unhealthy,
component,
message,
checked_at: chrono::Utc::now().to_rfc3339(),
duration_ms: duration.as_millis() as u64,
metadata: HashMap::new(),
}
}
pub fn with_metadata(mut self, key: String, value: String) -> Self {
self.metadata.insert(key, value);
self
}
}
#[async_trait]
pub trait HealthCheck: Send + Sync {
async fn check_liveness(&self) -> HealthCheckResult;
async fn check_readiness(&self) -> HealthCheckResult;
fn component_name(&self) -> String;
}
pub struct HealthChecker {
checks: Arc<parking_lot::RwLock<Vec<Arc<dyn HealthCheck>>>>,
}
impl HealthChecker {
pub fn new() -> Self {
Self {
checks: Arc::new(parking_lot::RwLock::new(Vec::new())),
}
}
pub fn register<H: HealthCheck + 'static>(&self, check: H) {
self.checks.write().push(Arc::new(check));
}
pub async fn check_liveness(&self) -> AggregateHealthResult {
let checks = self.checks.read().clone();
let mut results = Vec::new();
for check in checks {
results.push(check.check_liveness().await);
}
AggregateHealthResult::from_results(results)
}
pub async fn check_readiness(&self) -> AggregateHealthResult {
let checks = self.checks.read().clone();
let mut results = Vec::new();
for check in checks {
results.push(check.check_readiness().await);
}
AggregateHealthResult::from_results(results)
}
pub async fn detailed_status(&self) -> DetailedHealthStatus {
let checks = self.checks.read().clone();
let mut liveness_results = Vec::new();
let mut readiness_results = Vec::new();
for check in checks {
liveness_results.push(check.check_liveness().await);
readiness_results.push(check.check_readiness().await);
}
DetailedHealthStatus {
liveness: AggregateHealthResult::from_results(liveness_results),
readiness: AggregateHealthResult::from_results(readiness_results),
}
}
}
impl Default for HealthChecker {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AggregateHealthResult {
pub status: HealthStatus,
pub components: Vec<HealthCheckResult>,
pub total_components: usize,
pub healthy_count: usize,
pub degraded_count: usize,
pub unhealthy_count: usize,
}
impl AggregateHealthResult {
pub fn from_results(components: Vec<HealthCheckResult>) -> Self {
let total_components = components.len();
let mut healthy_count = 0;
let mut degraded_count = 0;
let mut unhealthy_count = 0;
for result in &components {
match result.status {
HealthStatus::Healthy => healthy_count += 1,
HealthStatus::Degraded => degraded_count += 1,
HealthStatus::Unhealthy => unhealthy_count += 1,
}
}
let status = if unhealthy_count > 0 {
HealthStatus::Unhealthy
} else if degraded_count > 0 {
HealthStatus::Degraded
} else {
HealthStatus::Healthy
};
Self {
status,
components,
total_components,
healthy_count,
degraded_count,
unhealthy_count,
}
}
pub fn all_healthy(&self) -> bool {
self.status == HealthStatus::Healthy
}
pub fn any_unhealthy(&self) -> bool {
self.unhealthy_count > 0
}
pub fn unhealthy_components(&self) -> Vec<&HealthCheckResult> {
self.components
.iter()
.filter(|r| r.status == HealthStatus::Unhealthy)
.collect()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetailedHealthStatus {
pub liveness: AggregateHealthResult,
pub readiness: AggregateHealthResult,
}
impl DetailedHealthStatus {
pub fn is_alive(&self) -> bool {
self.liveness.status != HealthStatus::Unhealthy
}
pub fn is_ready(&self) -> bool {
self.readiness.status != HealthStatus::Unhealthy
}
}
#[derive(Clone)]
pub struct SimpleHealthCheck {
name: String,
is_healthy: Arc<parking_lot::RwLock<bool>>,
}
impl SimpleHealthCheck {
pub fn new(name: String) -> Self {
Self {
name,
is_healthy: Arc::new(parking_lot::RwLock::new(true)),
}
}
pub fn set_healthy(&self, healthy: bool) {
*self.is_healthy.write() = healthy;
}
}
#[async_trait]
impl HealthCheck for SimpleHealthCheck {
async fn check_liveness(&self) -> HealthCheckResult {
let start = Instant::now();
let is_healthy = *self.is_healthy.read();
let duration = start.elapsed();
if is_healthy {
HealthCheckResult::healthy(
self.name.clone(),
"Component is alive".to_string(),
duration,
)
} else {
HealthCheckResult::unhealthy(
self.name.clone(),
"Component is not alive".to_string(),
duration,
)
}
}
async fn check_readiness(&self) -> HealthCheckResult {
let start = Instant::now();
let is_healthy = *self.is_healthy.read();
let duration = start.elapsed();
if is_healthy {
HealthCheckResult::healthy(
self.name.clone(),
"Component is ready".to_string(),
duration,
)
} else {
HealthCheckResult::unhealthy(
self.name.clone(),
"Component is not ready".to_string(),
duration,
)
}
}
fn component_name(&self) -> String {
self.name.clone()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_health_status() {
assert!(HealthStatus::Healthy.is_healthy());
assert!(!HealthStatus::Degraded.is_healthy());
assert!(!HealthStatus::Unhealthy.is_healthy());
assert!(HealthStatus::Healthy.is_ready());
assert!(HealthStatus::Degraded.is_ready());
assert!(!HealthStatus::Unhealthy.is_ready());
}
#[tokio::test]
async fn test_simple_health_check() {
let check = SimpleHealthCheck::new("test".to_string());
let result = check.check_liveness().await;
assert!(result.status.is_healthy());
check.set_healthy(false);
let result = check.check_liveness().await;
assert!(result.status.is_unhealthy());
}
#[tokio::test]
async fn test_health_checker_aggregate() {
let checker = HealthChecker::new();
let check1 = SimpleHealthCheck::new("component1".to_string());
let check2 = SimpleHealthCheck::new("component2".to_string());
checker.register(check1.clone());
checker.register(check2.clone());
let result = checker.check_liveness().await;
assert!(result.all_healthy());
assert_eq!(result.healthy_count, 2);
check1.set_healthy(false);
let result = checker.check_liveness().await;
assert!(!result.all_healthy());
assert!(result.any_unhealthy());
assert_eq!(result.healthy_count, 1);
assert_eq!(result.unhealthy_count, 1);
}
#[tokio::test]
async fn test_detailed_status() {
let checker = HealthChecker::new();
let check = SimpleHealthCheck::new("test".to_string());
checker.register(check);
let status = checker.detailed_status().await;
assert!(status.is_alive());
assert!(status.is_ready());
}
#[tokio::test]
async fn test_aggregate_health_result() {
let results = vec![
HealthCheckResult::healthy(
"comp1".to_string(),
"OK".to_string(),
Duration::from_millis(10),
),
HealthCheckResult::degraded(
"comp2".to_string(),
"Slow".to_string(),
Duration::from_millis(100),
),
];
let aggregate = AggregateHealthResult::from_results(results);
assert_eq!(aggregate.status, HealthStatus::Degraded);
assert_eq!(aggregate.healthy_count, 1);
assert_eq!(aggregate.degraded_count, 1);
assert_eq!(aggregate.unhealthy_count, 0);
}
}