use std::sync::Arc;
use serde::{Deserialize, Serialize};
use tracing::{debug, warn};
use crate::service::ServiceRegistry;
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum HealthStatus {
Healthy,
Degraded(String),
Unhealthy(String),
Unknown,
}
impl std::fmt::Display for HealthStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
HealthStatus::Healthy => write!(f, "healthy"),
HealthStatus::Degraded(msg) => write!(f, "degraded: {msg}"),
HealthStatus::Unhealthy(msg) => write!(f, "unhealthy: {msg}"),
HealthStatus::Unknown => write!(f, "unknown"),
}
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum OverallHealth {
Healthy,
Degraded {
unhealthy_services: Vec<String>,
},
Down,
}
impl std::fmt::Display for OverallHealth {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
OverallHealth::Healthy => write!(f, "healthy"),
OverallHealth::Degraded { unhealthy_services } => {
write!(f, "degraded ({})", unhealthy_services.join(", "))
}
OverallHealth::Down => write!(f, "down"),
}
}
}
pub struct HealthSystem {
check_interval_secs: u64,
}
impl HealthSystem {
pub fn new(check_interval_secs: u64) -> Self {
Self {
check_interval_secs,
}
}
pub fn check_interval_secs(&self) -> u64 {
self.check_interval_secs
}
pub async fn aggregate(
&self,
registry: &Arc<ServiceRegistry>,
) -> (OverallHealth, Vec<(String, HealthStatus)>) {
let results = registry.health_all().await;
if results.is_empty() {
return (OverallHealth::Down, results);
}
let mut unhealthy = Vec::new();
let mut all_unhealthy = true;
for (name, status) in &results {
match status {
HealthStatus::Healthy => {
debug!(service = %name, "health check: healthy");
all_unhealthy = false;
}
HealthStatus::Degraded(msg) => {
warn!(service = %name, reason = %msg, "health check: degraded");
unhealthy.push(name.clone());
all_unhealthy = false;
}
HealthStatus::Unhealthy(msg) => {
warn!(service = %name, reason = %msg, "health check: unhealthy");
unhealthy.push(name.clone());
}
HealthStatus::Unknown => {
warn!(service = %name, "health check: unknown");
unhealthy.push(name.clone());
}
}
}
let overall = if unhealthy.is_empty() {
OverallHealth::Healthy
} else if all_unhealthy {
OverallHealth::Down
} else {
OverallHealth::Degraded {
unhealthy_services: unhealthy,
}
};
(overall, results)
}
}
#[non_exhaustive]
#[cfg(feature = "os-patterns")]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProbeResult {
Live,
NotLive { reason: String },
Ready,
NotReady { reason: String },
}
#[cfg(feature = "os-patterns")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProbeConfig {
pub liveness_interval_secs: u64,
pub readiness_interval_secs: u64,
pub failure_threshold: u32,
pub success_threshold: u32,
}
#[cfg(feature = "os-patterns")]
impl Default for ProbeConfig {
fn default() -> Self {
Self {
liveness_interval_secs: 10,
readiness_interval_secs: 5,
failure_threshold: 3,
success_threshold: 1,
}
}
}
#[cfg(feature = "os-patterns")]
#[derive(Debug, Clone)]
pub struct ProbeState {
pub liveness_failures: u32,
pub readiness_failures: u32,
pub readiness_successes: u32,
pub is_live: bool,
pub is_ready: bool,
}
#[cfg(feature = "os-patterns")]
impl Default for ProbeState {
fn default() -> Self {
Self {
liveness_failures: 0,
readiness_failures: 0,
readiness_successes: 0,
is_live: true,
is_ready: true,
}
}
}
#[cfg(feature = "os-patterns")]
impl ProbeState {
pub fn record_liveness(&mut self, result: &ProbeResult, config: &ProbeConfig) -> bool {
match result {
ProbeResult::Live => {
self.liveness_failures = 0;
self.is_live = true;
false
}
ProbeResult::NotLive { .. } => {
self.liveness_failures += 1;
if self.liveness_failures >= config.failure_threshold {
self.is_live = false;
true
} else {
false
}
}
_ => false, }
}
pub fn record_readiness(&mut self, result: &ProbeResult, config: &ProbeConfig) -> Option<bool> {
match result {
ProbeResult::Ready => {
self.readiness_failures = 0;
self.readiness_successes += 1;
if !self.is_ready && self.readiness_successes >= config.success_threshold {
self.is_ready = true;
Some(true) } else {
None
}
}
ProbeResult::NotReady { .. } => {
self.readiness_successes = 0;
self.readiness_failures += 1;
if self.is_ready && self.readiness_failures >= config.failure_threshold {
self.is_ready = false;
Some(false) } else {
None
}
}
_ => None, }
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::service::{ServiceType, SystemService};
use async_trait::async_trait;
struct HealthyService;
#[async_trait]
impl SystemService for HealthyService {
fn name(&self) -> &str {
"healthy-svc"
}
fn service_type(&self) -> ServiceType {
ServiceType::Core
}
async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
Ok(())
}
async fn stop(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
Ok(())
}
async fn health_check(&self) -> HealthStatus {
HealthStatus::Healthy
}
}
struct UnhealthyService;
#[async_trait]
impl SystemService for UnhealthyService {
fn name(&self) -> &str {
"unhealthy-svc"
}
fn service_type(&self) -> ServiceType {
ServiceType::Core
}
async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
Ok(())
}
async fn stop(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
Ok(())
}
async fn health_check(&self) -> HealthStatus {
HealthStatus::Unhealthy("test failure".into())
}
}
#[tokio::test]
async fn aggregate_all_healthy() {
let registry = Arc::new(ServiceRegistry::new());
registry.register(Arc::new(HealthyService)).unwrap();
let health = HealthSystem::new(30);
let (overall, results) = health.aggregate(®istry).await;
assert!(matches!(overall, OverallHealth::Healthy));
assert_eq!(results.len(), 1);
}
#[tokio::test]
async fn aggregate_mixed() {
let registry = Arc::new(ServiceRegistry::new());
registry.register(Arc::new(HealthyService)).unwrap();
registry.register(Arc::new(UnhealthyService)).unwrap();
let health = HealthSystem::new(30);
let (overall, results) = health.aggregate(®istry).await;
assert!(matches!(overall, OverallHealth::Degraded { .. }));
assert_eq!(results.len(), 2);
}
#[tokio::test]
async fn aggregate_all_unhealthy() {
let registry = Arc::new(ServiceRegistry::new());
registry.register(Arc::new(UnhealthyService)).unwrap();
let health = HealthSystem::new(30);
let (overall, _) = health.aggregate(®istry).await;
assert!(matches!(overall, OverallHealth::Down));
}
#[tokio::test]
async fn aggregate_empty_registry() {
let registry = Arc::new(ServiceRegistry::new());
let health = HealthSystem::new(30);
let (overall, results) = health.aggregate(®istry).await;
assert!(matches!(overall, OverallHealth::Down));
assert!(results.is_empty());
}
#[test]
fn health_status_display() {
assert_eq!(HealthStatus::Healthy.to_string(), "healthy");
assert_eq!(
HealthStatus::Degraded("slow".into()).to_string(),
"degraded: slow"
);
assert_eq!(
HealthStatus::Unhealthy("crash".into()).to_string(),
"unhealthy: crash"
);
assert_eq!(HealthStatus::Unknown.to_string(), "unknown");
}
#[test]
fn overall_health_display() {
assert_eq!(OverallHealth::Healthy.to_string(), "healthy");
assert_eq!(OverallHealth::Down.to_string(), "down");
assert_eq!(
OverallHealth::Degraded {
unhealthy_services: vec!["svc-a".into(), "svc-b".into()]
}
.to_string(),
"degraded (svc-a, svc-b)"
);
}
#[test]
fn check_interval() {
let health = HealthSystem::new(15);
assert_eq!(health.check_interval_secs(), 15);
}
#[cfg(feature = "os-patterns")]
mod probe_tests {
use super::super::*;
#[test]
fn probe_config_default() {
let config = ProbeConfig::default();
assert_eq!(config.liveness_interval_secs, 10);
assert_eq!(config.readiness_interval_secs, 5);
assert_eq!(config.failure_threshold, 3);
assert_eq!(config.success_threshold, 1);
}
#[test]
fn probe_config_serde_roundtrip() {
let config = ProbeConfig {
liveness_interval_secs: 15,
readiness_interval_secs: 10,
failure_threshold: 5,
success_threshold: 2,
};
let json = serde_json::to_string(&config).unwrap();
let restored: ProbeConfig = serde_json::from_str(&json).unwrap();
assert_eq!(restored.liveness_interval_secs, 15);
assert_eq!(restored.failure_threshold, 5);
}
#[test]
fn probe_result_serde_roundtrip() {
let results = vec![
ProbeResult::Live,
ProbeResult::NotLive { reason: "oom".into() },
ProbeResult::Ready,
ProbeResult::NotReady { reason: "init".into() },
];
for result in results {
let json = serde_json::to_string(&result).unwrap();
let restored: ProbeResult = serde_json::from_str(&json).unwrap();
assert_eq!(restored, result);
}
}
#[test]
fn probe_state_default_is_live_and_ready() {
let state = ProbeState::default();
assert!(state.is_live);
assert!(state.is_ready);
assert_eq!(state.liveness_failures, 0);
assert_eq!(state.readiness_failures, 0);
}
#[test]
fn liveness_resets_on_success() {
let config = ProbeConfig {
failure_threshold: 3,
..Default::default()
};
let mut state = ProbeState::default();
state.liveness_failures = 2;
let restart = state.record_liveness(&ProbeResult::Live, &config);
assert!(!restart);
assert_eq!(state.liveness_failures, 0);
assert!(state.is_live);
}
#[test]
fn liveness_triggers_restart_at_threshold() {
let config = ProbeConfig {
failure_threshold: 3,
..Default::default()
};
let mut state = ProbeState::default();
assert!(!state.record_liveness(&ProbeResult::NotLive { reason: "hang".into() }, &config));
assert!(!state.record_liveness(&ProbeResult::NotLive { reason: "hang".into() }, &config));
assert!(state.is_live);
assert!(state.record_liveness(&ProbeResult::NotLive { reason: "hang".into() }, &config));
assert!(!state.is_live);
}
#[test]
fn readiness_removes_at_threshold() {
let config = ProbeConfig {
failure_threshold: 2,
..Default::default()
};
let mut state = ProbeState::default();
assert!(state.record_readiness(&ProbeResult::NotReady { reason: "init".into() }, &config).is_none());
let change = state.record_readiness(&ProbeResult::NotReady { reason: "init".into() }, &config);
assert_eq!(change, Some(false)); assert!(!state.is_ready);
}
#[test]
fn readiness_recovery_re_adds() {
let config = ProbeConfig {
failure_threshold: 1,
success_threshold: 1,
..Default::default()
};
let mut state = ProbeState::default();
state.record_readiness(&ProbeResult::NotReady { reason: "init".into() }, &config);
assert!(!state.is_ready);
let change = state.record_readiness(&ProbeResult::Ready, &config);
assert_eq!(change, Some(true));
assert!(state.is_ready);
}
#[test]
fn threshold_prevents_flapping() {
let config = ProbeConfig {
failure_threshold: 3,
success_threshold: 2,
..Default::default()
};
let mut state = ProbeState::default();
assert!(state.record_readiness(&ProbeResult::NotReady { reason: "x".into() }, &config).is_none());
assert!(state.is_ready);
assert!(state.record_readiness(&ProbeResult::Ready, &config).is_none());
assert!(state.is_ready);
}
#[test]
fn default_probe_returns_live_ready() {
let live = ProbeResult::Live;
let ready = ProbeResult::Ready;
assert_eq!(live, ProbeResult::Live);
assert_eq!(ready, ProbeResult::Ready);
}
}
}