diff --git a/README.md b/README.md index c7735a5..e8cdf8c 100644 --- a/README.md +++ b/README.md @@ -17,12 +17,32 @@ This crate is intentionally thin. It does not replace axum's router, extractors, 1. `on_init` in registration order, passing a `TypeMap` to build state. 2. `S::try_from(TypeMap)` to build the final typed app state. 3. `on_setup` in registration order, passing `Router` and `&S`. -4. `on_shutdown` consecutively in reverse registration order: the last registered plugin shuts down first, and each hook finishes before the next one starts. +4. `InitializedApp::shutdown()` runs `on_shutdown` consecutively in reverse registration order: the last registered plugin shuts down first, and each hook finishes before the next one starts. ## Example See the `examples` folder for full examples. +`App::init()` returns an initialized app handle. Use `router()` to get a ready-to-serve `Router<()>`, `state()` to inspect the finalized state, and `shutdown()` from your graceful shutdown path: + +```rust +let app = App::::new() + .store(config) + .route("/health", axum::routing::get(health)) + .register(metrics_plugin) + .init() + .await?; + +let router = app.router(); +let service_name = &app.state().config.service_name; + +axum::serve(listener, router) + .with_graceful_shutdown(async move { + tokio::signal::ctrl_c().await.expect("failed to listen for ctrl-c"); + app.shutdown().await.expect("failed to shut down"); + }) + .await?; +``` For `on_setup` closures that access typed state, construct the plugin as `AdHocPlugin::::new()`. That gives Rust enough context to infer the `state` parameter: @@ -41,15 +61,15 @@ Implement `AppPlugin` directly when setup should be reusable across apps: ```rust use axum::Router; -use axum_app_wrapper::{AppPlugin, TypeMap}; -use futures::{future::BoxFuture, FutureExt}; +use axum_app_wrapper::{AppPlugin, InitFuture, Result, TypeMap}; +use futures::FutureExt; struct ConfigPlugin { config: Config, } impl AppPlugin for ConfigPlugin { - fn on_init(&mut self, mut state: TypeMap) -> BoxFuture<'static, anyhow::Result> { + fn on_init(&mut self, mut state: TypeMap) -> InitFuture { let config = self.config.clone(); async move { state.insert(config); @@ -62,7 +82,7 @@ impl AppPlugin for ConfigPlugin { &mut self, router: Router, _state: &AppState, - ) -> anyhow::Result> { + ) -> Result> { Ok(router.route("/health", axum::routing::get(health))) } } diff --git a/crates/axum-app-wrapper-macros/src/lib.rs b/crates/axum-app-wrapper-macros/src/lib.rs index a0d325f..d12efb9 100644 --- a/crates/axum-app-wrapper-macros/src/lib.rs +++ b/crates/axum-app-wrapper-macros/src/lib.rs @@ -22,10 +22,10 @@ use syn::{Data, DeriveInput, Fields, parse_macro_input}; /// } /// // To wrap the whole state in `Arc`, implement `TryFrom` for `Arc`: /// impl TryFrom for Arc { -/// type Error = anyhow::Error; +/// type Error = axum_app_wrapper::Error; /// /// fn try_from(map: TypeMap) -> Result { -/// Ok(Self(Arc::new(AppState::try_from(map)?))) +/// Ok(Arc::new(AppState::try_from(map)?)) /// } /// } /// ``` @@ -57,10 +57,11 @@ pub fn derive_app_state(input: TokenStream) -> TokenStream { }; let field_names = fields.iter().map(|f| &f.ident); + let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl(); quote! { - impl ::std::convert::TryFrom<::axum_app_wrapper::TypeMap> for #name { - type Error = ::anyhow::Error; + impl #impl_generics ::std::convert::TryFrom<::axum_app_wrapper::TypeMap> for #name #ty_generics #where_clause { + type Error = ::axum_app_wrapper::Error; fn try_from(mut map: ::axum_app_wrapper::TypeMap) -> ::std::result::Result { Ok(#name { diff --git a/examples/app.rs b/examples/app.rs index a7cc3e2..3ace388 100644 --- a/examples/app.rs +++ b/examples/app.rs @@ -67,27 +67,30 @@ async fn main() -> anyhow::Result<()> { let metrics = Arc::clone(&state.metrics); async move { metrics.flush().await; + Ok(()) } }); - // Register your plugins in the desired order + // Register your plugins in the desired order, and initialize the app let app = App::::new() .register(config_plugin) - .register(metrics_plugin); - - let (router, state, on_shutdown) = app.init().await?; - let router = router.with_state(state); + .register(metrics_plugin) + .init() + .await?; + tracing::info!(service = %app.state().config.service_name, "starting server"); let addr: SocketAddr = "127.0.0.1:3000".parse()?; let listener = tokio::net::TcpListener::bind(addr).await?; // Start the axum server with graceful shutdown - axum::serve(listener, router) + axum::serve(listener, app.router()) .with_graceful_shutdown(async { tokio::signal::ctrl_c() .await .expect("failed to listen for ctrl-c"); - on_shutdown.await; // Run the on_shutdown future for graceful shutdown + app.shutdown() + .await + .expect("failed to run graceful shutdown"); }) .await?; diff --git a/src/lib.rs b/src/lib.rs index 87abfc5..2979112 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,3 @@ -#![allow(clippy::type_complexity)] //! A small plugin layer for [`axum`](https://docs.rs/axum/latest/axum/) applications. //! //! `axum-app-wrapper` lets plugins contribute startup state, router setup, and shutdown work @@ -9,7 +8,7 @@ //! 1. [`AppPlugin::on_init`] runs in registration order and builds a [`TypeMap`]. //! 2. The `TypeMap` is converted into the app state `S` with `TryFrom`. //! 3. [`AppPlugin::on_setup`] runs in registration order and receives `Router` plus `&S`. -//! 4. The returned shutdown future runs [`AppPlugin::on_shutdown`] consecutively in reverse +//! 4. [`InitializedApp::shutdown`] runs [`AppPlugin::on_shutdown`] consecutively in reverse //! registration order. //! //! For one-off plugins, use [`AdHocPlugin`]. If the setup closure needs access to typed state, @@ -27,30 +26,60 @@ //! }); //! ``` +extern crate self as axum_app_wrapper; + use std::{fmt::Display, sync::Arc}; -use anyhow::anyhow; -use axum::Router; +use axum::{Router, routing::MethodRouter}; use futures::{FutureExt, future::BoxFuture}; -// State extraction utilities - pub use axum_app_wrapper_macros::AppState; pub use type_map::concurrent::TypeMap; +pub use anyhow; + +/// Error type used by this crate. +pub type Error = anyhow::Error; + +/// Result type used by this crate. +pub type Result = std::result::Result; + +/// Startup initialization future returned by [`AppPlugin::on_init`]. +pub type InitFuture = BoxFuture<'static, Result>; + +/// Shutdown future returned by [`AppPlugin::on_shutdown`]. +pub type ShutdownFuture = BoxFuture<'static, Result<()>>; + +/// Default app state that keeps the startup [`TypeMap`] available behind an [`Arc`]. +#[derive(Clone)] +pub struct TypeMapState(Arc); + +impl TypeMapState { + /// Borrow the underlying [`TypeMap`]. + pub fn type_map(&self) -> &TypeMap { + &self.0 + } +} + +impl From for TypeMapState { + fn from(map: TypeMap) -> Self { + Self(Arc::new(map)) + } +} + /// Extracts a value of type `T` from a [`TypeMap`], returning an Anyhow error if the type is not present. /// /// This is used internally by the [`AppState`] macro but is also available for manual /// [`TryFrom`] implementations. -pub fn extract_type_field(map: &mut TypeMap) -> anyhow::Result { +pub fn extract_type_field(map: &mut TypeMap) -> Result { map.remove::() - .ok_or_else(|| anyhow!("Missing type in TypeMap: {}", std::any::type_name::())) + .ok_or_else(|| ::anyhow::anyhow!("Missing type in TypeMap: {}", std::any::type_name::())) } // Plugin system /// An axum app builder with plugin-managed state, router setup, and shutdown hooks. -pub struct App> { +pub struct App { base_router: Router, plugins: Vec>>, state: TypeMap, @@ -69,6 +98,27 @@ where } } + /// Add a route to the base router before plugins run setup. + pub fn route(mut self, path: &str, method_router: MethodRouter) -> Self { + self.base_router = self.base_router.route(path, method_router); + self + } + + /// Mount routes at the given path before plugins run setup. + pub fn mount(mut self, path: &str, router: Router) -> Self { + self.base_router = match path { + "" | "/" => self.base_router.merge(router), + _ => self.base_router.nest(path, router), + }; + self + } + + /// Store a startup value for the final state conversion step. + pub fn store(mut self, state: T) -> Self { + self.state.insert(state); + self + } + /// Register a plugin. /// /// `on_init` and `on_setup` run in registration order. Shutdown runs in reverse registration @@ -78,26 +128,8 @@ where self } - // /// Mount the routes at the given path. All layers / middleware from plugins - // /// will be applied to these routes. - // pub fn mount(mut self, path: &str, router: Router) -> Self { - // self.base_router = match path { - // "" | "/" => self.base_router.merge(router), - // _ => self.base_router.nest(path, router), - // }; - // self - // } - - // /// Store type T in state - // pub fn store(mut self, state: T) -> Self { - // self.state.insert(state); - // self - // } - - /// Build the router, finalized state, and graceful shutdown future. - /// - /// The returned shutdown future must be awaited by the caller when the server is shutting down. - pub async fn init(mut self) -> anyhow::Result<(Router, S, impl Future + Send)> + /// Build the router, finalized state, and graceful shutdown handle. + pub async fn init(mut self) -> Result> where S::Error: Display, { @@ -106,7 +138,8 @@ where state = plugin.on_init(state).await?; } - let state = S::try_from(state).map_err(|err| anyhow!("Error creating state: {err}"))?; + let state = + S::try_from(state).map_err(|err| ::anyhow::anyhow!("Error creating state: {err}"))?; let mut router = self.base_router; for plugin in self.plugins.iter_mut() { @@ -121,35 +154,82 @@ where .collect(); let on_shutdown = async move { for shutdown_fn in shutdown_fns { - shutdown_fn.await; + shutdown_fn.await?; } - }; + Ok(()) + } + .boxed(); - Ok((router, state, on_shutdown)) + Ok(InitializedApp { + router, + state, + on_shutdown, + }) + } +} + +impl Default for App +where + S: TryFrom + Clone + Send + Sync + 'static, +{ + fn default() -> Self { + Self::new() + } +} + +/// A fully initialized axum app. +pub struct InitializedApp { + router: Router, + state: S, + on_shutdown: ShutdownFuture, +} + +impl InitializedApp +where + S: Clone + Send + Sync + 'static, +{ + /// Build a ready-to-serve router by attaching the finalized state. + pub fn router(&self) -> Router<()> { + self.router.clone().with_state(self.state.clone()) + } + + /// Borrow the finalized app state. + pub fn state(&self) -> &S { + &self.state + } + + /// Consume the initialized app and return its raw parts. + pub fn into_parts(self) -> (Router, S, ShutdownFuture) { + (self.router, self.state, self.on_shutdown) + } + + /// Run graceful shutdown work for all plugins. + pub async fn shutdown(self) -> Result<()> { + self.on_shutdown.await } } /// A plugin that can participate in app initialization, router setup, and shutdown. #[allow(unused_variables, reason = "trait functions with default no-op")] -pub trait AppPlugin> { +pub trait AppPlugin { /// Run during startup before typed state exists. /// /// Use this hook to insert or transform values in the shared [`TypeMap`]. - fn on_init(&mut self, app_state: TypeMap) -> BoxFuture<'static, anyhow::Result> { + fn on_init(&mut self, app_state: TypeMap) -> InitFuture { async { Ok(app_state) }.boxed() } /// Run after typed state has been created. /// /// Use this hook to add routes, services, middleware, or router state. - fn on_setup(&mut self, router: Router, state: &S) -> anyhow::Result> { + fn on_setup(&mut self, router: Router, state: &S) -> Result> { Ok(router) } /// Return shutdown work for this plugin. /// /// Shutdown hooks are awaited consecutively in reverse registration order. - fn on_shutdown(&mut self, state: &S) -> Option> { + fn on_shutdown(&mut self, state: &S) -> Option { None } } @@ -158,12 +238,16 @@ pub trait AppPlugin> { /// /// `on_init` and `on_setup` accept capturing closures. If `on_setup` uses typed state, prefer /// `AdHocPlugin::::new()` so the closure parameter type can be inferred. -pub struct AdHocPlugin> { - on_init: Option BoxFuture<'static, anyhow::Result> + Send>>, - on_setup: Option, &S) -> anyhow::Result> + Send>>, - on_shutdown: Option BoxFuture<'static, ()>>>, +pub struct AdHocPlugin { + on_init: Option, + on_setup: Option>, + on_shutdown: Option>, } +type InitFn = Box InitFuture + Send>; +type SetupFn = Box, &S) -> Result> + Send>; +type ShutdownFn = Box ShutdownFuture + Send>; + impl AdHocPlugin { /// Create an empty ad-hoc plugin. pub fn new() -> Self { @@ -178,7 +262,7 @@ impl AdHocPlugin { pub fn on_init(mut self, on_init: F) -> Self where F: FnOnce(TypeMap) -> T + Send + 'static, - T: Future> + Send + 'static, + T: Future> + Send + 'static, { self.on_init = Some(Box::new(move |s| Box::pin(on_init(s)))); self @@ -187,38 +271,45 @@ impl AdHocPlugin { /// Set router setup for this plugin. pub fn on_setup(mut self, on_setup: F) -> Self where - F: FnOnce(Router, &S) -> anyhow::Result> + Send + 'static, + F: FnOnce(Router, &S) -> Result> + Send + 'static, { self.on_setup = Some(Box::new(on_setup)); self } /// Set shutdown work for this plugin. - pub fn on_shutdown(mut self, on_shutdown: fn(state: &S) -> T) -> Self + pub fn on_shutdown(mut self, on_shutdown: F) -> Self where - T: Future + Send + 'static, + F: FnOnce(&S) -> T + Send + 'static, + T: Future> + Send + 'static, { self.on_shutdown = Some(Box::new(move |s| Box::pin(on_shutdown(s)))); self } } +impl Default for AdHocPlugin { + fn default() -> Self { + Self::new() + } +} + impl AppPlugin for AdHocPlugin { - fn on_init(&mut self, app_state: TypeMap) -> BoxFuture<'static, anyhow::Result> { + fn on_init(&mut self, app_state: TypeMap) -> InitFuture { match self.on_init.take() { Some(init_fn) => async move { init_fn(app_state).await }.boxed(), None => async { Ok(app_state) }.boxed(), } } - fn on_setup(&mut self, router: Router, state: &S) -> anyhow::Result> { + fn on_setup(&mut self, router: Router, state: &S) -> Result> { match self.on_setup.take() { Some(setup_fn) => setup_fn(router, state), None => Ok(router), } } - fn on_shutdown(&mut self, state: &S) -> Option> { + fn on_shutdown(&mut self, state: &S) -> Option { match self.on_shutdown.take() { Some(shutdown_fn) => Some(shutdown_fn(state)), None => None, @@ -228,8 +319,9 @@ impl AppPlugin for AdHocPlugin { #[cfg(test)] mod tests { + use super::*; + use std::{ - convert::Infallible, sync::{ Mutex, atomic::{AtomicUsize, Ordering}, @@ -237,23 +329,11 @@ mod tests { task::Poll, }; - use super::*; - - #[derive(Clone)] + #[derive(Clone, self::AppState)] struct TestState { value: Arc, } - impl TryFrom for TestState { - type Error = Infallible; - - fn try_from(mut map: TypeMap) -> Result { - Ok(Self { - value: map.remove::>().expect("missing state value"), - }) - } - } - #[test] fn adhoc_plugin_with_basic_state() { let init_value = Arc::new(String::from("ready")); @@ -270,10 +350,75 @@ mod tests { }); let app = App::::new().register(plugin); - let (_router, state, _shutdown) = - futures::executor::block_on(app.init()).expect("app should initialize"); + let app = futures::executor::block_on(app.init()).expect("app should initialize"); - assert_eq!(state.value.as_str(), "ready"); + assert_eq!(app.state().value.as_str(), "ready"); + } + + #[test] + fn store_and_router_handle_expose_initialized_state() { + let app = App::::new().store(Arc::new(String::from("stored"))); + + let app = futures::executor::block_on(app.init()).expect("app should initialize"); + let _router = app.router(); + + assert_eq!(app.state().value.as_str(), "stored"); + } + + #[test] + fn default_state_keeps_type_map_available() { + let app: App = App::new().store(Arc::new(String::from("typemap"))); + + let app = futures::executor::block_on(app.init()).expect("app should initialize"); + + let value = app + .state() + .type_map() + .get::>() + .expect("stored value should remain in typemap"); + assert_eq!(value.as_str(), "typemap"); + } + + #[derive(Clone, AppState)] + struct GenericState { + value: Arc, + } + + #[test] + fn app_state_derive_supports_generics() { + let app = App::>::new().store(Arc::new(String::from("generic"))); + + let app = futures::executor::block_on(app.init()).expect("app should initialize"); + + assert_eq!(app.state().value.as_str(), "generic"); + } + + #[test] + fn adhoc_shutdown_accepts_capturing_fallible_closure() { + let events = Arc::new(Mutex::new(Vec::new())); + let shutdown_events = Arc::clone(&events); + + let app = App::::new() + .store(Arc::new(String::from("ready"))) + .register(AdHocPlugin::::new().on_shutdown(move |state| { + let events = Arc::clone(&shutdown_events); + let value = Arc::clone(&state.value); + async move { + events + .lock() + .expect("events lock poisoned") + .push(value.to_string()); + Ok(()) + } + })); + + let app = futures::executor::block_on(app.init()).expect("app should initialize"); + futures::executor::block_on(app.shutdown()).expect("shutdown should succeed"); + + assert_eq!( + *events.lock().expect("events lock poisoned"), + [String::from("ready")] + ); } struct ShutdownOrderPlugin { @@ -283,7 +428,7 @@ mod tests { } impl AppPlugin for ShutdownOrderPlugin { - fn on_shutdown(&mut self, _state: &TestState) -> Option> { + fn on_shutdown(&mut self, _state: &TestState) -> Option { let name = self.name; let events = Arc::clone(&self.events); let active_shutdowns = Arc::clone(&self.active_shutdowns); @@ -307,7 +452,7 @@ mod tests { .expect("events lock poisoned") .push(format!("{name}:finish")); active_shutdowns.fetch_sub(1, Ordering::SeqCst); - Poll::Ready(()) + Poll::Ready(Ok(())) }))) } } @@ -333,9 +478,8 @@ mod tests { active_shutdowns, }); - let (_router, _state, on_shutdown) = - futures::executor::block_on(app.init()).expect("app should initialize"); - futures::executor::block_on(on_shutdown); + let app = futures::executor::block_on(app.init()).expect("app should initialize"); + futures::executor::block_on(app.shutdown()).expect("shutdown should succeed"); assert_eq!( *events.lock().expect("events lock poisoned"),