Recall.ai のリングバッファの設計を考察する

あなたとわたしとロバストとパフォーマンス。どうも、かわしんです。

先日 Recall.ai というビデオ会議に関連するサービスのブログ記事を読みました。

www.recall.ai

インフラ費用を減らすために動画処理サーバーのプロファイルをとったところ、CPU 時間を一番使っていたのがビデオフレームを送信する際の Web Socket の通信のメモリコピーだったということがわかったので、共有メモリ上に実装したリングバッファを使うことで CPU 使用量を半分にしてサーバー代を半分にしたという豪快なお話です。

同じサーバー内のプロセス間通信に Web Socket を使うと、以下のオーバーヘッドがあります。

• Chromium プロセスからカーネル内のバッファへのメモリコピー
• カーネル内のバッファから動画処理プロセスへのメモリコピー
• Web Socket の改竄防止のために自動で行う全データのマスキング

これらのコストは、ロックフリーでゼロコピーの共有メモリのリングバッファを使うことで削減できます。

まさに技術で課題を解決している感じがしていい話なのですが、以下の要件を満たす共有メモリ上のリングバッファライブラリがなかったため自作したらしいです。

• ロックフリー
• マルチプロデューサ / シングルコンシューマ
• 可変長のフレームに対応
• ゼロコピー
• サンドボックス内から送信
• 低遅延のシグナリング

しかし、この記事にはざっくりとしたリングバッファの概要しか記述がなく、またそのリングバッファライブラリは OSS にはなっておらず、どのように実装したかは不明です。Recall.ai の Github にはそれっぽいリポジトリは見つかりませんでした。

