Databricksの高性能レートリミティング技術

Databricksのエンジニアとして、私たちは優れた仲間と共に、常に挑戦的な課題に取り組む機会に恵まれています。
本記事では、Databricksがどのように高性能なレートリミティング（Rate Limiting）システムを構築したか、その舞台裏を紹介します。

Databricksのユーザーであれば、この仕組みを理解していなくてもプラットフォームを最大限に活用できます。
しかし、もし「Databricksの内部で何が起きているのか」を少し覗いてみたい方は、ぜひ読み進めてください。私たちが取り組んできた興味深い技術の一端をお見せします。

背景

レートリミティング（Rate Limiting）とは、リソースの使用量を制御し、マルチテナントシステムにおける各テナント間の分離と過負荷防止を実現する仕組みです。

Databricksの文脈では、アカウント・ワークスペース・ユーザー間の分離を目的として使われ、
たとえば「1秒あたりに実行できるジョブ数」や「1秒あたりのAPIリクエスト数」など、時間単位での制限として外部から見える形で適用されます。

また、内部的にもレート制限は活用されており、サービスの各クライアント間でキャパシティを調整するために使用されるケースもあります。

このようなレート制限の適用（enforcement）は、Databricksの信頼性を支える重要な要素ですが、同時にシステム上のオーバーヘッドを生じさせるため、
その負荷をいかに最小化するかが極めて重要になります。

問題点

2023年初頭には、Databricksの既存のレート制限インフラストラクチャは、Envoyイングレスゲートウェイがレート制限サービスに対して呼び出しを行い、そのサービスをバックアップするための単一のRedisインスタンスで構成されていました（図1）。これは、地域内の任意のマシンクラスタが受け取ることが予想される既存のクエリー数（QPS）や、秒単位のカウントの一時的な性質に対して完全に適していました。しかし、会社が顧客基盤を拡大し、新たなユースケースを追加するにつれて、私たちをその地点にまで導いたものが、将来のニーズを満たすには十分でないことが明らかになりました。リアルタイムモデル提供や他の高qpsユースケースの導入により、一つの顧客がレート制限サービスが現在処理できる量よりもはるかに多くのトラフィックを送信する可能性があるため、いくつかの問題が浮上しました：

• 高いテールレイテンシー - 私たちのサービスのテールレイテンシーは、特に2つのネットワークホップが関与し、クラウドプロバイダーの1つで10ms-20msのP99ネットワークレイテンシーがあった場合、重いトラフィックの下で許容できないほど高かった。
• 制限されたスループット - ある点を超えると、マシンを追加したり、キャッシングなどのポイント最適化を行ったりしても、トラフィックをさらに処理することはできません。
• Redisが単一の障害点 - 私たちの単一のRedisインスタンスが単一の障害点であり、それに対処する必要がありました。私たちのサービスを再設計する時が来ました。

用語

Databricksでは、RatelimitGroup（RLG）という概念があり、これはAPIエンドポイントなど、保護する必要があるリソースまたはリソースのセットを表す文字列識別子です。これらのリソースは、ワークスペース/ユーザー/アカウントレベルで制限を設定するなど、特定の次元で保護されます。例えば、次元は「FooBarAPIをworkspaceIdで保護したい、そしてこのリクエストのworkspaceIdは12345です」と伝えます。ディメンションはこのように表現されます：

単一のshouldRateLimitリクエストには複数のディスクリプタが含まれている可能性があり、例えば特定のAPIに対して、ワークスペースレベルとユーザーレベルで制限を設定することがあります。

このようにDescriptorスキーマが見えるでしょう:

ソリューション

低レイテンシー応答

最初に取り組みたかった問題は、レート制限サービスのレイテンシを改善することでした。レート制限は結局のところカウント問題であり、理想的には、レート制限リクエストを常にメモリ内で回答できるモデルに移行したいと考えていました。なぜなら、それは超高速で、ほとんどのレート制限はQPSに基づいていたため、これらのカウントは一時的なものであり、サービスインスタンスが再起動またはクラッシュしても耐性を持つ必要はありませんでした。既存のセットアップでは、Envoyの一貫性のあるハッシングを使用してキャッシュヒット率を上げることにより、すでに限定的なメモリ内カウントを行っていました。これにより、同じリクエストを同じマシンに送信することができました。しかし、1) これは非Envoyサービスと共有することができず、2) サービスのリサイズや再起動時の割り当ての変動は、私たちが定期的にRedisと同期する必要があることを意味し、3) 一貫したハッシングはホットスポットになりやすく、負荷が均等に分散されないときには、負荷をより良く分散するためにインスタンスの数を増やすしかなく、サービスの利用が最適でないことが多かった。

