高速に metrics を記録するための仕組み

概要

アプリケーションを動かす際にその内部状態を外部から観測したいことが多く、その1つの方法にメトリクスを出力しそれを可視化するというやり方があります。こんな感じです。

今回は Rust のアプリケーションでメトリクスを記録するのに用いられる metrics-rs/metrics の実装を追います。

前提

この記事はこちらの commit に関する情報を元にしてます。

https://github.com/metrics-rs/metrics/commit/e6cf12457b4e9b81eb319f4aab70a573be7998b3

metrics-rs/metrics とは

https://metrics.rs/

metrics makes it easy to instrument your application to provide real-time insight into what's happening. It provides a number of practical features that make it easy for library and application authors to start collecting and exporting metrics from their codebase.

https://github.com/metrics-rs/metrics?tab=readme-ov-file#whats-it-all-about

要はメトリクスの収集と外部への出力を簡単にできるようにしたものです。

使い方は非常に簡単です。次のような関数があったとします。

pub fn process(query: &str) -> u64 {
    let row_count = run_query(query);
    row_count
}

これに metrics を導入し、run_query の所要測定と返却されたデータ数を記録したい場合は以下のようにできます。

+use metrics::{counter, histogram};
+
 pub fn process(query: &str) -> u64 {
+    let start = Instant::now();
     let row_count = run_query(query);
+    let delta = start.elapsed();
+
+    histogram!("process.query_time").record(delta);
+    counter!("process.query_row_count").increment(row_count);
+
     row_count
 }

元々のコードが小さいので大きい変更に見えますが、どれだけ処理が複雑になってもこれ以上 metrics の導入が煩雑になることはありません。少しだけ丁寧に見ておくと、例えば以下の行では

counter!("process.query_row_count").increment(row_count);

process.query_row_count という名前のカウンターの値を row_count だけ増やすという処理をしています。注目いただきたいのは、process.query_row_count という名前のカウンターを別の場所で初期化したり metrics 関連の変数を関数間で引き渡したりする必要はなしにその場でいきなり利用できるという手軽さです。
また、どこにメトリクスを出力するのかという部分は完全に切り離されているのでライブラリは意識する必要はありません。アプリケーション側（実行する側）で必要に応じて以下のように「recorder」を定義します。

fn main() {
    // Run a Prometheus scrape endpoint on 127.0.0.1:9000.
    let _ = PrometheusBuilder::new()
        .install()
        .expect("failed to install prometheus exporter");
}

Go だと以下のライブラリが相当しそうです。
https://github.com/hashicorp/go-metrics
自分はちゃんと使ったことないですが、examples を見ると

func SlowMethod() {
    // Profiling the runtime of a method
    defer metrics.MeasureSince([]string{"SlowMethod"}, time.Now())
}

// Configure a statsite sink as the global metrics sink
sink, _ := metrics.NewStatsiteSink("statsite:8125")
metrics.NewGlobal(metrics.DefaultConfig("service-name"), sink)

// Emit a Key/Value pair
metrics.EmitKey([]string{"questions", "meaning of life"}, 42)

となっており、使い心地は比較的 Rust のものと似ているように感じます。なので以降で話すことに関しても似たような形になっているかもしれません。そうでないかもしれません。教えてください。

導入

メトリクスはアプリケーションのあらゆる部分、さまざまなレイヤーで欲しくなります。メトリクスの数も動的に増加しますし、1つのメトリクスの更新が集中することもあるでしょう。つまり、同時に様々なところから並行でアクセスされ必要に応じて上書きされるということです。そのため、実装によってはアプリケーションとは無関係のメトリクスを処理するための部分で律速になってしまう可能があります。そのため、どうメトリクスを記録する部分の実装は高速化のためにさまざまな工夫が必要そうです。

皆さんはどう実装しますか？ 最もナイーブな方法から考えてみます。

以降では、メトリクスのうち、Counters のみを扱いますが、Gauges, Histograms なども同じ議論が当てはまります。むしろ、Histograms のような少し更新が複雑なものの方が効果は大きいかもしれません。

案1: Arc<Mutex<HashMap<Key, u64>>>

