Skip to main content

rightkit_process/
health.rs

1use std::{error::Error, fmt, time::Duration};
2
3#[derive(Clone, Copy, Debug, Eq, PartialEq)]
4pub enum HealthProbe {
5    Success,
6    Failure,
7}
8
9#[derive(Clone, Copy, Debug, Eq, PartialEq)]
10pub enum HealthOutcome {
11    Starting,
12    Healthy,
13    Unhealthy { consecutive_failures: u32 },
14    StartupTimedOut,
15}
16
17#[derive(Clone, Copy, Debug, Eq, PartialEq)]
18pub struct HealthPolicy {
19    startup_timeout: Duration,
20    unhealthy_threshold: u32,
21}
22
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24pub struct InvalidHealthPolicy;
25
26impl fmt::Display for InvalidHealthPolicy {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        f.write_str("startup timeout and unhealthy threshold must be non-zero")
29    }
30}
31
32impl Error for InvalidHealthPolicy {}
33
34impl HealthPolicy {
35    pub fn new(
36        startup_timeout: Duration,
37        unhealthy_threshold: u32,
38    ) -> Result<Self, InvalidHealthPolicy> {
39        if startup_timeout.is_zero() || unhealthy_threshold == 0 {
40            return Err(InvalidHealthPolicy);
41        }
42        Ok(Self {
43            startup_timeout,
44            unhealthy_threshold,
45        })
46    }
47}
48
49#[derive(Clone, Debug)]
50pub struct HealthTracker {
51    policy: HealthPolicy,
52    started_at: Duration,
53    ever_healthy: bool,
54    consecutive_failures: u32,
55}
56
57impl HealthTracker {
58    pub fn new(policy: HealthPolicy, started_at: Duration) -> Self {
59        Self {
60            policy,
61            started_at,
62            ever_healthy: false,
63            consecutive_failures: 0,
64        }
65    }
66
67    pub fn observe(&mut self, at: Duration, probe: HealthProbe) -> HealthOutcome {
68        match probe {
69            HealthProbe::Success => {
70                self.ever_healthy = true;
71                self.consecutive_failures = 0;
72            }
73            HealthProbe::Failure => {
74                self.consecutive_failures = self.consecutive_failures.saturating_add(1);
75            }
76        }
77        self.outcome(at)
78    }
79
80    pub fn outcome(&self, at: Duration) -> HealthOutcome {
81        if !self.ever_healthy && at.saturating_sub(self.started_at) >= self.policy.startup_timeout {
82            return HealthOutcome::StartupTimedOut;
83        }
84        if self.consecutive_failures >= self.policy.unhealthy_threshold {
85            return HealthOutcome::Unhealthy {
86                consecutive_failures: self.consecutive_failures,
87            };
88        }
89        if self.ever_healthy {
90            HealthOutcome::Healthy
91        } else {
92            HealthOutcome::Starting
93        }
94    }
95}