feat: Add lifecycle handlers (#360)

Add `AppLifecycleHandler` trait with methods that are called at various
stages of an app's lifecycle. The initial set of methods are the
following:

- `before_service_cli`
- `before_health_checks`
- `before_services`
- `on_shutdown`

Closes #350

src/app/mod.rs

@@ -15,6 +15,7 @@ use crate::config::app_config::AppConfig;
 use crate::config::environment::Environment;
 use crate::error::RoadsterResult;
 use crate::health_check::registry::HealthCheckRegistry;
+use crate::lifecycle::registry::LifecycleHandlerRegistry;
 use crate::service::registry::ServiceRegistry;
 use crate::tracing::init_tracing;
 use async_trait::async_trait;
@@ -30,7 +31,7 @@ use sea_orm_migration::MigratorTrait;
 use std::env;
 use std::future;
 use std::sync::Arc;
-use tracing::{instrument, warn};
+use tracing::{error, info, instrument, warn};
 
 pub async fn run<A, S>(app: A) -> RoadsterResult<()>
 where
@@ -59,6 +60,7 @@ where
     run_prepared_without_app_cli(prepare_from_cli_and_state(cli_and_state).await?).await
 }
 
+#[non_exhaustive]
 struct CliAndState<A, S>
 where
     A: App<S> + 'static,
@@ -132,6 +134,7 @@ where
     pub app_cli: A::Cli,
     pub state: S,
     pub service_registry: ServiceRegistry<A, S>,
+    pub lifecycle_handler_registry: LifecycleHandlerRegistry<A, S>,
 }
 
 /// Prepare the app. Does everything to prepare the app short of starting the app. Specifically,
@@ -166,6 +169,10 @@ where
     } = cli_and_state;
     let context = AppContext::from_ref(&state);
 
+    let mut lifecycle_handler_registry = LifecycleHandlerRegistry::new(&state);
+    app.lifecycle_handlers(&mut lifecycle_handler_registry, &state)
+        .await?;
+
     let mut health_check_registry = HealthCheckRegistry::new(&context);
     app.health_checks(&mut health_check_registry, &state)
         .await?;
@@ -182,6 +189,7 @@ where
         app_cli,
         state,
         service_registry,
+        lifecycle_handler_registry,
     })
 }
 
@@ -226,8 +234,16 @@ where
         return Ok(());
     }
 
+    let lifecycle_handlers = prepared_app.lifecycle_handler_registry.handlers(state);
+
     #[cfg(feature = "cli")]
     {
+        info!("Running AppLifecycleHandler::before_service_cli");
+        for handler in lifecycle_handlers.iter() {
+            info!(name=%handler.name(), "Running AppLifecycleHandler::before_service_cli");
+            handler.before_service_cli(state).await?;
+        }
+
         let PreparedApp {
             roadster_cli,
             app_cli,
@@ -240,16 +256,37 @@ where
         }
     }
 
-    #[cfg(feature = "db-sql")]
-    if context.config().database.auto_migrate {
-        A::M::up(context.db(), None).await?;
+    info!("Running AppLifecycleHandler::before_health_checks");
+    for handler in lifecycle_handlers.iter() {
+        info!(name=%handler.name(), "Running AppLifecycleHandler::before_health_checks");
+        handler.before_health_checks(state).await?;
     }
-
     crate::service::runner::health_checks(&context).await?;
 
+    info!("Running AppLifecycleHandler::before_services");
+    for handler in lifecycle_handlers.iter() {
+        info!(name=%handler.name(), "Running AppLifecycleHandler::before_services");
+        handler.before_services(state).await?
+    }
     crate::service::runner::before_run(service_registry, state).await?;
+    let result =
+        crate::service::runner::run(prepared_app.app, prepared_app.service_registry, state).await;
+    if let Err(err) = result {
+        error!("An error occurred in the app: {err}");
+    }
 
-    crate::service::runner::run(prepared_app.app, prepared_app.service_registry, state).await?;
+    info!("Shutting down");
+
+    info!("Running AppLifecycleHandler::before_shutdown");
+    for handler in lifecycle_handlers.iter() {
+        info!(name=%handler.name(), "Running AppLifecycleHandler::before_shutdown");
+        let result = handler.on_shutdown(state).await;
+        if let Err(err) = result {
+            error!(name=%handler.name(), "An error occurred when running AppLifecycleHandler::before_shutdown: {err}");
+        }
+    }
+
+    info!("Shutdown complete");
 
     Ok(())
 }
@@ -288,6 +325,14 @@ where
     /// See the following for more details regarding [FromRef]: <https://docs.rs/axum/0.7.5/axum/extract/trait.FromRef.html>
     async fn provide_state(&self, context: AppContext) -> RoadsterResult<S>;
 
