fa-sharp/axum-app-wrapper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Releases Activity

API improvements

Browse Source 
- Added public `Error`, `Result`, `InitFuture`, and `ShutdownFuture`
  aliases.
- Added `TypeMapState` as the default state instead of `Arc<TypeMap>`.
- Added `App::route`, `App::mount`, and `App::store`.
- Changed `App::init()` to return `InitializedApp<S>`.
- Added `InitializedApp::router()`, `state()`, `into_parts()`, and
  `shutdown()`.
- Made shutdown hooks fallible: `Result<()>`.
- Made `AdHocPlugin::on_shutdown` accept capturing closures.
- Added `Default` for `App` and `AdHocPlugin`.

Improved AppState macro:
- Uses `axum_app_wrapper::Error`.
- Supports generic state structs.
- Fixed the stale `Arc<AppState>` doc snippet.
This commit is contained in:
Farshad
2026-05-27 01:04:50 -04:00
parent 958060e538
commit 1d6708e23b
4 changed files with 254 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
README.md
View File

@@ -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<S>` 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::<AppState>::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::<AppState>::new()`. That gives Rust enough context to infer the `state` parameter:
@@ -41,15 +61,15 @@ Implement `AppPlugin<S>` 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<AppState> for ConfigPlugin {
fn on_init(&mut self, mut state: TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> {
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<AppState> for ConfigPlugin {
&mut self,
router: Router<AppState>,
_state: &AppState,
) -> anyhow::Result<Router<AppState>> {
) -> Result<Router<AppState>> {
Ok(router.route("/health", axum::routing::get(health)))
}
}

9
crates/axum-app-wrapper-macros/src/lib.rs
View File

@@ -22,10 +22,10 @@ use syn::{Data, DeriveInput, Fields, parse_macro_input};
/// }
/// // To wrap the whole state in `Arc`, implement `TryFrom<TypeMap>` for `Arc<AppState>`:
/// impl TryFrom<TypeMap> for Arc<AppState> {
/// type Error = anyhow::Error;
/// type Error = axum_app_wrapper::Error;
///
/// fn try_from(map: TypeMap) -> Result<Self, Self::Error> {
/// 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<Self, Self::Error> {
Ok(#name {

17
examples/app.rs
View File

@@ -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::<AppState>::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?;

284
src/lib.rs
View File

@@ -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<TypeMap>`.
//! 3. [`AppPlugin::on_setup`] runs in registration order and receives `Router<S>` 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<T> = std::result::Result<T, Error>;
/// Startup initialization future returned by [`AppPlugin::on_init`].
pub type InitFuture = BoxFuture<'static, Result<TypeMap>>;
/// 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<TypeMap>);
impl TypeMapState {
/// Borrow the underlying [`TypeMap`].
pub fn type_map(&self) -> &TypeMap {
&self.0
}
}
impl From<TypeMap> 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<TypeMap>`] implementations.
pub fn extract_type_field<T: Send + Sync + 'static>(map: &mut TypeMap) -> anyhow::Result<T> {
pub fn extract_type_field<T: Send + Sync + 'static>(map: &mut TypeMap) -> Result<T> {
map.remove::<T>()
.ok_or_else(|| anyhow!("Missing type in TypeMap: {}", std::any::type_name::<T>()))
.ok_or_else(|| ::anyhow::anyhow!("Missing type in TypeMap: {}", std::any::type_name::<T>()))
}
// Plugin system
/// An axum app builder with plugin-managed state, router setup, and shutdown hooks.
pub struct App<S = Arc<TypeMap>> {
pub struct App<S = TypeMapState> {
base_router: Router<S>,
plugins: Vec<Box<dyn AppPlugin<S>>>,
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<S>) -> 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<S>) -> 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<T: Send + Sync + 'static>(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<S>) -> 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<T: Send + Sync + 'static>(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>, S, impl Future + Send)>
/// Build the router, finalized state, and graceful shutdown handle.
pub async fn init(mut self) -> Result<InitializedApp<S>>
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<S> Default for App<S>
where
S: TryFrom<TypeMap> + Clone + Send + Sync + 'static,
{
fn default() -> Self {
Self::new()
}
}
/// A fully initialized axum app.
pub struct InitializedApp<S> {
router: Router<S>,
state: S,
on_shutdown: ShutdownFuture,
}
impl<S> InitializedApp<S>
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>, 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<S = Arc<TypeMap>> {
pub trait AppPlugin<S = TypeMapState> {
/// 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<TypeMap>> {
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<S>, state: &S) -> anyhow::Result<Router<S>> {
fn on_setup(&mut self, router: Router<S>, state: &S) -> Result<Router<S>> {
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<BoxFuture<'static, ()>> {
fn on_shutdown(&mut self, state: &S) -> Option<ShutdownFuture> {
None
}
}
@@ -158,12 +238,16 @@ pub trait AppPlugin<S = Arc<TypeMap>> {
///
/// `on_init` and `on_setup` accept capturing closures. If `on_setup` uses typed state, prefer
/// `AdHocPlugin::<State>::new()` so the closure parameter type can be inferred.
pub struct AdHocPlugin<S = Arc<TypeMap>> {
on_init: Option<Box<dyn FnOnce(TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> + Send>>,
on_setup: Option<Box<dyn FnOnce(Router<S>, &S) -> anyhow::Result<Router<S>> + Send>>,
on_shutdown: Option<Box<dyn FnOnce(&S) -> BoxFuture<'static, ()>>>,
pub struct AdHocPlugin<S = TypeMapState> {
on_init: Option<InitFn>,
on_setup: Option<SetupFn<S>>,
on_shutdown: Option<ShutdownFn<S>>,
}
type InitFn = Box<dyn FnOnce(TypeMap) -> InitFuture + Send>;
type SetupFn<S> = Box<dyn FnOnce(Router<S>, &S) -> Result<Router<S>> + Send>;
type ShutdownFn<S> = Box<dyn FnOnce(&S) -> ShutdownFuture + Send>;
impl<S: 'static> AdHocPlugin<S> {
/// Create an empty ad-hoc plugin.
pub fn new() -> Self {
@@ -178,7 +262,7 @@ impl<S: 'static> AdHocPlugin<S> {
pub fn on_init<F, T>(mut self, on_init: F) -> Self
where
F: FnOnce(TypeMap) -> T + Send + 'static,
T: Future<Output = anyhow::Result<TypeMap>> + Send + 'static,
T: Future<Output = Result<TypeMap>> + Send + 'static,
{
self.on_init = Some(Box::new(move |s| Box::pin(on_init(s))));
self
@@ -187,38 +271,45 @@ impl<S: 'static> AdHocPlugin<S> {
/// Set router setup for this plugin.
pub fn on_setup<F>(mut self, on_setup: F) -> Self
where
F: FnOnce(Router<S>, &S) -> anyhow::Result<Router<S>> + Send + 'static,
F: FnOnce(Router<S>, &S) -> Result<Router<S>> + Send + 'static,
{
self.on_setup = Some(Box::new(on_setup));
self
}
/// Set shutdown work for this plugin.
pub fn on_shutdown<T>(mut self, on_shutdown: fn(state: &S) -> T) -> Self
pub fn on_shutdown<F, T>(mut self, on_shutdown: F) -> Self
where
T: Future<Output = ()> + Send + 'static,
F: FnOnce(&S) -> T + Send + 'static,
T: Future<Output = Result<()>> + Send + 'static,
{
self.on_shutdown = Some(Box::new(move |s| Box::pin(on_shutdown(s))));
self
}
}
impl<S: 'static> Default for AdHocPlugin<S> {
fn default() -> Self {
Self::new()
}
}
impl<S> AppPlugin<S> for AdHocPlugin<S> {
fn on_init(&mut self, app_state: TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> {
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<S>, state: &S) -> anyhow::Result<Router<S>> {
fn on_setup(&mut self, router: Router<S>, state: &S) -> Result<Router<S>> {
match self.on_setup.take() {
Some(setup_fn) => setup_fn(router, state),
None => Ok(router),
}
}
fn on_shutdown(&mut self, state: &S) -> Option<BoxFuture<'static, ()>> {
fn on_shutdown(&mut self, state: &S) -> Option<ShutdownFuture> {
match self.on_shutdown.take() {
Some(shutdown_fn) => Some(shutdown_fn(state)),
None => None,
@@ -228,8 +319,9 @@ impl<S> AppPlugin<S> for AdHocPlugin<S> {
#[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<String>,
}
impl TryFrom<TypeMap> for TestState {
type Error = Infallible;
fn try_from(mut map: TypeMap) -> Result<Self, Self::Error> {
Ok(Self {
value: map.remove::<Arc<String>>().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::<TestState>::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::<TestState>::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::<Arc<String>>()
.expect("stored value should remain in typemap");
assert_eq!(value.as_str(), "typemap");
}
#[derive(Clone, AppState)]
struct GenericState<T: Send + Sync + 'static> {
value: Arc<T>,
}
#[test]
fn app_state_derive_supports_generics() {
let app = App::<GenericState<String>>::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::<TestState>::new()
.store(Arc::new(String::from("ready")))
.register(AdHocPlugin::<TestState>::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<TestState> for ShutdownOrderPlugin {
fn on_shutdown(&mut self, _state: &TestState) -> Option<BoxFuture<'static, ()>> {
fn on_shutdown(&mut self, _state: &TestState) -> Option<ShutdownFuture> {
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"),