rust_zero_core/
health.rs

1use std::{
2    collections::BTreeMap,
3    sync::{Arc, Mutex},
4};
5use tokio::sync::watch;
6
7/// Cloneable aggregate of named dependency readiness states.
8///
9/// An empty registry is healthy. Once dependencies are registered, the aggregate is ready only
10/// when every dependency is ready. Updates are watchable so transports can project health without
11/// polling application code.
12#[derive(Debug, Clone)]
13pub struct HealthRegistry {
14    state: Arc<Mutex<BTreeMap<String, bool>>>,
15    updates: watch::Sender<HealthSnapshot>,
16}
17
18#[derive(Debug, Clone, Default, PartialEq, Eq)]
19pub struct HealthSnapshot {
20    pub dependencies: BTreeMap<String, bool>,
21}
22
23impl HealthSnapshot {
24    pub fn is_ready(&self) -> bool {
25        self.dependencies.values().all(|ready| *ready)
26    }
27
28    pub fn unhealthy(&self) -> Vec<String> {
29        self.dependencies
30            .iter()
31            .filter_map(|(name, ready)| (!ready).then_some(name.clone()))
32            .collect()
33    }
34}
35
36impl Default for HealthRegistry {
37    fn default() -> Self {
38        let snapshot = HealthSnapshot::default();
39        let (updates, _) = watch::channel(snapshot);
40        Self {
41            state: Arc::new(Mutex::new(BTreeMap::new())),
42            updates,
43        }
44    }
45}
46
47impl HealthRegistry {
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    pub fn set(&self, name: impl Into<String>, ready: bool) {
53        let mut state = self.state.lock().expect("health registry mutex poisoned");
54        state.insert(name.into(), ready);
55        self.updates.send_replace(HealthSnapshot {
56            dependencies: state.clone(),
57        });
58    }
59
60    pub fn remove(&self, name: &str) {
61        let mut state = self.state.lock().expect("health registry mutex poisoned");
62        state.remove(name);
63        self.updates.send_replace(HealthSnapshot {
64            dependencies: state.clone(),
65        });
66    }
67
68    pub fn snapshot(&self) -> HealthSnapshot {
69        self.updates.borrow().clone()
70    }
71
72    pub fn subscribe(&self) -> watch::Receiver<HealthSnapshot> {
73        self.updates.subscribe()
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[tokio::test]
82    async fn aggregates_and_publishes_dependency_readiness() {
83        let registry = HealthRegistry::new();
84        let mut updates = registry.subscribe();
85        assert!(registry.snapshot().is_ready());
86
87        registry.set("users", false);
88        updates.changed().await.unwrap();
89        assert_eq!(updates.borrow().unhealthy(), vec!["users"]);
90
91        registry.set("users", true);
92        updates.changed().await.unwrap();
93        assert!(updates.borrow().is_ready());
94    }
95}
rust_zero_core/health.rs

rust_zero_core/
health.rs
