😺

Rust/AxumでStateをArcに包む

に公開

RustのhttpライブラリAxumで状態を共有するにはStateを使います。

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

#[derive(Clone)]
struct AppState {}

let routes = Router::new()
    .route("/", get(|State(state): State<AppState>| async {
        // use state
    }))
    .with_state(AppState {});

let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
axum::serve(listener, routes).await.unwrap();

(docs.rs/axum より)

しかし、StateにはCloneが要求されるため、負荷をかけるとAppStateのcloneが並列的に大量に走り、メモリ使用量のピークをすさまじく押し上げます。
また、この小さな大量のcloneで確保されたヒープを、glibcなどは往々にしてOSに返しません。
ドキュメントには素のAppStateで載っているし、exampleもまちまちなので、わたしも素で使っていましたが、Arcで包めば簡単に劇的にメモリ効率、パフォーマンスを改善できたので覚えを書いておきます。

環境

$ cat /etc/debian_version 
13.2

$ dpkg -l | grep "GNU C Library"
ii  libc-bin                                                 2.41-12                                  amd64        GNU C Library: Binaries
Cargo.toml
[package]
name = "axum_state_arc_test"
version = "0.1.0"
edition = "2024"

[dependencies]
tokio = { version = "1", features = ["full"] }
axum = "0.8"
rand = "0.9"

Arcなし

main.rs
use axum::{
    Router,
    extract::{Path, State},
    routing::get,
};
use rand::Rng;
use rand::distr::Alphanumeric;
use tokio::net::TcpListener;

#[derive(Clone)]
struct AppState {
    state_vec: Vec<String>,
}

fn build_vec() -> Vec<String> {
    let mut vec = Vec::new();
    let mut rng = rand::rng();

    for _i in 0..100000 {
        let s = (0..32).map(|_| rng.sample(Alphanumeric) as char).collect();
        vec.push(s);
    }

    vec
}

#[tokio::main]
async fn main() {
    let state = AppState {
        state_vec: build_vec(),
    };

    let app = Router::new().route("/{id}", get(handler)).with_state(state);

    let listener = TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

async fn handler(State(state): State<AppState>, Path(id): Path<usize>) -> String {
    format!("{}: {}", id, state.state_vec[id])
}

適当な大きい配列をStateで共有します。
起動時のメモリ使用量は22MB。人間のCtrl-R長押しで50MBまで増えます。

$ wrk -t12 -c400 -d10s http://localhost:8080/12345
Running 10s test @ http://localhost:8080/12345
  12 threads and 400 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   552.03ms  381.57ms   1.96s    69.69%
    Req/Sec    29.52     17.94    90.00     57.90%
  1918 requests in 10.11s, 292.20KB read
  Socket errors: connect 0, read 0, write 0, timeout 153
Requests/sec:    189.77
Transfer/sec:     28.91KB

wrkにより負荷をかけると、1.8GBまでのぼり、また負荷後も使用量は下がりません。

Arcで包む

main.rs
use axum::{
    Router,
    extract::{Path, State},
    routing::get,
};
use rand::Rng;
use rand::distr::Alphanumeric;
+use std::sync::Arc;
use tokio::net::TcpListener;

+type AppState = Arc<AppStateInner>;

#[derive(Clone)]
+struct AppStateInner {
    state_vec: Vec<String>,
}

fn build_vec() -> Vec<String> {
    let mut vec = Vec::new();
    let mut rng = rand::rng();

    for _i in 0..100000 {
        let s = (0..32).map(|_| rng.sample(Alphanumeric) as char).collect();
        vec.push(s);
    }

    vec
}

#[tokio::main]
async fn main() {
+    let state = Arc::new(AppStateInner {
+        state_vec: build_vec(),
+    });

    let app = Router::new().route("/{id}", get(handler)).with_state(state);

    let listener = TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

async fn handler(State(state): State<AppState>, Path(id): Path<usize>) -> String {
    format!("{}: {}", id, state.state_vec[id])
}

起動時は8.4MB、Ctrl-R長押しで8.5MB。

$ wrk -t12 -c400 -d10s http://localhost:8080/12345
Running 10s test @ http://localhost:8080/12345
  12 threads and 400 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency     0.86ms    0.97ms  34.52ms   92.55%
    Req/Sec    41.91k     5.37k   72.23k    72.67%
  5010019 requests in 10.06s, 745.36MB read
Requests/sec: 498187.24
Transfer/sec:     74.12MB

wrkによる負荷で18MBとなりました。
また、Requests/secが2000倍以上伸びていますね。
Arcのcloneはメモリを複製せず参照カウンタを増やすだけなので、メモリ効率もパフォーマンスもかなり上げられます。

type AppState を作ると、既存のコードを変えずに済みますし、毎回 Arc<> を書く必要もありません。

Discussion