1use std::pin::Pin;
2use std::sync::Arc;
3
4use crate::Result;
5
6pub trait HealthCheck: Send + Sync + 'static {
15 fn check(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>>;
21}
22
23struct 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
36pub struct HealthChecks {
52 checks: Vec<(String, Arc<dyn HealthCheck>)>,
53}
54
55impl HealthChecks {
56 pub fn new() -> Self {
58 Self { checks: Vec::new() }
59 }
60
61 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 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 pub(crate) fn as_slice(&self) -> &[(String, Arc<dyn HealthCheck>)] {
82 &self.checks
83 }
84}
85
86impl 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}