rws_server/health/
mod.rs

1//! Health checks for Kubernetes probes (18.35.0): register what "alive",
2//! "ready" and "started" mean for your app, and `/livez`, `/readyz` and
3//! `/startupz` run them.
4//!
5//! ```rust
6//! # extern crate rws_server as rust_web_server;
7//! use rust_web_server::health;
8//!
9//! health::readiness("cache", || {
10//!     // e.g. ping Redis; Err(message) takes the pod out of the Service.
11//!     Ok(())
12//! });
13//! health::startup("warm-up", || if warmed() { Ok(()) } else { Err("loading the index".into()) });
14//! # fn warmed() -> bool { true }
15//! // What /readyz?verbose answers (not ready here: no server is running).
16//! println!("{}", health::report(health::Probe::Readiness).body(true));
17//! assert!(health::report(health::Probe::Startup).ok);
18//! ```
19//!
20//! | Endpoint | Probe | `200` when |
21//! |---|---|---|
22//! | `/livez` (also `/healthz`) | liveness | every liveness check passes (none by default: the process answers) |
23//! | `/readyz` | readiness | the server is up and not shutting down, and every readiness check passes |
24//! | `/startupz` | startup | every startup check has passed once (then it stays `200`) |
25//!
26//! Each check runs with a timeout ([`set_timeout`], default 2 s; a check
27//! that takes longer fails). Add `?verbose` for one line per check
28//! (`[+]cache ok`, `[-]db failed: …`), as the Kubernetes API server does;
29//! otherwise the body is `OK` or the failing checks. Endpoints with checks
30//! run off the event loop, so a check may block (a database ping).
31//!
32//! Keep liveness checks to what a restart fixes (a deadlocked worker);
33//! a database outage belongs in readiness, or every pod restarts at once.
34
35#[cfg(test)]
36mod tests;
37
38use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
39use std::sync::{Arc, OnceLock, RwLock};
40use std::time::Duration;
41
42/// Which probe a check belongs to.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
44pub enum Probe {
45    Liveness,
46    Readiness,
47    Startup,
48}
49
50impl Probe {
51    /// The endpoint path.
52    pub fn path(self) -> &'static str {
53        match self {
54            Probe::Liveness => "/livez",
55            Probe::Readiness => "/readyz",
56            Probe::Startup => "/startupz",
57        }
58    }
59}
60
61type Check = Arc<dyn Fn() -> Result<(), String> + Send + Sync>;
62
63/// A set of checks. The module functions use the process-wide one the
64/// probe endpoints answer from; build your own for tests.
65pub struct Registry {
66    checks: RwLock<Vec<(Probe, String, Check)>>,
67    timeout_ms: AtomicU64,
68    started: AtomicBool,
69}
70
71impl Default for Registry {
72    fn default() -> Self {
73        Registry { checks: RwLock::new(Vec::new()), timeout_ms: AtomicU64::new(2000), started: AtomicBool::new(false) }
74    }
75}
76
77/// The process-wide registry.
78pub fn global() -> &'static Registry {
79    static R: OnceLock<Registry> = OnceLock::new();
80    R.get_or_init(Registry::default)
81}
82
83impl Registry {
84    pub fn new() -> Self {
85        Registry::default()
86    }
87
88    /// Add a check to `probe` (one with the same probe and name is replaced).
89    pub fn register(&self, probe: Probe, name: &str, check: impl Fn() -> Result<(), String> + Send + Sync + 'static) {
90        let mut checks = self.checks.write().unwrap_or_else(|e| e.into_inner());
91        checks.retain(|(p, n, _)| !(*p == probe && n == name));
92        checks.push((probe, name.to_string(), Arc::new(check)));
93    }
94
95    /// Remove a check; whether it was there.
96    pub fn unregister(&self, probe: Probe, name: &str) -> bool {
97        let mut checks = self.checks.write().unwrap_or_else(|e| e.into_inner());
98        let before = checks.len();
99        checks.retain(|(p, n, _)| !(*p == probe && n == name));
100        checks.len() != before
101    }
102
103    /// How long one check may take before it counts as failed (default 2 s).
104    pub fn set_timeout(&self, timeout: Duration) {
105        self.timeout_ms.store(timeout.as_millis().max(1) as u64, Ordering::Relaxed);
106    }
107
108    /// Whether `probe` has checks.
109    pub fn has_checks(&self, probe: Probe) -> bool {
110        self.checks.read().unwrap_or_else(|e| e.into_inner()).iter().any(|(p, _, _)| *p == probe)
111    }
112
113    /// Run `probe`'s checks now, concurrently, each within the timeout.
114    /// Readiness also fails while the server isn't ready (starting or
115    /// shutting down); startup, once passed, stays passed.
116    pub fn report(&self, probe: Probe) -> Report {
117        let checks: Vec<(String, Check)> = self
118            .checks
119            .read()
120            .unwrap_or_else(|e| e.into_inner())
121            .iter()
122            .filter(|(p, _, _)| *p == probe)
123            .map(|(_, n, c)| (n.clone(), Arc::clone(c)))
124            .collect();
125        if probe == Probe::Startup && self.started.load(Ordering::Relaxed) {
126            return Report { probe, ok: true, checks: checks.into_iter().map(|(name, _)| CheckResult { name, error: None }).collect(), note: None };
127        }
128        let timeout = Duration::from_millis(self.timeout_ms.load(Ordering::Relaxed));
129        let handles: Vec<(String, std::thread::JoinHandle<Result<(), String>>)> =
130            checks.into_iter().map(|(name, check)| (name, std::thread::spawn(move || run_one(check, timeout)))).collect();
131        let results: Vec<CheckResult> = handles
132            .into_iter()
133            .map(|(name, h)| CheckResult { name, error: h.join().unwrap_or_else(|_| Err("the check panicked".into())).err() })
134            .collect();
135        let mut note = None;
136        if probe == Probe::Readiness && !crate::metrics::SERVER_READY.load(Ordering::Relaxed) {
137            note = Some("server not ready (starting or shutting down)".to_string());
138        }
139        let ok = note.is_none() && results.iter().all(|c| c.error.is_none());
140        if probe == Probe::Startup && ok {
141            self.started.store(true, Ordering::Relaxed);
142        }
143        Report { probe, ok, checks: results, note }
144    }
145}
146
147/// Add a check to `probe` in the process-wide registry.
148pub fn register(probe: Probe, name: &str, check: impl Fn() -> Result<(), String> + Send + Sync + 'static) {
149    global().register(probe, name, check)
150}
151
152/// A liveness check (`/livez`, `/healthz`).
153pub fn liveness(name: &str, check: impl Fn() -> Result<(), String> + Send + Sync + 'static) {
154    register(Probe::Liveness, name, check)
155}
156
157/// A readiness check (`/readyz`).
158pub fn readiness(name: &str, check: impl Fn() -> Result<(), String> + Send + Sync + 'static) {
159    register(Probe::Readiness, name, check)
160}
161
162/// A startup check (`/startupz`).
163pub fn startup(name: &str, check: impl Fn() -> Result<(), String> + Send + Sync + 'static) {
164    register(Probe::Startup, name, check)
165}
166
167/// An async check, awaited on the shared runtime (`http2`).
168#[cfg(feature = "http2")]
169pub fn register_async<F, Fut>(probe: Probe, name: &str, check: F)
170where
171    F: Fn() -> Fut + Send + Sync + 'static,
172    Fut: std::future::Future<Output = Result<(), String>> + Send + 'static,
173{
174    register(probe, name, move || crate::async_bridge::block_on_shared(check()))
175}
176
177/// A readiness check that the database answers `SELECT 1` (`model-*`).
178#[cfg(any(feature = "model-sqlite", feature = "model-postgres", feature = "model-mysql"))]
179pub fn database(name: &str, pool: crate::model::DbPool) {
180    register_async(Probe::Readiness, name, move || {
181        let pool = pool.clone();
182        async move { pool.query_rows("SELECT 1", &[]).await.map(|_| ()).map_err(|e| e.to_string()) }
183    })
184}
185
186/// Remove a check from the process-wide registry.
187pub fn unregister(probe: Probe, name: &str) -> bool {
188    global().unregister(probe, name)
189}
190
191/// The process-wide check timeout (default 2 s).
192pub fn set_timeout(timeout: Duration) {
193    global().set_timeout(timeout)
194}
195
196/// Whether `probe` has checks in the process-wide registry.
197pub fn has_checks(probe: Probe) -> bool {
198    global().has_checks(probe)
199}
200
201/// One check's outcome.
202#[derive(Debug, Clone, PartialEq, Eq)]
203pub struct CheckResult {
204    pub name: String,
205    /// `None`: passed; `Some(message)`: failed.
206    pub error: Option<String>,
207}
208
209/// A probe's outcome.
210#[derive(Debug, Clone, PartialEq, Eq)]
211pub struct Report {
212    pub probe: Probe,
213    pub ok: bool,
214    pub checks: Vec<CheckResult>,
215    /// Why it failed besides its checks (`server not ready`).
216    pub note: Option<String>,
217}
218
219impl Report {
220    /// The probe's body: `OK`, or one line per failed check; with
221    /// `verbose`, one line per check and a summary.
222    pub fn body(&self, verbose: bool) -> String {
223        let mut lines = Vec::new();
224        if let Some(note) = &self.note {
225            lines.push(format!("[-]{note}"));
226        }
227        for c in &self.checks {
228            match &c.error {
229                None if verbose => lines.push(format!("[+]{} ok", c.name)),
230                None => {}
231                Some(e) => lines.push(format!("[-]{} failed: {e}", c.name)),
232            }
233        }
234        if verbose {
235            let what = &self.probe.path()[1..];
236            lines.push(if self.ok { format!("{what} check passed") } else { format!("{what} check failed") });
237            lines.join("\n")
238        } else if self.ok {
239            "OK".to_string()
240        } else {
241            lines.join("\n")
242        }
243    }
244}
245
246fn run_one(check: Check, timeout: Duration) -> Result<(), String> {
247    let (tx, rx) = std::sync::mpsc::channel();
248    let spawned = std::thread::Builder::new().name("rws-health".into()).spawn(move || {
249        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| check())).unwrap_or_else(|_| Err("the check panicked".into()));
250        let _ = tx.send(result);
251    });
252    if spawned.is_err() {
253        return Err("could not start the check".into());
254    }
255    match rx.recv_timeout(timeout) {
256        Ok(result) => result,
257        Err(_) => Err(format!("timed out after {} ms", timeout.as_millis())),
258    }
259}
260
261/// Run `probe`'s checks in the process-wide registry.
262pub fn report(probe: Probe) -> Report {
263    global().report(probe)
264}
265
266/// The probe a request path asks for (`/healthz` is liveness).
267pub fn probe_for_path(path: &str) -> Option<Probe> {
268    match path {
269        "/livez" | "/healthz" => Some(Probe::Liveness),
270        "/readyz" => Some(Probe::Readiness),
271        "/startupz" => Some(Probe::Startup),
272        _ => None,
273    }
274}
275
276/// The response for a probe request: `200` or `503`, text.
277pub fn respond(probe: Probe, verbose: bool) -> crate::response::Response {
278    use crate::response::{Response, STATUS_CODE_REASON_PHRASE};
279    let report = report(probe);
280    let status = if report.ok { STATUS_CODE_REASON_PHRASE.n200_ok } else { STATUS_CODE_REASON_PHRASE.n503_service_unavailable };
281    let body = report.body(verbose);
282    let body = if !report.ok && body.is_empty() { "not ready".to_string() } else { body };
283    let mut response = Response::text(status, &body);
284    response.headers.push(crate::header::Header::new(crate::header::Header::_CACHE_CONTROL, "no-store"));
285    response
286}
287
rws_server/health/mod.rs

rws_server/health/
mod.rs
