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

# 概要

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

# AxumのExtensionsとは

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

```rust
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の正式ドキュメントからの抜粋：

```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`を返してくれます。

```rust
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の型を定義しましょう。

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

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

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

```rust
struct Extensions {
    map: AnyMap,
}
```

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

```rust
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`まで生き延びるかどうか、わからないので、制約してあげないと怒るのです。

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

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

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

```rust
    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())
    }
```

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

```rust
#[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のものを使っていたのです。

```rust
    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みたいに、いろんないいコードがあって、常に勉強になります。

良いお年を！