失われたメトリクスのケース: Rustクロージャーの謎

2025年1月26日 20:39

問題 - 不正確なメトリクス

RisingWaveのようなストリーミングデータベースでは、正確なメトリクスは単なる便利なものではなく、パフォーマンスの監視、ボトルネックの特定、システムがリアルタイムでスムーズに動作することを確保するために不可欠です。これは、車のダッシュボードの計器のようなもので、エンジンの下で何が起こっているのかをリアルタイムでフィードバックしてくれます。私たちが監視している重要な項目の1つは、ストレージ層からのデータ取得速度で、これはリアルタイムシステムにとって非常に重要です。

最近、奇妙な状況に遭遇しました：ストレージメトリクスの1つがゼロを報告しており、データがシステム内で流れていることがわかっているにもかかわらず、ゼロのままでした。これは、壊れたスピードメーターで車を運転しているようなもので、動いていることはわかりますが、どれくらいの速度で進んでいるのかがわからないという状態でした。これがきっかけとなり、Rustコードの奥深くを調査することになり、Rust 2021がクロージャーをどのように扱うかについて興味深い発見をしました。

RAIIとは？そしてそれが重要な理由

Rustには、RAIIという概念があります。これは「リソース取得は初期化」という意味で、パーティーの後片付けを自動で行ってくれるクリーンアップチームのようなものだと考えてください。

データ（リソース）がもはや必要でなくなったとき、RAIIはそれを適切に廃棄することを保証します。これは、メモリリークを防ぐために重要であり、私たちの場合、メトリクスが正しく報告されることを保証します。

例えば、あるイベントが発生するたびにインクリメントする必要があるカウンターがあるとしましょう。RAIIを使用すれば、このカウンターを設定して、使用が終わった時に自動的に値を報告できるようにできます。これにより、コードが複雑になっても、更新が見逃されることはありません。

しかし、この強力なツールを使用しても、私たちはメトリクスにゼロの値が表示され続けました。これは理にかなりませんでした。

私たちのセットアップ: `MonitoredStateStoreIterStats` と `try_stream`

すべてのデータポイントをキャプチャするために、私たちは特別なツール `MonitoredStateStoreIterStats` を設計しました。これは、処理されたデータ項目の数など、メトリクスを追跡します。データストリームで何が起こっているかを厳密に監視するカウンターのようなものです。

データがシステムを通過するたびに、このカウンターはそれを注意深く記録します。以下は、この構造体の簡略版です：

struct MonitoredStateStoreIterStats {
    total_items: usize,
    total_size: usize,
    storage_metrics: Arc<MonitoredStorageMetrics>,
}

また、Rustの便利な機能である`try_stream`を使って、データのストリームを扱うプロセスを簡素化しています。これは効率的にデータを運ぶコンベアベルトのようなものです。`MonitoredStateStoreIterStats`はこのベルトの横に座り、通過するすべての項目をカウントします。

次のようにストリームを設定しました：

pub struct MonitoredStateStoreIter<S> {
    inner: S,
    stats: MonitoredStateStoreIterStats,
}

impl<S: StateStoreIterItemStream> MonitoredStateStoreIter<S> {
    #[try_stream(ok = StateStoreIterItem, error = StorageError)]
    async fn into_stream_inner(mut self) {
        let inner = self.inner;
        futures::pin_mut!(inner);
        while let Some((key, value)) = inner.try_next().await? {
            self.stats.total_items += 1;
            self.stats.total_size += key.encoded_len() + value.len();
            yield (key, value);
        }
    }
}

`MonitoredStateStoreIterStats`が終了すると、それは自動的に収集されたメトリクスを監視システムに報告します。これは、`Drop`の実装で行われます：

impl Drop for MonitoredStateStoreIterStats {
    fn drop(&mut self) {
        self.storage_metrics.iter_item.observe(self.total_items as f64);
        self.storage_metrics.iter_size.observe(self.total_size as f64);
    }
}

調査: コードの掘り下げ

何が問題なのかを解明するために、コードの簡略版を作成しました。これは、実験室でのミニ実験のようなものです。

私たちはRust Playgroundを使ってこのテストを実行しました。問題を再現するために使用したコードは以下の通りです：

