用 Axum 在 Rust 中编写 REST API

1. 概述

Axum 是 Tokio 团队的 web framework。它构建于 hyper 和 tower 之上,利用 Rust 的类型系统在编译时检查处理程序。

在本文中,你将编写一个 REST API 来管理订单(Order),包括所有 CRUD 操作。本文还将介绍如何在处理程序之间共享 state、返回结构化的 JSON 错误以及编写不需要打开网络端口的路由器测试。

本文中的代码使用 Axum 0.8 和 Rust 2021 edition 或更高版本。

2. 项目设置

创建新项目:

cargo new order-api
cd order-api

将依赖项添加到 Cargo.toml:

[dependencies]
axum = "0.8"
tokio = { version = "1", features = ["full"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"

[dev-dependencies]
tower = { version = "0.5", features = ["util"] }
http-body-util = "0.1"

每个 crate 的作用:

  • axum:router、extractor 和 response 类型。
  • tokio:async runtime。Axum 必须在 Tokio 上运行。
  • serde 和 serde_json:将 struct 转换为 JSON 和反向转换。
  • tower 和 http-body-util:仅在测试时使用。它们帮助将请求直接发送到路由器并读取响应体。

3. 第一个 Axum 应用程序

以一个检查服务健康状态的端点开始:

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

#[tokio::main]
async fn main() {
    let app = Router::new().route("/health", get(|| async { "OK" }));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    println!("Listening on {}", listener.local_addr().unwrap());
    axum::serve(listener, app).await.unwrap();
}

这段代码做三件事:

  • #[tokio::main] 将 main 转换为在 Tokio runtime 上运行的异步函数。
  • Router::new().route(...) 将路径 /health 绑定到 GET 方法。处理程序是一个返回 &'static str 的异步闭包。Axum 会自动将此字符串转换为 200 OK 响应,带有 Content-Type: text/plain。
  • axum::serve 接受一个 TcpListener 并开始服务请求。

尝试运行:

cargo run
curl -i http://localhost:3000/health
HTTP/1.1 200 OK
content-type: text/plain; charset=utf-8
content-length: 2

OK

4. 模型和共享 state

4.1. 定义模型

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Order {
    pub id: u64,
    pub customer: String,
    pub amount: u64,
}

#[derive(Debug, Deserialize)]
pub struct CreateOrder {
    pub customer: String,
    pub amount: u64,
}

将 CreateOrder 与 Order 分开是有意的。客户端不应该自己发送 id,因为 server 才是生成它的一方。如果使用同一个 struct,你必须处理客户端随意发送 id 的情况。

4.2. 共享 state

在本文中,数据存储在内存中。在实际应用中,这个位置通常是一个连接池,如 sqlx::PgPool。

use std::collections::HashMap;
use std::sync::atomic::AtomicU64;
use std::sync::{Arc, RwLock};

#[derive(Clone, Default)]
pub struct AppState {
    orders: Arc<RwLock<HashMap<u64, Order>>>,
    next_id: Arc<AtomicU64>,
}

Axum 为每个请求 clone state。因此 AppState 必须实现 Clone,并且克隆必须代价很小。将数据包装在 Arc 中满足这两个要求:克隆只是增加引用计数,而所有处理程序都指向同一个数据区域。

RwLock 允许多个线程并行读取,但一次只有一个线程可以写。AtomicU64 用于生成 id,无需加锁。

注意: 这里使用 std::sync::RwLock 是因为在获取和释放锁之间没有 .await。如果你需要在 .await 上保持锁,请切换到 tokio::sync::RwLock。否则,future 将不会是 Send,处理程序将无法编译。

5. 使用 IntoResponse 处理错误

在编写处理程序之前,定义一个通用的错误类型。任何实现 IntoResponse 的类型都可以由处理程序返回,包括在 Result 中。

use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::Json;
use serde_json::json;

#[derive(Debug)]
pub enum AppError {
    NotFound(u64),
    Validation(String),
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let (status, message) = match self {
            AppError::NotFound(id) => (StatusCode::NOT_FOUND, format!("Order {id} not found")),
            AppError::Validation(msg) => (StatusCode::UNPROCESSABLE_ENTITY, msg),
        };
        (status, Json(json!({ "error": message }))).into_response()
    }
}

元组 (StatusCode, Json<_>) 也实现了 IntoResponse。因此,你可以同时设置状态码和响应体。

从现在起,处理程序只需返回 Err(AppError::NotFound(id))。客户端将收到如下响应:

HTTP/1.1 404 Not Found
content-type: application/json

{"error":"Order 42 not found"}

6. 编写 CRUD 处理程序

6.1. 创建订单

use axum::extract::State;
use std::sync::atomic::Ordering;

async fn create_order(
    State(state): State<AppState>,
    Json(input): Json<CreateOrder>,
) -> Result<(StatusCode, Json<Order>), AppError> {
    if input.customer.trim().is_empty() {
        return Err(AppError::Validation("customer must not be empty".into()));
    }
    let id = state.next_id.fetch_add(1, Ordering::Relaxed) + 1;
    let order = Order { id, customer: input.customer, amount: input.amount };
    state.orders.write().unwrap().insert(id, order.clone());
    Ok((StatusCode::CREATED, Json(order)))
}

处理程序的每个参数都是一个 extractor:

  • State(state) 获取已绑定到路由器的 AppState。
  • Json(input) 读取请求体并反序列化为 CreateOrder。如果请求体格式错误,Axum 会在你的函数有机会运行之前自动返回错误。

fetch_add 返回旧值,所以需要加 1 以使 id 从 1 开始。处理程序返回 201 Created 以及新创建的订单。

注意: 读取请求体的 extractor,如 Json、String 或 Bytes,必须是 最后一个参数。请求体只能读取一次。将 Json 放在 State 之前会导致编译错误,提示“the trait Handler<_, _> is not implemented”,如果你不知道原因,这个错误消息会很难理解。

6.2. 按 id 获取一个订单

use axum::extract::Path;

async fn get_order(
    State(state): State<AppState>,
    Path(id): Path<u64>,
) -> Result<Json<Order>, AppError> {
    state
        .orders
        .read()
        .unwrap()
        .get(&id)
        .cloned()
        .map(Json)
        .ok_or(AppError::NotFound(id))
}

Path<u64> 将 URL 上的 {id} 部分解析为 u64。如果客户端调用 /orders/abc,Axum 会返回 400 Bad Request,无需额外代码。当找不到订单时,ok_or 会将 None 转换为 NotFound 错误。

6.3. 分页列表

use axum::extract::Query;

#[derive(Debug, Deserialize)]
pub struct Pagination {
    #[serde(default)]
    offset: usize,
    #[serde(default = "default_limit")]
    limit: usize,
}

fn default_limit() -> usize {
    20
}

async fn list_orders(
    State(state): State<AppState>,
    Query(page): Query<Pagination>,
) -> Json<Vec<Order>> {
    let orders = state.orders.read().unwrap();
    let mut result: Vec<Order> = orders.values().cloned().collect();
    result.sort_by_key(|o| o.id);
    let result = result.into_iter().skip(page.offset).take(page.limit).collect();
    Json(result)
}

Query<Pagination> 读取查询字符串,如 ?offset=10&limit=5。得益于 #[serde(default)],客户端可以省略这两个参数,仍然会获得有效结果。

代码必须按 id 排序,因为 HashMap 不保证顺序。如果跳过这一步,分页结果会在每次调用之间改变。

6.4. 删除订单

async fn delete_order(
    State(state): State<AppState>,
    Path(id): Path<u64>,
) -> Result<StatusCode, AppError> {
    match state.orders.write().unwrap().remove(&id) {
        Some(_) => Ok(StatusCode::NO_CONTENT),
        None => Err(AppError::NotFound(id)),
    }
}

StatusCode 本身也实现了 IntoResponse。对于成功的删除操作,返回 204 No Content 是常见的做法。

7. 组织路由器

pub fn app(state: AppState) -> Router {
    Router::new()
        .route("/health", get(|| async { "OK" }))
        .route("/orders", get(list_orders).post(create_order))
        .route("/orders/{id}", get(get_order).delete(delete_order))
        .with_state(state)
}

#[tokio::main]
async fn main() {
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app(AppState::default())).await.unwrap();
}