+    async fn lifecycle_handlers(
+        &self,
+        _registry: &mut LifecycleHandlerRegistry<Self, S>,
+        _state: &S,
+    ) -> RoadsterResult<()> {
+        Ok(())
+    }
+
     /// Provide the [crate::health_check::HealthCheck]s to use throughout the app.
     async fn health_checks(
         &self,
@@ -315,6 +360,10 @@ where
 
     /// Override to provide custom graceful shutdown logic to clean up any resources created by
     /// the app. Roadster will take care of cleaning up the resources it created.
+    ///
+    /// Alternatively, provide a [`crate::lifecycle::AppLifecycleHandler::on_shutdown`]
+    /// implementation and provide the handler to the [`LifecycleHandlerRegistry`] in
+    /// [`Self::lifecycle_handlers`].
     #[instrument(skip_all)]
     async fn graceful_shutdown(self: Arc<Self>, _state: &S) -> RoadsterResult<()> {
         Ok(())

src/config/app_config.rs

@@ -3,6 +3,7 @@ use crate::config::auth::Auth;
 use crate::config::database::Database;
 use crate::config::environment::{Environment, ENVIRONMENT_ENV_VAR_NAME};
 use crate::config::health_check::HealthCheck;
+use crate::config::lifecycle::LifecycleHandler;
 use crate::config::service::Service;
 use crate::config::tracing::Tracing;
 use crate::error::RoadsterResult;
@@ -28,6 +29,8 @@ pub struct AppConfig {
     #[validate(nested)]
     pub app: App,
     #[validate(nested)]
+    pub lifecycle_handler: LifecycleHandler,
+    #[validate(nested)]
     pub health_check: HealthCheck,
     #[validate(nested)]
     pub service: Service,
@@ -172,6 +175,8 @@ impl AppConfig {
         #[cfg(feature = "sidekiq")]
         let config = config.add_source(crate::config::service::worker::sidekiq::default_config());
 
+        let config = config.add_source(crate::config::lifecycle::default_config());
+
         let config = config.add_source(crate::config::health_check::default_config());
 
         config

src/config/default.toml

@@ -1,6 +1,9 @@
 [app]
 shutdown-on-error = true
 
+[lifecycle-handler]
+default-enable = true
+
 [service]
 default-enable = true
 

src/config/health_check/mod.rs

@@ -19,11 +19,17 @@ pub fn default_config() -> config::File<FileSourceString, FileFormat> {
 pub struct HealthCheck {
     #[serde(default = "default_true")]
     pub default_enable: bool,
+
     pub max_duration: MaxDuration,
+
     #[cfg(feature = "db-sql")]
+    #[validate(nested)]
     pub database: HealthCheckConfig<()>,
+
     #[cfg(feature = "sidekiq")]
+    #[validate(nested)]
     pub sidekiq: HealthCheckConfig<()>,
+
     /// Allows providing configs for custom health checks. Any configs that aren't pre-defined above
     /// will be collected here.
     ///

src/config/lifecycle/default.toml

@@ -0,0 +1,2 @@
+[lifecycle-handler.db-migration]
+priority = 0

src/config/lifecycle/mod.rs

@@ -0,0 +1,115 @@
+use crate::app::context::AppContext;
+use crate::config::app_config::CustomConfig;
+use crate::util::serde::default_true;
+use config::{FileFormat, FileSourceString};
+use serde_derive::{Deserialize, Serialize};
+use std::collections::BTreeMap;
+use validator::Validate;
+
+pub fn default_config() -> config::File<FileSourceString, FileFormat> {
+    config::File::from_str(include_str!("default.toml"), FileFormat::Toml)
+}
+
+#[derive(Debug, Clone, Validate, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+#[non_exhaustive]
+pub struct LifecycleHandler {
+    #[serde(default = "default_true")]
+    pub default_enable: bool,
+
+    #[cfg(feature = "db-sql")]
+    #[validate(nested)]
+    pub db_migration: LifecycleHandlerConfig<()>,
+
+    /// Allows providing configs for custom lifecycle handlers. Any configs that aren't pre-defined
+    /// above will be collected here.
+    ///
+    /// # Examples
+    ///
+    /// ```toml
+    /// [lifecycle-handler.foo]
+    /// enable = true
+    /// x = "y"
+    /// ```
+    ///
+    /// This will be parsed as:
+    /// ```raw
+    /// LifecycleHandler#custom: {
+    ///     "foo": {
+    ///         LifecycleHandlerConfig#common: {
+    ///             enable: true,
+    ///             priority: 10
+    ///         },
+    ///         LifecycleHandlerConfig<CustomConfig>#custom: {
+    ///             "x": "y"
+    ///         }
+    ///     }
+    /// }
+    /// ```
+    #[serde(flatten)]
+    pub custom: BTreeMap<String, LifecycleHandlerConfig<CustomConfig>>,
+}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case", default)]
+#[non_exhaustive]
+pub struct CommonConfig {
+    // Optional so we can tell the difference between a consumer explicitly enabling/disabling
+    // the lifecycle handler, vs the lifecycle handler being enabled/disabled by default.
+    // If this is `None`, the value will match the value of `LifecycleHandler#default_enable`.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub enable: Option<bool>,
+    pub priority: i32,
+}
+
+impl CommonConfig {
+    pub fn enabled(&self, context: &AppContext) -> bool {
+        self.enable
+            .unwrap_or(context.config().lifecycle_handler.default_enable)
+    }
+}
+
+#[derive(Debug, Clone, Validate, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+#[non_exhaustive]
+pub struct LifecycleHandlerConfig<T> {
+    #[serde(flatten, default)]
+    pub common: CommonConfig,
+    #[serde(flatten)]
+    pub custom: T,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::app_config::AppConfig;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case(true, None, true)]
+    #[case(true, Some(true), true)]
+    #[case(true, Some(false), false)]
+    #[case(false, None, false)]
+    #[case(false, Some(true), true)]
+    #[case(false, Some(false), false)]
+    #[cfg_attr(coverage_nightly, coverage(off))]
+    fn common_config_enabled(
+        #[case] default_enable: bool,
+        #[case] enable: Option<bool>,
+        #[case] expected_enabled: bool,
+    ) {
+        // Arrange
+        let mut config = AppConfig::test(None).unwrap();
+        config.lifecycle_handler.default_enable = default_enable;
+
+        let context = AppContext::test(Some(config), None, None).unwrap();
+
+        let common_config = CommonConfig {
+            enable,
+            priority: 0,
+        };
+
+        // Act/Assert
+        assert_eq!(common_config.enabled(&context), expected_enabled);
+    }
+}

src/config/mod.rs

@@ -4,5 +4,6 @@ pub mod auth;
 pub mod database;
 pub mod environment;
 pub mod health_check;
+mod lifecycle;
 pub mod service;
 pub mod tracing;

src/config/snapshots/roadster__config__app_config__tests__test.snap

@@ -8,6 +8,12 @@ environment = 'test'
 name = 'Test'
 shutdown-on-error = true
 
+[lifecycle-handler]
+default-enable = true
+
+[lifecycle-handler.db-migration]
+priority = 0
+
 [health-check]
 default-enable = true
 

src/health_check/default.rs

@@ -26,6 +26,7 @@ pub fn default_health_checks(
             context: context.clone(),
         }),
     ];
+
     health_checks
         .into_iter()
         .filter(|check| check.enabled())

src/health_check/mod.rs

@@ -61,10 +61,10 @@ pub struct ErrorData {
 
 /// Trait used to check the health of the app before its services start up.
 ///
-/// This is a separate trait, vs adding a "health check" method to `AppService`, to allow defining
-/// health checks that apply to multiple services. For example, most services would require
-/// the DB and Redis connections to be valid, so we would want to perform a check for these
-/// resources a single time before starting any service instead of once for every service that
+/// This is a separate trait, vs adding a "health check" method to [`crate::service::AppService`],
+/// to allow defining health checks that apply to multiple services. For example, most services
+/// would require the DB and Redis connections to be valid, so we would want to perform a check for
+/// these resources a single time before starting any service instead of once for every service that
 /// needs the resources.
 ///
 /// Another benefit of using a separate trait is, because the health checks are decoupled from
@@ -73,14 +73,14 @@ pub struct ErrorData {
 #[cfg_attr(test, mockall::automock)]
 #[async_trait]
 pub trait HealthCheck: Send + Sync {
-    /// The name of the health check.
+    /// The name of the [`HealthCheck`].
     fn name(&self) -> String;
 
     /// Whether the health check is enabled. If the health check is not enabled, Roadster will not
-    /// run it. However, if a consumer wants, they can certainly create a [HealthCheck] instance
-    /// and directly call `HealthCheck#check` even if `HealthCheck#enabled` returns `false`.
+    /// run it. However, if a consumer wants, they can certainly create a [`HealthCheck`] instance
+    /// and directly call [`HealthCheck::check`] even if [`HealthCheck::enabled`] returns `false`.
     fn enabled(&self) -> bool;
 
-    /// Run the health check.
+    /// Run the [`HealthCheck`].
     async fn check(&self) -> RoadsterResult<CheckResponse>;
 }

src/lib.rs

@@ -12,6 +12,7 @@ pub mod app;
 pub mod config;
 pub mod error;
 pub mod health_check;
+pub mod lifecycle;
 pub mod middleware;
 #[cfg(feature = "db-sql")]
 pub mod migration;