struct Stat {
    count: usize,
    vec: Vec<u8>,
}

impl Drop for Stat {
    fn drop(&mut self) {
        println!("count: {}", self.count);
    }
}

fn main() {
    let mut stat = Stat {
        count: 0,
        vec: Vec::new(),
    };

    let mut f = move || {
        stat.count += 1;
        1
    };

    println!("num: {}", f());
}

このコードでは、`MonitoredStateStoreIterStats`に似た`Stat`構造体を持っています。`count`フィールドは何かが発生するたびに回数を追跡し、`vec`フィールドは後で使用します。`Drop`の実装では、`Stat`がドロップされるときに`count`を表示します。

`main`関数では、`Stat`インスタンスとクロージャ`f`を作成します。このクロージャは呼び出されるたびに`count`を1増やすことになっています。

しかし、このコードを実行すると、次の出力が得られます：

num: 1
count: 0

`num: 1`はクロージャが呼ばれ、1を返したことを示しています。しかし、`count: 0`は`Stat`構造体の`count`フィールドが期待通りにインクリメントされなかったことを示しています。これは驚くべきことで、RisingWaveで見ていた問題と一致していました。私たちのメトリクスが更新されなかったのです。

罪の元凶：Rustのクロージャの挙動

私たちは、問題がRustがクロージャを処理する方法に関連していることを発見しました。クロージャはコード内の小さな関数のようなものです。通常、クロージャ内で変数を使用すると、その変数を「キャプチャ」します。これは、変数の所有権を取得するか、使用方法によってはその変数を借用することを意味します。

`move` キーワードを使用すると、面白いことが起こります。`move` を使用すると、クロージャはそれが使用するすべての変数の所有権を強制的に取得します。これで問題が解決すると思うかもしれませんが、実際には問題がさらに複雑になったのです。

アナロジー：選択的ショッパー

ショッピングリスト（データ）を持っていて、誰か（クロージャ）を店に送り出すことを想像してください。Rustの古いバージョンでは、ショッパーはリスト全体を持ち帰っていました、たとえ一つのアイテムだけが必要でも。

しかし、Rust 2021では、ショッパーはより効率的です。彼らは、タスクに関連する部分だけを持ち帰ります。

これは通常、メモリを節約し、速度を向上させる良いことです。しかし、私たちの場合、これにより「カウンター」が指標を更新するために必要な情報を取得できなかったのです。まるでショッパーが「リンゴ」リストだけを持ち帰り、「数量」部分を無視していたかのようです。この「数量」こそ、彼らに何個のリンゴを買うべきかを伝える部分です。

詳細の理解

これがなぜ起こっているのかを理解するために、いくつかのシナリオを見てみましょう。

構造体全体を移動

まず、クロージャに`stat`全体を明示的に移動するように変更してみました：

let mut f = move || {
    let mut stat = stat; // statをクロージャ内に明示的に移動
    stat.count += 1;
    1
};

この変更を加えると、出力は次のようになります：

num: 1
count: 1

今、`count`が正しく1にインクリメントされています。これは、`stat`を明示的にクロージャ内に移動させることで、クロージャがその完全な所有権を取得することを示しています。クロージャ内での変更は、元の`stat`に反映されます。

他のフィールドの使用

次に、クロージャ内で`stat`の他のフィールド`vec`を使用してみました：

let mut f = move || {
    let _ = stat.vec.len(); // vecフィールドを使用
    stat.count += 1;
    1
};

驚くべきことに、出力はまだ正しいです：

num: 1
count: 1

これは、`vec`のような`Copy`を実装していないフィールドを使用すると、クロージャが`stat`全体の所有権を強制的に取得することを示しています。単に`count`フィールドだけをコピーするわけではありません。

`Copy`型の問題

問題は、クロージャが`Copy`を実装しているフィールド（例えば、`usize`型の`count`フィールド）だけを使用する場合に発生します。この場合、`move`を使用しても、クロージャは構造体全体の所有権を取得せず、`Copy`フィールドだけがコピーされます。

私たちは、`stat`が`Drop`を実装しているため部分的に移動できないものの、`stat`内の`Copy`型フィールドのみがクロージャ内で使用される場合、これらのフィールドはクロージャ内にコピーされ、クロージャ内でこれらのフィールドに対する変更はコピーされた値にしか影響を与えず、元のフィールドは変更されないと仮定しました。

