データ構造
学生の頃は、試験の要件でもアルゴリズム問題でも、メモリ上のデータ構造を学ぶことに多くの時間を使っていました。フルタイムの開発者になってからはデータ構造を学ぶ機会が減り、標準ライブラリやコミュニティの hash map や btree map をそのまま使うことが増えて、メモリデータ構造の性能や計算量について考えることが少なくなりました。その意味では、仕事がシンプルになったと言えます。
TiDB v8.5 の開発中、偶然のきっかけで TiDB の非常に重要なパスにあるデータ構造を最適化するチャンスを得ました。学生時代ほど未熟ではなくなった今、これまでの知識を総動員して TiDB の MemBuffer に強力なエンジンを載せ替え。その結果、Sysbench oltp_insert の TiDB v8.5.0 でのスループットを 13% 向上させました。
MemBuffer とは
TiDB のアーキテクチャでは TiDB が計算ノード、TiKV がストレージノードであり、トランザクションが commit された後に書き込まれた内容が TiKV に保存されます。トランザクションが commit される前は、すべての内容が MemBuffer に保持されています。MemBuffer はメモリ上のデータ構造で、すべての DML が MemBuffer に書き込まれ、さらにいくつかの特性を備えています。
Read Your Own Write
ダーティリードを防ぐため、トランザクションが commit されるまではその DML 書き込みを他トランザクションが読むことはできませんが、自分自身からは見える必要があります。
次の例では、commit 文の前は現在のトランザクションだけが自分で書き込んだ (1, 1) を見ることができます。
create table t(id int primary key, v int);
begin;
insert into t values(1, 1);
select * from t where id = 1; -- (1, 1)
commit;
細粒度のロールバック
トランザクションは原子性を持つと聞いたことがあると思いますが、ほかにも多くの操作が原子性を必要とし、失敗したら完全にロールバックする必要があります。
次の例では、(1, 2) の書き込みは失敗していますが、それ以前に成功した文の結果は保持する必要があり、トランザクション全体をそのまま abort してはいけません。
create table t(id int primary key, v int);
begin;
insert into t values(1, 1);
insert into t values(1, 2); -- Duplicate entry '1' for key 't.PRIMARY'
commit;
select * from t where id = 1; -- (1, 1)
もう一つ例を挙げます。insert ignore 文では、ユニークインデックスや主キーの衝突で失敗した行があっても、成功した行の書き込みには影響しないはずです。
create table t(id int primary key, v int);
begin;
insert into t values(1, 1);
insert ignore into t values(1, 2), (2, 2); -- 1 row affected
commit;
select * from t; -- (1, 1), (2, 2)
TiDB では、これらの細粒度ロールバック機能を MemBuffer が提供しています。
Iterator
データ構造が Iterator をサポートするには追加のコストがかかります。一般に、iterator を実装する場合は書き込み時に順序を保つ必要があり、red-black tree、btree map、skiplist といったデータ構造がこれに当たります。一方、hash map は iterator をサポートしないのでデータを無秩序に格納できます。通常、無秩序な格納の方が set/get の性能が良好です。
TiDB のユースケースでは iterator が求められるため、iterator をサポートする MemBuffer 実装を選ばざるを得ません。
MemBuffer のアーキテクチャ
細粒度のロールバックを実現するため、MemBuffer をインデックス構造（index）とログ（vlog）の二つに分割しています。
index と vlog はポインタで互いを参照します。
ロールバックが必要なときは、まず vlog 上の value を取得し、ポインタで対応する index データ構造上の leaf と一つ前のバージョンの value を見つけ出します。そして leaf のデータを古いバージョンの value に向け替えます（存在しなければ空ポインタを書き込みます）。これで一回のロールバックが完了します。この操作を繰り返し、vlog を目標の位置までロールバックします。
TiDB v8.1 以前では index に red-black tree (RBT) を使っていましたが、TiDB v8.5 からは adaptive radix tree (ART) データ構造を使っています。
ART は adaptive radix tree の略で、従来の radix tree のメモリ膨張問題を解決し、積極的な検索パス圧縮戦略を採用しています。TiDB のテストでは、ART は RBT より多くのメモリを使ってはいません。
TiDB v8.5 の実装の最適化
ART vs RBT
比較ベースのデータ構造では、検索の計算量は容量に依存することが多く（一般に O(log n)）、ART のデータ構造は比較に基づかず、計算量は O(key length) です。容量増による性能劣化を避けられるため、大量書き込みのシナリオでは大きな利点になります。
強調したいのは、ART にはデータへの強い制約があることです。インデックスに使う key は []byte か []uint8 形式でなければなりません。これは強い制約なので、比較ベースのデータ構造と性能比較をするのは不公平です。
もし利用シナリオの key が []byte または []uint8 形式であれば、ART はたいてい最良の選択になります。TiDB もこのシナリオに当たります。
Memory Arena
TiDB は Golang で実装されており、Golang の GC は小さなオブジェクトの頻繁な割り当てに敏感です。そのため MemBuffer 向けに Memory Arena を開発しました。Memory Arena は一度に大きめのメモリブロック（[]byte）を確保し、小さなオブジェクトが必要なときにはその大きなブロックから切り出して使います。
type MemdbArenaAddr struct {
idx uint32
off uint32
}
func (f *artAllocator) allocNode4() (arena.MemdbArenaAddr, *node4) {
addr, data = f.nodeAllocator.Alloc(node4size, true)
n4 := (*node4)(unsafe.Pointer(&data[0]))
n4.init()
return addr, n4
}
お気づきかもしれませんが、小さなオブジェクト（n4）を確保するとき、オブジェクト自体に加えて addr も生成しています。Memory Arena を使うと通常のポインタは意味を失うため、確保したオブジェクトの位置を示す独自のポインタが必要になります。idx は大きなメモリブロックのシーケンス番号、off はそのブロック内のアドレスを表します。この addr のサイズは uint64 と同じです。
Chunk Bytes Comparason
ART では、二つの []byte の最長共通 prefix を探す処理が頻繁に行われます。思いつきやすい実装は次のとおりです。
func longestCommonPrefix(key1, key2 []byte) int {
limit := min(len(key1), len(key2))
idx := 0
for ; idx < limit; idx++ {
if key1[idx] != key2[idx] {
return idx
}
}
return limit
}
しかしこの実装は性能が十分ではありません。共通の prefix が長い場合は比較に時間がかかりますし、各 loop に現れる if 命令が CPU の分岐予測の阻害になります。
TiDB の実装では、64bit リトルエンディアンの CPU アーキテクチャ（amd64, arm64）に対して 8 バイトを uint64 に変換してビット演算を行うことで、一度の比較で 8 バイトを判定できるようにしています。
func longestCommonPrefixByChunk(key1, key2 []byte) int {
limit := min(len(key1), len(key2))
idx := 0
p1 := unsafe.Pointer(&key1)
p2 := unsafe.Pointer(&key2)
// Compare 8 bytes at a time
remaining := limit
for remaining >= 8 {
if *(*uint64)(p1) != *(*uint64)(p2) {
// Find first different byte using trailing zeros
xor := *(*uint64)(p1) ^ *(*uint64)(p2)
return limit - remaining + (bits.TrailingZeros64(xor) >> 3)
}
p1 = unsafe.Add(p1, 8)
p2 = unsafe.Add(p2, 8)
remaining -= 8
}
// Compare rest bytes
idx = limit - remaining
for ; idx < limit; idx++ {
if key1[idx] != key2[idx] {
break
}
}
return idx
}
まとめ
ART は素晴らしいデータ構造であり、それを Golang 上でさらに高速に動かすために多くのエンジニアリング最適化を施しました。RBT と比べて性能が大きく向上していることが分かります。
私たちのデータ構造の研究に興味があれば、ぜひメールをください。