こんにちは。データサイエンス部の石川です。
先日 Go で書かれていた広告スコアリングサーバを Rust でリプレイスし、その後レイテンシの改善に取り組みました。 この記事では具体的にどのような実装の変更によって Rust サーバのレイテンシを改善したかについて紹介したいと思います。
広告スコアリングサーバのリプレイスに関する詳しい背景については以下の記事をご覧ください。
背景
広告スコアリングサーバとは
はじめに広告スコアリングサーバとは何かについて説明します。
弊社の広告配信システムは、配信サーバとスコアリングサーバで構成されています。 スコアリングサーバは配信サーバから広告候補のリクエストを受け、CTR と CVR の推定により各広告候補の入札単価と広告スコアを算出し、配信サーバにレスポンスを返します。配信サーバは算出されたスコアに基づいて広告候補を並び替え、広告枠の上から順に広告を表示します*1。
広告スコアリングサーバに高速化が必要な理由
次に、なぜ広告スコアリングサーバにおいてレイテンシの改善が重要なのかについて書きます。
広告効果や収益性の観点から、配信サーバはアプリクライアントから受けた広告リクエストに迅速に応答する必要があり、スコアリングサーバも配信サーバの応答速度に合わせて応答速度が制限されています*2。 この応答速度の要件を満たすために、配信サーバはスコアリングサーバのレスポンスが遅いとリクエストをキャンセルし、最適化されていない状態で広告表示を行う設定になっています。
このような背景から、スコアリングサーバは広告スコアの最適化だけでなく、安定した応答速度の担保がビジネス的な価値を持ちます。
課題
最初に紹介した記事の通り、最近 Go で書かれた広告スコアリングサーバを Rust で書き直しました。
しかし、移行直後に Go と Rust のスコアリングサーバを比較すると、Rust で書き直した新規のサーバの方が既存の Go サーバよりもレイテンシが悪化していたことが分かりました。
そのため、トレーシングを仕込み、アプリケーションコード内のどの部分がレイテンシの悪化につながっているかを理解し、ボトルネックを解消していきました。その具体的な内容について、背景を説明しつつ、問題点と解決策、そして結果を紹介します。
レイテンシ改善施策
1. 有効期限付きキャッシュの導入
背景
広告情報の集計値を利用するロジックがあり、これは 1 時間ごとに更新されるログを元に集計が行われています。そのため、データストアに保存されても 1 時間程度は更新がありません。そこで、高速化のために有効期限付きのインメモリキャッシュを導入しています。
課題
広告リクエストには複数の広告 ID があり、各広告 ID に対してスコアリングが行われていますが、キャッシュの有効期限が広告リクエストごとに設定されているため、一斉にキャッシュが切れる現象が発生し、データストアへのアクセスが集中してレイテンシが大きくなる原因となりました。
解決策
この問題を発見したため、有効期限付きキャッシュの有効期限を、一定の範囲内でのランダム値でキーごとに設定できるようにしました。 これにより、キャッシュの有効期限切れができるだけバラつくことを意図しました。
結果
この改善により、レイテンシの 99 パーセンタイルが約 40% 改善しました。 特にレイテンシが大きく悪化するケースに対して有効な改善策となりました。
2. Array を Vec に
背景
スコアリングサーバでは、広告やユーザーの特徴量のベクトルをそれぞれ別のテーブルに保存し、サーバ側で取得・結合してモデルの入力に使用しています。これらのベクトルは、 ndarray クレートの Array 型で保持されています*3。
課題
しかし、ベクトルの結合処理において、ndarray クレートの concatenate 関数を使用していたことが予想外に時間を消費していることが判明しました。この結合処理において、実装をたどってみると内部的に入力の Array が一件ずつ Clone されていることを特定できました。
解決策
これを解決するため、 Array は行列処理で便利な関数を持っているのですが、行列処理が必要にならない限り Array を使用せず、標準で提供される Vec を使用して値を保持することにしました。
結果
この変更により、レイテンシの 50 パーセンタイルが約 20% 改善しました。 ベクトルの前処理に対する改善であり、リクエスト全体のレイテンシに有効な施策となりました
3. Redis コネクションの設定
背景
特徴量ベクトルは基本的にインメモリにキャッシュされていますが、ユーザーの特徴量ベクトルのみレコード数が多いため、Redis にキャッシュしています。 Redis は外部データストアであるため、コネクションの確立に時間がかかります。そのため、あらかじめ構築したコネクションプールからコネクションを取り出してデータを取得しています。 また、コネクションプールで管理されるコネクションは一定時間利用されないとアイドル状態として破棄されます。
課題
Rust の Redis クライアントライブラリでは、アイドル状態と見なされる時間を設定する idle_timeout や、アイドル状態でも保持する最小コネクション数を設定する min_idle を指定できます。 しかし、これらの値が非常に短く設定されていたため、コネクションの破棄と確立が不必要に繰り返され、レスポンスの遅延が発生しました。
解決策
コネクションの破棄・確立が不必要に繰り返されないように、 idle_timeout を長く設定しました。
結果
これにより、レイテンシの 50 パーセンタイルが約 10% 改善しました。 多くのリクエストで Redis からのユーザーの特徴量ベクトルの取得が行われるため、リクエスト全体のレイテンシが改善されました。
4. CPU バウンドな処理を Tokio のブロッキングスレッドに移す
背景
Rust では非同期処理のランタイムが外部クレートとして用意されていて、今回のプロジェクトでは Tokio クレートを利用しています。 Tokio は各スレッドの行き来しながら多数の非同期タスクを少数のスレッドで捌けるので I/O バウンドな処理に適しているとされています。
課題
しかし、各スレッドの行き来は .await
の時点のみで行われるため、ブロッキングな処理によってスレッドが専有され、CPU バウンドな処理に対しては工夫が必要でした。
解決策
実際、 Tokio のメンテナのブログ記事*4の説明では、 10μs ~ 100μs 以上の処理はブロッキングスレッドに移すことが推奨されています。 そこで、 CPU バウンドな処理(例:モデルによる推論)はブロッキングスレッドで処理するようにしました。
結果
この変更により、レイテンシの 50 パーセンタイルが約 5% 改善しました。 CPU バウンドな処理がスコアリングロジック内に点在しているため、リクエスト全体のレイテンシが改善されました。
Tokio の並列実行とブロッキング処理について
余談ですが、Tokio がどのようにして並列実行をしているか、なぜブロッキングな処理がリクエスト全体のスループットを低下させてしまうかについて書きます。
Tokio は、各スレッドで現在実行中のタスクを繰り返しスワップすることで、少数のスレッドで多数の非同期タスクを同時に実行することができます*5。しかし、このようなスワップは前述の通り .await
時点でしか行えないため .await
に到達しないまま長時間を過ごすコードは他のタスクの実行を妨げることになります。
これに対処するため、Tokio はコアスレッドとブロッキングスレッドの 2 種類のスレッドを提供しています。コアスレッドは非同期コードが実行される場所で、各 CPU コアに 1 つずつスレッドが生成されます。一方、ブロッキングスレッドは他のタスクの実行をブロックするようなブロッキングコードを実行するのに使われます。ブロッキングコードをコアスレッドで実行すると、そのスレッドが処理に専有されてしまい、他のタスクの実行が妨げられて全体のパフォーマンスが低下します。そのため、CPU バウンドな重い計算処理はブロッキングスレッドで実行されることが推奨されています。
5. RwLock を ShardedLock に
背景
広告スコアリングサーバで使用する機械学習モデルはモデル管理サービスとスコアリングサービスで状態が共有されているため、読み込みや書き込み時にロックが必要です。 モデルの更新による書き込み時のロックは 1 時間に数回である一方、スコアリングサービスはリクエスト回数分の読み込み時のロックが発生します。
課題
読み込み時のロックはリクエストのたびに発生する処理でボトルネックとなりうる状態でした。
解決策
crossbeam クレートの ShardedLock を導入しました。 ShardedLock は標準の RwLock と同等の機能を持ちつつ、書き込みは遅いが読み込みは早いという特性があります。この特性を利用して、書き込みよりも読み込みが多いような今回のユースケースでレイテンシ改善が見込めると判断しました。
結果
リクエストごとに必ず通る処理のため、モジュールの置き換えだけでレイテンシが数ミリ秒改善しました。
6. HashMap を HashBrown に
背景
複数の広告候補に対する値を取り扱うために、多くの場所で HashMap が使用され、Get や Set が頻繁に呼び出されていました。
課題
HashMap の亜種が多く存在し、速度面で改善可能な外部クレートがあることが分かりました。
解決策
hashbrown クレートを採用しました。このクレートは Google の SwissTable と呼ばれるハッシュテーブルを利用しています。std::collections::HashMap
から hashbrown::HashMap
にするだけの実装で容易に変更できました。
ただし、hashbrown はデフォルトで aHash というハッシュ関数を使用していますが、これはハッシュ DoS 攻撃に対する耐性が十分ではないため注意が必要です。
ちなみに insert_unique_unchecked
メソッドを使用すると、キーの存在確認を行わずにインサートできるため、キーがマップに確実に存在しない場合に若干高速になるらしいです。
結果
わずかにレイテンシが改善しました。
まとめ
以上のようにレイテンシ改善を行い、ピーク帯でも安定したレスポンスを返せるようになり、ビジネス KPI にも既存のサーバと大きな差異がなかったため、徐々に適用割合を増やし、現在では完全に置き換えることができました。
今後の展望としては、配信サーバとスコアリングサーバの AZ 間の Pod の台数の相違により*6、CPU 負荷が高まるときにレイテンシが悪化する問題があるため、AZ ごとに Pod の台数が偏らないように調整することが課題です。
また、今回のサーバのフルリプレイスの最終的なゴールはディープラーニングモデルをデプロイしてビジネス KPI を改善することでしたので、ディープラーニングモデルも高速に推論できるように精進していきたいと考えています。
*1:実際にはもう少し細かい表示ルールがありますがここでは割愛します
*2:リアルタイム入札の世界では 50ms or die という言葉がありますが、弊社の場合はリアルタイム入札をしていないため、タイムアウトの設定は 50ms よりは若干緩いです
*3:Python の numpy の ndarray に相当するものだと考えてもらうと分かりやすいかもしれません
*4:https://ryhl.io/blog/async-what-is-blocking/
*5:https://docs.rs/tokio/1.29.1/tokio/#cpu-bound-tasks-and-blocking-code
*6:主にアクセスのピーク時間帯に備えて事前に一気にスケールする場合に生じています