幸運にも、私たちのDatabricksには素晴らしい人々が参加し、彼らはDicerというオートシャーディング技術を設計していました。これにより、ステートフルなサービスを管理しやすくし、ステートレスなサービスデプロイメントの利点を保持することができました。これにより、レート制限のカウントをすべてメモリ内に保持することで、サーバーサイドのレイテンシを抑制することができました。なぜなら、クライアントはDicerにリクエストをマップして目的地のサーバーに問い合わせることができ、サーバーはDicerでリクエストの適切な所有者であることを確認することができたからです。メモリ内でのカウントは明らかに他のソースからこの情報を探すよりも簡単で速いですし、Dicerは私たちがサーバーサイドのテールレイテンシーを改善し、ストレージソリューションを心配せずに水平にスケールすることを可能にしました。つまり、これにより私たちの単一の障害点（Redis）が取り除かれ、同時にリクエストが速くなりました！

効率的なスケーリング

私たちは問題の一部を解決する方法を理解していましたが、予想される大量のリクエストを処理する本当に良い方法はまだありませんでした。私たちは、大量のサーバーを問題に投げつけるのではなく、これについてより効率的で賢い方法を見つけなければなりませんでした。最終的には、一つのクライアントリクエストがレート制限サービスへの一つのリクエストに変換されることを望んでいませんでした。なぜなら、スケールすると、レート制限サービスへの数百万のリクエストは高コストになるからです。

選択肢は何だったのでしょうか？私たちは多くの選択肢を考えましたが、考慮した選択肢の一部は以下の通りです。

• クライアント上でトークンをプリフェッチし、ローカルでリクエストに応答しようとする。
• リクエストのセットをバッチ化し、送信し、レスポンスを待ってからトラフィックを通す。
• リクエストの一部だけを送信する（つまり、サンプリング）。

これらのオプションのどれも特に魅力的ではありませんでした。プリフェッチング（a）は初期化時やクライアントのトークンが切れたり期限切れになったりしたときに多くのエッジケースがあります。バッチング（b）は不必要な遅延とメモリ圧力を加えます。そして、サンプリング（c）は高qpsのケースには適していますが、一般的には、実際には低レート制限がある可能性があるため、適していません。

私たちが最終的に設計したのは、バッチレポーティングと呼ぶ仕組みで、次の2つの原則を組み合わせています：1）クライアントは、重要なレート制限パスでリモート呼び出しを行わない、2）クライアントは、デフォルトではリクエストを通過させ、特定のリクエストを拒否したいとすでに知っている場合に限り、楽観的なレート制限を行う。私たちは、スケーラビリティと引き換えにレート制限に厳格な保証を持たないことを受け入れました。なぜなら、バックエンドサービスは一定の割合の超過を許容できるからです。高レベルでは、バッチ報告はクライアント側でローカルカウントを行い、定期的に（例えば、100ms）がサーバーにカウントを報告します。サーバーはエントリのどれがレート制限を必要とするかをクライアントに伝えます。

バッチレポーティングのフローは次のようになります：

• クライアントは、どれだけのリクエストを通過させたか（outstandingHits）と、どれだけのリクエストを拒否したか（rejectedHits）を記録します。
• 定期的に、クライアント上のプロセスが収集したカウントをサーバーに報告します。
  • 例えば、KeyABC_SubKeyXYZ: outstandingHits=624, rejectedHits=857;
KeyQWD_SubKeyJHP: outstandingHits=876, rejectedHits=0
• サーバーはレスポンスの配列を返します
  • KeyABC_SubKeyXYZ: rejectTilTimestamp=..., rejectionRate=...
    KeyQWD_SubKeyJHP: rejectTilTimestamp=..., rejectionRate=...

このアプローチの利点は大きく、実質的にゼロレイテンシーのレート制限呼び出しが可能になり、一部のテールレイテンシー呼び出しと比較して10倍の改善が見られ、スパイキーなレート制限トラフィックを（比較的）一定のqpsトラフィックに変えることができました！メモリ内レート制限のためのDicerソリューションと組み合わせると、ここからはすべてスムーズに進むはずですよね？

ディテールに魔物がいる

最終目標については良い考えがあったものの、それを実現するためには多くの困難なエンジニアリング作業が必要でした。ここでは、途中で遭遇したいくつかの課題と、それらをどのように解決したかについて説明します。

高いファンアウト

レート制限グループと次元に基づいてシャード化する必要があったため、これは以前の単一のRateLimitRequestがN次元を持つことができ、つまり典型的なファンアウトリクエストになる可能性があることを意味します。これは特にバッチ報告と組み合わせると問題になる可能性があります。なぜなら、単一のバッチリクエストが多数（500+）の異なるリモート呼び出しにファンアウトする可能性があるからです。対処しないと、クライアント側のテールレイテンシーが大幅に増加する可能性があります（1つのリモート呼び出しの待機から500+のリモート呼び出しの待機へ）、そしてサーバー側の負荷も増加します（全体のリモートリクエストが1つから500+へ）。これを最適化するために、ディスクリプタをそのDicerの割り当てによってグループ化しました - 同じレプリカに割り当てられたディスクリプタは単一のレート制限バッチリクエストにグループ化され、それに対応する宛先サーバーに送信されました。これにより、クライアント側のテールレイテンシーの増加を最小限に抑え（バッチリクエストはクリティカルパス上ではなく、バックグラウンドスレッドで処理されるため、テールレイテンシーの一部の増加は許容可能）、サーバーへの負荷増加も最小限に抑えることができました（各サーバーレプリカは、バッチサイクルごとにクライアントレプリカからの1つのリモートリクエストを最大で処理します）。

