1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
//! Health-check abstraction for operational readiness probes.
//!
//! Implement [`HealthCheck`] on your session stores, registries, or any other
//! backend to expose a `/healthz`-style endpoint. Use [`CompositeHealthCheck`]
//! to aggregate multiple backends into a single probe.
//!
//! # Example
//!
//! ```rust,ignore
//! use axess_core::health::{CompositeHealthCheck, HealthCheck};
//!
//! let health = CompositeHealthCheck::new()
//! .add("session_store", store.clone())
//! .add("session_registry", registry.clone());
//!
//! // In an Axum handler:
//! let status = health.check_all().await;
//! if status.is_healthy() { /* 200 */ } else { /* 503 */ }
//! ```
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
/// Result of a single health check.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HealthStatus {
/// The component is operating normally.
Healthy,
/// The component is functional but experiencing issues (e.g. high latency).
Degraded(String),
/// The component is not functional.
Unhealthy(String),
}
impl HealthStatus {
/// Returns `true` if the status is [`Healthy`](HealthStatus::Healthy).
pub fn is_healthy(&self) -> bool {
matches!(self, Self::Healthy)
}
/// Returns `true` if the status is [`Unhealthy`](HealthStatus::Unhealthy).
pub fn is_unhealthy(&self) -> bool {
matches!(self, Self::Unhealthy(_))
}
}
/// Aggregated result of checking multiple components.
#[derive(Debug, Clone)]
pub struct CompositeStatus {
/// Per-component results, in the order they were registered.
pub components: Vec<(String, HealthStatus)>,
}
impl CompositeStatus {
/// Returns `true` if every component is [`Healthy`](HealthStatus::Healthy).
pub fn is_healthy(&self) -> bool {
self.components.iter().all(|(_, s)| s.is_healthy())
}
/// Returns `true` if any component is [`Unhealthy`](HealthStatus::Unhealthy).
pub fn has_unhealthy(&self) -> bool {
self.components.iter().any(|(_, s)| s.is_unhealthy())
}
/// Returns only the aggregate status without component details.
///
/// Use this for external-facing endpoints to avoid leaking infrastructure
/// information (component names, error messages) to unauthenticated callers.
///
/// The message is fixed text: it never names a component or repeats a
/// component's own message.
pub fn for_external(&self) -> HealthStatus {
match self.worst() {
HealthStatus::Healthy => HealthStatus::Healthy,
HealthStatus::Degraded(_) => {
HealthStatus::Degraded("one or more components degraded".to_string())
}
HealthStatus::Unhealthy(_) => {
HealthStatus::Unhealthy("one or more components unhealthy".to_string())
}
}
}
/// Returns the worst status across all components.
pub fn worst(&self) -> HealthStatus {
if self.has_unhealthy() {
let msgs: Vec<&str> = self
.components
.iter()
.filter_map(|(name, s)| match s {
HealthStatus::Unhealthy(_) => Some(name.as_str()),
_ => None,
})
.collect();
HealthStatus::Unhealthy(format!("unhealthy: {}", msgs.join(", ")))
} else if self
.components
.iter()
.any(|(_, s)| matches!(s, HealthStatus::Degraded(_)))
{
HealthStatus::Degraded("one or more components degraded".to_string())
} else {
HealthStatus::Healthy
}
}
}
/// Async health check for a backend component.
///
/// Implement this on session stores, registries, or any other
/// infrastructure dependency that should be monitored.
pub trait HealthCheck: Send + Sync {
/// Probe the component and return its current health.
fn check(&self) -> Pin<Box<dyn Future<Output = HealthStatus> + Send + '_>>;
}
/// Aggregates multiple [`HealthCheck`] implementations into a single probe.
///
/// Use this to build a composite `/healthz` endpoint that checks all backends.
#[derive(Clone, Default)]
pub struct CompositeHealthCheck {
checks: Vec<(String, Arc<dyn HealthCheck>)>,
}
impl CompositeHealthCheck {
/// Create an empty composite health check.
pub fn new() -> Self {
Self::default()
}
/// Register a named health check.
///
/// If a check with the same name already exists, it is replaced.
pub fn add(mut self, name: impl Into<String>, check: impl HealthCheck + 'static) -> Self {
let name = name.into();
self.checks.retain(|(n, _)| n != &name);
self.checks.push((name, Arc::new(check)));
self
}
/// Run all registered checks and return the aggregate status.
///
/// Checks are run sequentially to avoid pulling in a `futures` dependency.
/// Health probes are typically cheap (PING, SELECT 1) so this is fine for
/// a small number of backends.
pub async fn check_all(&self) -> CompositeStatus {
let mut components = Vec::with_capacity(self.checks.len());
for (name, check) in &self.checks {
let status = check.check().await;
components.push((name.clone(), status));
}
CompositeStatus { components }
}
}
#[cfg(test)]
mod tests {
use super::*;
struct AlwaysHealthy;
impl HealthCheck for AlwaysHealthy {
fn check(&self) -> Pin<Box<dyn Future<Output = HealthStatus> + Send + '_>> {
Box::pin(async { HealthStatus::Healthy })
}
}
struct AlwaysUnhealthy;
impl HealthCheck for AlwaysUnhealthy {
fn check(&self) -> Pin<Box<dyn Future<Output = HealthStatus> + Send + '_>> {
Box::pin(async { HealthStatus::Unhealthy("down".into()) })
}
}
#[test]
fn health_status_predicates() {
assert!(HealthStatus::Healthy.is_healthy());
assert!(!HealthStatus::Healthy.is_unhealthy());
assert!(!HealthStatus::Degraded("slow".into()).is_healthy());
assert!(HealthStatus::Unhealthy("down".into()).is_unhealthy());
}
#[tokio::test]
async fn composite_all_healthy() {
let health = CompositeHealthCheck::new()
.add("a", AlwaysHealthy)
.add("b", AlwaysHealthy);
let status = health.check_all().await;
assert!(status.is_healthy());
assert!(status.worst().is_healthy());
}
#[tokio::test]
async fn composite_one_unhealthy() {
let health = CompositeHealthCheck::new()
.add("ok", AlwaysHealthy)
.add("down", AlwaysUnhealthy);
let status = health.check_all().await;
assert!(!status.is_healthy());
assert!(status.has_unhealthy());
assert!(status.worst().is_unhealthy());
}
#[tokio::test]
async fn for_external_names_no_component() {
let status = CompositeHealthCheck::new()
.add("postgres-primary", AlwaysUnhealthy)
.check_all()
.await;
assert_eq!(
status.for_external(),
HealthStatus::Unhealthy("one or more components unhealthy".to_string())
);
let shown = format!("{:?}", status.for_external());
assert!(!shown.contains("postgres") && !shown.contains("down"));
}
#[tokio::test]
async fn composite_empty_is_healthy() {
let health = CompositeHealthCheck::new();
let status = health.check_all().await;
assert!(status.is_healthy());
}
/// `worst()` rolls the names of all `Unhealthy` components
/// into the aggregate message so operators can pinpoint which
/// backend is down. Pins `delete match arm Unhealthy(_)` in the
/// `filter_map`; without that arm, every component would fall
/// through to the `_ => None` branch and the rolled-up message
/// would contain no names ("unhealthy: ").
#[tokio::test]
async fn composite_worst_message_lists_unhealthy_component_names() {
let health = CompositeHealthCheck::new()
.add("primary-db", AlwaysUnhealthy)
.add("cache", AlwaysHealthy)
.add("queue", AlwaysUnhealthy);
let status = health.check_all().await;
let worst = status.worst();
let HealthStatus::Unhealthy(msg) = worst else {
panic!("worst() must return Unhealthy when components are down");
};
assert!(
msg.contains("primary-db"),
"worst() message must list 'primary-db' (got: {msg})"
);
assert!(
msg.contains("queue"),
"worst() message must list 'queue' (got: {msg})"
);
// Healthy components must NOT be in the rolled-up message.
assert!(
!msg.contains("cache"),
"worst() must not name healthy components (got: {msg})"
);
}
/// Registering a check with a name that already exists
/// replaces the prior registration. Pins `!=` → `==` on the
/// `retain` predicate at line 126; the mutant flips dedup so the
/// new entry would *replace everything else* (keep only entries
/// whose name matches the new one), making the per-name semantics
/// silently break for any composite that registers more than one
/// check.
#[tokio::test]
async fn composite_add_replaces_existing_name_and_keeps_others() {
let health = CompositeHealthCheck::new()
.add("db", AlwaysHealthy)
.add("cache", AlwaysHealthy)
// Re-register `db` with an Unhealthy stub; must replace the
// prior `db` entry, not other entries.
.add("db", AlwaysUnhealthy);
let status = health.check_all().await;
// Original behavior: ["db" -> Unhealthy, "cache" -> Healthy].
// Mutant `==` would keep only entries whose name == new name,
// so adding "cache" would wipe out "db", and re-adding "db"
// would wipe out "cache", leaving exactly ["db" -> Unhealthy].
assert_eq!(
status.components.len(),
2,
"dedup must replace same-named entry, not wipe others (got: {:?})",
status.components
);
let names: Vec<_> = status.components.iter().map(|(n, _)| n.as_str()).collect();
assert!(names.contains(&"db"), "'db' must be present");
assert!(names.contains(&"cache"), "'cache' must remain");
// The 'db' entry must be the new (Unhealthy) registration, not stale.
let db_status = status
.components
.iter()
.find(|(n, _)| n == "db")
.map(|(_, s)| s.clone())
.unwrap();
assert!(
db_status.is_unhealthy(),
"'db' must reflect the most recent registration (Unhealthy)"
);
}
}