これは Rust Advent Calendar 2022 のカレンダー2の22日目の記事です。
YJITとは
私は今年の7月からShopifyという会社でRubyのJITコンパイラであるYJITを開発している。このJITは今年CからRustに書き直されたため、現職では業務としてフルタイムでRustを書いている。
実用段階になったYJIT
おそらく世界最大規模でRubyを使っている弊社では、お客様のお店のサイトをレンダリングするアプリには社内最大のトラフィックが来ていて、実は最近そのアプリほぼ全台で最新のYJITが有効化されたことが昨日公開情報になった。あまりにも大量のトラフィックが来るので、YJITによって行なわれているRubyコード実行量はYJITが使われてないものよりも世界全体で見て多くなったのではないかとCEOが言っていた。
この規模のトラフィックを捌くためにこのアプリは比較的よくチューニングされていて、多くのリクエストでIOが大体キャッシュを読みまくるだけになるのでYJITがなくても結構速いのだが、それでもYJITを有効化するだけでRuby 3.2のインタプリタと比較して最大10%の高速化が観測できている。
メモリ使用量の改善
そのチューニングの結果かどうかはわからないが、このアプリのJITによる生成コードはずっと走らせ続けてもせいぜい30~40MiBで収まっている。というか、JITするかどうかの閾値を下げていくと本来一時的にしか実行されないものもJITしてしまいコード量が結構増えるのだが、このアプリの本番の挙動を見ながらRuby 3.2の設定パラメータをチューニングしたため、いい感じになっている。
こういった背景から、YJITが生成した機械語によるメモリ消費はRubyインタプリタ本体が (10~20倍くらい) 消費するメモリに比べても比較的小さいのだが、我々がRustのメモリ使用量の改善を行なう前、Rustによって機械語のメタデータを管理するのに消費しているメモリが機械語のサイズの5倍くらいある状態だった。
このRustによるメモリ消費が現在では、保存している情報を本質的にそれほど変更しない細かなテクニックだけで機械語の3~4倍程度まで削減させられた。次のバージョンではデータ構造そのものを変えることでもっと大きく削減していく予定だが、比較的短い期間でどのようにメモリ削減を実現していったかについて紹介したい。
計測テクニック
何がどのように性能に影響しているかプロファイリングをしたり、実際に施策を打った後にどれくらいインパクトがあったか計測することは、効率よく確実にパフォーマンスチューニングをしていく上で非常に重要である。この辺も私が実験や整備を行なったので、先に紹介しておく。
ヒーププロファイラ
The Rust Performance Book の Profiling の章 を見ると、結構色々なプロファイラが紹介されている。この中で valgrind (cachegrind & callgrind)、DHAT、dhat-rs、heaptrack、bytehound あたりを試したのだが、我々がYJITのRustのコードをCで書かれているRubyインタプリタと同居させて使っている関係でCの世界のものが混ざってしまい結果が見づらかったり、そもそもクラッシュして全然使えないものが多かった。
Cのところは無視してRustの部分だけヒーププロファイリングをしたいという目的ではdhat-rsが一番使いやすかった。これはmasterにはコミットしてないが、こんな感じで割と簡単にセットアップできる。
実行すると dhat-heap.json というファイルが生成され、これをDHAT Viewerにアップロードするとヒープのメモリ消費のブレークダウンをWeb UIで見ることができる。
最初はこれをいい感じに見る方法もよくわかってなかったのだが、ベンチマーク時の最後まで使われ続けたメモリの消費を抑えたいという私の用途だと Sort metric を At t-end (bytes) にすると大分読みやすくなった。
このSort metricを変えるテクを知らなかった時はやたら情報が欠落していると思ってtatsuya6502さんにrust-lang-jpで相談していたのだが、実際、ある閾値を下回るものは落とされてはいるので、それを直すパッチも書いていただいたので、もしそれが必要になった方はそちらもどうぞ。
Rustのグローバルアロケーションのサイズ計測
メモリ改善を始める前の時点ではRustが使用するメモリの量を計測できておらず、JITが生成するコードサイズくらいしか自前で計測できていなかったため、何故生成コードサイズよりも多くメモリ消費量が増えているのかよくわかっていない状態だった。Rustだろうな、という疑いは持っていたが、これはRustのグローバルアロケーションのサイズを計測することで確信を持つことができた。
それを入れたPull Requestがこれで、これはmasterにマージされている。Rubyをconfigure時に --enable-yjit=stats や --enable-yjit=dev を指定してビルドすると使えるようになっており、その場合はstats_alloc crate (使い方は GitHub を参照) を使ってアロケータに計測用のミドルウェアを差し込む形で動いている。そのようにRubyをビルドし --yjit-stats を使うと、yjit_alloc_size というメトリックでRustが消費するメモリ量を見られるようになった。
inline_code_size: 4855668
outlined_code_size: 4852936
freed_code_size: 0
yjit_alloc_size: 25114317
以後の省メモリテクニックの適用においては、前後のバージョンでこの yjit_alloc_size を比較することで効果検証を行なった。
省メモリ化のためにやったこと
手元で高速に改善サイクルを回す上で、本番にデプロイして検証するよりは本番に近い構造のベンチマークを使う方が便利であったため、以後、Railsのベンチマークであるrailsbenchを使って効果検証をしている。
改善を始める前の時点では --enable-yjit=stats ビルドで yjit_alloc_size が上記の通り25.1MBだった。
参考: Heap Allocations - The Rust Performance Book
#[allow(unused)] で放置していたフィールドの削除
Pull Request: YJIT: Remove unused src_ctx from Block
テクニックというにはあまりにもしょうもなさすぎる話だが、過去のリファクタリングの過程で将来的に使う予定で残されたフィールドがそのまま使われずに時が経ってしまっていて、それを削除した、というのが最初の一歩だった。
これで 25.1MB → 23.3MB まで削減。
2通りの方法で保存していた重複フィールドの削除
Pull Request: YJIT: Invalidate redefined methods only through cme
これもなんというかRustテクニックというよりはYJIT特有の話なのだが、メソッドが再定義された時にJITコードをinvalidateするための方法を2通り用意していて、これはRubyインタプリタ側の実装をいじらずに実装する都合上そうなっていたのだが、単にインタプリタ側の実装を簡略化していくことでYJIT側もinvalidation用のフィールドの重複削除ができた。
これで 22.3MB → 22.2MB まで削減。
Vec を shrink_to_fit で縮める
Pull Request: YJIT: Shrink the vectors of Block after mutation
やっと純粋なRustのテクニックぽい話。Vecのドキュメントを読むと、こう書いてある:
Vec will never automatically shrink itself, even if completely empty. This ensures no unnecessary allocations or deallocations occur. Emptying a Vec and then filling it back up to the same len should incur no calls to the allocator. If you wish to free up unused memory, use shrink_to_fit or shrink_to.
Vecが伸縮する度に必ずリアロケーションをしてると当然遅くなるので、速度のためにRustはリアロケーションを少しサボっているというような話が書いてある。
従って、本来必要な領域よりも少し多くVecがメモリを使う可能性があり、YJITが使うVecが変更された後にshrink_to_fitを呼んで明示的にリアロケーションを発生させることでこれを改善した。YJITのワークロードでは、一度あるコンパイルしてしまえばあとはそこで生成されたコードがひたすら使われるだけなので、コンパイルがshrink_to_fitで多少遅くなっても無視できる性能特性であることが背景にある。
これで 22.2MB → 20.9MB まで削減。
#[repr(packed)] で構造体のフィールドを詰める
Pull Request: YJIT: Pack BlockId and CodePtr
一般に、CやRustで構造体を書いた時、それをコンパイルしたバイナリでは必ずしも各フィールドに必要な最低限のサイズで並んで敷き詰められているわけではなく、アーキテクチャに応じて適当に何もない空間(パディング)を間に挟んでアラインメントを取ることでフィールドのアクセスが速くなるようにコンパイラが最適化をしたりする。
YJITのワークロードでは、生成コードの性能に比べるとコンパイルにかかる (#[repr(packed)] をつけることで発生する程度の) オーバーヘッドは重要度が低いため、そこでアラインメントを取って速くするよりは、多少コンパイル速度を犠牲にしてガチガチに敷き詰めてでもメモリ消費量を抑えたいという状態だった。そのため、structに #[repr(packed)] をつけていくことでこれを実現した。
これをやると使えなくなる機能もあるので、必ずしも使えるわけじゃないが、struct CodePtr(*const u8) みたいなフィールド1つだけのstructでも何故か効果があったりと、使える時は使っておくと便利。
これで 20.9MB → 18.2MB まで削減。
Vec を shrink_to_fit で縮める (2)
Pull Request: YJIT: Shrink version lists after mutation
段々手詰りになってきて、本当にここ入れてええんかみたいな微妙な場所にshrink_to_fitをいれて、やや顰蹙を買う。正直今回はそれほどインパクトもなかったし、ほどほどにした方が良さそう。
これで 18.2MB → 17.8MB まで削減。
ポインタのOptionにNonNullを使う
Pull Request: YJIT: Use NonNull pointer for CodePtr
これはrust-lang-jpでmonochromeさんから教わった話なのだが、ポインタのOptionはポインタのかわりにNonNull<T>を使うことでメモリ消費を抑えられるようになっている:
https://doc.rust-lang.org/stable/std/ptr/struct.NonNull.html
Unlike *mut T, the pointer must always be non-null, even if the pointer is never dereferenced. This is so that enums may use this forbidden value as a discriminant – Option has the same size as *mut T.
まあ要するに、ポインタの値が0ではないことが分かっていればSomeの中身が0なケースを考慮する必要がなくなるので、0の値をNoneとして使うことで、ポインタのOptionをポインタ1つ分の幅で表現可能になるという最適化だと思われる。
CodePtrというのがその名の通りポインタを1つ持つだけのstructなのだが、Option<CodePtr>が他のstructからよく使われていて、そのためCodePtrが使うポインタを*const u8からNonNull<u8>に変えたことでメモリ消費量が減らせた。
これで 17.8MB → 17.2MB まで削減。
メモリリークの修正
Pull Request: Plug YJIT memory leaks
ここで紹介しているPull Requestだけは私ではなく同僚のAlanがやったものなのだが、YJITにRuby 3.2で導入したCode GCの(生成コードではなく、Rustの)メモリリークを修正したというもので、これは使わなくなった時にfreeしていたつもりのブロックが循環参照などの理由で開放されていなかった等の問題を解決している。この変更の過程で、Code GCが使われないケースでもより多くのブロックが開放されるようになり、メモリ消費が減った。
KOBA789さんがちょっと前に書いていた話だが、Rustを書いているからといってメモリリークがないことが自動的に保証されるわけではないという問題がある。メモリを複数箇所から共有しないようなコードを書いていれば話は別だが、既存のCのコードベースでヒープに割り付けられているオブジェクトに関連するメタデータをRustで管理する、といった使い方をしていると、必然的にRust側もヒープアロケーションみたいな感じになるし、(Rustなのに)実質GCみたいなものを自分で実装する必要がある。なので、そういう状況で使う場合に関しては、Rustを使っているからといってC++やZigのような言語でやる場合よりメモリ管理が楽になっているわけではなさそうだなと最近感じている。
同僚の環境/計測なのでベースが前回の17.2MBとはずれてしまっているが、15.2MB → 11.7MB になるくらいのインパクトがあった模様。
まとめ
railsbench上でYJITのRustのコードが消費するメモリを 25.1MB から少なくとも 17.2MBまで、更にそこからおそらく3~4MBくらいは削減できた。本番環境では手元のrailsbenchほどの割合で差は出なかったが、それでもメモリ消費がちゃんと減ったことが観測できた。
マルチプロセスサーバーが主流なRubyの運用においてはメモリの消費量が割と重要になるので、今後も改善を続けていきたい。