From eee057f09df9f0b7a624f25c953832d0675672e7 Mon Sep 17 00:00:00 2001 From: Spencer Ferris <3319370+spencewenski@users.noreply.github.com> Date: Mon, 1 Jul 2024 11:43:51 -0700 Subject: [PATCH] feat!: Remove interior mutability of `HealthCheckRegistry` (#258) Switch to using `OnceLock` to allow setting the health checks on the `AppContext` after it is created. The method to set the health checks is private to the crate, so we have control over when the checks are set on the context. This has some impacts on public-facing APIs: - `App::health_checks` takes a `&mut HealthCheckRegistry` instead of a non-mutable ref - `AppContext::health_checks` returns the list of health checks. (Alternatively, we could have returned an Option) - We also made the `HealthCheckRegistry::new` method non-public. This aligns with the `ServiceRegistry::new` visibility. Closes https://github.com/roadster-rs/roadster/issues/257 --- src/api/core/health.rs | 2 +- src/app/context.rs | 32 ++++++-- src/app/mod.rs | 16 +--- src/health_check/default.rs | 24 ++++-- src/health_check/registry.rs | 75 +++++++------------ ...ult__tests__default_middleware@case_1.snap | 6 +- 6 files changed, 75 insertions(+), 80 deletions(-) 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', -] +[]