AxumのExtensionsを自前で作りましょう - Rust

概要

Axum Webサーバーのクレートで使われているExtensionsの仕組みを説明し、コードで簡略版を再現します。

AxumのExtensionsとは

Axumを愛用している筆者ですが、Axumの主な便利な機能として、Extractorsがあります。以下、Axumのドキュメントからの抜粋です：

use axum::{
    extract::{Request, Json, Path, Extension, Query},
    routing::post,
    http::header::HeaderMap,
    body::{Bytes, Body},
    Router,
};
use serde_json::Value;
use std::collections::HashMap;

// `Path` gives you the path parameters and deserializes them. See its docs for
// more details
async fn path(Path(user_id): Path<u32>) {}

// `Query` gives you the query parameters and deserializes them.
async fn query(Query(params): Query<HashMap<String, String>>) {}

// `HeaderMap` gives you all the headers
async fn headers(headers: HeaderMap) {}

// `String` consumes the request body and ensures it is valid utf-8
async fn string(body: String) {}

// `Bytes` gives you the raw request body
async fn bytes(body: Bytes) {}

// We've already seen `Json` for parsing the request body as json
async fn json(Json(payload): Json<Value>) {}

// `Extension` extracts data from "request extensions"
// This is commonly used to share state with handlers
async fn extension(Extension(state): Extension<State>) {}

#[derive(Clone)]
struct State { /* ... */ }

let app = Router::new()
    .route("/path/:user_id", post(path))
    .route("/query", post(query))
    .route("/string", post(string))
    .route("/bytes", post(bytes))
    .route("/json", post(json))
    .route("/request", post(request))
    .route("/extension", post(extension));

引用：https://docs.rs/axum/latest/axum/extract/index.html

さまざまなExtensionがありますが、今回の記事で注目したいのは、Extensionです。筆者は、何かのサイズを知らないと鞭打ちにされるRustの世界に慣れているので、「何を入れてもOK,何を入れても出せるよ」というExtensionのニュアンスが不思議に思えてなりませんでした；なぜこのようなことがRustで可能でしょうか？

このExtensionの仕組みが興味深くて面白いです。

RustのAny型について

Extensionの仕組みの説明に入る前に、必要な知識があります。それはRustのAnyタイプについてです。

trait Anyは、動的タイプdynと組み合わせて使って、どの型でもその型のユニークな識別番号を教えてくれるメソッドです。それは.type_idのメソッドで、TypeIdのstructを返してくれます。

以下、Rustの正式ドキュメントからの抜粋：

use std::any::{Any, TypeId};

let boxed: Box<dyn Any> = Box::new(3_i32);

// You're more likely to want this:
let actual_id = (&*boxed).type_id();
// ... than this:
let boxed_id = boxed.type_id();

assert_eq!(actual_id, TypeId::of::<i32>());
assert_eq!(boxed_id, TypeId::of::<Box<dyn Any>>());

引用：https://doc.rust-lang.org/std/any/index.html

このAny型にはもう一つ大事な機能があります：downcast。

downcastメソッドは、trait Anyの動的オブジェクトを任意のstructに変換することを試みるものです。とある&dyn Anyが&Tだったら、&Tを返してくれます。

use std::any::Any;

fn print_if_string(value: Box<dyn Any>) {
    if let Ok(string) = value.downcast::<String>() {
        println!("String ({}): {}", string.len(), string);
    }
}

let my_string = "Hello World".to_string();
print_if_string(Box::new(my_string));
print_if_string(Box::new(0i8));

_引用：https://doc.rust-lang.org/std/any/trait.Any.html#method.downcast

これがとても面白いのです。そして、AxumのExtensionの抽出詞の秘訣もここにあるのです。

AxumのExentionsの仕組み

AxumのExtensionsの仕組みは、上記のtrait Anyを生かしたHashMapなのです。

具体的に言えば、HashMap<TypeId, Box<dyn Any>>なのです。そして、Axumがdowncastを使ってそのHashMapから取り出しているのです。

実際、同じものを作ってみましょう！

