use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Health {
Healthy,
Busy,
Stuck,
Dead,
DeadlineExceeded,
}
impl Health {
pub fn needs_teardown(self) -> bool {
matches!(
self,
Health::Stuck | Health::Dead | Health::DeadlineExceeded
)
}
}
#[derive(Debug, Clone, Copy)]
pub struct LivenessConfig {
pub progress_timeout: Duration,
pub pong_timeout: Duration,
pub ping_interval: Duration,
}
impl Default for LivenessConfig {
fn default() -> Self {
Self::new(Duration::from_secs(120), Duration::from_secs(10))
}
}
impl LivenessConfig {
pub fn new(progress_timeout: Duration, pong_timeout: Duration) -> LivenessConfig {
let ping_interval =
(pong_timeout / 3).clamp(Duration::from_millis(50), Duration::from_secs(5));
LivenessConfig {
progress_timeout,
pong_timeout,
ping_interval,
}
}
pub fn from_env() -> LivenessConfig {
let d = LivenessConfig::default();
let ms = |k: &str, fallback: Duration| {
std::env::var(k)
.ok()
.and_then(|v| v.parse::<u64>().ok())
.map(Duration::from_millis)
.unwrap_or(fallback)
};
LivenessConfig::new(
ms("AGENTD_PROGRESS_TIMEOUT_MS", d.progress_timeout),
ms("AGENTD_PONG_TIMEOUT_MS", d.pong_timeout),
)
}
}
#[derive(Debug)]
pub struct Liveness {
deadline: Instant,
cfg: LivenessConfig,
last_event_at: Instant,
last_pong_at: Instant,
eof: bool,
}
impl Liveness {
pub fn new(now: Instant, deadline: Instant, cfg: LivenessConfig) -> Liveness {
Liveness {
deadline,
cfg,
last_event_at: now,
last_pong_at: now,
eof: false,
}
}
pub fn on_event(&mut self, now: Instant) {
self.last_event_at = now;
self.last_pong_at = now;
}
pub fn on_pong(&mut self, now: Instant) {
self.last_pong_at = now;
}
pub fn on_eof(&mut self) {
self.eof = true;
}
pub fn deadline(&self) -> Instant {
self.deadline
}
pub fn classify(&self, now: Instant) -> Health {
if self.eof {
return Health::Dead;
}
if now >= self.deadline {
return Health::DeadlineExceeded;
}
if now.duration_since(self.last_event_at) <= self.cfg.progress_timeout {
Health::Healthy
} else if now.duration_since(self.last_pong_at) <= self.cfg.pong_timeout {
Health::Busy
} else {
Health::Stuck
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn cfg() -> LivenessConfig {
LivenessConfig::new(Duration::from_secs(100), Duration::from_secs(10))
}
#[test]
fn ping_interval_is_derived_below_the_pong_window_and_clamped() {
assert_eq!(
LivenessConfig::new(Duration::from_secs(100), Duration::from_secs(9)).ping_interval,
Duration::from_secs(3)
);
assert_eq!(
LivenessConfig::new(Duration::from_secs(1), Duration::from_millis(30)).ping_interval,
Duration::from_millis(50)
);
assert_eq!(
LivenessConfig::new(Duration::from_secs(600), Duration::from_secs(60)).ping_interval,
Duration::from_secs(5)
);
let d = LivenessConfig::default();
assert!(d.ping_interval < d.pong_timeout);
}
#[test]
fn recent_events_are_healthy() {
let t0 = Instant::now();
let l = Liveness::new(t0, t0 + Duration::from_secs(3600), cfg());
assert_eq!(l.classify(t0 + Duration::from_secs(50)), Health::Healthy);
}
#[test]
fn no_events_but_pongs_is_busy() {
let t0 = Instant::now();
let mut l = Liveness::new(t0, t0 + Duration::from_secs(3600), cfg());
l.on_pong(t0 + Duration::from_secs(145));
assert_eq!(l.classify(t0 + Duration::from_secs(150)), Health::Busy);
}
#[test]
fn no_events_no_pongs_is_stuck() {
let t0 = Instant::now();
let l = Liveness::new(t0, t0 + Duration::from_secs(3600), cfg());
assert_eq!(l.classify(t0 + Duration::from_secs(200)), Health::Stuck);
assert!(l.classify(t0 + Duration::from_secs(200)).needs_teardown());
}
#[test]
fn eof_is_dead_and_dominates() {
let t0 = Instant::now();
let mut l = Liveness::new(t0, t0 + Duration::from_secs(3600), cfg());
l.on_event(t0 + Duration::from_secs(1)); l.on_eof();
assert_eq!(l.classify(t0 + Duration::from_secs(2)), Health::Dead); }
#[test]
fn deadline_exceeded() {
let t0 = Instant::now();
let mut l = Liveness::new(t0, t0 + Duration::from_secs(60), cfg());
l.on_event(t0 + Duration::from_secs(59)); assert_eq!(
l.classify(t0 + Duration::from_secs(61)),
Health::DeadlineExceeded
);
}
#[test]
fn on_event_refreshes_both_clocks() {
let t0 = Instant::now();
let mut l = Liveness::new(t0, t0 + Duration::from_secs(3600), cfg());
l.on_event(t0 + Duration::from_secs(300));
assert_eq!(l.classify(t0 + Duration::from_secs(301)), Health::Healthy);
}
}