Rust で並行処理をする際に必ずお世話になるのが、Arc<Mutex<_>> です。メトリクスを記録する上でも Arc<Mutex<_>> を用いた実装が一番初めに思いつくと思います。その場合、

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

fn main() {
    let counters = Arc::new(Mutex::new(HashMap::<String, u64>::new()));

    {
        // lock then get inner HashMap
        let mut hmap = counters.lock().unwrap();
        // get or create counter for metric "metric-name"
        let counter = hmap.entry("metric-a".to_owned()).or_insert(0);
        // increment the counter
        *counter += 1;
        // lock is released here
    }
}

のような実装ができます。（あくまでイメージです）

ちなみに、自分もこの方と基本的に同じ感覚を持っています。とりあえず、Arc<Mutex<_>> で囲みます。
https://x.com/0918nobita/status/1841071225646563401

案2: Arc<Mutex<HashMap<Key, Arc<AtomicU64>>>>

なるべく全体へのロックを獲得する時間を短くしたいなら、HashMap のバリューを Arc<AtomicU64> にすることもできるでしょう。

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

fn main() {
    let counters = Arc::new(Mutex::new(HashMap::<String, Arc<AtomicU64>>::new()));

    {
        // lock then get or create counter for metric "metric-name"
        let counter = counters
            .lock()
            .unwrap()
            .entry("metric-a".to_owned())
            .or_insert(Arc::new(AtomicU64::new(0)))
            .clone();
        // LOCK IS RELEASED HERE!

        // increment the counter
        counter.fetch_add(1, Ordering::SeqCst);
    }
}

この counter は可変参照ではないことに注意してください。つまり、既にキーが存在する場合は特に counters に書き込みを行なってはいないです。

ここは内部可変性を用いています。

内部可変性は、そのデータへの不変参照がある時でさえもデータを可変化できるRustでのデザインパターンです

https://doc.rust-jp.rs/book-ja/ch15-05-interior-mutability.html

Counters の場合、メトリクスの更新処理が非常に単純（fetch_add）であるため、Arc<AtomicU64>::clone() が必要であるを考えるとそこまでパフォーマンスは向上しないかもしれません。

きっとまだまだ改善できる部分があるはず！そう信じて metrics-rs/metrics のコードを読んでみましょう。

コード

couter! マクロの実装が以下になります。
https://github.com/metrics-rs/metrics/blob/ce9084b948b48cb087b3fa973a348f8ead315690/metrics/src/macros.rs#L106-L123

これを見ると、おそらく、再帰の部分ではないところ、つまり以下の部分で必要な処理がされていそうだと推測できます。つまり、recorder.register_counter の部分です。
https://github.com/metrics-rs/metrics/blob/ce9084b948b48cb087b3fa973a348f8ead315690/metrics/src/macros.rs#L112

では、recorder とはなんでしょう？名前から記録のための仕組みのようだとわかりますが、これ自体は単なる Record trait ですね。つまり、実際にどうメトリクスを記録するかという部分は Strategy パターンが用いられており柔軟に変更できるようになっています。
https://github.com/metrics-rs/metrics/blob/d97f80146b0f5e0ea4f8f0809fdfa76bb49357f5/metrics/src/recorder/mod.rs#L21-L55

今回は、Prometheus 向けの実装について深ぼってみます。Prometheus のための Recorder の実装は以下にあります。
https://github.com/metrics-rs/metrics/blob/ed64fb60152cdc91b241d22d2facb33058a1fdad/metrics-exporter-prometheus/src/recorder.rs#L245-L253

フィールドの inner は複数スレッドからアクセスする必要があるので Arc が必要です。Arc は内部に参照カウントをもつのでアトミックなインクリメント処理があるのですが、ここが律速になることはないでしょう。

recorder.register_counter の実装は以下。self.inner.registry がメトリックの名前とその値を管理していることがわかります。
https://github.com/metrics-rs/metrics/blob/ed64fb60152cdc91b241d22d2facb33058a1fdad/metrics-exporter-prometheus/src/recorder.rs#L293-L295

registry は以下ですね。まさにこの Registry は HashMap とシグネチャが非常に似ており、ここに高速に書き込むための仕組みがありそうです。
https://github.com/metrics-rs/metrics/blob/ed64fb60152cdc91b241d22d2facb33058a1fdad/metrics-exporter-prometheus/src/recorder.rs#L19-L27