仮説の確認

仮説を確認するために、Rustコンパイラが生成するコードを確認しました。私たちは、Rustコンパイラの内部コードを確認するためにMIR（中間表現）というツールを使用しました。

問題のあるコードに対して、MIRは次のように表示されました：

      _3 = {closure@src/main.rs:19:17: 19:24} { stat: (_1.0: usize) };

しかし、`stat`を明示的に移動させた修正後のコードでは、MIRは次のように表示されました：

      _3 = {closure@src/main.rs:19:17: 19:24} { stat: move _1 };

これにより、私たちの仮説が確認されました。問題のあるコードでは、クロージャは`stat`の所有権をキャプチャせず、`Copy`フィールドのみをコピーしていることがわかりました。

解決策

修正は驚くほど簡単でした。クロージャに`stats`オブジェクト全体の所有権を明示的に渡す必要がありました。

以下は、元の問題のあるコードです：

#[try_stream(ok = StateStoreIterItem, error = StorageError)]
async fn into_stream_inner(mut self) {
    let inner = self.inner;
    ...
    self.stats.total_items += 1;
    self.stats.total_size += key.encoded_len() + value.len();
    ...
}

そして、修正されたコードは次の通りです：

#[try_stream(ok = StateStoreIterItem, error = StorageError)]
async fn into_stream_inner(self) {
    let inner = self.inner;
    let mut stats = self.stats; // self.statsの所有権を取得
    ...
    stats.total_items += 1;
    stats.total_size += key.encoded_len() + value.len();
    ...
}

`let mut stats = self.stats;` を追加することで、クロージャが`stats`オブジェクトの所有権を取得するようになります。これで、`stats.total_items`や`stats.total_size`をインクリメントする際に、実際の`stats`オブジェクトが変更されることが保証されます。

より広い視点：Rust開発者への影響

私たちの調査は、Rust 2021の`move`クロージャの挙動に関する微妙な点を浮き彫りにしました。クロージャは必要な最小限の構造体の部分だけをキャプチャします。効率的である一方で、`Copy`型のフィールドを扱う構造体内でこの動作を理解していないと、予期しない結果を引き起こす可能性があります。

私たちの場合、この挙動が指標の損失を引き起こしました。クロージャを内部で使用する`try_stream`マクロのカプセル化が、問題をさらに複雑にしました。これにより、パラメータの所有権をストリームに移動する意図が、期待通りに移動されない状況が生まれました。

私たちは、以下の2つの改善点を特定しました：

マクロのカプセル化：`try_stream`マクロは、所有権の扱いについてもっと明示的であるべきです。これにより、開発者が所有権の移動が行われていないと誤解することを防げます。
Rustのクロージャの挙動：Rustの挙動は技術的には正しいものの、混乱を招く可能性があります。この問題について、Rustコミュニティに問題報告を行い、明確さやドキュメントの改善について議論しています。

この発見はRisingWaveのコードベースにとっても広範な影響を及ぼします。私たちは、同様のパターンに依存している他の部分を慎重にレビューし、データ損失や予期しない動作を防ぐ必要があります。

結論：学び取った教训

この経験は、Rustの所有権およびクロージャの挙動の微妙な違いを理解する重要性を再認識させてくれました。特に、Rust 2021で導入された変更に関しては、その理解が重要です。また、マクロAPIの設計が曖昧さを避けるために慎重である必要があることも再確認できました。

RisingWaveにとって、これは単なるバグ修正以上のものでした。それは、より堅牢なシステムに向けた貴重な学びの経験であり、今後はより信頼性の高いコードを書くための助けとなるでしょう。

私たちの経験を共有することで、Rustコミュニティに貢献できることを願っています。他の開発者がクロージャ、`move`のセマンティクス、およびマクロカプセル化に関する類似の落とし穴を避けられることを期待しています。RisingWaveと、Rustを使って次世代のストリーミングデータベースを構築している方法についてさらに学びたい方は、GitHubリポジトリをご覧いただくか、私たちのコミュニティに参加してください！

いいなと思ったら応援しよう！

この記事が参加している募集

#note書き初め

18,059件