YJIT・pitchforkのreforkingを導入し、パフォーマンスが3割向上しました

はじめに

こんにちは。アカツキゲームスの松下です。

弊社ではRailsを用いた大規模なゲームサーバを運用しており、お客様に良い体験を届けるべく、パフォーマンス改善が常に重要事項であり続けています。そのため、RubyにおけるJITコンパイラ・高速化技術であるYJITについても、登場時から注目していました。弊社でもこれまでに何度か導入を試みてきたものの、YJITを有効化すると負荷試験中に segmentation fault（SEGV）がごく稀かつランダムな箇所で発生するという現象に悩まされており、原因の特定すら難航する、歯痒い状態が続いている状況でした。

そのため、長らく本番導入を見送らざるを得ない状況が続いていましたが、一念発起し試行錯誤を重ねた結果、YJITとpitchforkのreforkingを本番環境に導入することに成功し、大きなパフォーマンス改善を得ることができました。本記事では、無事に導入に至るまでの経緯と、その過程で得られた知見について紹介します。

前提

今回YJITを導入したサービスはゲームサーバであり、マスタデータ*1をプロセス上に保持する構成をとっています。このような特性から、メモリ効率の最適化が重要です。

将来的なYJIT導入を見据えて、昨年夏時点にUnicornからpitchforkへ移行していました。後述するSimpleMasterなど、reforkingによるCoW（Copy-on-Write）の恩恵を十分に得ることが狙いです。これはYJITに使うメモリ使用量削減にも寄与します。*2

マスタデータのロードと管理には、弊社の趙が開発したSimpleMasterを利用しています。SimpleMasterは、マスタデータをプロセス起動時に読み込み、不変データとして扱うことでCoWを活用し、ワーカプロセス間でメモリを効率よく共有できる点が特徴です。ゲームサーバのように大規模なマスタデータを扱う環境において、メモリ使用量の削減に大きく寄与しています。

また、JSONレンダリングには、同じく趙が開発したSimpleJson を利用しています。SimpleJsonは、Lambda記法により直感的にviewを記述可能でありながら、軽量かつ高速に動作する点が特徴です。

なお、負荷試験環境および本番環境はAWS上でGraviton 4のインスタンスを利用しています。

導入のきっかけ

これまでRubyの新しいバージョンがリリースされるたびにYJITの有効化の検討・負荷試験などを行ってきましたが、負荷試験中に高い確率でSEGVが発生してしまい、導入を断念せざるを得ませんでした。原因の特定も難しく、調査が進まない状況が続いていました。なお、このSEGVはUnicorn・pitchforkいずれの構成においても発生していました。

昨年12月頃、Ruby v3.4.8 へのアップグレードを検討するにあたりリリースノートを確認したところ、YJITに関する気になるチケット Bug #21266: YJIT GC safety crash with proc objects as block argument がありました。内容的に、これまで遭遇していたSEGVと関係がありそうだったため、再度本格的にYJIT導入に挑戦しました。ただし、後にRuby v3.4.8にアップグレードするだけではSEGVは解消しないことがわかります。

導入に向けた流れ

負荷試験環境でRuby v3.4.8, Rails v8.0を使って負荷試験を進めていました。yjit_exec_mem_size: 128MiB（default）ではSEGVは起きず、yjit_exec_mem_size: 256MiB（本番想定サイズ）ではSEGVが発生しました。このため、Ruby v3.4.8の導入のみだと、SEGVになる原因は取り除けていません。

SEGVのスタックトレースを見ていると、SimpleJson Gemが怪しいことに気がつきました。SimpleJsonはテンプレートファイルを読み込んでLambdaとして評価し、instance_exec(&lambda)で実行します。SimpleJsonには起動時にテンプレートキャッシュを保持する仕組みがあります。しかしながら、改めて設定を見返すと、正しくテンプレートキャッシュが機能していないことがわかりました。そのため、大量にLambdaを生成・実行しており、これがYJITと相性が良くない可能性があると推測しました。

