🚀

axum の State という新しい概念について

2022/12/19に公開

こんにちは，なっふぃです．

凍てつくような厳しい寒さが続く中，個人的にもっとも心躍った発見を共有します．

いつも axum という Rust 製のライブラリにお世話になっているのですが，今までは v0.4.8 を使っていました．
それで現在（2022-12-19）の最新バージョン v0.6 のドキュメントを眺めていたところ，State という新しい概念を見つけました．

詳しく読んでみると，今までの v0.4.8 で悩まされていたアレヤコレヤが華麗に解決できそうだったので，興奮冷めやらぬままにこの記事を書くことにしました．

axum とは

axum は Rust で Web サーバを書くためのライブラリで，Tower 系統のライブラリとの互換性が特徴です．

また型システムに優れていて，他の Web フレームワークよりも安心してコーディングできます．
ここで紹介する State も，axum の型を活かした機能になっています．

ちなみにアイキャッチは競合（？）のフレームワーク Rocket です．

サンプル

State が導入されたのは v0.6.0 からなので，v0.5.17 までは Extensions を用いて状態を扱っていました：

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

#[derive(Clone)]
struct Count(i64);

async fn get_state(Extension(state): Extension<Count>) -> String {
    state.0.to_string()
}

#[tokio::main]
async fn main() {
    let state = Count(0);

    let app = Router::new()
        .route("/", get(get_state))
        .layer(Extension(state));

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

v0.6.0 以降では次のように書けます：

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

#[derive(Clone)]
struct Count(i64);

async fn get_state(State(state): State<Count>) -> String {
    state.0.to_string()
}

#[tokio::main]
async fn main() {
    let state = Count(0);

    let app = Router::new().route("/", get(get_state)).with_state(state);

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

見た目はほとんど変わりませんが，型レベルでは大きく変わっています．
それを以下で説明していきます．

`Router` の型パラメータ

以前の Router は

pub struct Router<B = Body> { /* private fields */ }

という形でした．B にはほとんどの場合デフォルトの hyper::Body が入ります．

この Router の impl Service は

impl<B> Service<Request<B>> for Router<B>
where
    B: HttpBody + Send + 'static,
{
    type Response = Response;
    // 省略
}

となっているので，単純に「http::Request<Body> を受け取って axum::response::Response を返す」Service です．

一方で新しい Router は

pub struct Router<S = (), B = Body> { /* private fields */ }

と宣言されていて，impl Service は

impl<B> Service<Request<B>> for Router<(), B>
where
    B: HttpBody + Send + 'static,
{
    // 省略
}

と変わりました．

状態無し（S = ()）の場合にのみ impl されていますね．
同様に Router::into_make_service() も S = () でのみ定義されています．

ここが最初は戸惑うところで，「() 以外の状態を使いたいときには into_make_service() できないの？？？」と思ってしまいますね．私は思いました．

`Router::with_state` の仕組み

このトリックは Router::with_state() にあります．
with_state() の定義箇所は

impl<S, B> Router<S, B>
where
    B: HttpBody + Send + 'static,
    S: Clone + Send + Sync + 'static,
{
    pub fn with_state<S2>(self, state: S) -> Router<S2, B> {
        // 省略
    }
}

のようになっています．

あれ？ state: S を受け取って Router<S, B> を返すんじゃないの？ S2 って何？？

と言いたいところですが，すぐ下にちゃんと説明されていました．

Router<S> means a router that is missing a state of type S to be able to handle requests. It does not mean a Router that has a state of type S.

Router<S> は「リクエストを処理できるようになるためには，状態 S が足りないよ」という意味であって，「状態 S を今持っている」ことではない．

Router<S> は with_state() を呼び出して初めて，リクエストが処理できるようになるわけです．

これは実際は Service の作用順序を考えれば自然な仕様で，たとえば上の例

Router::new()
    .route("/", get(get_state))
    .with_state(state)

であれば，

        requests
           |
           v
+----- with_state -----+
|                      |
|    get(get_state)    |
|                      |
+----- with_state -----+
           |
           v
        responses

with_state のサービスが get_state ハンドラーへ state を与える形になっています．

これがもし

Router::new()
    .with_state(state)
    .route("/", get(get_state))

        requests
           |
           v
+--- get(get_state) ---+
|                      |
|      with_state      |
|                      |
+--- get(get_state) ---+
           |
           v
        responses

だったら「誰が get_state に state をあげるの？」となり破綻してしまいます．

実はこれが起こり得たのが v0.5 までの世界です．後ほど S の利点として説明します．

続いて S2 というパラメータについてですが，

pub fn with_state<S2>(self, state: S) -> Router<S2, B>

これは基本的に型推論によって決定されます．

つまり with_state を呼んだ時点では「どの S2 を受け取れば，リクエスト処理可能な状態になるのか」が分からないので，

続く self.with_state(state2) があれば，state2 の型に，
self.into_make_service() が続けば () に

推論されるわけです．

よくできていますね．

`S` パラメータの利点

さて，こんなに複雑な仕組みにしたのには当然意味があります．
実際の axum チームのディスカッションを見ていないので私の推測にはなりますが，いくつかの利点を挙げてみます．

まず上でもほんの少し触れましたが，route() と with_state() の順番を間違えてもコンパイルエラーになってくれます．
v0.5 であれば，

#[derive(Clone)]
struct Count(i64);

async fn get_state(Extension(state): Extension<Count>) -> String {
    state.0.to_string()
}

let state = Count(0);

let app = Router::new()
    .layer(Extension(state))
    .route("/", get(get_state));

これはランタイムエラーになります．get_state() が layer(Extension(state)) の範囲外なので，「Extension<Count> が存在しないよ」みたいに言われます．

コンパイル時ではなくランタイムにエラーを吐くというのは，Rust っぽくありません．

v0.6 以降なら，

#[derive(Clone)]
struct Count(i64);

async fn get_state(State(state): State<Count>) -> String {
    state.0.to_string()
}

let state = Count(0);

let app = Router::new()
    .with_state(state)
    .route("/", get(get_state));

はコンパイルエラーになるので安心できます．

第二の利点は，ハンドラー（get_state）の引数がより柔軟に設定できます．

たとえば状態が複数のフィールドを持つとき，

struct AuthToken {}
struct DatabasePool {}

struct AppState {
    auth_token: AuthToken,
    database_pool: DatabasePool,
}

Extension だと「auth_token だけ欲しい」という場合でも Extension<AppState> を使わなければいけません．

#[derive(Clone)]
struct AuthToken {}
#[derive(Clone)]
struct DatabasePool {}

#[derive(Clone)]
struct AppState {
    auth_token: AuthToken,
    database_pool: DatabasePool,
}

async fn get_token(Extension(state): Extension<AppState>) -> String {
    let token = &state.auth_token;
    //
}

let state = AppState::default();

let app = Router::new()
    .route("/", get(get_token))
    .layer(Extension(state));

State なら必要な部分のみで大丈夫です：

use axum::extract::FromRef;

#[derive(Clone)]
struct AuthToken {}
#[derive(Clone)]
struct DatabasePool {}

#[derive(FromRef, Clone)]
struct AppState {
    auth_token: AuthToken,
    database_pool: DatabasePool,
}

async fn get_token(State(token): State<AuthToken>) -> String {
    //
}

let state = AppState::default();

let app = Router::new()
    .route("/", get(get_token));
    .with_state(State(state))

ただし AuthToken は clone されるので，&AuthToken が欲しい場合は工夫を要します．

第三の利点はパフォーマンスです．

v0.5 までは Extensions という TypeId => Extension のマップ

HashMap<TypeId, Box<dyn Any + Send + Sync>, BuildHasherDefault<IdHasher>>;

を用いていたのに対し，State は

pub struct HandlerService<H, T, S, B> {
    handler: H,
    state: S,
    _marker: PhantomData<fn() -> (T, B)>,
}

のフィールドとして提供されるので，より高速です．

~~まあ小さいサーバなら，そこまでメリットにはならない気もします．~~

欠点

一方で，この仕組みには当然欠点もあり，それが許容できないなら従来の Extension で妥協するしかありません．

私的に一番大きな欠点は，複数の状態を扱えないことです．

たとえば

Router::new()
    .route("/token", get(get_token))
    .route("/pool", get(get_pool))

において

get_pool() は DatabasePool のみ必要
get_token() は DatabasePool と AuthToken の両方が必要

であれば，Extension だと次のように書くことができます：

Router::new()
    .route("/pool", get(get_pool))
    .layer(Extension(AuthToken {}))
    .route("/token", get(get_token))
    .layer(Extension(DatabasePool {}))

これならそれぞれの Extension の影響範囲がはっきりして，分かりやすくなります．

しかし State を2つの with_state() に分割することはできず，上の例のような AppState でまとめなければなりません．
State の影響範囲はコンパイラにしか判然としません．

こちらの方が「影響範囲が分かりやすくて気持ちいい！！」と思う人もいるでしょうが，私は Extension 派です．

他にも，FromRef を impl するのが面倒とか，型パラメータ S を付けるのが面倒とかありますが，本質的な問題にはなりませんね．

あと地味に面倒なのが，v0.5 以前からの移行です．
もともと Router<B> などと書いていたら，これは Router<S = B, B = Body> と解釈されてしまうので，適切に型を付け直す手間がかかります．

まとめ

axum の State の仕組みはややこしいですが，Rust の型の堅牢性を活かした非常に便利な機能です．
この State に気づいてから，私は毎日が幸せです．

皆さんもぜひ幸せになりましょう．

axum とは

サンプル

Router の型パラメータ

Router::with_state の仕組み

S パラメータの利点

欠点

まとめ

Discussion

`Router` の型パラメータ

`Router::with_state` の仕組み

`S` パラメータの利点