施行の精度

バッチレポーティングアルゴリズムは非同期であり、更新されたカウントをRatelimitサービスに報告するための時間ベースの間隔を使用するため、レート制限を実施する前に多くのリクエストを許可してしまう可能性がありました。これらの制限を曖昧に定義することができたとしても、私たちはX％（例えば、5％）以上の制限を超えることを保証したいと思っていました。制限を大幅に超えることが起こる主な理由は2つあります：

• 一つのバッチングウィンドウ（例えば、100ms）がレート制限ポリシーを超える可能性があります。
• 私たちの多くのユースケースでは、固定窓アルゴリズムと秒単位のレート制限を使用していました。固定窓アルゴリズムの特性は、各「窓」が新鮮に始まることです（つまり、リセットしてゼロから始まるため、一定（しかし高い）トラフィックがあっても、毎秒レート制限を超える可能性があります！

これを修正した方法は3つあります：

• レート制限サービスのレスポンスに拒否率を追加しました。これにより、過去の履歴を使用して、いつどのくらいのトラフィックをクライアント上で拒否するかを予測することができました。
  rejectionRate=(estimatedQps-rateLimitPolicy)/estimatedQps これは、次の秒のトラフィックが過去の秒のトラフィックと同じになるという仮定を使用しています。
• 明らかに過度なトラフィックを即座にカットオフするために、クライアント側のローカルレートリミッターを追加することで防御の深さを追加しました。
• オートシャーディングを実装した後、私たちはメモリ内のトークンバケットレート制限アルゴリズムを実装しました。これにはいくつかの大きな利点がありました：
  1. これで私たちは制御されたトラフィックのバーストを許可することができました
  2. さらに重要なことは、トークンバケットは固定ウィンドウアルゴリズムのように毎回リセットするのではなく、連続してカウントし、さらにマイナスになることさえあるため、時間間隔をまたいで情報を「覚えて」います。したがって、顧客が多すぎるリクエストを送信した場合、制限を超えた分を「覚えて」おき、バケットが少なくともゼロに補充されるまでリクエストを拒否することができます。以前にはRedisでこのトークンバケットをサポートできなかったのは、トークンバケットがRedisでかなり複雑な操作を必要とし、それが私たちのRedisのレイテンシーを増加させるからです。今、トークンバケットが毎回の時間間隔で記憶喪失にならないため、拒否率メカニズムを廃止することができました。
  3. トークンバケットは、余分なバースト機能を有効にしないで、スライディングウィンドウアルゴリズムを近似することができます。これは、「リセット」問題を持たない固定ウィンドウのより良いバージョンです。

トークンバケットアプローチの利点は非常に大きかったため、すべてのレート制限をトークンバケットに変換しました。

飛行中の飛行機を作り直すような挑戦

私たちは到達したい最終状態を知っていましたが、それには重要なサービスに対して2つの独立した大きな変更を行う必要があり、そのどちらも単独でうまく機能することが保証されていませんでした。そして、これら二つの変更を一緒にロールアウトすることは、技術的な理由とリスク管理の理由からも選択肢ではありませんでした。途中で行った興味深い作業の一部は以下の通りです：

• 私たちはenvoyイングレスにローカルホストのサイドカーを構築しました。これにより、バッチ報告とオートシャーディングの両方を適用できました。なぜなら、envoyは私たちが変更できないサードパーティのコードだからです。
• メモリ内レート制限がなかった前は、バッチ報告リクエストのテールレイテンシーを下げるために、Luaスクリプトを介してRedisへの書き込みをバッチ化しなければなりませんでした。なぜなら、ディスクリプタを一つずつRedisに送るのは、すべてのネットワークの往復で遅すぎたからです。それであってもバッチ実行に切り替えていたとしても。
• 私たちは、多様なトラフィックパターンとレート制限ポリシーを持つトラフィックシミュレーションフレームワークを構築しました。これにより、この遷移の全体を通じて、私たちの精度、パフォーマンス、スケーラビリティを評価することができました。

現状と今後の作業

バッチ報告とメモリ内トークンバケットレート制限の両方の成功したロールアウトにより、私たちはサーバーサイドのトラフィックの成長がサブリニアルである一方で、テールレイテンシーの大幅な改善（一部のケースでは最大10倍！）を見ました。レート制限呼び出しを行う際にリモート呼び出しがないこと、そしてレート制限サービスとは独立してスケールする自由があることに、私たちの内部サービスクライアントは特に満足しています。

チームはまた、サービスメッシュルーティングやゼロ設定オーバーロード保護など、他のエキサイティングなエリアにも取り組んでいるので、これからのブログ投稿をお楽しみに！そして、Databricksは常に探しています。あなたが違いを生むことができる素晴らしいエンジニアを、私たちはあなたが参加するのを待っています！