Registry の実装は、Vec<RwLock<RegistryHashMap<K, S::Counter>>> を内部で持っているようです。つまり、RegistryHashMap という HashMap 相当のものに Vec<RwLock<_>> が被せてあります。一体これはどういうことなのでしょう。
https://github.com/metrics-rs/metrics/blob/fa95f3a9b1c23056de5b7feb8238323b98d86c85/metrics-util/src/registry/mod.rs#L49-L58

RwLock について

まず RwLock についてですが、これは同時に、複数の reader か1つの writer を許可するための型です。一方、Mutex は同時に1つの reader か writer を許可します。
https://doc.rust-lang.org/std/sync/struct.RwLock.html

一般的に、RwLock は read 頻度 >> write 頻度の場合に、 Mutex はそれ以外の場合に好まれることが多いです。^[1]
metrics では、メトリクスは作成された後は内部可変性を用いて更新されるので writer が必要なのは作成時のみ。圧倒的に read する頻度の方が大きいです。 つまり、Mutex よりも RwLock の方が好ましいです。Mutex を用いることで必要最低限だけ writer がロックを獲得し、read に関しては並列で可能ということが可能になりました。

RwLock は read 頻度 >> write 頻度の場合に、

上でこのように述べましたが、少し前までは、std::sync::RwLock の場合、ずっと reader がロックをとり続けていつまで経っても writer が書き込みをできない（writer starvation）問題がありました。そのため、parking_lot というライブラリの使用が推奨されてました。しかし、Rust 1.62.0 以降ではこの問題は解決されているようです。
https://blog.rust-lang.org/2022/06/30/Rust-1.62.0.html

実際、metrics の issue を見ると Rust 1.62.0 のリリース以降 parking_lot から std::sync::RwLock になったことがわかります。

https://github.com/metrics-rs/metrics/issues/343

Vec<_> について

一つの場所へのアクセスが集中するなら分散させれないかを考えることはよくあります。この Vec<_> はまさにそのためにあります。つまり、write ロックが必要な場所を分散させることでボトルネックになるのを回避するということです。
念の為、コードを追ってみます。まず、Registry の初期化の部分です。
https://github.com/metrics-rs/metrics/blob/fa95f3a9b1c23056de5b7feb8238323b98d86c85/metrics-util/src/registry/mod.rs#L80-L96

shard_count の長さの Vec を用意してます。この shard_count は利用可能な CPU の数で初期化されます。次に、これを実際に用いる部分です。

https://github.com/metrics-rs/metrics/blob/fa95f3a9b1c23056de5b7feb8238323b98d86c85/metrics-util/src/registry/mod.rs#L346-L374

この部分で key の hash と対応する shard を取得します。これによって shard の数だけアクセスを分散することができます。

let (hash, shard) = self.get_hash_and_shard_for_counter(key);

それ以降の部分では、key に対応する entry がすでに存在すれば reader を返し、なければ write ロックを取得して entry を作成してます。

Redis Cluster も key のハッシュに基づいてデータの所属するノードを決定しているので、仕組みは同じと言えます。
https://redis.io/docs/latest/operate/oss_and_stack/reference/cluster-spec/#key-distribution-model

まとめ

• read 頻度 >> write 頻度なら RwLock を検討しよう
• アクセスが集中する部分は分散させて並列性を高めよう

今回のテクニックは、動的に変化する集合に対して操作を行う際に、その操作が要素ごとに独立に実行することが可能であれば、汎用的に用いれるものです。

追記

X でフィードバックいただきました。感謝です。Dashmap なるものがあり、内部実装も上記の仕組みに近いようです。みなさんぜひ積極的に使っていきましょう。
https://x.com/_iy4/status/1874792194307232172
https://github.com/xacrimon/dashmap

CPU に寄り添って最適化をしたい場合はキャッシュを意識してメモリの配置を工夫するなども非常に効果があるようです。
https://www.youtube.com/watch?v=_LQ6jvB7sq8

参考文献

https://www.oreilly.co.jp//books/9784873119595/

脚注

1. 詳しくは並行プログラミング入門に詳細な説明があります。