1use 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
29pub type CheckError = String;
31
32pub type CheckFuture = BoxFuture<'static, Result<(), CheckError>>;
34
35type CheckFn = Arc<dyn Fn() -> CheckFuture + Send + Sync>;
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
39#[serde(rename_all = "lowercase")]
40pub enum CheckStatus {
41 Ok,
43 Failed,
45 Timeout,
47}
48
49#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
51pub struct ReadyReport {
52 pub status: &'static str,
54 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#[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 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 #[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 #[must_use]
145 pub fn timeout(mut self, timeout: Duration) -> Self {
146 self.timeout = timeout;
147 self
148 }
149
150 #[must_use]
152 pub fn concurrency(mut self, limit: usize) -> Self {
153 self.concurrency = limit.max(1);
154 self
155 }
156
157 #[must_use]
160 pub fn cache(mut self, period: Duration) -> Self {
161 self.cache = period;
162 self
163 }
164
165 pub async fn report(&self, ready: bool) -> ReadyReport {
168 if !ready {
169 return ReadyReport {
170 status: "unavailable",
171 checks: BTreeMap::new(),
172 };
173 }
174 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 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
206fn 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
233pub 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
258pub fn readiness(path: &'static str, checks: ReadinessChecks) -> Route {
267 Route::new(path, readiness_router(checks))
268}
269
270pub 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}