This commit is contained in:
2026-05-26 02:14:41 -04:00
parent a5b42ab55a
commit 995f3cc61a
2 changed files with 184 additions and 16 deletions

130
README.md Normal file
View File

@@ -0,0 +1,130 @@
# axum-app-wrapper
A small plugin layer for `axum` applications, inspired by [fastify](https://fastify.dev/) from the Node/JS ecossytem.
Plugins can:
- add startup state before the final axum state type exists
- add routes, services, and middleware after state is initialized
- run graceful shutdown work in reverse registration order
This crate is intentionally thin. It does not replace axum's router, extractors, or middleware. Rather, it organizes your server setup and teardown into plugins.
## Lifecycle
`App::init()` runs plugins in this order:
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
## Example
```rust
use std::{net::SocketAddr, sync::Arc};
use axum::{
extract::State,
routing::get,
};
use axum_app_wrapper::{AdHocPlugin, App, AppState};
#[derive(Clone, AppState)]
struct AppState {
config: Arc<Config>,
}
#[derive(Clone)]
struct Config {
service_name: String,
}
async fn health(State(state): State<AppState>) -> String {
format!("{}:ok", state.config.service_name)
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let config = Arc::new(Config {
service_name: "api".to_owned(),
});
let app = App::<AppState>::new().register(
AdHocPlugin::<AppState>::new()
.on_init(async move |mut state| {
state.insert(config);
Ok(state)
})
.on_setup(|router, _state| Ok(router.route("/health", get(health))))
.on_shutdown(|state| {
let service_name = state.config.service_name.clone();
async move {
tracing::info!(%service_name, "shutting down");
}
}),
);
let (router, state, on_shutdown) = app.init().await?;
let router = router.with_state(state);
let addr: SocketAddr = "127.0.0.1:3000".parse()?;
let listener = tokio::net::TcpListener::bind(addr).await?;
axum::serve(listener, router)
.with_graceful_shutdown(async {
tokio::signal::ctrl_c().await.expect("failed to listen for ctrl-c");
on_shutdown.await;
})
.await?;
Ok(())
}
```
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:
```rust
let plugin = AdHocPlugin::<AppState>::new()
.on_setup(|router, state| {
let config = Arc::clone(&state.config);
Ok(router.layer(axum::Extension(config)))
});
```
## Reusable Plugins
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};
struct ConfigPlugin {
config: Config,
}
impl AppPlugin<AppState> for ConfigPlugin {
fn on_init(&mut self, mut state: TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> {
let config = self.config.clone();
async move {
state.insert(config);
Ok(state)
}
.boxed()
}
fn on_setup(
&mut self,
router: Router<AppState>,
_state: &AppState,
) -> anyhow::Result<Router<AppState>> {
Ok(router.route("/health", axum::routing::get(health)))
}
}
```
Shutdown hooks run like stack unwinding: the last registered plugin shuts down first, and each hook
finishes before the next one starts.

View File

@@ -1,4 +1,31 @@
#![allow(clippy::type_complexity)] #![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
//! around an `axum::Router`.
//!
//! Lifecycle:
//!
//! 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
//! registration order.
//!
//! For one-off plugins, use [`AdHocPlugin`]. If the setup closure needs access to typed state,
//! spell the state type at construction time so Rust can infer the closure parameters:
//!
//! ```ignore
//! let plugin = AdHocPlugin::<AppState>::new()
//! .on_init(async |mut state| {
//! state.insert(config);
//! Ok(state)
//! })
//! .on_setup(|router, state| {
//! let config = state.config.clone();
//! Ok(router.layer(axum::Extension(config)))
//! });
//! ```
use std::{fmt::Display, sync::Arc}; use std::{fmt::Display, sync::Arc};
@@ -22,7 +49,7 @@ pub fn extract_type_field<T: Send + Sync + 'static>(map: &mut TypeMap) -> anyhow
// Plugin system // Plugin system
/// A lightweight wrapper around axum that enables building plugins around the router /// An axum app builder with plugin-managed state, router setup, and shutdown hooks.
pub struct App<S = Arc<TypeMap>> { pub struct App<S = Arc<TypeMap>> {
base_router: Router<S>, base_router: Router<S>,
plugins: Vec<Box<dyn AppPlugin<S>>>, plugins: Vec<Box<dyn AppPlugin<S>>>,
@@ -33,7 +60,7 @@ impl<S> App<S>
where where
S: TryFrom<TypeMap> + Clone + Send + Sync + 'static, S: TryFrom<TypeMap> + Clone + Send + Sync + 'static,
{ {
/// Create a new server with the given routes. /// Create an empty app.
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
base_router: Router::new(), base_router: Router::new(),
@@ -42,8 +69,10 @@ where
} }
} }
/// Register a plugin. This can be considered analogous to axum/tower's `layer` function /// Register a plugin.
/// and follows the same ordering (each plugin "wraps" the previous one). ///
/// `on_init` and `on_setup` run in registration order. Shutdown runs in reverse registration
/// order, awaiting each hook before starting the next.
pub fn register(mut self, plugin: impl AppPlugin<S> + 'static) -> Self { pub fn register(mut self, plugin: impl AppPlugin<S> + 'static) -> Self {
self.plugins.push(Box::new(plugin)); self.plugins.push(Box::new(plugin));
self self
@@ -65,8 +94,9 @@ where
// self // self
// } // }
/// Build and initialize the server. Returns the base router, finalized state, and a future to run /// Build the router, finalized state, and graceful shutdown future.
/// on graceful shutdown. ///
/// 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)> pub async fn init(mut self) -> anyhow::Result<(Router<S>, S, impl Future + Send)>
where where
S::Error: Display, S::Error: Display,
@@ -99,27 +129,35 @@ where
} }
} }
/// Trait for a plugin that can be attached to the server /// A plugin that can participate in app initialization, router setup, and shutdown.
#[allow(unused_variables, reason = "trait functions with default no-op")] #[allow(unused_variables, reason = "trait functions with default no-op")]
pub trait AppPlugin<S = Arc<TypeMap>> { pub trait AppPlugin<S = Arc<TypeMap>> {
/// Init function that will run on server startup. Can add and manipulate state. /// 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) -> BoxFuture<'static, anyhow::Result<TypeMap>> {
async { Ok(app_state) }.boxed() async { Ok(app_state) }.boxed()
} }
/// Setup function that will run _after_ state is initialized (e.g. routes, middleware, /// Run after typed state has been created.
/// and services should be added here) ///
/// 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) -> anyhow::Result<Router<S>> {
Ok(router) Ok(router)
} }
/// Teardown function that will run on server shutdown /// 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<BoxFuture<'static, ()>> {
None None
} }
} }
/// Utility to build a plugin on the fly /// A closure-based plugin for application-local setup.
///
/// `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>> { pub struct AdHocPlugin<S = Arc<TypeMap>> {
on_init: Option<Box<dyn FnOnce(TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> + Send>>, 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_setup: Option<Box<dyn FnOnce(Router<S>, &S) -> anyhow::Result<Router<S>> + Send>>,
@@ -127,6 +165,7 @@ pub struct AdHocPlugin<S = Arc<TypeMap>> {
} }
impl<S: 'static> AdHocPlugin<S> { impl<S: 'static> AdHocPlugin<S> {
/// Create an empty ad-hoc plugin.
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
on_init: None, on_init: None,
@@ -135,7 +174,7 @@ impl<S: 'static> AdHocPlugin<S> {
} }
} }
/// Init function that will run on server startup. Can add and manipulate state. /// Set startup state initialization for this plugin.
pub fn on_init<F, T>(mut self, on_init: F) -> Self pub fn on_init<F, T>(mut self, on_init: F) -> Self
where where
F: FnOnce(TypeMap) -> T + Send + 'static, F: FnOnce(TypeMap) -> T + Send + 'static,
@@ -145,8 +184,7 @@ impl<S: 'static> AdHocPlugin<S> {
self self
} }
/// Setup function that will run _after_ state is initialized. (e.g. routes, middleware, /// Set router setup for this plugin.
/// and services should be added here)
pub fn on_setup<F>(mut self, on_setup: F) -> Self pub fn on_setup<F>(mut self, on_setup: F) -> Self
where where
F: FnOnce(Router<S>, &S) -> anyhow::Result<Router<S>> + Send + 'static, F: FnOnce(Router<S>, &S) -> anyhow::Result<Router<S>> + Send + 'static,
@@ -155,7 +193,7 @@ impl<S: 'static> AdHocPlugin<S> {
self self
} }
/// Teardown function that will run on server shutdown /// Set shutdown work for this plugin.
pub fn on_shutdown<T>(mut self, on_shutdown: fn(state: &S) -> T) -> Self pub fn on_shutdown<T>(mut self, on_shutdown: fn(state: &S) -> T) -> Self
where where
T: Future<Output = ()> + Send + 'static, T: Future<Output = ()> + Send + 'static,