🚀

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

2022/12/19に公開約8,500字

こんにちは,なっふぃです.

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

いつも 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 が入ります.

この Routerimpl 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_statestate をあげるの?」となり破綻してしまいます.

実はこれが起こり得たのが 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()DatabasePoolAuthToken の両方が必要

であれば,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 に気づいてから,私は毎日が幸せです.

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

Discussion

ログインするとコメントできます