diff --git a/src/api/core/health.rs b/src/api/core/health.rs index 6f8158cb..7ddbf7ec 100644 --- a/src/api/core/health.rs +++ b/src/api/core/health.rs @@ -53,7 +53,7 @@ where let context = AppContext::from_ref(state); let timer = Instant::now(); - let check_futures = context.health_checks().checks()?.into_iter().map(|check| { + let check_futures = context.health_checks().into_iter().map(|check| { Box::pin(async move { let name = check.name(); info!(name=%name, "Running check"); diff --git a/src/app/context.rs b/src/app/context.rs index 560d3f27..35f72786 100644 --- a/src/app/context.rs +++ b/src/app/context.rs @@ -3,10 +3,12 @@ use crate::app::App; use crate::config::app_config::AppConfig; use crate::error::RoadsterResult; use crate::health_check::registry::HealthCheckRegistry; +use crate::health_check::HealthCheck; +use anyhow::anyhow; use axum::extract::FromRef; #[cfg(feature = "db-sql")] use sea_orm::DatabaseConnection; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; #[cfg(not(test))] type Inner = AppContextInner; @@ -75,7 +77,7 @@ impl AppContext { let inner = AppContextInner { config, metadata, - health_checks: HealthCheckRegistry::new(), + health_checks: OnceLock::new(), #[cfg(feature = "db-sql")] db, #[cfg(feature = "sidekiq")] @@ -127,10 +129,17 @@ impl AppContext { self.inner.metadata() } - pub fn health_checks(&self) -> &HealthCheckRegistry { + pub fn health_checks(&self) -> Vec> { self.inner.health_checks() } + pub(crate) fn set_health_checks( + &self, + health_checks: HealthCheckRegistry, + ) -> RoadsterResult<()> { + self.inner.set_health_checks(health_checks) + } + #[cfg(feature = "db-sql")] pub fn db(&self) -> &DatabaseConnection { self.inner.db() @@ -150,7 +159,7 @@ impl AppContext { struct AppContextInner { config: AppConfig, metadata: AppMetadata, - health_checks: HealthCheckRegistry, + health_checks: OnceLock, #[cfg(feature = "db-sql")] db: DatabaseConnection, #[cfg(feature = "sidekiq")] @@ -173,8 +182,19 @@ impl AppContextInner { &self.metadata } - fn health_checks(&self) -> &HealthCheckRegistry { - &self.health_checks + fn health_checks(&self) -> Vec> { + self.health_checks + .get() + .map(|health_checks| health_checks.checks()) + .unwrap_or_default() + } + + fn set_health_checks(&self, health_checks: HealthCheckRegistry) -> RoadsterResult<()> { + self.health_checks + .set(health_checks) + .map_err(|_| anyhow!("Unable to set health check registry"))?; + + Ok(()) } #[cfg(feature = "db-sql")] diff --git a/src/app/mod.rs b/src/app/mod.rs index 935826d6..6e1172fe 100644 --- a/src/app/mod.rs +++ b/src/app/mod.rs @@ -12,7 +12,6 @@ use crate::config::app_config::AppConfig; #[cfg(not(feature = "cli"))] use crate::config::environment::Environment; use crate::error::RoadsterResult; -use crate::health_check::default::default_health_checks; use crate::health_check::registry::HealthCheckRegistry; use crate::service::registry::ServiceRegistry; use crate::tracing::init_tracing; @@ -68,10 +67,9 @@ where let state = A::provide_state(context.clone()).await?; - default_health_checks(&context) - .into_iter() - .try_for_each(|check| context.health_checks().register_arc(check))?; - A::health_checks(context.health_checks(), &state).await?; + let mut health_checks = HealthCheckRegistry::new(&context); + A::health_checks(&mut health_checks, &state).await?; + context.set_health_checks(health_checks)?; #[cfg(feature = "cli")] if crate::api::cli::handle_cli(&app, &roadster_cli, &app_cli, &state).await? { @@ -142,13 +140,7 @@ where async fn provide_state(context: AppContext) -> RoadsterResult; /// Provide the [crate::health_check::HealthCheck]s to use throughout the app. - /// - /// Note that a non-mutable reference to the [HealthCheckRegistry] is provided. This is because - /// [HealthCheckRegistry] implements the - /// [interior mutability](https://doc.rust-lang.org/reference/interior-mutability.html) pattern. - /// As such, ___it is not recommended to register additional health checks outside of - /// this method___ -- doing so may result in a panic. - async fn health_checks(_registry: &HealthCheckRegistry, _state: &S) -> RoadsterResult<()> { + async fn health_checks(_registry: &mut HealthCheckRegistry, _state: &S) -> RoadsterResult<()> { Ok(()) } diff --git a/src/health_check/default.rs b/src/health_check/default.rs index eedccfec..2a47f96f 100644 --- a/src/health_check/default.rs +++ b/src/health_check/default.rs @@ -6,12 +6,13 @@ use crate::health_check::sidekiq_enqueue::SidekiqEnqueueHealthCheck; #[cfg(feature = "sidekiq")] use crate::health_check::sidekiq_fetch::SidekiqFetchHealthCheck; use crate::health_check::HealthCheck; +use std::collections::BTreeMap; use std::sync::Arc; pub fn default_health_checks( #[allow(unused_variables)] context: &AppContext, -) -> Vec> { - vec![ +) -> BTreeMap> { + let health_checks: Vec> = vec![ #[cfg(feature = "db-sql")] Arc::new(DatabaseHealthCheck { context: context.clone(), @@ -24,7 +25,12 @@ pub fn default_health_checks( Arc::new(SidekiqFetchHealthCheck { context: context.clone(), }), - ] + ]; + health_checks + .into_iter() + .filter(|check| check.enabled()) + .map(|check| (check.name(), check)) + .collect() } #[cfg(all(test, feature = "sidekiq", feature = "db-sql",))] @@ -32,9 +38,11 @@ mod tests { use crate::app::context::AppContext; use crate::config::app_config::AppConfig; use crate::util::test_util::TestCase; + use bb8::Pool; use insta::assert_toml_snapshot; use itertools::Itertools; use rstest::{fixture, rstest}; + use sidekiq::RedisConnectionManager; #[fixture] #[cfg_attr(coverage_nightly, coverage(off))] @@ -45,17 +53,21 @@ mod tests { #[rstest] #[case(false)] #[case(true)] + #[tokio::test] #[cfg_attr(coverage_nightly, coverage(off))] - fn default_middleware(_case: TestCase, #[case] default_enable: bool) { + async fn default_middleware(_case: TestCase, #[case] default_enable: bool) { // Arrange let mut config = AppConfig::test(None).unwrap(); config.health_check.default_enable = default_enable; - let context = AppContext::test(Some(config), None, None).unwrap(); + let redis_fetch = RedisConnectionManager::new("redis://invalid_host:1234").unwrap(); + let pool = Pool::builder().build_unchecked(redis_fetch); + + let context = AppContext::test(Some(config), None, Some(pool)).unwrap(); // Act let health_checks = super::default_health_checks(&context); - let health_checks = health_checks.iter().map(|check| check.name()).collect_vec(); + let health_checks = health_checks.keys().collect_vec(); // Assert assert_toml_snapshot!(health_checks); diff --git a/src/health_check/registry.rs b/src/health_check/registry.rs index a7214374..2d755785 100644 --- a/src/health_check/registry.rs +++ b/src/health_check/registry.rs @@ -1,8 +1,10 @@ +use crate::app::context::AppContext; use crate::error::RoadsterResult; +use crate::health_check::default::default_health_checks; use crate::health_check::HealthCheck; use anyhow::anyhow; use std::collections::BTreeMap; -use std::sync::{Arc, RwLock}; +use std::sync::Arc; use tracing::info; /// Registry for [HealthCheck]s that will be run in the app. @@ -11,46 +13,21 @@ use tracing::info; /// 1. As pre-boot checks to ensure the app's resource dependencies are healthy. /// 2. As a "core" API that can be used from multiple components, e.g. the `_health` HTTP endpoint /// and the health CLI command. -/// -/// # Internal mutability -/// In order to make this registry available to multiple parts of the app, this is included -/// as part of the [AppContext][crate::app::context::AppContext]. This is not strictly necessary -/// for the Axum handlers (the registry could be provided via an [Extension][axum::Extension]), -/// but it is (currently) required for other things, such as the CLI handlers. -/// -/// In order to include the registry as part of the context, but also allow checks to be added -/// to the registry after the context is created, the registry implements the -/// [interior mutability](https://doc.rust-lang.org/reference/interior-mutability.html) pattern -/// using a [RwLock]. As such, ___it is not recommended to register additional health checks -/// outside of the app initialization process___ -- doing so may result in a panic. -/// -/// Because of the internal mutability, methods that modify the internal state can accept `&self` -/// instead of `&mut self`. pub struct HealthCheckRegistry { - health_checks: Arc>>>, -} - -impl Default for HealthCheckRegistry { - fn default() -> Self { - HealthCheckRegistry::new() - } + health_checks: BTreeMap>, } impl HealthCheckRegistry { - pub fn new() -> Self { + pub(crate) fn new(context: &AppContext) -> Self { Self { - health_checks: Arc::new(RwLock::new(Default::default())), + health_checks: default_health_checks(context), } } - pub fn register(&self, health_check: H) -> RoadsterResult<()> + pub fn register(&mut self, health_check: H) -> RoadsterResult<()> where H: HealthCheck + 'static, { - self.register_arc(Arc::new(health_check)) - } - - pub(crate) fn register_arc(&self, health_check: Arc) -> RoadsterResult<()> { let name = health_check.name(); if !health_check.enabled() { @@ -60,27 +37,25 @@ impl HealthCheckRegistry { info!(name=%name, "Registering health check"); - let mut health_checks = self.health_checks.write().map_err(|err| { - anyhow!("Unable to acquire write lock on health check registry: {err}") - })?; - if health_checks.insert(name.clone(), health_check).is_some() { + if self + .health_checks + .insert(name.clone(), Arc::new(health_check)) + .is_some() + { return Err(anyhow!("Health check `{}` was already registered", name).into()); } Ok(()) } - pub fn checks(&self) -> RoadsterResult>> { - let health_checks = self - .health_checks - .read() - .map_err(|err| anyhow!("Unable to acquire read lock on heath check registry: {err}"))?; - Ok(health_checks.values().cloned().collect()) + pub fn checks(&self) -> Vec> { + self.health_checks.values().cloned().collect() } } #[cfg(test)] mod tests { use super::*; + use crate::config::app_config::AppConfig; use crate::health_check::MockHealthCheck; use rstest::rstest; @@ -88,25 +63,25 @@ mod tests { #[case(true, 1)] #[case(false, 0)] #[cfg_attr(coverage_nightly, coverage(off))] - fn register_check(#[case] service_enabled: bool, #[case] expected_count: usize) { + fn register_check(#[case] check_enabled: bool, #[case] expected_count: usize) { // Arrange + let mut config = AppConfig::test(None).unwrap(); + config.health_check.default_enable = false; + let context = AppContext::test(Some(config), None, None).unwrap(); + let mut check: MockHealthCheck = MockHealthCheck::default(); - check.expect_enabled().return_const(service_enabled); + check.expect_enabled().return_const(check_enabled); check.expect_name().return_const("test".to_string()); // Act - let subject: HealthCheckRegistry = HealthCheckRegistry::new(); + let mut subject: HealthCheckRegistry = HealthCheckRegistry::new(&context); subject.register(check).unwrap(); // Assert - assert_eq!(subject.checks().unwrap().len(), expected_count); + assert_eq!(subject.checks().len(), expected_count); assert_eq!( - subject - .checks() - .unwrap() - .iter() - .any(|check| check.name() == "test"), - service_enabled + subject.checks().iter().any(|check| check.name() == "test"), + check_enabled ); } } diff --git a/src/health_check/snapshots/roadster__health_check__default__tests__default_middleware@case_1.snap b/src/health_check/snapshots/roadster__health_check__default__tests__default_middleware@case_1.snap index 1bd9b8b9..c3b7bad8 100644 --- a/src/health_check/snapshots/roadster__health_check__default__tests__default_middleware@case_1.snap +++ b/src/health_check/snapshots/roadster__health_check__default__tests__default_middleware@case_1.snap @@ -2,8 +2,4 @@ source: src/health_check/default.rs expression: health_checks --- -[ - 'db', - 'sidekiq-enqueue', - 'sidekiq-fetch', -] +[]