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 #257

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");

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<Arc<dyn HealthCheck>> {
         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<HealthCheckRegistry>,
     #[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<Arc<dyn HealthCheck>> {
+        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")]

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<S>;
 
     /// 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(())
     }
 

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<Arc<dyn HealthCheck>> {
-    vec![
+) -> BTreeMap<String, Arc<dyn HealthCheck>> {
+    let health_checks: Vec<Arc<dyn HealthCheck>> = vec![
         #[cfg(feature = "db-sql")]
         Arc::new(DatabaseHealthCheck {
             context: context.clone(),
@@ -24,17 +25,24 @@ 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",))]
 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);

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<RwLock<BTreeMap<String, Arc<dyn HealthCheck>>>>,
-}
-
-impl Default for HealthCheckRegistry {
-    fn default() -> Self {
-        HealthCheckRegistry::new()
-    }
+    health_checks: BTreeMap<String, Arc<dyn HealthCheck>>,
 }
 
 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<H>(&self, health_check: H) -> RoadsterResult<()>
+    pub fn register<H>(&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<dyn HealthCheck>) -> RoadsterResult<()> {
         let name = health_check.name();
 
         if !health_check.enabled() {
@@ -60,53 +37,51 @@ 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<Vec<Arc<dyn HealthCheck>>> {
-        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<Arc<dyn HealthCheck>> {
+        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;
 
     #[rstest]
     #[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
         );
     }
 }

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',
-]
+[]