そこで、SimpleJsonに変更を加えることにしました。PRはこちらです Replace lambda with method to avoid YJIT crash by AZQ1994 · Pull Request #12 · aktsk/simple_json。既存のテンプレートファイルの置き換えをしないために、テンプレートのLambdaを正規表現でメソッド定義文字列に変換し、class_evalでクラスに直接メソッドとして定義するようにしました。あわせて設定を見直しSimpleJsonのテンプレートキャッシュを有効にしました。

これらのSimpleJsonの変更を加え、Ruby v3.4.8 / yjit_exec_mem_size: 256MiBの構成で負荷試験を進めましたが、SimpleJsonのLambda実行とは関係ない箇所でSEGVが発生しました。該当コードを確認しても不自然に思える点はなく、対応に行き詰まっていました。ちょうどこの検証を進めていた時期がクリスマス頃でRuby v4.0がリリースされており、リリースノートを確認するとGC周りの改善やYJITのバグ修正が含まれていることがわかりました。これらの変更が影響する可能性も考え、Ruby v4.0.0でも負荷試験を進めることにしました。

Ruby v4.0.0 / yjit_exec_mem_size: 256MiBで負荷試験をすると、これまで遭遇していたタイミングでSEGVが起きなくなりました。Ruby v3.4.8とv4.0.0の差分を調査しましたが、どの修正が直接効いたのかまでは特定できませんでした（推測ではありますが、GCやメモリ周りの改善によるものかと考えています）。念のため、十分な時間をかけて繰り返し負荷試験を回しましたが、SEGVは起きませんでした。負荷試験を十分な時間回してSEGVが起きないことを確認できたため、本番環境へ導入を進めました。一方で、明確な原因が不明だったこともあり、本番環境で、SEGVが起きないという保証・確証があるわけではありませんでした。そこで本番環境へ導入するにあたり、SEGVを検知できるようにアラートを仕込み、すぐに対応ができるようにしておきました。

本番導入による効果

導入前後での構成の違い

Ruby v3.4.7 → v4.0.0
YJIT無効 → 有効(yjit_exec_mem_size: 256MiB)
reforking無効 → 有効
SimpleJsonのテンプレートキャッシュ無効 → 有効

YJITの有効化およびreforkingの導入が大きく寄与していると考えていますが、個別の寄与割合については計測できていません。効果は以下になり、大きなパフォーマンス改善が確認できました。

レスポンスタイム
- 平均: 約28%改善
- P90: 約32%改善
メモリ使用量
- 約30%削減

SEGV再び

本番環境導入後、特定のサーバ環境のみにおいて稀ですが何度かSEGVが発生しました。スタックトレースを見ると、SimpleJsonとは関係ない箇所のLambdaの実行でSEGVが発生していました。SEGVが起きるエラー箇所はランダムで異なり、再現性が低く、再現手順の確立が難しい状況です。

SEGVが起きるのが稀であることから、SEGVが発生したときにpitchforkのhookを利用して自動的に復旧できる仕組みを実装しました。具体的には、SEGV発生を after_worker_exit で検知し、 after_mold_fork においてRubyVM::YJIT.code_gc()（YJITのコードを捨てる）とreforkingを行って、SEGVを引き起こす原因を取り除くというものです。コンテナの再起動などでの対応も可能ですが、再起動時に初期化処理でマスタデータのロードなどで幾分か時間がかかるワークロードであることと、CoWによって他のpitchforkのワーカでもSEGVが発生しやすい状況になっており、ダウンタイムを小さくしたいためこの手段を選択しています。これにより、ごく稀に発生したとしても、即座に自動復旧・正常化することが確認できています。

現状、Ruby v4.0.2へアップグレードしていますが、依然として稀にSEGVが発生しています。関連しそうな修正として YJIT: Fix remaining version_map aliasing UB from get_iseq_payload by rwstauner · Pull Request #16510 · ruby/ruby があり、これによって改善される可能性も考えていますが、現時点では影響は不明です。