Skip to main content

modo/health/
check.rs

1use std::pin::Pin;
2use std::sync::Arc;
3
4use crate::Result;
5
6/// A health check that can verify the readiness of a service.
7///
8/// Implement this trait for types that can verify their own health (e.g.,
9/// database pools, cache connections). The check should be fast and
10/// non-destructive.
11///
12/// When the `db` feature is enabled, [`crate::db::Database`] implements this
13/// trait automatically.
14pub trait HealthCheck: Send + Sync + 'static {
15    /// Run the health check.
16    ///
17    /// # Errors
18    ///
19    /// Returns [`crate::Error`] if the service is unhealthy or unreachable.
20    fn check(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>>;
21}
22
23/// Internal adapter that wraps a closure into a [`HealthCheck`].
24struct FnHealthCheck<F>(F);
25
26impl<F, Fut> HealthCheck for FnHealthCheck<F>
27where
28    F: Fn() -> Fut + Send + Sync + 'static,
29    Fut: Future<Output = Result<()>> + Send + 'static,
30{
31    fn check(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
32        Box::pin((self.0)())
33    }
34}
35
36/// A collection of named health checks.
37///
38/// Built with a fluent API and registered in the
39/// [`service::Registry`](crate::service::Registry). The readiness endpoint
40/// runs all checks concurrently and reports failures.
41///
42/// # Example
43///
44/// ```
45/// use modo::HealthChecks;
46///
47/// let checks = HealthChecks::new()
48///     .check_fn("database", || async { Ok(()) })
49///     .check_fn("redis", || async { Ok(()) });
50/// ```
51pub struct HealthChecks {
52    checks: Vec<(String, Arc<dyn HealthCheck>)>,
53}
54
55impl HealthChecks {
56    /// Creates an empty collection.
57    pub fn new() -> Self {
58        Self { checks: Vec::new() }
59    }
60
61    /// Register a named [`HealthCheck`] implementation under the given name.
62    pub fn check(mut self, name: &str, c: impl HealthCheck) -> Self {
63        self.checks.push((name.to_owned(), Arc::new(c)));
64        self
65    }
66
67    /// Register a named health check from an async closure.
68    ///
69    /// The closure must return [`crate::Result<()>`].
70    pub fn check_fn<F, Fut>(mut self, name: &str, f: F) -> Self
71    where
72        F: Fn() -> Fut + Send + Sync + 'static,
73        Fut: Future<Output = Result<()>> + Send + 'static,
74    {
75        self.checks
76            .push((name.to_owned(), Arc::new(FnHealthCheck(f))));
77        self
78    }
79
80    /// Returns a slice of all registered checks.
81    pub(crate) fn as_slice(&self) -> &[(String, Arc<dyn HealthCheck>)] {
82        &self.checks
83    }
84}
85
86/// Returns an empty [`HealthChecks`] collection.
87impl Default for HealthChecks {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93#[cfg(feature = "db")]
94impl HealthCheck for crate::db::Database {
95    fn check(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
96        Box::pin(async {
97            self.conn()
98                .query("SELECT 1", ())
99                .await
100                .map_err(|e| crate::Error::internal("db health check failed").chain(e))?;
101            Ok(())
102        })
103    }
104}
105
106#[cfg(all(test, feature = "db"))]
107mod tests {
108    use super::*;
109
110    #[tokio::test]
111    async fn database_health_check() {
112        let config = crate::db::Config {
113            path: ":memory:".to_string(),
114            ..Default::default()
115        };
116        let db = crate::db::connect(&config).await.unwrap();
117        db.check().await.unwrap();
118    }
119
120    #[tokio::test]
121    async fn fn_health_check_succeeds() {
122        let checks = HealthChecks::new().check_fn("ok", || async { Ok(()) });
123        let (_, c) = &checks.as_slice()[0];
124        c.check().await.unwrap();
125    }
126
127    #[tokio::test]
128    async fn fn_health_check_fails() {
129        let checks =
130            HealthChecks::new().check_fn("fail", || async { Err(crate::Error::internal("down")) });
131        let (_, c) = &checks.as_slice()[0];
132        assert!(c.check().await.is_err());
133    }
134
135    #[tokio::test]
136    async fn chaining_preserves_order() {
137        let checks = HealthChecks::new()
138            .check_fn("a", || async { Ok(()) })
139            .check_fn("b", || async { Ok(()) })
140            .check_fn("c", || async { Ok(()) });
141        let names: Vec<&str> = checks.as_slice().iter().map(|(n, _)| n.as_str()).collect();
142        assert_eq!(names, vec!["a", "b", "c"]);
143    }
144
145    #[tokio::test]
146    async fn health_checks_default_is_empty() {
147        let checks = HealthChecks::default();
148        assert!(checks.as_slice().is_empty());
149    }
150}