diff --git a/foundations/src/dataloader/mod.rs b/foundations/src/dataloader/mod.rs index dc048fb3..23658cbf 100644 --- a/foundations/src/dataloader/mod.rs +++ b/foundations/src/dataloader/mod.rs @@ -7,8 +7,6 @@ use std::sync::atomic::AtomicU64; use std::sync::Arc; use std::time::Duration; -use futures::FutureExt; - use self::batch_loader::BatchLoader; pub use self::types::LoaderOutput; use self::types::{BatchState, DataLoaderInner}; diff --git a/foundations/src/telementry/server.rs b/foundations/src/telementry/server.rs index 088d81ae..ece4bacd 100644 --- a/foundations/src/telementry/server.rs +++ b/foundations/src/telementry/server.rs @@ -160,46 +160,46 @@ async fn metrics( } #[cfg(feature = "health-check")] -pub fn register_health_check(check: F) -> usize -where - F: Fn() -> Fut + Send + Sync + 'static, - Fut: std::future::Future + Send + Sync + 'static, -{ - register_health_check_boxed(Box::new(move || Box::pin(check()))) -} - -#[cfg(feature = "health-check")] -pub fn register_health_check_boxed(check: health_check::BoxHealthCheck) -> usize { - health_check::register(check) -} - -#[cfg(feature = "health-check")] -pub fn unregister_health_check(id: usize) { - health_check::unregister(id); -} +pub use health_check::{register as register_health_check, unregister as unregister_health_check, HealthCheck}; #[cfg(feature = "health-check")] mod health_check { use std::pin::Pin; use std::sync::atomic::AtomicUsize; + use futures::Future; use scc::HashMap; - type BoxFut = Pin + Send + Sync>>; - pub(super) type BoxHealthCheck = Box BoxFut + Send + Sync>; + pub trait HealthCheck: Send + Sync + 'static { + fn check(&self) -> Pin + Send + '_>>; + } + + impl HealthCheck for F + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future + Send + Sync + 'static, + { + fn check(&self) -> Pin + Send + '_>> { + Box::pin((self)()) + } + } #[derive(Default)] - struct HealthCheck { + struct HealthChecker { id: AtomicUsize, - health_checks: HashMap, + health_checks: HashMap>, } - static HEALTH_CHECK: once_cell::sync::Lazy = - once_cell::sync::Lazy::::new(|| HealthCheck::default()); + static HEALTH_CHECK: once_cell::sync::Lazy = + once_cell::sync::Lazy::::new(|| HealthChecker::default()); - pub fn register(check: BoxHealthCheck) -> usize { + pub fn register(check: impl HealthCheck) -> usize { let id = HEALTH_CHECK.id.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - HEALTH_CHECK.health_checks.insert(id, check).ok().expect("id already exists"); + HEALTH_CHECK + .health_checks + .insert(id, Box::new(check)) + .ok() + .expect("id already exists"); id } @@ -211,7 +211,7 @@ mod health_check { let mut o_entry = HEALTH_CHECK.health_checks.first_entry_async().await; while let Some(entry) = o_entry { - if (entry.get())().await { + if (entry.get()).check().await { return false; }