Axum 设计与实现

第18章 State 管理与 FromRef 子状态提取

作者 杨艺韬 · 12,406 字 · 发布于 · 更新于

一个真实的 Web 服务需要各种全局资源:数据库连接池、HTTP 客户端、配置对象、缓存、metrics registry、feature flag 系统。这些资源在 Router 生命周期里只初始化一次、被所有请求共享。axum 的 State 机制就是管这个——给 Router 一个 state 对象、handler 通过 State<T> 提取器取用。

第 7 章简单介绍过 State<T> 的用法。第 18 章深入——讲 Router<S> 的类型状态设计、FromRef trait 如何让 handler 只声明自己需要的子状态、#[derive(FromRef)] 宏如何让多字段 state 变成可组合的小块。这一章理解了,你就知道怎么组织一个生产 axum 项目的依赖图。

为什么 state 这件事需要一整章

其他 Web 框架里 state / context / app config 通常是个默认就有、不太需要细讲的东西——Flask 的 current_app、Express 的 app.locals、Spring 的 @Autowired 都是"你要用时就能用"。axum 不一样——state 是类型系统的一部分。这带来几个根本差异:

编译期检查:handler 里用了什么 state 字段、在编译时由类型签名决定。加字段不破坏代码、删字段让依赖该字段的 handler 编译失败。

没有运行时容器:state 不走任何容器、没有哈希表查找、没有 Arc 计数(除了你主动用 Arc 的字段)。每次 handler 调用的 state 访问是一次 clone——ZST 0 成本、Arc 常数成本。

类型驱动的依赖关系:handler 签名 fn h(State(db): State<PgPool>, State(auth): State<AuthService>) 直接表达"这个 handler 依赖 PgPool 和 AuthService"——signature 即 dependency graph。

这三点让 axum 的 state 机制值得一整章展开——既深入源码、又涉及生产模式。读完你应该能给任意规模的 axum 项目设计 state 架构——小项目的单 AppState、大项目的多模块分离、复杂项目的多租户解析——都有对应的 pattern。

Router<S> 的类型状态模式

Router<S> 的 S 不是装饰——它是类型状态。S 代表"这个 Router 期望 state 类型是 S":

  • Router<S> 未绑定状态:像一张草图——声明要 S 类型的 state 才能工作
  • Router<()> 已绑定空状态:最终形态——state 已经塞到内部 handler 里、对外不再需要 state
  • Router 无泛型参数:默认 Router<()>——即"不需要 state"

用 with_state 从前者转后者——这个 API 改变 Router 的类型:

// axum/src/routing/mod.rs:444-450
pub fn with_state<S2>(self, state: S) -> Router<S2> {
    map_inner!(self, this => RouterInner {
        path_router: this.path_router.with_state(state.clone()),
        default_fallback: this.default_fallback,
        catch_all_fallback: this.catch_all_fallback.with_state(state),
    })
}

签名很特殊:Router<S>::with_state(self, state: S) -> Router<S2>——输入期望 S 的 Router、接收一个 S 值、产出期望 S2 的 Router。S2 完全不受约束——可以是任何类型。

典型用法:

let app: Router = Router::new()   // 注意标注的是 Router(即 Router<()>),不是 Router<AppState>
    .route("/", get(handler))     // handler 里声明 State<AppState>
    .with_state(AppState {...});  // 传入 AppState 值

// with_state 的返回类型 S2 没有任何约束, 这里由左侧的类型标注推导成 ()。
// 如果写成 `let app: Router<AppState> = ...`, 那 S2 就被推成 AppState, 一样能编译——
// 只是这样得到的 Router 还需要外部再喂一次 state, 不能直接交给 axum::serve

with_state(state) 在内部把 state 注入到所有 handler 的 Service 包装里——handler 执行时能通过 State 提取器拿到。注入完成后 Router 自己不再需要外部提供 state——所以上面的例子把 S2 推成 ()。

为什么 S2 不被约束?这是精妙的设计:S2 由下游代码决定。如果你还要继续组装 Router(比如 nest、merge),下游可能需要不同的 state 类型——S2 让灵活性保留到最后一刻。

stateDiagram-v2
    [*] --> Router_S: Router new
    Router_S --> Router_S: add routes, handler 可声明 State S
    Router_S --> Router_S2: with_state s
    Router_S2 --> Service: into_make_service
    Service --> Running: axum serve
    note right of Router_S
        Router 带 AppState 类型参数
        等待绑定 AppState
    end note
    note right of Router_S2
        Router 以 unit 作为 state 类型
        已绑定, 对外不需要 state
    end note

into_make_service 隐式绑定 state

axum/src/routing/mod.rs:558-562:

pub fn into_make_service(self) -> IntoMakeService<Self> {
    IntoMakeService::new(self.with_state(()))
}

into_make_service 隐式调用 with_state(())——把 Router 最终锁定成 Router<()>。不过直接交给 axum::serve 时并不需要显式调它,见下文。

注释里的原因:

call Router::with_state such that everything is turned into Route eagerly rather than doing that per request

每次请求才注入 state 会慢——提前注入让请求路径上不再有注入 state 的工作(handler 里每个 State<T> 提取的那次 clone 仍然在,见后文「state 的开销来自哪里」)。这是构建期的优化,不是编译期的——运行时路由里看不到 state 单独存在,因为它早就被 clone 进各个 handler 的 Route 里了。

但要注意:into_make_service 只适用于 Router<()>。如果你的 Router 是 Router<AppState>,调 into_make_service 编译错——说明还没绑定 state。正确顺序是 .with_state(state).into_make_service()。

那为什么 axum::serve(listener, app) 直接传 Router 也能编译?因为 axum/src/routing/mod.rs:577-596 单独给 Router<()> 实现了 Service<serve::IncomingStream<'_, L>>——它本身就是一个 MakeService,call 里同样做 self.clone().with_state(())(注释和 into_make_service 里那句一模一样)。所以实际代码几乎不用显式写 .into_make_service(),只有需要 ConnectInfo 时才改用 into_make_service_with_connect_info::<SocketAddr>()。

FromRef:子状态抽取的核心

第 6 章讲过的 FromRef trait(axum-core/src/extract/from_ref.rs:14-26):

pub trait FromRef<T> {
    fn from_ref(input: &T) -> Self;
}

impl<T> FromRef<T> for T where T: Clone {
    fn from_ref(input: &T) -> Self {
        input.clone()
    }
}

两部分:

trait 本身:FromRef<T> 意思是"我能从 &T 产出一个 Self"。典型是 clone。

blanket impl:任何 Clone 类型都自动 FromRef<Self>——T::from_ref(&t) == t.clone()。这让 T: FromRef<T> 在 T: Clone 时自动成立。

State 提取器:通过 FromRef 取子状态

第 7 章讲过 State<T> 的实现(axum/src/extract/state.rs:303-317):

impl<OuterState, InnerState> FromRequestParts<OuterState> for State<InnerState>
where
    InnerState: FromRef<OuterState>,
    OuterState: Send + Sync,
{
    type Rejection = Infallible;

    async fn from_request_parts(
        _parts: &mut Parts,
        state: &OuterState,
    ) -> Result<Self, Self::Rejection> {
        let inner_state = InnerState::from_ref(state);
        Ok(Self(inner_state))
    }
}

核心 bound:InnerState: FromRef<OuterState>——让 State<InnerState> 能从 Router 的 OuterState 里抽出来。

OuterState = InnerState 的简单情形:handler 声明 State<AppState>、Router 是 Router<AppState>——AppState: FromRef<AppState> 成立(blanket impl for Clone),state 提取是一次 clone。

OuterState ≠ InnerState 的子状态情形:handler 声明 State<PgPool>、Router 是 Router<AppState>——需要 PgPool: FromRef<AppState>,用户自己实现(或 derive)这个 impl。

这个 trait bound 在编译期就检查——handler 声明 State<T>、编译器看 T: FromRef<S> 是否成立。成立编译过、不成立报错。这就是 axum 依赖注入的类型安全:handler 没法假设 state 里有什么——必须通过 FromRef 声明。

