Rust 异步 Web 框架 Axum 的深入原理与高级用法

自定义请求提取器示例：

use axum::{async_trait, extract::FromRequestParts, http::request::Parts, response::IntoResponse, routing::get, Router};

struct UserAgent(String);

#[async_trait]
impl FromRequestParts<()> for UserAgent {
    type Rejection = ();

    async fn from_request_parts(parts: &mut Parts, _state: &()) -> Result<Self, Self::Rejection> {
        parts.headers.get("user-agent")
            .and_then(|value| value.to_str().ok())
            .map(|s| UserAgent(s.to_string()))
            .ok_or(())
    }
}

async fn get_user_agent(agent: UserAgent) -> impl IntoResponse {
    format!("User-Agent: {}", agent.0)
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/user-agent", get(get_user_agent));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

1.2.2 响应映射器

响应映射器负责将请求处理函数的返回值转换为 HTTP 响应。Axum 提供了多种内置的响应映射器，并允许开发者自定义响应映射器。

内置响应映射器示例：

use axum::{http::StatusCode, response::IntoResponse, routing::get, Router};
use serde_json::json;

async fn get_user() -> impl IntoResponse {
    (StatusCode::OK, json!({"id": 1, "name": "张三", "email": "[email protected]"}))
}

async fn create_user() -> impl IntoResponse {
    (StatusCode::CREATED, "User created successfully")
}

async fn delete_user() -> impl IntoResponse {
    StatusCode::NO_CONTENT
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/users/1", get(get_user))
        .route("/users", get(create_user))
        .route("/users/1", get(delete_user));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

自定义响应映射器示例：

use axum::{http::StatusCode, response::IntoResponse, routing::get, Router};
use serde_json::json;

struct ApiResponse {
    code: i32,
    message: String,
    data: serde_json::Value,
}

impl IntoResponse for ApiResponse {
    fn into_response(self) -> axum::response::Response {
        let status = if self.code == 200 { StatusCode::OK } else { StatusCode::BAD_REQUEST };
        (status, json!({"code": self.code, "message": self.message, "data": self.data})).into_response()
    }
}

async fn get_user() -> ApiResponse {
    ApiResponse {
        code: 200,
        message: "Success".to_string(),
        data: json!({"id": 1, "name": "张三", "email": "[email protected]"}),
    }
}

async fn create_user() -> ApiResponse {
    ApiResponse {
        code: 400,
        message: "Invalid request".to_string(),
        data: serde_json::Value::Null,
    }
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/users/1", get(get_user))
        .route("/users", get(create_user));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

1.2.3 中间件

中间件是 Axum 框架中用于请求和响应处理的通用组件，可以在请求到达路由之前或响应返回客户端之前执行特定的逻辑，如身份验证、日志记录、CORS 处理等。

内置中间件示例：

use axum::{middleware, routing::get, Router};
use tower_http::trace::TraceLayer;

async fn handler() -> &'static str {
    "Hello, World!"
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(handler))
        .layer(TraceLayer::new_for_http());
    // 日志记录中间件
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

自定义中间件示例：

use axum::{async_trait, extract::FromRequestParts, http::request::Parts, middleware::Next, response::IntoResponse, routing::get, Router};
use std::time::Duration;
use tokio::time::Instant;

struct RequestTime(Duration);

#[async_trait]
impl FromRequestParts<()> for RequestTime {
    type Rejection = ();

    async fn from_request_parts(parts: &mut Parts, _state: &()) -> Result<Self, Self::Rejection> {
        parts.extensions.get::<RequestTime>().copied().ok_or(())
    }
}

async fn timing_middleware<B>(request: axum::http::Request<B>, next: Next<B>) -> impl IntoResponse {
    let start = Instant::now();
    let response = next.run(request).await;
    let duration = start.elapsed();
    let mut response = response.into_response();
    response.headers().insert(
        "X-Response-Time",
        format!("{}ms", duration.as_millis()).parse().unwrap(),
    );
    response
}

async fn handler(time: RequestTime) -> impl IntoResponse {
    format!("Response time: {}ms", time.0.as_millis())
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(handler))
        .layer(middleware::from_fn(timing_middleware));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

二、Axum 框架的路由系统

2.1 路由的定义与匹配

Axum 的路由系统基于路径匹配，支持静态路径、动态路径参数和通配符路径。

静态路径与动态路径参数：

use axum::{extract::Path, response::IntoResponse, routing::get, Router};

async fn get_user(Path(user_id): Path<i32>) -> impl IntoResponse {
    format!("Get user with ID: {}", user_id)
}

async fn get_product(Path((category_id, product_id)): Path<(i32, i32)>) -> impl IntoResponse {
    format!("Get product {} in category {}", product_id, category_id)
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/users/:id", get(get_user)) // 单路径参数
        .route("/categories/:category_id/products/:product_id", get(get_product)); // 多路径参数
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

通配符路径：

use axum::{extract::Path, response::IntoResponse, routing::get, Router};

async fn catch_all(Path(path): Path<String>) -> impl IntoResponse {
    format!("Not found: {}", path)
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/:rest..", get(catch_all)); // 通配符路径
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

2.2 路由的嵌套与组合

Axum 支持路由的嵌套与组合，允许开发者将相关的路由组织成模块，提高代码的可读性和可维护性。

路由的嵌套：

use axum::{extract::Path, response::IntoResponse, routing::{get, post}, Router};

async fn get_user(Path(user_id): Path<i32>) -> impl IntoResponse {
    format!("Get user with ID: {}", user_id)
}

async fn create_user() -> impl IntoResponse {
    "User created successfully"
}

async fn delete_user(Path(user_id): Path<i32>) -> impl IntoResponse {
    format!("Delete user with ID: {}", user_id)
}

#[tokio::main]
async fn main() {
    let user_routes = Router::new()
        .route("/:id", get(get_user).delete(delete_user))
        .route("/", post(create_user));
    let app = Router::new().nest("/users", user_routes); // 嵌套路由
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

路由的组合：

use axum::{response::IntoResponse, routing::get, Router};

async fn home() -> impl IntoResponse {
    "Home page"
}

async fn about() -> impl IntoResponse {
    "About page"
}

async fn contact() -> impl IntoResponse {
    "Contact page"
}

#[tokio::main]
async fn main() {
    let public_routes = Router::new()
        .route("/", get(home))
        .route("/about", get(about));
    let contact_routes = Router::new().route("/contact", get(contact));
    let app = public_routes.merge(contact_routes); // 组合路由
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

2.3 路由的状态共享

Axum 支持通过路由状态共享数据，如数据库连接池、Redis 连接、配置信息等。

使用 Router.with_state 共享状态：

use axum::{extract::State, response::IntoResponse, routing::get, Router};
use sqlx::PgPool;
use std::sync::Arc;

#[derive(Clone)]
struct AppState {
    db_pool: Arc<PgPool>,
    config: crate::config::Config,
}

async fn get_user_count(State(state): State<AppState>) -> impl IntoResponse {
    let count = sqlx::query_scalar!("SELECT COUNT(*) FROM users")
        .fetch_one(&*state.db_pool)
        .await
        .unwrap();
    format!("Total users: {}", count)
}

#[tokio::main]
async fn main() {
    let config = crate::config::Config::from_env().unwrap();
    let db_pool = Arc::new(sqlx::PgPool::connect(&config.db.url).await.unwrap());
    let state = AppState { db_pool, config };
    let app = Router::new()
        .route("/users/count", get(get_user_count))
        .with_state(state);
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

三、Axum 框架的高级功能

3.1 WebSocket 支持

Axum 提供了对 WebSocket 的原生支持，允许开发者实现实时通信功能。

WebSocket 服务器示例：

use axum::{extract::WebSocketUpgrade, response::IntoResponse, routing::get, Router};
use tokio_tungstenite::tungstenite::Message;

async fn websocket_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(|mut socket| async move {
        println!("WebSocket connection established");
        while let Some(msg) = socket.next().await {
            match msg {
                Ok(Message::Text(text)) => {
                    println!("Received text message: {}", text);
                    socket.send(Message::Text(format!("Echo: {}", text))).await.unwrap();
                }
                Ok(Message::Binary(data)) => {
                    println!("Received binary message with length: {}", data.len());
                    socket.send(Message::Binary(data)).await.unwrap();
                }
                Ok(Message::Ping(ping)) => {
                    socket.send(Message::Pong(ping)).await.unwrap();
                }
                Ok(Message::Pong(_)) => {}
                Ok(Message::Close(frame)) => {
                    println!("WebSocket connection closing: {:?}", frame);
                }
                Err(e) => {
                    println!("WebSocket error: {:?}", e);
                }
            }
        }
        println!("WebSocket connection closed");
    })
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/ws", get(websocket_handler));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

3.2 流式请求与响应

Axum 支持流式处理请求体和响应体，适用于处理大量数据的场景。

流式响应示例：

use axum::{body::Body, response::IntoResponse, routing::get, Router};
use futures::stream::{self, StreamExt};
use http_body_util::Full;

async fn stream_response() -> impl IntoResponse {
    let items = vec!["First item", "Second item", "Third item"];
    let stream = stream::iter(items).map(|item| Ok::<_, std::io::Error>(Full::new(item.as_bytes())));
    Body::wrap_stream(stream)
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/stream", get(stream_response));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

流式请求示例：

use axum::{body::Body, extract::Request, response::IntoResponse, routing::post, Router};
use futures::StreamExt;

async fn stream_request(request: Request<Body>) -> impl IntoResponse {
    let mut body = request.into_body();
    let mut buffer = Vec::new();
    while let Some(chunk) = body.next().await {
        buffer.extend_from_slice(&chunk.unwrap());
    }
    format!("Received {} bytes", buffer.len())
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/upload", post(stream_request));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

3.3 错误处理与响应

Axum 提供了灵活的错误处理机制，允许开发者自定义错误类型和错误响应。

自定义错误类型与响应：

use axum::{extract::Path, http::StatusCode, response::{IntoResponse, Response}, routing::get, Router};
use serde_json::json;
use thiserror::Error;

#[derive(Error, Debug)]
enum AppError {
    #[error("User not found")]
    UserNotFound,
    #[error("Invalid request")]
    InvalidRequest,
    #[error(transparent)]
    Unexpected(#[from] anyhow::Error),
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let (status, message) = match self {
            AppError::UserNotFound => (StatusCode::NOT_FOUND, "User not found"),
            AppError::InvalidRequest => (StatusCode::BAD_REQUEST, "Invalid request"),
            AppError::Unexpected(_) => (StatusCode::INTERNAL_SERVER_ERROR, "Internal server error"),
        };
        (status, json!({"code": status.as_u16(), "message": message})).into_response()
    }
}

async fn get_user(Path(user_id): Path<i32>) -> Result<impl IntoResponse, AppError> {
    if user_id == 0 {
        return Err(AppError::InvalidRequest);
    }
    if user_id == 999 {
        return Err(AppError::UserNotFound);
    }
    Ok(json!({"id": user_id, "name": "张三", "email": "[email protected]"}))
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/users/:id", get(get_user));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

3.4 CORS 支持

Axum 通过 cors 中间件提供了对 CORS 的支持，允许开发者配置跨域请求的规则。

CORS 中间件示例：

use axum::{response::IntoResponse, routing::get, Router};
use tower_http::cors::{Any, CorsLayer};

async fn handler() -> impl IntoResponse {
    "CORS enabled"
}

#[tokio::main]
async fn main() {
    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods(Any)
        .allow_headers(Any);
    let app = Router::new()
        .route("/", get(handler))
        .layer(cors);
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

3.5 身份验证与授权

Axum 支持多种身份验证与授权方法，如 JWT、API 密钥、OAuth2 等。

JWT 身份验证示例：

use axum::{async_trait, extract::FromRequestParts, http::request::Parts, response::IntoResponse, routing::{get, post}, Router};
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use serde::{Deserialize, Serialize};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

#[derive(Debug, Serialize, Deserialize)]
struct Claims {
    sub: String,
    exp: usize,
}

impl Claims {
    fn new(sub: &str) -> Self {
        let expiration = SystemTime::now()
            .checked_add(Duration::from_secs(3600))
            .unwrap()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs() as usize;
        Claims { sub: sub.to_string(), exp: expiration }
    }
}

struct JwtSecret(String);

#[async_trait]
impl FromRequestParts<JwtSecret> for Claims {
    type Rejection = ();

    async fn from_request_parts(parts: &mut Parts, state: &JwtSecret) -> Result<Self, Self::Rejection> {
        parts.headers.get("authorization")
            .and_then(|value| value.to_str().ok())
            .and_then(|s| s.strip_prefix("Bearer ").map(|s| s.to_string()))
            .and_then(|token| {
                decode::<Claims>(&token, &DecodingKey::from_secret(state.0.as_bytes()), &Validation::new(Algorithm::HS256)).ok()
            })
            .map(|data| data.claims)
            .ok_or(())
    }
}

async fn login() -> impl IntoResponse {
    let claims = Claims::new("user123");
    let token = encode(&Header::new(Algorithm::HS256), &claims, &EncodingKey::from_secret(b"secret")).unwrap();
    token
}

async fn protected(claims: Claims) -> impl IntoResponse {
    format!("Welcome, {}", claims.sub)
}

#[tokio::main]
async fn main() {
    let secret = JwtSecret("secret".to_string());
    let public_routes = Router::new().route("/login", post(login));
    let protected_routes = Router::new()
        .route("/protected", get(protected))
        .with_state(secret.clone());
    let app = public_routes.merge(protected_routes);
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

四、Axum 框架的性能优化

4.1 工作线程数配置

Axum 使用 Tokio 异步运行时，工作线程数的配置会影响系统的并发能力。

use axum::{response::IntoResponse, routing::get, Router};
use num_cpus;
use tokio::runtime::Builder;

async fn handler() -> impl IntoResponse {
    "Hello, World!"
}

fn main() {
    let runtime = Builder::new_multi_thread()
        .worker_threads(num_cpus::get()) // 使用 CPU 核心数作为工作线程数
        .max_blocking_threads(10)
        .build()
        .unwrap();
    runtime.block_on(async {
        let app = Router::new().route("/", get(handler));
        axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
            .serve(app.into_make_service())
            .await
            .unwrap();
    });
}

4.2 请求提取器与响应映射器的优化

请求提取器与响应映射器的优化可以提高请求处理的效率。

避免重复解析：

use axum::{async_trait, extract::{FromRef, FromRequestParts}, http::request::Parts, response::IntoResponse, routing::get, Router};
use serde_json::json;
use std::sync::Arc;

#[derive(Clone)]
struct AppState {
    db_pool: Arc<sqlx::PgPool>,
    config: crate::config::Config,
}

struct UserExtractor(i32);

#[async_trait]
impl FromRef<AppState> for crate::db::DbPool {
    fn from_ref(state: &AppState) -> Self {
        state.db_pool.clone()
    }
}

#[async_trait]
impl FromRequestParts<AppState> for UserExtractor {
    type Rejection = ();

    async fn from_request_parts(parts: &mut Parts, state: &AppState) -> Result<Self, Self::Rejection> {
        parts.headers.get("user-id")
            .and_then(|value| value.to_str().ok())
            .and_then(|s| s.parse().ok())
            .map(|id| UserExtractor(id))
            .ok_or(())
    }
}

async fn get_user(extractor: UserExtractor) -> impl IntoResponse {
    // 使用 extractor.0 直接访问用户 ID，避免重复解析
    json!({"id": extractor.0, "name": "张三", "email": "[email protected]"})
}

#[tokio::main]
async fn main() {
    let config = crate::config::Config::from_env().unwrap();
    let db_pool = Arc::new(sqlx::PgPool::connect(&config.db.url).await.unwrap());
    let state = AppState { db_pool, config };
    let app = Router::new()
        .route("/users", get(get_user))
        .with_state(state);
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

4.3 中间件的优化

中间件的优化可以减少请求处理的开销。

使用 tower_http 的内置中间件：

use axum::{response::IntoResponse, routing::get, Router};
use tower_http::trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer};

async fn handler() -> impl IntoResponse {
    "Hello, World!"
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(handler))
        .layer(TraceLayer::new_for_http()
            .make_span_with(DefaultMakeSpan::new().include_headers(true))
            .on_response(DefaultOnResponse::new().include_headers(true)),
        );
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

五、实战项目的 Axum 应用

5.1 用户同步服务的 Axum 集成

// user-sync-service/src/main.rs
use axum::{http::StatusCode, response::IntoResponse, routing::{get, post}, Router};
use user_sync_service::sync;
use user_sync_service::config::Config;

async fn health() -> impl IntoResponse {
    StatusCode::OK
}

async fn sync_users() -> impl IntoResponse {
    match sync::sync_users().await {
        Ok(_) => StatusCode::ACCEPTED,
        Err(e) => {
            tracing::error!("User sync failed: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        }
    }
}

#[tokio::main]
async fn main() {
    let config = Config::from_env().unwrap();
    let app = Router::new()
        .route("/health", get(health))
        .route("/api/users/sync", post(sync_users));
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

5.2 订单处理服务的 Axum 集成

// order-processing-service/src/main.rs
use axum::{http::StatusCode, response::IntoResponse, routing::{get, post}, Router};
use order_processing_service::process;
use order_processing_service::config::Config;

async fn health() -> impl IntoResponse {
    StatusCode::OK
}

async fn process_order() -> impl IntoResponse {
    match process::process_orders().await {
        Ok(_) => StatusCode::ACCEPTED,
        Err(e) => {
            tracing::error!("Order processing failed: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        }
    }
}

#[tokio::main]
async fn main() {
    let config = Config::from_env().unwrap();
    let app = Router::new()
        .route("/health", get(health))
        .route("/api/orders/process", post(process_order));
    axum::Server::bind(&"0.0.0.0:3001".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

5.3 监控服务的 Axum 集成

// monitoring-service/src/main.rs
use axum::{extract::WebSocketUpgrade, http::StatusCode, response::IntoResponse, routing::{get, post}, Router};
use monitoring_service::monitor;
use monitoring_service::config::Config;

async fn health() -> impl IntoResponse {
    StatusCode::OK
}

async fn websocket_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    monitor::handle_websocket_connection(ws).await
}

#[tokio::main]
async fn main() {
    let config = Config::from_env().unwrap();
    let app = Router::new()
        .route("/health", get(health))
        .route("/ws", get(websocket_handler));
    axum::Server::bind(&"0.0.0.0:3002".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

六、Axum 框架的常见问题与解决方案

6.1 常见问题 1：请求提取器的类型不匹配

问题描述：当请求提取器的类型与请求中的数据类型不匹配时，会导致编译错误或运行时错误。

解决方案：确保请求提取器的类型与请求中的数据类型一致。例如，路径参数 :id 的类型应该是 i32 或 String，而不是 f64。

6.2 常见问题 2：响应映射器的返回值类型不匹配

问题描述：当响应映射器的返回值类型与 IntoResponse trait 的要求不匹配时，会导致编译错误。

解决方案：确保响应映射器的返回值类型实现了 IntoResponse trait。例如，使用 (StatusCode, json) 或自定义类型。

6.3 常见问题 3：中间件的顺序问题

问题描述：中间件的顺序会影响请求处理的流程，错误的顺序会导致预期的逻辑无法执行。

解决方案：按照从外到内的顺序配置中间件。例如，身份验证中间件应该放在路由处理函数之前，而 CORS 中间件应该放在最外层。

6.4 常见问题 4：状态共享的生命周期问题

问题描述：当状态共享的生命周期不正确时，会导致编译错误或运行时错误。

解决方案：确保状态共享的类型实现了 Clone trait，并使用 Arc 或 Rc 管理共享数据的生命周期。

七、总结

Axum 是一个功能强大、简单易用的异步 Web 框架，基于 Tokio 异步运行时，具有高度模块化的架构和类型安全的特点。通过本文的介绍，我们学习了 Axum 框架的架构与核心组件、路由系统、高级功能、性能优化方法、实战项目的应用以及常见问题的解决方案。希望这些内容能够帮助我们深入掌握 Axum 框架的使用，并在实际项目中开发高质量的异步 Web 应用。