豪鬼メモ

一瞬千撃

スピンロックによる並列処理の高速化

C++標準のstd::shared_mutexを自前のスピンロックによる実装に置き換えたところ、スループットが激烈に向上した。その経緯と実装と性能評価について説明する。
f:id:fridaynight:20210806012345p:plain


Tkrzwはマルチスレッドによる並列処理性能を重視したデータベースライブラリである。ところで、GitHubにて、TkrzwのGoインターフェイスを6倍高速化するというアイデアを報告いただいた。ボトルネックC++のコアライブラリ内のstd::shared_timed_mutexにあるということを突き止め、それをサードパーティの実装に置き換えると高速化するという話であった。Goだから高速化するというより、ファイル層でスレッド同士がブロックしやすいような条件でテストをすると差が出やすいということらしい。

そのサードパーティの実装とはobject_threadsafeというパッケージにあるcontention_free_shared_mutexクラスである。私もそれを使って自分の環境で実行してみたら、確かに激烈に早くなった。Macbook AirApple M1 3.2GHz)で動かしたところ、4スレッドで1億レコードを格納するテストが3倍ほど早くなった。実際に実行したテストコマンドは以下のものである。

$ tkrzw_dbm_perf sequence --path casket.tkh --buckets 200000000 --iter 25000000 --threads 4

MacOS上では、91万QPSが287万QPSに改善するという結果であった。3.12倍とは驚きだ。ただ、こういう時はだいたい新しいものに革新性があるというよりは、古いものに非効率があることが多い。なので、LinuxCore i7 1.8GHz)上でも同じ測定をしたところ、156万QPSが181万QPSに改善するという結果が得られた。つまり1.15倍だ。ここから考えると、contention_free_shared_mutexが凄いというよりは、MacOS上のstd::shared_timed_mutexの実装に非効率があると言えそうだ。

話は脱線する。標準のstd::shared_timed_mutexクラスは、以前はstd::shared_mutexという名前であった。以前のstd::shared_mutexクラスは、timedというラベルがついていないにもかかわらず、タイムアウト付きのtry_lockを実装していた。よって、他のmutexクラスの命名規則に合わせるために、std::shared_timed_mutexに改名された。そして、std::shared_mutexという名前のクラスがタイムアウト機能を省いた最適化を施す余地を与えられ、C++17から存在している。ただし、Linux上のlibstdc++などの多くの標準ライブラリでは、std::shared_timed_mutexとstd::shared_mutexは単なる別名の関係である。よって、Tkrzwの従来のバージョンでは、タイムアウト機能を使っていないにもかかわらず、std::shared_timed_mutexの方を使っていた。新バージョンからはstd::shared_mutexを使うように変えたが、それだけで性能は変わらない。

話を戻す。いわゆる野良実装であるところのcontention_free_shared_mutexがなぜ早いのかと思ってソースを読んでみたところ、スピンロックだから早いというのが結論だ。スピンロックとは、ロックが獲得できない場合にビジーループを回しながら獲得できるまで待機するという手法である。完全なビジーループにしてしまうとCPUパワーを使いすぎるので、適当な頻度でstd::this_thread::yieldを呼んで実行権を手放すようにはしているが、それでもブロックされた際にはCPUのロードアベレージを劇的に上げてしまうという副作用がある。ロック可能になった時点ですぐに処理が行われるという性能上の利点はあるが、ブロックされる確率が高かったり、ブロックされる時間が長かったりする場合には向かない。つまり、比較的短時間で終わるであろうCPU集中の処理を待つのには良いが、比較的長時間かかるであろうストレージやネットワークの応答を待つのには向いていない。

Tkrzwがデフォルトで使うMemoryMapParallelクラスは、メモリマップを使ってファイル入出力を行う。これはファイルのサイズが実メモリのサイズを下回ることを前提としていて、その場合にはストレージデバイスの応答待ちは最小化されて、CPU集中の処理が多くの時間を占めることになる。このような場合には、スピンロックを使っても差し支えない。よって、MemoryMapParallelクラスの内部で必要な排他制御にスピンロックを使うというのは合理的な提案と言えそうだ。

contention_free_shared_mutexは性能が良いので、CPU集中な処理で排他制御が必要な他のクラスでも導入してみようかと考えた。そこで、HashDBMやTreeDBMやSkipDBMなどに組み込んでみた。しかし、テストケースを動かしてみたところ、assertで落ちたりデッドロックしたりといった不具合が発生するようになってしまった。Tkrzwが悪いのか、contention_free_shared_mutexが悪いのか、両方が悪いのか、それはわからない。わからないが、よそ様のコードを完全に理解するのも大変そうだ。

ところで、shared_mutexとかreader-writerロックとか言われる機能をスピンロックで実装するのはKyoto Cabinetでやっていた。条件変数を使ってビジーループを排除した通常バージョンもある。それらを再利用した方が楽そうだ。よって、スピンロックによる自前のshared_mutexをTkrzwに組み込んでみた。もちろん全てのテストは通った。

性能評価をしてみよう。従来のバージョンと、ファイルクラスにcontention_free_shared_mutexを組み込んだバージョンと、自前のスピンロックを組み込んだ実装を比較する。実行したテストコマンドは冒頭に上げたものと同じで、1億レコードを4スレッドで分担して同時に書き込む。各レコードは8バイトのキーと8バイトの値を持ち、ファイルサイズは3GBほどになる。

OLD Third-party NEW
Linux 1,596,393 2,085,587 2,071,034
MacOS 917,057 3,129,749 3,031,818

