Skip to main content

siderite_core/health/
ready.rs

1//! Opt-in readiness with application checks:
2//! `App::routes([readiness("/readyz", checks)])`.
3//!
4//! The route answers `200` only while the server is Ready and every
5//! registered check passes, `503` otherwise. Checks are bounded so probe
6//! traffic cannot become load: each has a timeout, at most
7//! [`ReadinessChecks::concurrency`] run at once, and one evaluation is
8//! shared by every request until it is older than the cache period
9//! (concurrent probes wait for the evaluation in progress instead of
10//! starting their own). Checks reuse the application's existing clients
11//! and pools. The public body names each check with `ok`, `failed` or
12//! `timeout` only; failure details go to the log at `warn` level.
13
14use crate::http::StatusCode;
15use crate::responses::WithHeaders;
16use crate::{Json, MethodRouter, Readiness, Route, WithStatus, get};
17use futures_util::StreamExt;
18use futures_util::future::BoxFuture;
19use serde::Serialize;
20use serde_json::json;
21use siderite_orm::Db;
22use siderite_validation::{Dump, Schema, SchemaObject, SchemaRegistry};
23use std::collections::BTreeMap;
24use std::future::Future;
25use std::sync::Arc;
26use std::time::{Duration, Instant};
27use tokio::sync::Mutex;
28
29/// What one check failing means for the log; never sent to clients.
30pub type CheckError = String;
31
32/// What a check returns.
33pub type CheckFuture = BoxFuture<'static, Result<(), CheckError>>;
34
35type CheckFn = Arc<dyn Fn() -> CheckFuture + Send + Sync>;
36
37/// The outcome of one check.
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
39#[serde(rename_all = "lowercase")]
40pub enum CheckStatus {
41    /// It passed.
42    Ok,
43    /// It returned an error.
44    Failed,
45    /// It did not finish within its timeout.
46    Timeout,
47}
48
49/// The readiness answer.
50#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
51pub struct ReadyReport {
52    /// `ready` or `unavailable`.
53    pub status: &'static str,
54    /// Each check's outcome, by name; empty while the server is not
55    /// Ready (no check runs then).
56    pub checks: BTreeMap<String, CheckStatus>,
57}
58
59impl ReadyReport {
60    fn is_ready(&self) -> bool {
61        self.status == "ready"
62    }
63}
64
65impl Dump for ReadyReport {}
66
67impl Schema for ReadyReport {
68    fn schema_name() -> Option<&'static str> {
69        Some("ReadyReport")
70    }
71
72    fn schema(_: &mut SchemaRegistry) -> SchemaObject {
73        let outcome = json!({ "enum": ["ok", "failed", "timeout"] });
74        let properties = json!({
75            "status": { "enum": ["ready", "unavailable"] },
76            "checks": { "type": "object", "additionalProperties": outcome },
77        });
78        SchemaObject::of_type("object")
79            .with("properties", properties)
80            .with("required", json!(["status", "checks"]))
81    }
82}
83
84/// Named readiness checks and their bounds.
85#[derive(Clone)]
86pub struct ReadinessChecks {
87    checks: Vec<(String, CheckFn)>,
88    timeout: Duration,
89    concurrency: usize,
90    cache: Duration,
91    last: Arc<Mutex<Option<(Instant, ReadyReport)>>>,
92}
93
94impl std::fmt::Debug for ReadinessChecks {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        let names = self.checks.iter().map(|(n, _)| n.as_str());
97        let names: Vec<&str> = names.collect();
98        f.debug_struct("ReadinessChecks")
99            .field("checks", &names)
100            .field("timeout", &self.timeout)
101            .field("concurrency", &self.concurrency)
102            .field("cache", &self.cache)
103            .finish_non_exhaustive()
104    }
105}
106
107impl Default for ReadinessChecks {
108    fn default() -> Self {
109        Self::new()
110    }
111}
112
113impl ReadinessChecks {
114    /// No checks; a 2 s timeout, 4 at a time, results kept for 1 s.
115    pub fn new() -> Self {
116        Self {
117            checks: Vec::new(),
118            timeout: Duration::from_secs(2),
119            concurrency: 4,
120            cache: Duration::from_secs(1),
121            last: Arc::default(),
122        }
123    }
124
125    /// Add a check named `name` (shown to clients, so keep it generic:
126    /// `database`, not a host name).
127    ///
128    /// Args:
129    ///   name: The check's public name.
130    ///   check: Returns `Ok(())` when the dependency is usable; the error
131    ///     text is logged, never returned.
132    #[must_use]
133    pub fn check<F, Fut>(mut self, name: impl Into<String>, check: F) -> Self
134    where
135        F: Fn() -> Fut + Send + Sync + 'static,
136        Fut: Future<Output = Result<(), CheckError>> + Send + 'static,
137    {
138        let run: CheckFn = Arc::new(move || Box::pin(check()));
139        self.checks.push((name.into(), run));
140        self
141    }
142
143    /// Give each check at most `timeout` (default 2 s).
144    #[must_use]
145    pub fn timeout(mut self, timeout: Duration) -> Self {
146        self.timeout = timeout;
147        self
148    }
149
150    /// Run at most `limit` checks at once (default 4, at least 1).
151    #[must_use]
152    pub fn concurrency(mut self, limit: usize) -> Self {
153        self.concurrency = limit.max(1);
154        self
155    }
156
157    /// Reuse an evaluation for `period` (default 1 s); zero evaluates on
158    /// every request, still one at a time.
159    #[must_use]
160    pub fn cache(mut self, period: Duration) -> Self {
161        self.cache = period;
162        self
163    }
164
165    /// The report for a server in `ready` state, evaluating the checks
166    /// unless a fresh enough evaluation exists.
167    pub async fn report(&self, ready: bool) -> ReadyReport {
168        if !ready {
169            return ReadyReport {
170                status: "unavailable",
171                checks: BTreeMap::new(),
172            };
173        }
174        // Holding the lock while evaluating makes concurrent probes wait
175        // for this evaluation instead of starting their own.
176        let mut last = self.last.lock().await;
177        if let Some((at, report)) = last.as_ref()
178            && at.elapsed() < self.cache
179        {
180            return report.clone();
181        }
182        let report = self.evaluate().await;
183        *last = Some((Instant::now(), report.clone()));
184        report
185    }
186
187    async fn evaluate(&self) -> ReadyReport {
188        let timeout = self.timeout;
189        // Owned, boxed futures: the stream holds no borrow of `self`.
190        let runs: Vec<BoxFuture<'static, (String, CheckStatus)>> = self
191            .checks
192            .iter()
193            .map(|(name, check)| run(name.clone(), Arc::clone(check), timeout))
194            .collect();
195        let stream = futures_util::stream::iter(runs);
196        let running = stream.buffer_unordered(self.concurrency);
197        let checks: BTreeMap<String, CheckStatus> = running.collect().await;
198        let ok = checks.values().all(|s| *s == CheckStatus::Ok);
199        ReadyReport {
200            status: if ok { "ready" } else { "unavailable" },
201            checks,
202        }
203    }
204}
205
206/// Run one check within `timeout`, logging why it did not pass.
207fn run(
208    name: String,
209    check: CheckFn,
210    timeout: Duration,
211) -> BoxFuture<'static, (String, CheckStatus)> {
212    Box::pin(async move {
213        let status = match tokio::time::timeout(timeout, check()).await {
214            Ok(Ok(())) => CheckStatus::Ok,
215            Ok(Err(error)) => {
216                tracing::warn!(check = %name, %error, "check failed");
217                CheckStatus::Failed
218            }
219            Err(_) => {
220                tracing::warn!(check = %name, "check timed out");
221                CheckStatus::Timeout
222            }
223        };
224        (name, status)
225    })
226}
227
228type Answer = WithHeaders<WithStatus<Json<ReadyReport>>>;
229
230const CACHE: &str = "cache-control";
231const NO_STORE: &str = "no-store";
232
233/// The readiness handler for `checks`, documented under the `health` tag
234/// with the operation ID `readiness`.
235pub fn readiness_router(checks: ReadinessChecks) -> MethodRouter {
236    let checks = Arc::new(checks);
237    let handler = move |server: Readiness| {
238        let checks = Arc::clone(&checks);
239        async move {
240            let report = checks.report(server.is_ready()).await;
241            let status = if report.is_ready() {
242                StatusCode::OK
243            } else {
244                StatusCode::SERVICE_UNAVAILABLE
245            };
246            let body = WithStatus(status, Json(report));
247            let answer: Answer = WithHeaders::new(body);
248            answer.header(CACHE, NO_STORE)
249        }
250    };
251    get(handler)
252        .operation_id("readiness")
253        .tag("health")
254        .summary("Whether the app is ready to receive traffic.")
255        .response_model::<ReadyReport>()
256}
257
258/// A `GET path` readiness route running `checks`.
259///
260/// Args:
261///   path: Where to answer, for example `/readyz`.
262///   checks: The checks and their bounds.
263///
264/// Returns:
265///   The route; nothing is registered until the app adds it.
266pub fn readiness(path: &'static str, checks: ReadinessChecks) -> Route {
267    Route::new(path, readiness_router(checks))
268}
269
270/// A check that the database answers `SELECT 1` through `db`'s existing
271/// pool. SQL backends only: on MongoDB or Redis it fails, so write a
272/// check with the client the app already has.
273pub fn sql_check(db: Db) -> impl Fn() -> CheckFuture + Send + Sync {
274    move || {
275        let db = db.clone();
276        Box::pin(async move {
277            db.raw_sql("SELECT 1", Vec::new())
278                .await
279                .map(drop)
280                .map_err(|e| e.to_string())
281        })
282    }
283}