(https://www.recall.ai/post/how-websockets-cost-us-1m-on-our-aws-bill より)

リングバッファを具体的にどのように作ったのかが気になったので、自分だったらどのように設計するかを考察してみました。

上で列挙した以外には以下の要素がブログ記事中でわかっています。これらがパズルのように後で設計を絞るのに役立ちます。

• read ptr, peek ptr, write ptr の3つのポインタをアトミックに使う
• named semaphore をプロセス間のシグナリングに使う

デザインドキュメントを書く時、検討したがうまくいかない案は普通は後ろの方に "Alternative Considered" の章にまとめますが、この記事ではどのように考察したのかも読んでほしいためうまくいかない案を考察してからうまくいく案を導き出す方式で書いています。

データフォーマットとデータ構造

さて、ロックフリー + マルチプロデューサ + 可変長フレーム のリングバッファであるためリングバッファに書き込む際の操作はざっくり以下のステップになるはずです。

1. write pointer をデータ長分 atomic に増やしてリングバッファ内の領域を確保する
2. 確保した領域にデータを書き込む
3. データが読み込み可であることを通知する

リングバッファにヘッダとデータを詰める

可変長のリングバッファのデータフォーマットとして簡単に思いつくのは、固定長のヘッダと可変長のデータ分の領域を (1) のステップで確保するものですが、実はうまくいきません。

             read ptr                        write ptr
              ∨                               ∨
 | free space | header | data | header | data | free space |

なぜかというと、ヘッダ内の書き込み可フラグが (1) のステップの時点で初期化されていないためからです。ロックフリーで並列にアクセスするためには、リングバッファ内の空き領域の先頭の次のヘッダに該当する部分が (1) を行う前に初期化されている必要がありますが、事前に初期化しようにも (1) が終わるまで空き領域の先頭の位置はわからないため、ヘッダの初期化はできません。コンシューマ側で読み込み完了時に空き領域に戻される領域を全てゼロクリアするなどして初期化はできますが、毎回全てをゼロクリアするのはメモリアクセスの観点からも非効率的です。

メタデータ専用のリングバッファ

それを解決するために、2つのリングバッファを用意してみます。

• 固定長のメタデータのためのリングバッファ
• 可変長のデータのためのリングバッファ

           read index            write index
            ∨                     ∨
| free slot | metadata | metadata | free slot | free slot |
                    
| free space | data | data | free space |
             ^             ^
           read ptr   write ptr

固定長のメタデータのリングバッファを読み込み側が解放する時に初期化することでメタデータの初期化の問題が解決されます。しかし、メタデータリングバッファの write index とデータリングバッファの write ptr をロックなしに操作するため、2つのリングバッファの順序が入れ替わってしまう可能性があるなどロジックが複雑になってしまいます。さらに、データリングバッファの空き領域が無くなった時のシグナリングにセマフォを使うと、1 バイト単位でセマフォに登録しなければならないため、セマフォ操作が大量に発生し非効率です。

リングバッファを固定長のブロック で管理する。

元記事によると 1080p の動画の 1 フレームの大きさは、3110.4 KB だそうです。結構でかいです。最初の案では、空き領域の先頭がどこになるかわからないため領域の解放時に戻される領域を全てゼロクリアしないといけないことが問題でした。リングバッファを 1 MB 単位で確保したり解放したりするようにすると、空き領域の先頭になる可能性のある部分は 1 MB 毎の先頭に限定され、空き領域に戻すときにヘッダの初期化のためにゼロクリアする部分が大幅に少なくなります（1 MB あたり数バイト）。

また、リングバッファの空き領域が無くなった後に空き領域が増えたことを通知するセマフォも、<リングバッファのサイズ> / 1 MB 個だけ登録すればいいのでセマフォ操作も少なくなります。

| free block | header | data   | unused area | free block |
             | multiple of block size        |

データのサイズによっては確保したブロックの後ろの部分は使われない無駄な領域になりますが、ブロックの大きさを調整することで無駄になる領域の割合を下げることができます。全体的な CPU 時間のオーバーヘッドとのトレードオフで決めることになります。

プロデューサの処理

1. データサイズとヘッダサイズから必要なブロック数を計算する
2. ブロック確保用のセマフォを必要なブロック数 sem_wait して必要なブロック数を予約する
3. write ptr を確保したブロックサイズ分 atomic に増やして領域を確保する
4. ヘッダにデータの大きさなどのメタデータを書き込む
5. データをリングバッファの確保した領域にコピーする
6. ヘッダ内の読み込み可のフラグを有効にする
7. データ完了通知用のセマフォを 1 回 sem_post する

read/write ptr はリングバッファの先頭からのオフセットで表されます。

セマフォの獲得は 1 ブロック単位で行うので、全体のリングバッファのブロック数が少ないと複数のプロデューサが不完全な個数のブロックをセマフォから予約してデッドロックする可能性があります。そのためリングバッファのサイズは、最大のプロデューサの数と最大フレームのサイズの積以上に設定する必要があります。

コンシューマの処理

まず、peek ptr がある意味を考えます。もし 1 フレームごとに処理をするのであれば read ptr から先頭のフレームを peek して、処理が終わってから read ptr を移動させればいいはずなので peek ptr は必要ありません。つまり、peek ptr があるということは、複数フレームを並行に処理する可能性があるということがわかります。

これを元に、コンシューマ側での読み込み処理を考えてみます。

Peek 処理

1. peek ptr が指す先頭のヘッダの読み込み可フラグが有効になっているかを確認する
2. もし、読み込み可でない場合はデータ完了通知用のセマフォを sem_wait してステップ 1 に戻る。ただし、sem_wait が unblock された場合でも先頭ではなくその先のフレームの読み込みが可能になっただけの場合があることに注意。いずれにせよその場合は sem_wait しなおすことになる。
3. 先頭のフレームが読み込み可であった場合は、peek ptr をフレームを含む領域のサイズ分増やしてからデータのポインタを呼び出し元に返す。

この読み込み処理はシングルコンシューマなのでシングルスレッドで行われますが、返されたそれぞれのデータの読み込み自体は別スレッドから並列に行えます。

Pop 処理

フレームの処理が終わった後の解放処理は以下のようになります。read ptr の先頭ではなく途中のフレームが先に処理済になった場合に対応するために、ヘッダに処理済フラグを用意します。

1. 処理済フレームデータのポインタからヘッダの位置を逆算する
2. ヘッダの処理済フラグを立てる
3. もし、処理済フレームが read ptr の先頭でなかった場合は (read ptr != ヘッダの先頭) ここで解放処理は終了。
4. もし、処理済フレームが read ptr の先頭であった場合は、フレームデータに含まれる全てのブロックの先頭のヘッダの読み込み可フラグに相当する位置をゼロクリアする
5. read ptr をフレームを含む領域のサイズ分増やして、ブロック確保用のセマフォを解放されるブロック数 sem_post する
6. step 3/4 に戻って次のフレームがすでに処理済であった場合は引き続き解放処理をする

死活監視

もしプロデューサが処理の途中でクラッシュするなどして止まった場合、リングバッファの処理途中のフレーム以降全てが読み込みできずにシステムが止まってしまいます。

もしかしたら Recall.ai ではここまでやってないかもしれないですが、システムの壊滅的な停止を防ぐためにも死活監視の仕組みをリングバッファに入れる必要があります。

プロデューサの死の判定

プロデューサプロセスの死は Unix ドメインソケットのコネクションを事前に貼っておくことで検知することができます。もしプロデューサが死ぬと自動的にコネクションが切断状態になり、ソケットを epoll などで読み込み待ちしているコンシューマに通知されます。プロデューサが終了メッセージをソケットに書き込むことなくコネクションが切断された場合はコンシューマは異常状態からの復帰モードに入ることができます。コンシューマは、異常死を検知した時点での write ptr の値を atomic に取得して、peek ptr がその値に到達するまで読み切るまで異常状態を続けます。

コピー途中の死からの復帰

プロデューサがデータをコピーしている途中で死んだ場合、ヘッダにフレームの大きさが書いてあるのでそのフレームをコンシューマは捨てることができます。ただし、そのフレームを書いているプロデューサが突然死したプロデューサかどうかを判定するために、ヘッダにプロデューサの ID を書き込むことにします。

また、プロデューサの ID は前述のコネクションを接続するときにコンシューマから割り振ってプロデューサに伝えることで一意性が保たれます。

ヘッダ更新前の死からの復帰

プロデューサが write ptr を更新した直後のヘッダを更新する前に死んでしまった場合、そのフレームの大きさをコンシューマは知ることができません。それ以降のフレームは正常なフレームも含めて残念ながら捨てることになります。コンシューマはヘッダのサイズが長時間更新されないフレームを検出した時、それ以降のフレームを諦めて peek ptr を異常状態に入った時の write ptr の値に書き換えて処理を再開します。フレームを捨てるのは peek 済のフレームが全て処理済になった状態 (read ptr == peek ptr) で行い、peek ptr を動かすブロック数分、ブロック確保用のセマフォを sem_post します。

もし、死んだプロデューサ以外のプロデューサからのフレームを捨てることが受け入れられない場合はブロックごとの状態を管理することになりますが、データがブロックの境界で連続しなくなってしまうので設計を 1 から見直すことになると思います。

コンシューマの死

コンシューマが死んで復帰した場合も、プロデューサ側はコンシューマの死を Unix ドメインソケットによって検出できるのでコンシューマとのコネクションが貼り直されるまでリングバッファの書き込みを中断することで対応できます。

コンシューマが復帰した後はプロデューサの ID が全て新しくなるため、コンシューマはデータの書き込みを再開させる前にリングバッファ内にあるデータを全て処理します。その後、新しいプロデューサの ID をソケット経由で送信してプロデューサにデータの書き込みを再開させます。

不明な点

ブログの元記事からはどうするのかが不明な点はこんな感じだと思います。

• どうやって共有メモリを Chromium の JS 環境に繋ぎ込んでいるのか
  • 元記事では "our Chromium" と言っているので、Chromium の C++ のコードに手を入れて JS から渡されたビデオフレームをリングバッファに書き込んでいるのだと思われます。
• 共有メモリのデータがライブラリ外から壊されないのか
  • 共有メモリ自体を JS 環境に見えないようにすれば Third-party コードを実行する JS から共有メモリを正しく使うことを保証できます。
• ひとつのプロデューサが遅すぎる場合全体を律速してしまう
  • 全てのプロデューサが同じデバイス上での同質な Chromium プロセスなので速度の違いは想定しなくてもいい気もします。

まとめ

こう考えてみると、考察するのに必要な情報はあのブログ記事にまとまっていたので Overview Design としてはかなりよく書かれた記事だったのだなと思いました。

ライブラリを作る時はすべての場合に対応しないといけないので設計って大変です。今回は死活監視によってロバストなリングバッファに仕上がりました。

皆さんも、ロバストとパフォーマンスを両立したプロダクトを作っていきましょう。