use super::HealthCheck;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, error, warn};
type HealthFailureCallback = Arc<dyn Fn(&str) + Send + Sync>;
#[derive(Debug, Clone)]
pub struct HealthCheckResult {
pub name: String,
pub healthy: bool,
pub error: Option<String>,
}
pub struct HealthChecker {
checks: Vec<Arc<dyn HealthCheck>>,
failure_counts: Arc<RwLock<std::collections::HashMap<String, u32>>>,
failure_threshold: u32,
on_failure: Option<HealthFailureCallback>,
}
impl HealthChecker {
pub fn new() -> Self {
Self {
checks: Vec::new(),
failure_counts: Arc::new(RwLock::new(std::collections::HashMap::new())),
failure_threshold: 3,
on_failure: None,
}
}
pub fn with_failure_threshold(mut self, threshold: u32) -> Self {
self.failure_threshold = threshold;
self
}
pub fn with_on_failure(mut self, callback: HealthFailureCallback) -> Self {
self.on_failure = Some(callback);
self
}
pub fn add_check(&mut self, check: Arc<dyn HealthCheck>) {
let name = check.name().to_string();
self.checks.push(check);
debug!(check_name = %name, "Health check added");
}
pub async fn check_all(&self) -> Vec<HealthCheckResult> {
let mut results = Vec::new();
for check in &self.checks {
let result = self.check_one(check.as_ref()).await;
results.push(result);
}
results
}
async fn check_one(&self, check: &dyn HealthCheck) -> HealthCheckResult {
let name = check.name().to_string();
match check.check().await {
Ok(_) => {
let mut counts = self.failure_counts.write().await;
counts.insert(name.clone(), 0);
HealthCheckResult {
name: name.clone(),
healthy: true,
error: None,
}
}
Err(e) => {
let mut counts = self.failure_counts.write().await;
let count = counts.entry(name.clone()).or_insert(0);
*count += 1;
error!(
check_name = %name,
failure_count = *count,
threshold = self.failure_threshold,
error = %e,
"Health check failed"
);
if *count >= self.failure_threshold {
warn!(
check_name = %name,
"Health check failure threshold exceeded"
);
if let Some(callback) = &self.on_failure {
callback(&name);
}
}
HealthCheckResult {
name: name.clone(),
healthy: false,
error: Some(e.to_string()),
}
}
}
}
pub async fn is_healthy(&self) -> bool {
let results = self.check_all().await;
results.iter().all(|r| r.healthy)
}
pub fn check_count(&self) -> usize {
self.checks.len()
}
}
impl Default for HealthChecker {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::HealthError;
use std::future::Future;
use std::pin::Pin;
struct TestHealthCheck {
name: String,
healthy: bool,
}
impl HealthCheck for TestHealthCheck {
fn check(&self) -> Pin<Box<dyn Future<Output = Result<(), HealthError>> + Send + '_>> {
let healthy = self.healthy;
let name = self.name.clone();
Box::pin(async move {
if healthy {
Ok(())
} else {
Err(HealthError::CheckFailed {
name,
reason: "Test failure".to_string(),
})
}
})
}
fn name(&self) -> &str {
&self.name
}
}
#[tokio::test]
async fn test_health_checker_add() {
let mut checker = HealthChecker::new();
checker.add_check(Arc::new(TestHealthCheck {
name: "test-check".to_string(),
healthy: true,
}));
assert_eq!(checker.check_count(), 1);
}
#[tokio::test]
async fn test_health_checker_healthy() {
let mut checker = HealthChecker::new();
checker.add_check(Arc::new(TestHealthCheck {
name: "test-check".to_string(),
healthy: true,
}));
let is_healthy = checker.is_healthy().await;
assert!(is_healthy);
}
#[tokio::test]
async fn test_health_checker_unhealthy() {
let mut checker = HealthChecker::new();
checker.add_check(Arc::new(TestHealthCheck {
name: "test-check".to_string(),
healthy: false,
}));
let is_healthy = checker.is_healthy().await;
assert!(!is_healthy);
}
}