AxumのExtensionsを再現する

まず、HashMapの型を定義しましょう。

use std::{
    any::{Any, TypeId},
    collections::HashMap,
};

type AnyMap = HashMap<TypeId, Box<dyn Any>>;

そして、structも定義しましょう。

struct Extensions {
    map: AnyMap,
}

次に、このExtensionsにリクエストについていた何かFromRequestPartsを実装しているものを分解して入れていけるようにしたいです。そのためにはinsertメソッドが必要です。

impl Extensions {
    fn insert<T>(&mut self, obj: T)
    {
        let key = obj.type_id().to_owned();
        let obj_box = Box::new(obj);

        self.map.insert(key, obj_box);
    }
}

上記のコードでTypeIdを抽出して、Boxに梱包して、入れるという単純な流れができていますが、問題があります：

error[E0310]: the parameter type `T` may not live long enough
  --> axum_extension_map/src/main.rs:15:19
   |
15 |         let key = obj.type_id().to_owned();
   |                   ^^^^^^^^^^^^^
   |                   |
   |                   the parameter type `T` must be valid for the static lifetime...
   |                   ...so that the type `T` will meet its required lifetime bounds
   |
help: consider adding an explicit lifetime bound
   |
13 |     fn insert<T: 'static>(&mut self, obj: T)
   |                +++++++++

コンパイラーのご指摘の通りですが、一旦ここで、なぜTが'staticまで生き延びられなければならないのか、それについて考えたいのです。まず、'staticを要求しているのはtrait Anyです。trait Any: 'staticになっているのです。なぜAnyが'staticになるのかというと、そもそも、全ての型は、コンパイル時に生成され、コンパイルされたコードの中にあるからです。ソースコードに"Hello World"を書くときに、&'static strになるのと同様に、全てのstructそのものが'staticのライフタイムになります。

なので、AnyがAny: 'staticは自明なのですが、一応RustのコンパイラーはTが'staticまで生き延びるかどうか、わからないので、制約してあげないと怒るのです。

なので、以下のコードを追加します。

    fn insert<T>(&mut self, obj: T)
    where
        T: 'static,

そして、入れっぱなしじゃ役に立たないので、Refを取り出せるようにもしておきましょう：

    fn get<T>(&self) -> Option<&T>
    where
        T: 'static,
    {
        let key = TypeId::of::<T>();
        let obj = self.map.get(&key);

        obj.and_then(|any_obj| any_obj.downcast_ref())
    }

試してみると以下のようなことができます！

#[derive(Debug)]
struct Foo {
    s: usize
}

#[derive(Debug)]
struct Bar {
    s: String,
}

fn main() {
    let mut extensions = Extensions {
        map: HashMap::new()
    };

    extensions.insert(Foo { s: 48 });
    extensions.insert(Bar { s: String::from("Hello World") });

    let foo = extensions.get::<Foo>();
    assert!(foo.is_some());
    dbg!(foo);
}

おめでとう、Axumの関心な部分の秘密を解明できたのです！

実際のAxumのコード

実際は、Axumではなく、Hyperiumのhttpのものを使っていたのです。

    pub fn insert<T: Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
        self.map
            .get_or_insert_with(|| Box::new(HashMap::default()))
            .insert(TypeId::of::<T>(), Box::new(val))
            .and_then(|boxed| {
                (boxed as Box<dyn Any + 'static>)
                    .downcast()
                    .ok()
                    .map(|boxed| *boxed)
            })
    }

引用：https://github.com/hyperium/http/blob/bda93204b3da1a776cf471ed39e8e374cec652e7/src/extensions.rs#L62

Send + Syncも入れて、Asyncにも対応できるようにしていますね。

まとめ

以上、Axum（httpというべきか）のExtensionsの仕組みについて解説しましたが、いかがでしょうか？

Rustでも多少の魔法はありますが、きちんと読んでみると、解読できるというのも、Rustの素晴らしいところです。また、Axumみたいに、いろんないいコードがあって、常に勉強になります。

良いお年を！