ということで、自前のスピンロック実装で、contention_free_shared_mutexとほぼ同等の性能が出ることが確かめられた。


スピンロックによるshared_mutexは、SpinSharedMutexというクラスで実装されている。ちょっとした工夫をすると、とても単純化した実装で高性能を実現できるので、それについて解説したい。

shared_mutexは、共有ロックと排他ロックを提供する機能だ。共有ロックは複数のスレッドが同時に獲得することができる。排他ロックを獲得したスレッドがいる間は、他のスレッドは共有ロックも排他ロックも獲得することができない。それを実現するための素直な実装では、共有ロックを獲得したスレッドの数と、排他ロックを獲得したスレッドの数を別個に数えて、各要求の際に適切な応答をすればよい。それぞれの数をアトミックに参照および更新するためにmutexを使い、更新された旨を条件変数を介して待機中のスレッドに通知する。

スピンロックによる実装の場合、単一のアトミック整数(atomic_uint32_t)を使うだけで同じことができる。共有ロックを獲得したスレッドの数をアトミック変数で数えるとともに、排他ロックを獲得したスレッドがいた場合にはその値をINT32MAXにする。

排他ロックを獲得しようとした場合、カウントが0の場合のみ、その値をINT32MAXにする。これはcompare_exchangeメソッドでアトミックに実現できる。失敗した場合は共有ロックを獲得しているスレッドがいるということだ。失敗の場合、成功するまでビジーループを繰り返せば、いつかはカウントがINT32MAXになる。この状態は、排他ロックを獲得したスレッドがいることを示す。

void lock() {
  uint32_t old_value = 0;
  while (!count_.compare_exchange_weak(old_value, INT32MAX)) {
    std::this_thread::yield();
    old_value = 0;
  }
}

排他ロックをアンロックする場合、単にカウンタの値を0にすればよい。排他ロックを獲得しているスレッドがいるなら、共有ロックを獲得しているスレッドはいないはずなので、常に0を設定してOKだ。

void unlock() {
  return count_.store(0);
}

共有ロックを獲得しようとした場合、カウントにfetch_addでアトミックに1を足してみて、戻り値(更新前のカウント)がINT32MAX未満である場合には、成功とみなす。そうでない場合、失敗だ。成功するまでビジーループして足し算を続ければ、いつかIN32MAX未満の値が得られる。カウンタはunsignedなので、INT32MAXに足し算をしても、UINT32MAXになるまではオーバーフローしない。とはいえ、共有ロックを獲得しようとする試みが21億回ほど連続して失敗すると困るので、失敗した場合には、compare_exchangeでカウンタの値をINT32MAXまで戻す試みを行う。これも失敗する可能性があるのだが、21億回連続して失敗することはないだろう。

void lock_shared() {
  while (count_.fetch_add(1) >= INT32MAX) {
    uint32_t old_value = count_.load();
    if (old_value > INT32MAX) {
      count_.compare_exchange_weak(old_value, INT32MAX);
    }
    std::this_thread::yield();
  }
}

共有ロックをアンロックする場合、fetch_subでカウンタの値を1減らせば良い。

void unlock_shared() {
  count_.fetch_sub(1);
}

このように、アトミック変数をたった1つだけ使い、mutexすら使わずに、共有ロックを実装することができる。理論的には完全ではないが。実際には完全とみなして問題ない前提条件を置くことにより、実装を単純化、効率化することができるのだ。この手のヒューリスティクスに萌える御仁は私の友達だ。コロンブスの卵。atomic_int64_tに変えれば確率的安全性は向上するのだが、問題はそこじゃない。ロック失敗時の回復処理がビジーループの待ち時間を活用しているところが面白いのだ。


メモリマップI/Oにスピンロックを使うことは既に述べたが、通常I/Oにはスピンロックは使わない。通常I/Oでは毎回の入出力にシステムコールを呼ぶので、スレッドがブロックする確率が高いし、ストレージの応答待ちになるとブロック時間も長い。よって、スピンロックではなく標準のstd::shared_mutexを使い続けた方がよい。

一方、データベース層で使っているshared_mutexの多くはスピンロックに置き換えて差し支えない。ハッシュ表のデータベースではバケット単位で、B+木のデータベースではページ単位で排他制御を行っているので、データベース層のshared_mutexがスレッドをブロックする確率は小さい。オンメモリデータベースはそもそもシステムコールを呼ばない(mallocの裏で呼ばれるmmapを除いて)ので、長時間のブロックは起こらない。ファイルデータベースの場合にはファイル層でブロックするが、その場合はファイル層のshared_mutexがそれを制御する。

ジーループ中にstd::this_thread::yieldを呼んでネイティブスレッドの実行権を手放すのはよいが、GoのN:M対応スレッドモデルに最適とは言い難い気もする。コールバックを注入できるようにして、Goインターフェイスから呼ばれた場合にはruntime.Goschedを呼ぶとかすると良さそうな気がするが、まだ実験はしていない。


まとめ。Tkrzwにスピンロックを導入することで、マルチスレッドのスループットを、Macで3.1倍、Linuxで1.3倍に高めることができた。スピンロックには副作用もあるので、全てのユースケースで有効なわけではないが、Tkrzwの典型的なユースケースでは有用と言えるだろう。

それにしても、Mac上のstd::shared_mutexは遅すぎる。もちろん何か理由があってそうならざるを得ないということなんだろうけども、それがなぜなのか知りたいところだ。標準実装を単に受け入れるのではなく、代替を探るというのも時に有用なのだなというのに気付かされた。報告者様に感謝。