mod server;
mod types;
use std::{
future::Future,
pin::Pin,
sync::Arc,
time::{Duration, Instant},
};
pub use server::HealthServer;
pub use types::*;
pub trait Dependency: Send + Sync {
fn name(&self) -> &str;
fn check(&self) -> Pin<Box<dyn Future<Output = DependencyStatus> + Send + '_>>;
fn is_critical(&self) -> bool {
true
}
fn timeout(&self) -> Duration {
Duration::from_secs(5)
}
}
pub trait HealthCheck: Send + Sync {
fn dependencies(&self) -> Vec<Arc<dyn Dependency>>;
fn check_all(&self) -> Pin<Box<dyn Future<Output = HealthReport> + Send + '_>> {
let deps = self.dependencies();
Box::pin(async move {
let start = Instant::now();
let mut reports = Vec::with_capacity(deps.len());
let mut has_critical_failure = false;
let mut has_degradation = false;
for dep in deps {
let dep_start = Instant::now();
let timeout = dep.timeout();
let is_critical = dep.is_critical();
let name = dep.name().to_string();
let status = match tokio::time::timeout(timeout, dep.check()).await {
Ok(status) => status,
Err(_) => DependencyStatus::Unhealthy(format!(
"Health check timed out after {:?}",
timeout
)),
};
let duration = dep_start.elapsed();
match &status {
DependencyStatus::Unhealthy(_) if is_critical => {
has_critical_failure = true;
}
DependencyStatus::Unhealthy(_) | DependencyStatus::Degraded(_) => {
has_degradation = true;
}
_ => {}
}
reports.push(DependencyReport {
name,
status,
duration,
critical: is_critical,
});
}
let overall_status = if has_critical_failure {
OverallStatus::Unhealthy
} else if has_degradation {
OverallStatus::Degraded
} else {
OverallStatus::Healthy
};
HealthReport {
status: overall_status,
dependencies: reports,
total_duration: start.elapsed(),
timestamp: std::time::SystemTime::now(),
}
})
}
}
pub struct SimpleHealthCheck {
dependencies: Vec<Arc<dyn Dependency>>,
}
impl SimpleHealthCheck {
pub fn new() -> Self {
Self {
dependencies: Vec::new(),
}
}
pub fn add_dependency<D: Dependency + 'static>(mut self, dep: D) -> Self {
self.dependencies.push(Arc::new(dep));
self
}
pub fn add_arc_dependency(mut self, dep: Arc<dyn Dependency>) -> Self {
self.dependencies.push(dep);
self
}
}
impl Default for SimpleHealthCheck {
fn default() -> Self {
Self::new()
}
}
impl HealthCheck for SimpleHealthCheck {
fn dependencies(&self) -> Vec<Arc<dyn Dependency>> {
self.dependencies.clone()
}
}
pub struct AlwaysHealthy {
name: String,
}
impl AlwaysHealthy {
pub fn new(name: impl Into<String>) -> Self {
Self { name: name.into() }
}
}
impl Dependency for AlwaysHealthy {
fn name(&self) -> &str {
&self.name
}
fn check(&self) -> Pin<Box<dyn Future<Output = DependencyStatus> + Send + '_>> {
Box::pin(async { DependencyStatus::Healthy })
}
}
pub struct AlwaysUnhealthy {
name: String,
message: String,
}
impl AlwaysUnhealthy {
pub fn new(name: impl Into<String>, message: impl Into<String>) -> Self {
Self {
name: name.into(),
message: message.into(),
}
}
}
impl Dependency for AlwaysUnhealthy {
fn name(&self) -> &str {
&self.name
}
fn check(&self) -> Pin<Box<dyn Future<Output = DependencyStatus> + Send + '_>> {
let msg = self.message.clone();
Box::pin(async move { DependencyStatus::Unhealthy(msg) })
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_always_healthy() {
let dep = AlwaysHealthy::new("test");
assert_eq!(dep.name(), "test");
assert!(matches!(dep.check().await, DependencyStatus::Healthy));
}
#[tokio::test]
async fn test_always_unhealthy() {
let dep = AlwaysUnhealthy::new("test", "connection failed");
assert_eq!(dep.name(), "test");
match dep.check().await {
DependencyStatus::Unhealthy(msg) => assert_eq!(msg, "connection failed"),
_ => panic!("Expected unhealthy status"),
}
}
#[tokio::test]
async fn test_simple_health_check_empty() {
let checker = SimpleHealthCheck::new();
let report = checker.check_all().await;
assert_eq!(report.status, OverallStatus::Healthy);
assert!(report.dependencies.is_empty());
}
#[tokio::test]
async fn test_simple_health_check_all_healthy() {
let checker = SimpleHealthCheck::new()
.add_dependency(AlwaysHealthy::new("dep1"))
.add_dependency(AlwaysHealthy::new("dep2"));
let report = checker.check_all().await;
assert_eq!(report.status, OverallStatus::Healthy);
assert_eq!(report.dependencies.len(), 2);
}
#[tokio::test]
async fn test_simple_health_check_critical_failure() {
let checker = SimpleHealthCheck::new()
.add_dependency(AlwaysHealthy::new("healthy"))
.add_dependency(AlwaysUnhealthy::new("critical", "down"));
let report = checker.check_all().await;
assert_eq!(report.status, OverallStatus::Unhealthy);
}
#[tokio::test]
async fn test_default_timeout() {
let dep = AlwaysHealthy::new("test");
assert_eq!(dep.timeout(), Duration::from_secs(5));
}
#[tokio::test]
async fn test_default_critical() {
let dep = AlwaysHealthy::new("test");
assert!(dep.is_critical());
}
}