Rustの非同期タスクをリアルタイムで可視化するツール「await-tree」を作ってみた

*[Await-Tree](https://github.com/risingwavelabs/await-tree)：非同期タスクの正確で有益なツリー構造ダンプを生成します。このツールは Apache License（バージョン 2.0）のもとで配布されています。*

Async Rust における `Future` は、様々な制御フローを実現するために、自由に合成・ネストすることが可能です。それぞれの `Future` の実行をノードとして表現するならば、非同期タスクの実行全体は、**論理的なツリー構造**として整理できます。そしてこのツリーは、`Future` のポーリング、完了、キャンセルの過程で絶えず変化していきます。

`await-tree` は、各 `Future` に `instrument_await` を付与することで、この実行ツリーをランタイム中にダンプできるようにします。以下に基本的な例を示します。より複雑な制御フローの例については、[examples ディレクトリ](https://github.com/risingwavelabs/await-tree/blob/main/examples)をご参照ください。

````rust
async fn bar(i: i32) {
    // `&'static str` のスパン
    baz(i).instrument_await("baz in bar").await
}

async fn baz(i: i32) {
    // ランタイムで生成される `String` のスパンにも対応
    work().instrument_await(span!("working in baz {i}")).await
}

async fn foo() {
    // joinされたFutureのスパンはツリー上で兄弟ノードになります
    join(
        bar(3).instrument_await("bar"),
        baz(2).instrument_await("baz"),
    )
    .await;
}

// タスクのトレースを開始するためにグローバルレジストリを初期化
await_tree::init_global_registry(Default::default());
// ルートスパン "foo" とキー "foo" を用いてタスクを spawn
// 注意：`spawn` 関数の使用には `tokio` 機能の有効化が必要です
await_tree::spawn("foo", "foo", foo());
// タスクをしばらく実行させる
sleep(Duration::from_secs(1)).await;
// キー "foo" を持つタスクのツリーを取得
let tree = Registry::current().get("foo").unwrap();

// foo [1.006s]
//   bar [1.006s]
//     baz in bar [1.006s]
//       working in baz 3 [1.006s]
//   baz [1.006s]
//     working in baz 2 [1.006s]
println!("{tree}");
````

## **機能（Features）**

`await-tree` は以下のオプション機能を提供します：

- `serde`: serde を用いたツリー構造のシリアライズを可能にします。これにより、[serde の例](https://github.com/risingwavelabs/await-tree/blob/main/examples/serde.rs)で示されているように、ツリーを JSON などの形式でシリアライズできます。
    
    ````rust
    // Cargo.toml で serde 機能を有効化
    // await-tree = { version = "<version>", features = ["serde"] }
    
    // その後、ツリーをシリアライズできます
    let tree = Registry::current().get("foo").unwrap();
    let json = serde_json::to_string_pretty(&tree).unwrap();
    println!("{json}");
    ````
    
- `tokio`: Tokio ランタイムとの統合を可能にし、`spawn` や `spawn_anonymous` 関数によるタスク生成機能を提供します。この機能は、タスク生成を示す例で必要となります。
    
    ````rust
    // Cargo.toml で tokio 機能を有効化
    // await-tree = { version = "<version>", features = ["tokio"] }
    
    // その後、await-tree による計測付きでタスクを生成可能
    await_tree::spawn("task-key", "root_span", async {
        // ここに非同期処理を記述
        work().instrument_await("work_span").await;
    });
    ````


## **`async-backtrace` との比較**

[`tokio-rs/async-backtrace`](https://github.com/tokio-rs/async-backtrace) は、非同期タスクの実行ツリーをダンプできる類似のクレートです。以下に、`await-tree` と `async-backtrace` の主な違いを示します。

### **`await-tree` の利点：**

- `await-tree` は実行時の `String` を用いたスパンのカスタマイズが可能ですが、`async-backtrace` は関数名や行番号のみのサポートです。
  
    これは、共有リソース（例：ロック）の識別子のような動的情報でスパンを注釈したい場合に非常に有用です。これにより、異なるタスク間での競合の発生状況を可視化できます。

- `await-tree` は、任意の `Future` トポロジーを持つあらゆる種類の非同期制御フローに対応していますが、`async-backtrace` は一部のケースに対応できません。

    例えば、キャンセルに対する安全性の問題を避けるため、`&mut impl Future` を `select` の分岐として使うことは一般的です。`select` の完了後にこの `Future` を別の場所に移動して `await` し直すことがありますが、`async-backtrace` は親の変更によりこの `Future` を再度追跡できなくなります。詳細は [`examples/detach.rs`](https://github.com/risingwavelabs/await-tree/blob/main/examples/detach.rs) を参照してください。

- `await-tree` は [arena ベースのデータ構造](https://crates.io/crates/indextree) を使ってツリー構造を保持しており、追加の `unsafe` コードは一切使用していません。比較として、`async-backtrace` はこの構造を手作業で構築しており、前述の未対応トポロジーに対してメモリ安全性の問題を抱える可能性があります。

    ちなみに、`await-tree` は [RisingWave](https://github.com/risingwavelabs/risingwave)（分散型ストリーミングデータベース）のプロダクション環境で長期間運用されています。

- `await-tree` は `Future` 自体とは独立にツリー構造を保持しているため、`Future` がアクティブにポーリング中であっても、ペンディング状態であっても、競合なくいつでもツリーをダンプすることができます。一方で、`async-backtrace` はツリーをダンプする前にポーリングの完了を[待つ](https://docs.rs/async-backtrace/0.2.5/async_backtrace/fn.taskdump_tree.html)必要があり、それが長い遅延につながる可能性があります。

### **`async-backtrace` の利点：**

- `async-backtrace` は Tokio 組織の公式プロジェクトです。