Zenn
😸

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

に公開
GitHub
Rust
プログラミング
debug
async
tech

Await-Tree：非同期タスクの正確で有益なツリー構造ダンプを生成します。このツールは Apache License（バージョン 2.0）のもとで配布されています。

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

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

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 の例で示されているように、ツリーを JSON などの形式でシリアライズできます。

    // 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 関数によるタスク生成機能を提供します。この機能は、タスク生成を示す例で必要となります。

    // 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 は、非同期タスクの実行ツリーをダンプできる類似のクレートです。以下に、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 を参照してください。

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

    ちなみに、await-tree は RisingWave（分散型ストリーミングデータベース）のプロダクション環境で長期間運用されています。

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

async-backtrace の利点：

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

Discussion