多字段 state 与 #[derive(FromRef)]

一个真实项目的 state 通常是多字段结构:

#[derive(Clone)]
struct AppState {
    db: PgPool,
    redis: RedisClient,
    config: Arc<Config>,
    metrics: Arc<Metrics>,
}

想让 handler 只声明自己需要的字段:

async fn health_check(State(db): State<PgPool>) -> impl IntoResponse { /* ... */ }
async fn get_config(State(cfg): State<Arc<Config>>) -> impl IntoResponse { /* ... */ }
async fn full(State(s): State<AppState>) -> impl IntoResponse { /* ... */ }

需要给每个字段写 FromRef<AppState> impl:

impl FromRef<AppState> for PgPool {
    fn from_ref(s: &AppState) -> Self { s.db.clone() }
}
impl FromRef<AppState> for RedisClient {
    fn from_ref(s: &AppState) -> Self { s.redis.clone() }
}
impl FromRef<AppState> for Arc<Config> {
    fn from_ref(s: &AppState) -> Self { s.config.clone() }
}
impl FromRef<AppState> for Arc<Metrics> {
    fn from_ref(s: &AppState) -> Self { s.metrics.clone() }
}

4 个字段 4 个 impl——模板化。axum-macros::FromRef 宏把这些自动生成:

#[derive(Clone, FromRef)]
struct AppState {
    db: PgPool,
    redis: RedisClient,
    config: Arc<Config>,
    metrics: Arc<Metrics>,
}

一行 derive 替代 4 个 impl。

宏展开的精确逻辑

axum-macros/src/from_ref.rs:29-64 是展开代码:

fn expand_field(state: &Ident, idx: usize, field: &Field) -> TokenStream {
    let FieldAttrs { skip } = match parse_attrs("from_ref", &field.attrs) { /* ... */ };
    if skip.is_some() {
        return TokenStream::default();   // 该字段一行都不生成
    }

    let field_ty = &field.ty;
    let span = field.ty.span();

    let body = if let Some(field_ident) = &field.ident {
        if matches!(field_ty, Type::Reference(_)) {
            quote_spanned! {span=> state.#field_ident }          // 引用类型直接取
        } else {
            quote_spanned! {span=> state.#field_ident.clone() }  // 值类型 clone
        }
    } else {
        let idx = syn::Index { index: idx as _, span: field.span() };
        quote_spanned! {span=> state.#idx.clone() }              // tuple struct
    };

    quote_spanned! {span=>
        #[allow(clippy::clone_on_copy, clippy::clone_on_ref_ptr)]
        impl ::axum::extract::FromRef<#state> for #field_ty {
            fn from_ref(state: &#state) -> Self {
                #body
            }
        }
    }
}

两个细节值得留意:宏用的是 quote_spanned! 而不是 quote!,span 取自 field.ty——这样类型不满足 Clone 时,编译错误会精确指到那个字段的类型上,而不是整个 #[derive]。另外生成的 impl 上挂了 #[allow(clippy::clone_on_copy, clippy::clone_on_ref_ptr)]——因为对 Copy 字段和 Arc 字段无脑 .clone() 会触发这两条 clippy lint,宏统一压掉。

每字段生成一个 FromRef impl——字段类型作为目标、state clone 字段作为实现。

#[from_ref(skip)] 属性:某些字段不想自动生成(可能手动有复杂 impl、或者与别的字段同类型、会生成冲突的 impl,见文末陷阱三)——用 #[from_ref(skip)] 跳过:

#[derive(Clone, FromRef)]
struct AppState {
    db: PgPool,
    #[from_ref(skip)]
    internal: Arc<Mutex<InternalState>>,  // 不生成 FromRef, 用户自己写
}

限制:#[derive(FromRef)] 不支持 generics(from_ref.rs:11-17 明确拒绝)。derive 只覆盖"具体类型的 struct"这一种场景,需要泛型 state 要手写 impl。

嵌套 FromRef:从子到孙

FromRef 是 trait——可以 chained。比如:

#[derive(Clone, FromRef)]
struct AppState {
    db_state: DatabaseState,
    // ...
}

#[derive(Clone, FromRef)]
struct DatabaseState {
    pool: PgPool,
    // ...
}

derive 生成:

  • FromRef<AppState> for DatabaseState
  • FromRef<DatabaseState> for PgPool

但不自动生成 FromRef<AppState> for PgPool——Rust trait impl 默认不传递。要让 handler 直接用 State<PgPool> 而不是 State<DatabaseState>,需要手写:

impl FromRef<AppState> for PgPool {
    fn from_ref(app: &AppState) -> Self {
        PgPool::from_ref(&app.db_state)  // 两层 from_ref
    }
}

每一层嵌套的类型都需要这种"跨层"转发。大项目里可能有 3-4 层——手写 FromRef 很烦。实际架构选择:

  • 浅嵌套:顶层 AppState + 一层 derive——handler 只能用顶层字段类型
  • 深嵌套:多层 derive + 手写跨层转发——代码多但组织清晰

平衡:state 结构浅、字段扁平——让 derive 覆盖所有情况。子结构只在有强分组理由(比如同主题、同生命周期)时才引入。

derive(FromRef) 的 skip 和 ref type

宏还支持 skip 属性和引用类型——axum-macros/src/from_ref.rs:35-54 有实现:

skip:跳过某字段的 FromRef 生成。

下面是读者示例,不是 axum-macros 的源码原文——AppState 与 private_key 都是示例自取的名字:

#[derive(Clone, FromRef)]
struct AppState {
    db: PgPool,
    #[from_ref(skip)]
    private_key: [u8; 32],  // 不想让 handler 直接提取
}

skip 的字段用户要手写 FromRef impl(或者就不让 handler 提取这类字段)。

引用类型字段:如果字段是 &'static T——宏生成 state.field(不 clone、直接拷贝引用)而不是 state.field.clone():

#[derive(Clone, FromRef)]
struct AppState {
    static_config: &'static Config,  // clone 的是引用(Copy), 不是指向的数据
}

这不是性能上的微优化,而是为了生成的代码正确、干净:对 &'static Config 写 .clone(),若 Config: Clone,方法解析会先匹配到 Config::clone,返回 Config 而不是引用,类型对不上(手写 impl 实测 E0308);若指向的类型不是 Clone(如 &'static str),则是对引用的空操作 clone,会触发 noop_method_call lint。这个分支来自修复生成代码警告的 #1676。注意它只对具名字段生效,tuple struct 的字段一律走 .clone()(axum-macros/src/from_ref.rs:48-54)。

AppState 作为依赖注入容器

把 state 想成"依赖注入容器"——handler 通过类型声明"我需要什么依赖":

#[derive(Clone, FromRef)]
struct AppState {
    db: PgPool,
    auth: Arc<AuthService>,
    cache: Arc<CacheService>,
    email: Arc<EmailSender>,
}

// 每个 handler 签名描述自己的依赖集
async fn login(
    State(db): State<PgPool>,
    State(auth): State<Arc<AuthService>>,
    Json(creds): Json<Credentials>,
) -> Result<Json<LoginResponse>, AppError> {
    let user = auth.verify(&creds).await?;
    let session = db.create_session(user.id).await?;
    Ok(Json(LoginResponse { token: session.token }))
}

async fn send_welcome(
    State(email): State<Arc<EmailSender>>,
    Json(user): Json<NewUser>,
) -> Result<(), AppError> {
    email.send_template("welcome", &user.email, &user).await?;
    Ok(())
}

几个工程价值:

一、签名即文档:读 handler 签名就知道用了哪些依赖——login 用 db 和 auth、send_welcome 用 email。测试和重构时这些依赖一目了然。

二、单元测试简单:想测 login,构造一个 mock AuthService + 简单 PgPool——不需要完整 AppState、可以传 mock。

三、添加字段无副作用:给 AppState 加个字段不影响任何现有 handler——因为 handler 只声明需要的子集。不需要的字段被 handler 忽略。

四、删除字段有编译保证:移除 AppState 某字段——所有依赖它的 handler 编译失败——回溯出"哪些地方用了这个依赖"。这比 runtime 的 "null pointer exception" 好太多。

这是 static DI(静态依赖注入)——和 Java Spring 的 @Autowired、NestJS 的 provider 不同——没有运行时容器、所有依赖关系在编译期解决。Rust 类型系统让这成为可能。

flowchart TB
    classDef state fill:#efe,stroke:#080
    classDef handler fill:#fef,stroke:#808
    AS["AppState { db, auth, cache, email }"]:::state
    AS --> L["login handler<br/>需要 db + auth"]:::handler
    AS --> W["send_welcome handler<br/>需要 email"]:::handler
    AS --> P["profile handler<br/>需要 db + cache"]:::handler
    AS --> H["health handler<br/>需要 db"]:::handler
    L -.提取.-> Ldb["State&lt;PgPool&gt;"]
    L -.提取.-> Lauth["State&lt;Arc&lt;AuthService&gt;&gt;"]
    W -.提取.-> We["State&lt;Arc&lt;EmailSender&gt;&gt;"]
    P -.提取.-> Pdb["State&lt;PgPool&gt;"]
    P -.提取.-> Pcache["State&lt;Arc&lt;CacheService&gt;&gt;"]
    H -.提取.-> Hdb["State&lt;PgPool&gt;"]

每个 handler 的箭头只指向自己需要的 state——signature 里就能看到。添加新 handler 需要新依赖、就在 AppState 里加字段、derive 自动更新。

共享可变状态:三种锁的选择

AppState 里的数据有时需要可变——比如缓存、计数器、动态配置。Rust 的所有权让可变共享有三种选择:

std::sync::Mutex:同步锁

use std::sync::{Arc, Mutex};

#[derive(Clone, FromRef)]
struct AppState {
    counter: Arc<Mutex<u64>>,
}

async fn increment(State(counter): State<Arc<Mutex<u64>>>) -> String {
    let mut guard = counter.lock().expect("poisoned");
    *guard += 1;
    format!("count = {}", *guard)
}

适用:锁持有时间短、不跨 .await。

关键约束:std::sync::MutexGuard 不是 Send——持有 guard 跨 .await 会让 future 变成 !Send——axum 要求 handler 的 future 是 Send,因此编译失败(见文末陷阱二)。

tokio::sync::Mutex:异步锁

use tokio::sync::Mutex;

async fn increment(State(counter): State<Arc<Mutex<u64>>>) {
    let mut guard = counter.lock().await;
    *guard += 1;
    expensive_async_op().await;  // 锁持有期间做异步 IO 也可以
}

适用:需要跨 .await 持有锁的场景。

代价:lock 本身是 async、比同步锁开销大;锁等待时让出 task 给 runtime 调度。

经验:能用同步 Mutex 尽量用——lock 后快速做完同步工作释放、别跨 await。只有真的需要在锁内做 I/O 才用 tokio Mutex。

Arc<RwLock>:读多写少

use tokio::sync::RwLock;

#[derive(Clone, FromRef)]
struct AppState {
    config: Arc<RwLock<AppConfig>>,
}

async fn get_setting(State(cfg): State<Arc<RwLock<AppConfig>>>) -> String {
    let guard = cfg.read().await;
    guard.feature_x.clone()  // 多个 read 并发
}

async fn update_setting(State(cfg): State<Arc<RwLock<AppConfig>>>, Json(new): Json<AppConfig>) {
    let mut guard = cfg.write().await;
    *guard = new;  // write 独占
}

适用:读远多于写的场景(配置、路由表、特性 flag)。

要知道的语义:tokio::sync::RwLock 的策略是公平的(write-preferring)——有 writer 在排队时,后来的 reader 即便当前没人持写锁也要排队等。这恰恰是为了避免持续的读把 writer 饿死,代价是读的延迟没那么"想拿就拿"。如果你的场景是读远多于写、且能接受"写时短暂停顿",arc_swap::ArcSwap(下一节的 config reload 例子)通常比 RwLock 更合适——读路径完全无锁。

选型表

场景 推荐
锁持有短、无 await std::sync::Mutex
锁内需要 await tokio::sync::Mutex
读多写少 Arc<tokio::sync::RwLock>
无锁数据结构可用 dashmap::DashMap、arc_swap::ArcSwap
原子操作够 std::sync::atomic::*

轻量场景(计数器、flag)用 atomic 最省——一条原子指令,没有加锁/让出的路径。需要复杂数据结构时才上 Mutex / RwLock。生产里通常混着用——不同 state 字段按访问模式各选各的策略。

实战:动态 config reload

生产中配置常需要不重启服务就更新——读文件、订阅配置中心推送、SIGHUP 信号重载。怎么和 state 系统配合?

经典模式:ArcSwap 原子替换:

use arc_swap::ArcSwap;
use std::sync::Arc;

#[derive(Clone)]
struct AppConfig {
    rate_limit: u32,
    feature_flags: HashMap<String, bool>,
    upstream_url: String,
}

#[derive(Clone, FromRef)]
struct AppState {
    config: Arc<ArcSwap<AppConfig>>,
    // 其他字段
}

async fn handler(State(cfg): State<Arc<ArcSwap<AppConfig>>>) -> impl IntoResponse {
    // 读当前 config——无锁
    let current = cfg.load();
    format!("upstream: {}", current.upstream_url)
}

// 后台任务订阅配置变更
async fn config_reload_task(state: Arc<ArcSwap<AppConfig>>) {
    let mut rx = config_center::subscribe().await;
    while let Some(new_config) = rx.recv().await {
        state.store(Arc::new(new_config));
        tracing::info!("config reloaded");
    }
}

#[tokio::main]
async fn main() {
    let config = Arc::new(ArcSwap::new(Arc::new(load_initial_config().await)));
    tokio::spawn(config_reload_task(config.clone()));

    let state = AppState { config };
    axum::serve(listener, Router::new().with_state(state)).await;
}

几个关键:

一、ArcSwap<T>:读路径(load)不加锁,只做原子操作;写(store)整体原子替换。相比 Arc<RwLock<T>>,读者之间、读者与写者之间都不再有锁竞争——具体收益取决于读写比例和并发度,值得自己压一把。

二、两层 Arc:外层 Arc<ArcSwap<_>>(让 state 可 Clone)、内层 Arc<AppConfig>(ArcSwap 存的是 Arc)。读取时 cfg.load() 返回 Guard<Arc<AppConfig>>——按 arc-swap 文档,它通常比 load_full(等于拿一份 Arc 克隆)更便宜;每个线程的廉价 Guard 数量有限,超出后才退化成等价于 load_full 的做法。

三、后台 reload 任务:独立 tokio task 订阅变更——每次新配置来就 store——live 用户立即看到新值。

四、平滑切换:读者读当前 config、reload 原子替换、读者看不到中间状态。比 Arc<RwLock> 的 reader-writer 竞争更少。

这是生产 axum 的 live config reload 推荐 pattern——和 envoy、nginx 的 config reload 类似——不需要重启。

避免锁的选择:atomic + 无锁结构

很多看似需要锁的 state 其实能用无锁数据结构:

计数器 → AtomicU64 / AtomicUsize:

use std::sync::atomic::{AtomicU64, Ordering};

#[derive(Clone)]
struct Metrics {
    request_count: Arc<AtomicU64>,
}

// handler
async fn h(State(m): State<Metrics>) {
    m.request_count.fetch_add(1, Ordering::Relaxed);
}

一条原子指令,没有加锁和让出调度的路径。适合 monotonic counter 这类场景。

键值 map → DashMap:

use dashmap::DashMap;

#[derive(Clone)]
struct Sessions {
    map: Arc<DashMap<SessionId, Session>>,
}

async fn get_session(State(s): State<Sessions>, Path(id): Path<SessionId>) -> Option<Session> {
    s.map.get(&id).map(|e| e.clone())
}

DashMap 内部是分片锁——落在不同分片上的并发读写互不阻塞,而 Mutex<HashMap> 是整表一把锁。收益随并发度和 key 分布而变,热点 key 集中时优势会缩小。

单值 swap → ArcSwap:

前面讲的动态 config reload 场景。也适用任何"偶尔整体替换、频繁读"的数据。

订阅/通知 → tokio::sync::broadcast、tokio::sync::watch:

broadcast 适合"一条消息发给多个订阅者"(SSE 广播);watch 适合"最新值推送"(config reload 也可以用)。

工程直觉:能不用锁就不用锁——Rust 的并发 primitive 很丰富——针对场景选合适的工具。Mutex 是"最通用但最慢"、atomic/ArcSwap/DashMap 是"针对场景的无锁"——前者覆盖广、后者在特定场景快。

嵌套 Router 的 state 约束

多模块项目里 Router 常被拆分成子 Router 组合:

fn api_routes() -> Router<AppState> {
    Router::new()
        .route("/users", get(list_users))
        .route("/users/{id}", get(get_user))
}

fn admin_routes() -> Router<AppState> {
    Router::new()
        .route("/dashboard", get(dashboard))
        .route("/settings", post(update_settings))
}

fn build_app(state: AppState) -> Router {
    Router::new()
        .nest("/api", api_routes())
        .nest("/admin", admin_routes())
        .with_state(state)
}

关键:子 Router 声明 Router<AppState>——子 Router 的 handler 也用 State<SubField>。嵌套时 state 类型必须一致——nest 要求 inner Router 的 S 等于 outer Router 的 S。

如果子 Router 用的是另一套 state(不想和 AppState 耦合),做法是先 with_state 消化掉自己的 state——关键在于 with_state<S2> 的 S2 完全不受约束,编译器会把它推成 outer 需要的类型:

// 注意返回类型是泛型的 Router<S>,不是 Router(= Router<()>)
fn public_api<S>() -> Router<S>
where
    S: Clone + Send + Sync + 'static,
{
    Router::new()
        .route("/health", get(|| async { "ok" }))
        .with_state(())  // 自己的 state 就地绑定; S2 留给调用方推导
}

fn build_app() -> Router {
    let app_state = AppState { /* ... */ };
    Router::new()
        .nest("/public", public_api())         // S 推导为 AppState
        .nest("/api", api_routes())            // 本来就是 Router<AppState>
        .with_state(app_state)
}

这里有个非常容易踩的坑:Router::nest 的签名是 fn nest(self, path: &str, router: Self) -> Self——router 的类型是 Self,也就是必须和 outer 完全同一个 Router<S>。所以如果把上面的辅助函数写成 fn public_api() -> Router(等价于 Router<()>),再 nest 进 Router<AppState> 就直接编译不过。要么像上面那样让返回类型对 S 保持泛型、要么不给 .with_state(...) 的结果写死类型(写在 nest 调用的表达式里,或赋给一个不标类型的 let 绑定),让 S2 由 nest 推导。axum 自己的 nest 文档(axum/src/docs/routing/nest.md 的 "Nesting routers with state" 一节)给的正是后一种写法(不标类型的 let inner_router)。

这条规则让模块化更灵活——自带 state 的子 Router 靠 with_state 把类型"抹平"到 outer 的 S 上;没自带 state 的子 Router 直接和 outer state 类型对齐。两种都有用。

merge 的 state 一致性

Router::merge 合并两个 Router——要求它们 state 类型一致:

fn app_part1() -> Router<AppState> { /* ... */ }
fn app_part2() -> Router<AppState> { /* ... */ }

fn build_app(state: AppState) -> Router {
    app_part1()
        .merge(app_part2())  // 两者都是 Router<AppState>, 合并后仍然是 Router<AppState>
        .with_state(state)
}

不同 state 类型不能 merge——编译失败。这条约束让"大项目拆成多个 router 模块"成为可行方案:所有子模块共享同一个顶层 state 类型、各自用 State<子字段> 提取。

生产级 state 组织模式

模式一:单 AppState + FromRef derive

最常见——一个 AppState、所有字段都 FromRef 派生、handler 按需提取。

#[derive(Clone, FromRef)]
struct AppState {
    db: PgPool,
    redis: RedisClient,
    config: Arc<Config>,
    // 其他字段
}

优点:简单、熟悉、 handler 签名直接反映依赖关系。

缺点:AppState 膨胀(可能几十字段)、derive 生成的 impl 多、编译时间稍长。

模式二:分组 state + sub-state

大型项目把 state 分成几个逻辑组、每组一个 struct:

#[derive(Clone, FromRef)]
struct AppState {
    storage: StorageState,
    services: ServicesState,
    infra: InfraState,
}

#[derive(Clone, FromRef)]
struct StorageState {
    db: PgPool,
    redis: RedisClient,
    s3: S3Client,
}

#[derive(Clone, FromRef)]
struct ServicesState {
    auth: Arc<AuthService>,
    billing: Arc<BillingService>,
    notify: Arc<NotifyService>,
}

handler 可以提取 State<StorageState>(整组)或 State<PgPool>(通过嵌套 FromRef)——但后者需要额外一层 FromRef 实现(FromRef<StorageState> for PgPool、然后手动 FromRef<AppState> for PgPool 转发)。

优点:大项目里结构清晰——"这是存储、这是服务、这是基础设施"。

缺点:额外的层级和 impl——derive 只生成顶层。

模式三:Arc<AppState> 避免 clone 大结构

如果 AppState 很大(每字段本身是 Arc 但 struct 本身多字段、每次 clone 所有 Arc):

// 每次 FromRef 都 clone 一遍所有字段
let state = AppState { db, redis, s3, auth, /* ... 10+ fields */ };

虽然每字段是 Arc::clone(原子加一、常数成本)、多字段加起来有开销。优化:把整个 AppState 包 Arc:

struct AppStateInner {
    db: PgPool,
    // 多字段
}

// 用 newtype 包一层。若写成 `type AppState = Arc<AppStateInner>;`,
// 下面的 impl 会被孤儿规则拒绝(E0117):FromRef、PgPool、Arc 都不是本 crate 定义的
#[derive(Clone)]
struct AppState(Arc<AppStateInner>);

impl FromRef<AppState> for PgPool {
    fn from_ref(s: &AppState) -> Self { s.0.db.clone() }
}
// 其他字段手写

每次 State 提取是一次 Arc::clone(一次引用计数递增)——不管字段多少。代价是每个 handler 访问字段时多一次 Arc 的 Deref——一次指针解引用,编译后通常还会被内联掉。

模式四:每模块自己的 state

超大项目把 state 完全分开:

// 返回 Router(= Router<()>):自己的 state 已就地绑定,
// 不能写成 Router<UserState>,否则各模块类型不同、没法 nest 进同一个 app
fn user_module() -> Router {
    Router::new()
        .route("/users", get(list_users))
        .with_state(UserState::new())
}

fn order_module() -> Router { /* ... */ }

每个模块独立 state、独立 Router、通过 nest 挂到 app。模块间不共享——如果需要跨模块通信走 message passing 或共享的 Arc<T>。

适合大型多团队项目——各团队的 state 完全隔离、不互相污染。

state 与生命周期

一个常见困惑:state 能持有引用吗?——答案是不能('static bound)。

Router<S> 的完整 bound 是 S: Clone + Send + Sync + 'static,写在 axum/src/routing/mod.rs:154-156——因为 state 要跨 handler 调用、跨 tokio task 传递,必须是 'static 的。

别漏掉里面那个 Clone。它才是你往 AppState 里塞进一个非 Clone 字段时第一个报错的 bound,而不少人只记住了 Send + Sync + 'static 三条。

回到 'static:所以 AppState 的字段必须不带非 static 引用:

// ❌ config: &Config 不是 'static
struct AppState<'a> {
    db: PgPool,
    config: &'a Config,
}

// ✅ 用 Arc 或 &'static
struct AppState {
    db: PgPool,
    config: Arc<Config>,  // 或 config: &'static Config
}

Arc<Config> 是 'static(Arc 不带生命周期)——常用。&'static Config 适合编译期常量。

这条约束让 state 的生命周期管理简单——没有复杂的 lifetime 追踪——但也限制了某些 "state 借用 Router 内部资源" 的设计。实际 Rust 风格不推荐那种设计——更干净的做法是把共享数据 Arc 化。

state 的线程安全

Router<S> 的 bound 还要求 S: Send + Sync——state 可能从任意 tokio worker thread 访问。

  • Send:能跨线程 move。多数类型满足(Arc、PgPool、一般数据结构)
  • Sync:能跨线程共享引用(&S 是 Send)。Mutex、RwLock 本身是 Sync;但 Cell / RefCell 不是
// ❌ RefCell<T> 不是 Sync
struct AppState {
    counter: RefCell<u64>,
}

// ✅ tokio::sync::Mutex 是 Sync;再包一层 Arc 才满足 Clone
#[derive(Clone)]
struct AppState {
    counter: Arc<tokio::sync::Mutex<u64>>,
}

这两个 bound 是 Rust 并发安全的基石——如果你的 state 类型不满足、编译就会失败——实测(把 RefCell 包进 Arc 并 derive Clone,只剩 Sync 不满足)Router::new() 那行直接报 "RefCell<u64> cannot be shared between threads safely",信息准确。

测试:mock state

state 设计对测试极重要——mock 越容易测试越简单:

#[cfg(test)]
fn test_state() -> AppState {
    AppState {
        db: test_pool(),         // 测试数据库或内存 SQLite
        redis: mock_redis(),     // redis mock
        config: Arc::new(Config::default()),
        auth: Arc::new(mock_auth()),
    }
}

#[tokio::test]
async fn test_login() {
    let app = Router::new()
        .route("/login", post(login))
        .with_state(test_state());

    // oneshot 测试
    let response = app.oneshot(make_login_request()).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}

生产 state 和测试 state 可以完全不同——只要都是 AppState 类型就行。FromRef 保证的是 handler 需要的字段在 state 里找得到——不关心字段是真的还是 mock 的。

工程建议:

  • 接口抽象:AuthService、EmailSender 等用 trait 定义、Arc<dyn Trait> 作为字段——测试时替换 mock impl
  • feature flag:#[cfg(test)] 让 test_state 构造函数只在测试时编译
  • 测试数据库:用 sqlx::Pool 的 test database(每个 test 一个独立 schema)或者内存数据库

State 深度 FAQ

Q:State 和 Extension 的区别?

维度 State Extension
作用域 Router 级(构造时绑定) Request 级(中间件注入)
获取 State<T>,编译期保证存在 Extension<T>,运行时查 extensions
生命周期 整个 Router 期间 单个请求期间
性能 一次 clone HashMap 查找 + clone
类型约束 编译期 FromRef 运行期 TypeId
典型用途 DB pool、config、全局资源 request id、auth user、trace span

选择:全局不变的资源用 State、请求级动态数据用 Extension。

Q:State 和 Arc 是什么关系?

State 提取器本身就是一次 clone——所以字段如果是 Arc<T> 就是 Arc::clone(原子加一、零拷贝);如果字段是 T(非 Arc),就真正 clone 值(可能昂贵)。生产上state 的字段几乎都该是 Arc 或其他轻量 clone 类型——避免重 clone。

Q:State 泛型不好用,能用 dyn Trait 吗?

可以——Arc<dyn AuthService + Send + Sync> 作为字段就行。但这让字段失去具体类型的编译时优化(vtable 调用、不能 monomorphize)——只在需要 runtime 多态时用(比如 auth service 想运行时切换 impl)。

Q:AppState 太大、handler 都提取整个——怎么减少?

用 #[derive(FromRef)] 让 handler 按需提取字段。如果还嫌 clone 开销,把 AppState 包 Arc、每个字段也 Arc——两次 clone 都是 Arc::clone、常数开销。

Q:测试不想重新搭 AppState 的所有字段——怎么做?

把 AppState 的字段写成 Option<T> 或 Arc<dyn Trait>——test 时 None 或 mock。但这会让生产代码多 unwrap——不推荐。推荐做法是每个字段用接口类型(Arc<dyn Trait>)——测试用 mock impl。

Q:AppState 的字段能是泛型吗?

理论上可以——但会让整个 Router 类型都带上那个泛型参数——代码复杂度剧增。通常避免——用 trait object 替代(Arc<dyn Trait>)。

Q:能在 state 里保存 tokio::sync::mpsc::Sender 发送后台任务吗?

可以——Sender 是 Clone + Send + Sync + 'static(满足 state bound)。handler 里拿到 Sender、send 一个消息给后台 worker——worker task 消费。这是生产里的常见 pattern——handler 不做耗时工作、转给 background pipeline。注意 Receiver 不放 state(需要 exclusive ownership),启动时 spawn 一个 task 持有 Receiver。

state 的开销来自哪里

state 访问的开销可以完全从源码推出来——State<InnerState> 的 from_request_parts 就一行 InnerState::from_ref(state),而 FromRef 的 blanket impl 就是 input.clone():

操作 开销
State<T> 提取 一次 T::from_ref(&outer)——ZST 编译后什么都不剩;Arc<T> 是一次引用计数递增;裸的 Clone struct 则是真正的深拷贝
with_state(state) → Router<()> 构建期把 state 预先 clone 进每个 Route,运行时不再有 state 相关工作
handler 内部 state.field 直接字段访问

每次 handler 调用:每个 State<T> 参数各做一次 from_ref——声明 3 个就是 3 次 clone。所以真正的成本完全取决于字段类型:字段是 Arc / PgPool(内部也是 Arc)时是常数级的引用计数操作,字段是一个大 Vec/String 时就是实打实的复制。

构建期开销:with_state 会把 state 预先烤进所有 route(源码注释原话是 "turned into Route eagerly rather than doing that per request")——路由越多,这一步越久,但只发生一次、不影响请求路径。

优化方向:如果 state clone 是 profile 里的热点(极少见),把大 struct 包 Arc——clone 变成原子加一。

生产模式:library crate 的 state 要求

如果你写一个可复用的 axum-based library crate(比如一个 auth middleware library、或者提供特定路由的扩展包),怎么处理 state?

问题:你不知道 library 用户的 AppState 长什么样。不能硬编码 State<UserAppState>——用户的 AppState 类型你不知道。

解决:用 FromRef trait bound 让用户自己提供转换:

// 你的 library
#[derive(Clone)]
pub struct MyAuthState {
    pub jwt_key: String,
    pub session_store: Arc<dyn SessionStore>,
}

pub fn auth_routes<S>() -> Router<S>
where
    MyAuthState: FromRef<S>,   // 要求用户的 S 能产出 MyAuthState
    S: Clone + Send + Sync + 'static,
{
    Router::new()
        .route("/login", post(login::<S>))
        .route("/logout", post(logout::<S>))
}

async fn login<S>(State(auth): State<MyAuthState>, /* ... */)
where
    MyAuthState: FromRef<S>,
    S: Clone + Send + Sync + 'static,
{ /* ... */ }

用户怎么用:

#[derive(Clone)]
struct AppState {
    db: PgPool,
    auth_state: MyAuthState,
}

impl FromRef<AppState> for MyAuthState {
    fn from_ref(app: &AppState) -> Self { app.auth_state.clone() }
}

let app = Router::new()
    .nest("/auth", auth_routes::<AppState>())
    .with_state(AppState { /* ... */ });

library 不假设 user 的 AppState——只要用户把 MyAuthState 作为某种 FromRef 转换在,就能用。这是 DI 的库侧设计——library 声明"我需要什么"、user 负责"怎么从自己的 state 给我"。

axum 文档把这种写法列为库作者的推荐做法(axum/src/extract/state.rs:219 起的 "For library authors" 一节);axum-extra 的签名 cookie jar 与加密 cookie jar 两个提取器都要求 K: FromRef<S>(axum-extra/src/extract/cookie/signed.rs:145、axum-extra/src/extract/cookie/private.rs:128),读这两处能看到 state bound 的实际用法。并非所有生态库都走这条路——比如 axum-login 的 AuthSession 是由 layer 放进 request extensions、再从 extensions 取出的。

实战:handler 同时需要多个子状态

生产中一个 handler 常需要多个依赖——都从 AppState 按 FromRef 取:

async fn process_payment(
    State(db): State<PgPool>,
    State(stripe): State<Arc<StripeClient>>,
    State(email): State<Arc<EmailSender>>,
    State(metrics): State<Arc<Metrics>>,
    Json(req): Json<PaymentRequest>,
) -> Result<Json<PaymentResponse>, AppError> {
    let charge = stripe.create_charge(req.amount, &req.token).await?;
    let payment = db.record_payment(&charge).await?;
    let _ = email.send_receipt(&req.email, &payment).await;  // fire-and-forget
    metrics.counter("payment.success").increment(1);
    Ok(Json(PaymentResponse { id: payment.id }))
}

4 个 State 参数——每个独立 FromRef、都从 AppState 取。handler 签名告诉读者"这个 handler 依赖 db + stripe + email + metrics"——一眼看清依赖图。

这种"一个 handler 声明多个 State"基本不用担心开销——只要字段都是 Arc 一类的轻量 clone,每次提取就是一次引用计数递增,4 次加起来相对于 handler 里的 DB 查询和 Stripe API 调用完全不在一个量级。真要小心的是字段本身 clone 很贵(比如直接放了一个大 Vec)那种情况。

子状态的接口抽象

让 State<T> 的 T 是 trait object 支持运行时切换:

#[derive(Clone)]
struct AppState {
    // 用 Arc<dyn Trait> 作为字段——测试/生产可切换 impl
    // trait 须 dyn 兼容:方法写原生 async fn 会报 E0038,要加 #[async_trait](见后文 mockall 一节)
    auth: Arc<dyn AuthService + Send + Sync>,
    email: Arc<dyn EmailSender + Send + Sync>,
}

impl FromRef<AppState> for Arc<dyn AuthService + Send + Sync> {
    fn from_ref(app: &AppState) -> Self { app.auth.clone() }
}

// handler
async fn login(
    State(auth): State<Arc<dyn AuthService + Send + Sync>>,
    Json(creds): Json<Credentials>,
) -> Result<Json<Token>, AppError> {
    let token = auth.authenticate(&creds).await?;  // 调 trait 方法
    Ok(Json(token))
}

生产用真实 ProductionAuthService、测试用 MockAuthService——AppState 构造时不同、handler 代码零变化。

这是 Rust 里比较少见的运行时多态 state——代价是方法调用走 vtable、不能内联和单态化,好处是测试极简。

设计 state 字段的三条经验

三条 state 字段设计原则:

一、Arc<T> 是字段的默认包装——除非 T 本身已经是 Arc-based(比如 PgPool 内部就是 Arc)或 Copy/ZST。这让 state clone(FromRef 派生 impl 的默认行为)只是常数开销——原子加一、不复制实际数据。

二、用 trait object 抽象复杂服务——Arc<dyn Service + Send + Sync> 而不是 Arc<ConcreteService>——测试和运行时切换都容易。代价是方法调用走 vtable、失去内联机会——对一个内部还要做 I/O 的服务来说可以忽略。

三、静态字段 vs 动态字段——静态字段(config、api keys、常量)构造时初始化、之后不变——直接放 state;动态字段(cache、counter)包 ArcSwap 或 Atomic——支持热更。两种的混用是 state 的典型形态。

这三条都是类型选择层面的建议——比具体字段数量限制更有指导意义。

state 的架构选型指南

给真实项目选 state 架构的 rubric:

项目规模 state 字段数 推荐模式
小型(< 10 handlers) 1-5 单 AppState + derive FromRef
中型(10-50 handlers) 5-15 单 AppState + derive + 少数 Arc
大型(50-200 handlers) 15-30 单 AppState + Arc<Inner> wrap + trait objects
超大(200+ handlers) 30+ 分组 state or 每模块 state

选型考虑:

  • 编译时间:每字段一个 FromRef impl——字段过多让 axum 的 Router<S> 类型推导变慢
  • 测试难度:state 越大 mock 越烦——考虑接口化(trait object)
  • 依赖可见性:handler 签名里的 State 参数数——< 5 个合适、过多说明 handler 职责太多应该拆

这些都是经验值——没有硬规则。项目演进时逐步重构——初期 AppState 可能几十字段、随着 handler 职责清晰再拆分。

一张全景图

axum state 系统的完整依赖图:

flowchart TB
    classDef framework fill:#efe,stroke:#080
    classDef user fill:#fef,stroke:#808
    classDef router fill:#eef,stroke:#066
    Router["Router&lt;S&gt;"]:::router --> WS["with_state(state: S)"]:::router
    WS --> Router2["Router&lt;()&gt;"]:::router
    Router2 --> Serve["axum::serve"]
    AppState["AppState (用户定义)"]:::user
    AppState -.derive.-> FromRefImpls["FromRef impls<br/>(每字段一个)"]:::framework
    AppState --Clone--> State["State&lt;AppState&gt; 提取"]:::framework
    FromRefImpls --blanket impl--> StateT["State&lt;T&gt; where T: FromRef&lt;AppState&gt;"]:::framework
    Handlers["handler (声明 State&lt;T&gt;)"]:::user
    StateT --注入--> Handlers
    State --注入--> Handlers
    WS --注入--> Handlers

用户定义 AppState、derive 生成 FromRef impl、handler 声明自己要的子字段 State<T>、Router 在 with_state 时把 state 注入到每个 handler 的 Service 包装里——运行时 handler 提取时用 FromRef 从 AppState 取出子字段。

整个过程编译期解析——不像 Spring 的 runtime DI、没有反射、没有容器。这是 Rust 类型系统的典型应用——把运行时的动态依赖问题变成编译时的类型检查。

State 的 drop 与资源清理

state 的资源(DB pool、HTTP client、后台任务 handle)什么时候释放?

时机:Router 的最后一个 Arc 引用被 drop 时。

典型路径:

  1. with_state(state) 时 state 被 clone 进所有 route 的 Service 里——main 里如果还留着一份,也只是其中一个引用
  2. Router 在 serve 退出后 drop——所有 route 的 Service 被 drop——各字段 Arc 引用计数递减
  3. main 函数 scope 结束——main 里留着的那份(如果有)被 drop
  4. 最后一份引用归零——字段 Drop 触发(PgPool::drop 关连接池、Arc<T> 的 T 的 drop)

陷阱:如果某个 state 字段的 Drop 需要异步操作(比如 async 关数据库连接)——Rust 的 Drop 是同步的、不能 await。sqlx 0.8.6 里 PoolInner 的 Drop 只把池标记为关闭,连接随之被直接丢弃,不走优雅关闭;Pool::close().await 才会逐个优雅关闭连接、让数据库端尽早回收资源。需要 flush 到远端的资源(比如 metrics 缓冲)只靠同步 Drop 则可能丢数据。

生产做法:在 serve 退出后、main 函数返回前,显式调异步清理:

#[tokio::main]
async fn main() {
    let state = AppState::new().await;

    axum::serve(listener, Router::new().with_state(state.clone()))
        .with_graceful_shutdown(shutdown_signal())
        .await
        .unwrap();

    // serve 返回后——主动关闭
    state.db.close().await;  // 异步 close——PgPool 的 graceful close
    state.metrics.flush().await;
    // state 在函数结束时自然 drop——drop 只做同步清理(已经是 close 后的状态)
}

这种 "async cleanup then drop" pattern 让异步资源正确关闭。

跨专栏关联:Rust 依赖注入的哲学

很多后端开发者来自 Spring / Angular 背景——熟悉 "DI container 运行时管理对象图"。Rust + axum 的做法完全不同:

  • 没有 DI 容器:AppState 就是一个普通 struct
  • 没有 lifecycle 管理:RAII 处理创建和销毁(drop 时自然清理)
  • 没有 annotation 魔法:#[derive(FromRef)] 只是生成 trait impl、没 runtime 反射
  • 编译期解析:handler 需要什么依赖在编译时确定——不像 Spring 的 runtime 报 "bean not found"

这种"静态 DI"是 Rust 生态的 style——tonic、sqlx、其他 Rust 框架都类似。优点是性能好(无 runtime 反射)、类型安全(编译期捕获错误)、易测试(state 就是 struct,mock 任意字段)。缺点是灵活性差(没办法运行时注入不同的 impl——需要 trait object 静态声明)。

大多数 Web 服务场景 Rust 的 static DI 已经够用——复杂的运行时组合可以通过 trait object + Arc 实现。只有需要插件系统、运行时加载的场景才会觉得受限。

Spring DI 和 axum State 的概念映射

来自 Spring 世界的开发者看 axum state 会觉得熟悉又陌生。概念对照:

Spring axum
@Component / @Service AppState 的一个字段(Arc<T>)
@Autowired handler 的 State<T> 提取器
@Configuration + @Bean AppState::new() 构造函数
ApplicationContext Router<AppState>
@Qualifier 给相同类型的字段各套一个 newtype,再分别 FromRef(同类型不能有两份 impl,见文末陷阱三)
@Scope("prototype") 手动每次 new(axum 没有默认 prototype)
Profile / Env Rust 的 #[cfg(feature = "...")] 或 env-based 构造
runtime 报 "bean not found" 编译期 "FromRef not implemented"

对比下来:axum 的 state 更手动、更类型安全、性能更好——但灵活性稍差(runtime 动态组合不如 Spring 方便)。大多数业务场景这个 tradeoff 值得——生产稳定性和开发速度的平衡。

Nest / Angular 风格对比

NestJS / Angular 的 DI 容器是"按需注入"——构造时声明依赖、框架运行时查容器。axum 没有容器、state 就是普通 struct、字段访问就是字段访问。

优点:读代码直接看 struct 字段就懂、不需要搞清"哪个 module 提供哪个 provider"。

缺点:循环依赖处理要自己想(Arc + Weak 或延迟初始化)——容器框架能自动处理。

axum 的态度:循环依赖通常是设计问题——不提供自动处理,逼用户重新思考架构。

生产模式:多租户 state

一个进阶场景——SaaS 多租户:每个租户有独立的 db schema / redis namespace / 配置。怎么组织?

方案 A:每租户一个 AppState——不现实,1000 租户要 1000 Router。

方案 B:AppState 里持有 tenant resolver、每请求按 tenant 解析:

#[derive(Clone, FromRef)]
struct AppState {
    tenants: Arc<TenantResolver>,
    global_db: PgPool,  // 租户 registry 用
}

struct TenantResolver {
    pools: DashMap<TenantId, PgPool>,
}

impl TenantResolver {
    async fn pool_for(&self, id: TenantId) -> PgPool {
        self.pools.entry(id).or_insert_with(|| /* 按 tenant 创建 */).clone()
    }
}

async fn handler(
    State(tenants): State<Arc<TenantResolver>>,
    Extension(tenant): Extension<TenantId>,  // middleware 解析出 tenant
) -> impl IntoResponse {
    let pool = tenants.pool_for(tenant).await;
    // 用 pool 处理业务
}

middleware 在每请求解析 tenant ID(从子域名、header 或 session)、通过 Extension 传给 handler、handler 从 tenants resolver 拿对应 pool。这种 pattern 让 axum 单 Router 支持多租户——state 只保留 resolver(元信息)、具体资源按请求延迟绑定。

state 演化的重构路径

实际项目的 state 通常从小长到大——下面是一条典型演化路径:

flowchart LR
    S1["阶段1<br/>单个 PgPool<br/>Router<PgPool>"] -->|加一个 config| S2
    S2["阶段2<br/>AppState struct<br/>{ db, config }"] -->|用字段多| S3
    S3["阶段3<br/>derive FromRef<br/>handler 提子字段"] -->|项目更大| S4
    S4["阶段4<br/>Arc 包 AppStateInner<br/>减少 clone"] -->|复杂业务| S5
    S5["阶段5<br/>接口化字段<br/>Arc&lt;dyn Trait&gt;<br/>测试 mock"]
    style S1 fill:#efe
    style S5 fill:#fef

每一步都是自然演进——项目需求触发重构:

  • 阶段 1 → 2:加第二个 state(config)——从单 type 切到 struct
  • 阶段 2 → 3:handler 多了——用 derive 省 boilerplate
  • 阶段 3 → 4:clone 开销显现——包 Arc 优化
  • 阶段 4 → 5:测试 mock 复杂——接口化

这个路径不一定所有项目都走完——小项目停在阶段 2-3 够。大型企业级应用可能到阶段 5。关键是按需演进——不要从一开始就复杂化。axum 的 state 机制支持每个阶段——迁移成本都不高。

state 相关的常见代码 review 点

code review 时 state 相关的注意事项:

  • 新字段是否 Clone? 不 Clone 的类型加进 AppState 会让 derive 失败——要么 Clone、要么用 Arc 包
  • 字段类型是否过于具体? Arc<ProductionDb> 不好测——改成 Arc<dyn Database + Send + Sync> 能 mock
  • 有没有循环依赖? AppState 字段里直接包含 AppState——编译器直接报 "recursive type has infinite size";间接的 Arc 环则不报错但会内存泄漏——用 Weak<T> 打破
  • 是否重复保存同一数据? 两个字段都是 Arc<Config>——应该共享一份
  • 可变字段的选锁是否合适? Mutex vs RwLock vs ArcSwap——review 看用对没
  • 测试 state 和生产 state 的差异? 测试用不同字段——接口化让差异可管理

这些点在 PR 时一眼检查——比 run 时出问题再修好。

state 和可观测性的结合

生产项目 state 里常加入 observability 相关字段,配合 metrics / tracing。下面两段是示意代码:span / histogram / in_span 等方法名只为说明结构,并不是 prometheus / opentelemetry crate 的真实 API,照抄编译不过:

#[derive(Clone, FromRef)]
struct AppState {
    db: PgPool,
    // observability——注入到 state
    metrics: Arc<prometheus::Registry>,
    tracer: Arc<opentelemetry::Tracer>,
    sentry: Arc<SentryClient>,
}

async fn handler(
    State(metrics): State<Arc<prometheus::Registry>>,
    State(tracer): State<Arc<opentelemetry::Tracer>>,
    Json(req): Json<Request>,
) -> impl IntoResponse {
    let span = tracer.span("handler.logic");
    let timer = metrics.histogram("handler.duration").start_timer();
    let result = business_logic(&req).in_span(span).await;
    timer.observe_duration();
    result
}

或者把 observability 包成 middleware + Extension——handler 里不直接引用:

// middleware 层
async fn tracing_mw(State(tracer): State<Arc<opentelemetry::Tracer>>, request: Request, next: Next) -> Response {
    let span = tracer.span_for(&request);
    let response = next.run(request).in_span(span).await;
    response
}

选择:

  • 在 handler 里手动引用 metrics / tracer:每个 handler 可定制——但多了样板代码
  • 中间件统一处理:signature 清爽——但自定义受限

生产通常混用——全局 trace/metric 走 middleware,特定 handler 的业务指标(order_placed、payment_failed)手动打。state 是这两种模式共同的"依赖载体"。

演进:state 在 axum 版本间的变化

axum 的 state 机制经过几次关键演进:

axum 0.5 及以前:没有类型化 State——全局共享靠 .layer(Extension(AppState {})) + handler 里的 Extension<T>。类型安全差:忘了加 layer 只会在运行时报错。

axum 0.6.0(2022-11-25):一次到位地引入了这套机制。CHANGELOG 里可以逐条对上——State 提取器 + Router::with_state 来自 #1155(0.6.0-rc.1),子状态用的 FromRef 同在 0.6 落地(#1532 把 with_state 从构造函数改成链式方法时,明确写了 "Use FromRef for substates"),#[derive(axum::extract::FromRef)] 是 #1430。原文推荐语是 "We recommend migrating from Extension to State since that is more type safe and faster"。

axum 0.7 / 0.8:state 模型本身没有再变,State + FromRef + #[derive(FromRef)] 一直是当前形态。0.7 那次大改动集中在 hyper 1.0 / http-body 1.0 的迁移上(见第 17 章),跟 state 无关。

也就是说:关于 state 的重大版本断点只有 0.5 → 0.6 那一次。从 0.6 写到 0.8 的 state 代码基本不用动。

本章总结

state 是 axum 的依赖注入机制。核心思想一句话:把依赖关系编译期化。六个要点:

一、Router<S> 是类型状态——S 在编译期携带"这个 Router 还需要什么 state"的信息。with_state 做状态迁移。

二、FromRef<T> 是子状态抽取的 trait——State<U> 要求 U: FromRef<OuterState>。

三、#[derive(FromRef)] 自动化——多字段 state 场景下免手写 impl、一行搞定。

四、类型安全来自编译期——加字段不破坏现有 handler、删字段编译失败暴露依赖、类型改变被编译器捕获。

五、并发选型丰富——同步锁、异步锁、RwLock、ArcSwap、DashMap、Atomic 等各有适用场景。

六、state 是 handler 签名的依赖文档——读 State<T> 参数就知道 handler 依赖什么。

和运行时 DI container 相比没那么灵活、但性能好、错误更早被发现。生产 axum 项目的架构讨论经常围绕 state 组织——几字段、单 AppState 还是分组、怎么 mock——本章已给出完整指导。

state 在 AI 应用里的典型结构

LLM 应用的 state 有些特殊考虑:

#[derive(Clone, FromRef)]
struct AppState {
    db: PgPool,
    llm: Arc<LlmClient>,
    prompt_cache: Arc<PromptCache>,
    rate_limiter: Arc<TokenRateLimiter>,
    vector_db: Arc<QdrantClient>,
    embedder: Arc<EmbeddingClient>,
    feature_flags: Arc<ArcSwap<FeatureFlags>>,
    sentry: Arc<SentryClient>,
}

几个特有字段:

  • llm:调用 Anthropic / OpenAI 的客户端——通常 Arc 共享一个 client 复用连接池
  • prompt_cache:已计算过的 prompt 缓存(Arc<DashMap>)
  • rate_limiter:按 user 或 API key 限制 token 用量——DashMap 或 Redis
  • vector_db + embedder:RAG 需要——查询 embedding + 向量搜索
  • feature_flags:AI 服务经常用 feature flag 做灰度——ArcSwap 支持热更

LLM 应用典型的 handler 签名:

async fn chat(
    State(llm): State<Arc<LlmClient>>,
    State(cache): State<Arc<PromptCache>>,
    State(rate): State<Arc<TokenRateLimiter>>,
    Extension(user): Extension<UserContext>,
    Json(req): Json<ChatRequest>,
) -> Result<Sse<impl Stream<...>>, AppError> {
    rate.check(user.id).await?;
    if let Some(cached) = cache.get(&req.prompt) { /* ... */ }
    let stream = llm.stream_completion(req).await?;
    Ok(Sse::new(stream))
}

3 个 State + 1 个 Extension + 1 个 Json——handler 签名即文档。state 的结构让 "AI 相关依赖"一目了然。

这种 state 在生产规模下需要 careful design——LLM 客户端通常要配 retry、circuit breaker、metrics——这些全在 Arc<LlmClient> 的实现里处理。handler 只管"调一次"。

测试 state 的高级技巧

测试专用 builder

#[cfg(test)]
pub struct TestAppStateBuilder {
    db: Option<PgPool>,
    auth: Option<Arc<dyn AuthService + Send + Sync>>,
    // ... 更多字段
}

#[cfg(test)]
impl TestAppStateBuilder {
    pub fn new() -> Self {
        Self { db: None, auth: None }
    }

    pub fn with_db(mut self, db: PgPool) -> Self {
        self.db = Some(db);
        self
    }

    pub fn with_mock_auth(mut self, mock: MockAuthService) -> Self {
        self.auth = Some(Arc::new(mock));
        self
    }

    pub fn build(self) -> AppState {
        AppState {
            db: self.db.unwrap_or_else(|| test_pool()),
            auth: self.auth.unwrap_or_else(|| Arc::new(default_mock_auth())),
        }
    }
}

测试 handler 时:

#[tokio::test]
async fn login_succeeds_with_valid_creds() {
    let mut mock_auth = MockAuthService::new();
    mock_auth.expect_verify().returning(|_| Ok(test_user()));

    let state = TestAppStateBuilder::new().with_mock_auth(mock_auth).build();
    let app = Router::new().route("/login", post(login)).with_state(state);

    let response = app.oneshot(login_request()).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}

builder pattern 让测试能只覆盖自己关心的字段——其他用默认 mock。这比每个测试都完整构造 AppState 省事。

Mockall 与 trait object

mockall crate 自动生成 mock impl——和前面讲的 Arc<dyn Trait> 结合:

// 原生 async fn in trait 不能做成 dyn(实测 E0038),要用 async-trait;
// automock 须写在 async_trait 之上
#[mockall::automock]
#[async_trait::async_trait]
trait AuthService: Send + Sync {
    async fn verify(&self, creds: &Credentials) -> Result<User, AuthError>;
}

// 生产
let auth: Arc<dyn AuthService + Send + Sync> = Arc::new(RealAuthService::new(jwt_key));

// 测试
let mut mock = MockAuthService::new();
mock.expect_verify().returning(|_| Ok(test_user()));
let auth: Arc<dyn AuthService + Send + Sync> = Arc::new(mock);

一套代码、两种场景——state 字段的 trait object 让切换无缝。

最后一点:state 不是"越多越好"

state 字段越多 handler 可注入的东西越多、但项目复杂性也增加。设计时倾向最小可行 state——只放真正跨 handler 共享的资源。一次性使用的数据不要放 state——构造时传参就行。

几个不应该放 state 的东西:

  • 单个 handler 用的临时数据——参数传
  • 请求级的数据(request id、user context)——用 Extension
  • 构造时局部的中间值——函数内部变量

state 应该是"整个 Router 共享、所有 handler 潜在可用"的资源。这个约束让 state 保持干净、不膨胀。

state 共享的三个陷阱

陷阱一:Clone 过重。AppState 里塞了非 Arc 的大 struct——每个请求提取 State<AppState> 时整份克隆——CPU 热点出现在 memcpy。修复:所有非原子字段都裹 Arc——Arc 的 clone 只是引用计数 +1。

陷阱二:同步锁的 guard 跨 .await。std::sync::MutexGuard 不是 Send,一旦它活过一个 .await 点,整个 handler 的 future 就变成 !Send。好消息是这不会变成运行时问题——axum 对 handler 要求 Future: Send,所以这是一个编译错误,只是错误信息不直观:0.8.9 实测默认只报 "the trait Handler<_, _> is not implemented for fn item ...",看不出是 guard 的问题;给 handler 加上 #[axum::debug_handler] 才会报出 "future cannot be sent between threads safely" 并指明 std::sync::MutexGuard<'_, u64> 不是 Send。修复:改用 tokio::sync::Mutex,或者把临界区收缩到一个不含 await 的块里(let v = { let g = m.lock().unwrap(); g.clone() };——guard 在块结束就 drop)。

陷阱三:两个字段类型相同。#[derive(FromRef)] 是逐字段生成 impl FromRef<AppState> for 字段类型,所以 AppState 里出现两个同类型字段(比如两个 String、两个 PgPool)会生成两份同样的 impl——编译器报 "conflicting implementations"。这也说明 FromRef 的寻址方式是按类型而不是按字段名:State<PgPool> 问的是"给我一个 PgPool",不是"给我 db 那个字段"。修复:给其中一个套 newtype(struct ReadPool(PgPool);),或者给该字段加 #[from_ref(skip)] 再手写 impl。

顺带一提,#[derive(FromRef)] 也不支持泛型(axum-macros/src/from_ref.rs:12-17 直接报 "#[derive(FromRef)] doesn't support generics")。

这三条是 code review axum 项目时的必检项——新人容易踩。

下一章讲 axum-macros——#[derive(FromRef)]、#[derive(FromRequest)]、#[debug_handler] 等宏的内部实现。理解了本章的 state 机制,读下一章的 FromRef derive 宏源码就会发现它就是本章讨论的几个 impl 的自动生成——宏只是把重复代码化成模板。axum 的宏哲学是"不发明新语法、只减少 boilerplate"——第 19 章会详细讨论这套风格。