这个路由器的组织方式如下:

  • 一个路径可以通过链接多个方法来绑定多个 HTTP 方法:get(...).post(...)。
  • with_state 将 state 绑定到路由器。之后,路由器的类型为 Router<()> 并准备好服务请求。
  • 函数 app 被分离出来,以便测试可以构建路由器而无需打开网络端口。

尝试创建并读取一个订单:

curl -i -X POST http://localhost:3000/orders \
  -H "Content-Type: application/json" \
  -d '{"customer":"Nguyen Van An","amount":150000}'

curl http://localhost:3000/orders/1
HTTP/1.1 201 Created
content-type: application/json

{"id":1,"customer":"Nguyen Van An","amount":150000}

{"id":1,"customer":"Nguyen Van An","amount":150000}

注意: 从 Axum 0.8 开始,路径参数使用 {id} 语法,而不是以前版本的 /:id。如果你复制使用 /:id 的旧代码,应用程序会在构建路由器时 panic。通配符语法也从 *rest 改为 {*rest}。

8. 测试路由器

Router 实现了 tower 的 Service trait。因此,你可以直接将请求发送到路由器,使用 oneshot 而无需运行真实的 server。在 main.rs 的末尾添加测试模块:

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::Request;
    use http_body_util::BodyExt;
    use tower::ServiceExt;

    fn post_order(json: &str) -> Request<Body> {
        Request::builder()
            .method("POST")
            .uri("/orders")
            .header("content-type", "application/json")
            .body(Body::from(json.to_owned()))
            .unwrap()
    }

    #[tokio::test]
    async fn create_then_get_order() {
        let app = app(AppState::default());

        let response = app
            .clone()
            .oneshot(post_order(r#"{"customer":"An","amount":150000}"#))
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::CREATED);

        let response = app
            .oneshot(Request::get("/orders/1").body(Body::empty()).unwrap())
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        let bytes = response.into_body().collect().await.unwrap().to_bytes();
        let order: Order = serde_json::from_slice(&bytes).unwrap();
        assert_eq!(order, Order { id: 1, customer: "An".into(), amount: 150000 });
    }

    #[tokio::test]
    async fn empty_customer_is_rejected() {
        let response = app(AppState::default())
            .oneshot(post_order(r#"{"customer":"  ","amount":1}"#))
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
    }

    #[tokio::test]
    async fn unknown_order_returns_404() {
        let response = app(AppState::default())
            .oneshot(Request::get("/orders/42").body(Body::empty()).unwrap())
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::NOT_FOUND);
    }
}

测试中需要注意三点:

  • oneshot 消耗路由器。要向同一个 state 发送多个请求,你需要在每次调用前 clone() 路由器。State 仍然被共享,因为它在 Arc 中。
  • BodyExt::collect 将整个响应体收集为 Bytes,然后 serde_json 解析回 Order。
  • 每个测试创建一个新的 AppState,所以测试之间是独立的,可以并行运行。

运行测试:

cargo test
running 3 tests
test tests::unknown_order_returns_404 ... ok
test tests::empty_customer_is_rejected ... ok
test tests::create_then_get_order ... ok

test result: ok. 3 passed; 0 failed

9. 常见错误

缺少 Content-Type header。 Extractor Json 需要 application/json header。缺少此 header,客户端会收到 415 Unsupported Media Type,尽管请求体完全有效。这是使用 curl 进行测试时最常见的错误。

区分 Json 的错误代码。 JSON 语法错误返回 400。JSON 语法正确但缺少字段或类型错误返回 422。默认情况下,这些响应的响应体是纯文本,而不是 JSON。如果 API 需要统一的错误格式,应将 Json 包装在自定义 extractor 中,并将 JsonRejection 转换为 AppError。

在生产环境中对锁使用 unwrap()。 如果一个线程在持有锁时 panic,锁会被“poisoned”,之后所有的 unwrap() 调用也会 panic。在实际代码中,应处理这种情况或使用真正的数据库。

将所有内容都放入一个大的 AppState。 随着应用程序的增长,你可以使用 FromRef 让处理程序仅获取所需的 state 部分,例如 State<PgPool>,而不是整个 AppState。

什么时候不应该使用 Axum? 如果你需要一个“完整”的框架,其中包含集成的 ORM、模板、admin 和 migration,Axum 会要求你自己组合很多东西。如果团队对异步 Rust 还不熟悉,处理程序的 trait 界限错误消息也会在早期阶段成为重大障碍。

10. 总结

你刚刚用 Axum 0.8 构建了一个完整的 REST API CRUD。主要概念包括:

  • Router: 将路径和 HTTP 方法绑定到处理程序。
  • Extractor: State、Path、Query 和 Json 以类型安全的方式从请求中提取数据。
  • 共享 state: 共享数据被包装在 Arc 中。
  • 错误处理: 实现 IntoResponse 以返回结构化的 JSON 错误。
  • 测试: 使用 tower::ServiceExt::oneshot 直接向路由器发送请求。

合理的下一步是将 HashMap 替换为 sqlx 以连接 PostgreSQL。你还应该从 tower-http 添加中间件,如 TraceLayer 用于记录请求和 CorsLayer 用于配置 CORS。大多数处理程序无需更改,因为 HTTP 逻辑已从存储层分离。

rustaxumrest-apitokiobackend