Rust で LSM ツリーを実装した話
概要
データ指向アプリケーションデザインを読んでいて、出てきた LSM ツリーを理解するために、
簡単な LSM ツリーを実装してみました。
この記事では実装した LSM ツリーの読み書き処理に着目した説明を記載します。
LSM ツリーとは
LSM ツリー（Log-Structured Merge Tree）は主に書き込み操作のパフォーマンスを最適化するために使用されるデータ構造です。
LSM ツリーは、メモリ層とディスク層で構成されます。
書き込みを受けると以下の動作をします。
- メモリに書き込む
- メモリが一定のサイズを超えたとき、ストレージに書き込む。
ストレージへの書き込みを少なくすることで、効率的な書き込み操作を実現します。
この記事では書き込み対象について、それぞれ
- メモリ上のコンポーネントを
Memtable - ストレージ上のコンポーネントを
SSTable(Sorted String Table の略)
と呼びます。
実装の説明
ソースコードは以下のリポジトリにあります。
実装の説明に出てくるErrorは、std::io::Errorです。
LSM ツリーを表す構造体
LSM ツリーを表す構造体の実装は以下です。
Memtable, SSTable をそれぞれ以下の形で実現しています。
- Memtable は BTreeMap で実現
- SStable は SSTable の構造体を定義して実現 (後述)
pub struct LsmTree {
data_dir: PathBuf, // データを格納するディレクトリ
pub memtable: BTreeMap<String, String>, // Memtable (メモリ上のデータ格納場所)
limit: usize, // Memtable に格納できる行数の上限
sstable_list: Vec<SSTable>, // SSTable のリスト
}
LsmTreeを生成する関数(new)の実装は以下です。
limitのみ引数で受け取るようにしています。
load_sstable_files関数を実行して、既存の SSTable からsstable_listを作っています。
こちらの関数の処理は、LSM ツリーの読み書き処理と直接関係が無いので、説明を省かせていただきます。
impl LsmTree {
pub fn new(limit: usize) -> Result<Self, Error> { // Error は std::io::Error
let data_dir: PathBuf = PathBuf::from("./data");
let sstable_list: Vec<SSTable> = load_sstable_files(&data_dir)?;
Ok(LsmTree {
data_dir,
memtable: BTreeMap::new(),
limit,
sstable_list
})
}
※ データを格納するディレクトリも引数で受け取る形がいいと思いますが、
今回はあくまで動作させることが目的であるため、./dataで固定としています。
SSTable を表す構造体
SSTable を表す構造体の実装は以下です。
pub struct SSTable {
path: PathBuf,
index: HashMap<String, usize>
}
SSTableはcreate関数で生成します。
create関数の実装は以下です。
ファイル名はタイムスタンプ.datしています(例: 1729868505.dat)
impl SSTable {
pub fn create(data_dir: &PathBuf, memtable: &BTreeMap<String, String>, timestamp: &str) -> Result<Self, Error> {
// ファイルパスを生成
let mut filepath: PathBuf = data_dir.clone();
filepath.push(format!("{}.dat", timestamp));
// ファイルを作成し、BufWriterを生成
let mut writer: BufWriter<File> = match File::create(&filepath) {
Ok(f) => BufWriter::new(f),
Err(e) => return Err(e)
};
// memtable のデータをファイルに書き込む
let mut pointer: usize = 0;
let mut index: HashMap<String, usize> = HashMap::new();
for (k, v) in memtable.iter() {
index.insert(k.to_string(), pointer);
pointer = write_key_value(&mut writer, k, v)?;
}
Ok(SSTable {
path: filepath,
index
})
}
データをファイルに書き込む処理はwrite_key_value関数で行っています。
write_key_value関数の実装は以下です。
この関数は BufWriter とkey, valueの文字列を受け取ります。
受け取った文字列をバイト列に変換し、ファイルに書き込みます。
pub fn write_key_value(buf_writer: &mut BufWriter<File>, key: &str, value: &str) -> Result<usize, Error>{
// key, value をバイト列に変換する
let key_bytes: Vec<u8> = [&key.len().to_be_bytes(), key.as_bytes()].concat();
let value_bytes: Vec<u8> = [&value.len().to_be_bytes(), value.as_bytes()].concat();
let bytes: Vec<u8> = [key_bytes, value_bytes].concat();
buf_writer.write(&bytes)
}
データを書き込む際、以下の形式で書き込みます。
[キーの長さ(8バイト), キー, 値の長さ(8バイト), 値]
Memtable のデータが以下の2行の場合、
{
"k1": "hello",
"k2": "world"
}
ファイルには以下のデータが書き込まれます。
0~22バイトまでがk1: helloを表すデータ、残りがk2: worldを表すデータです。
[0, 0, 0, 0, 0, 0, 0, 2, 107, 49, 0, 0, 0, 0, 0, 0, 0, 5, 104, 101, 108, 108, 111, 0, 0, 0, 0, 0, 0, 0, 2, 107, 50, 0, 0, 0, 0, 0, 0, 0, 5, 119, 111, 114, 108, 100]
上記のデータは以下を示しています。
* 0, 0, 0, 0, 0, 0, 0, 2 => キーの長さが2文字であることを表す
* 107, 49 => キー本体 (107=k, 49=1 で k1)
* 0, 0, 0, 0, 0, 0, 0, 5 => 値が5文字であることを表す
* 104, 101, 108, 108, 111 => 値本体 (104=h, 101=e, 108=l. 108=l, 111=o で hello)
* 0, 0, 0, 0, 0, 0, 0, 2 => キーの長さが2文字であることを表す
* 107, 50 => キー本体 (107=k, 50=2 で k2)
* 0, 0, 0, 0, 0, 0, 0, 5 => 値が5文字であることを表す
* 119, 111, 114, 108, 100 => 値本体 (119=w, 111=o, 114=r. 108=l, 100=d で world)
SSTable のindexには、それぞれの行の開始バイトの位置を格納します。
上記の例ではk1の開始位置は0、k2の開始位置は23です。
その場合、indexは以下のようになります。
{
"k1": 0,
"k2": 23
}
indexは、読み取りの際に使用されます。
LSM ツリーへの読み書き
LSM ツリーへのデータの読み書きについて記載します。
書き込み処理 (put, delete)
書き込み処理はput関数で行います。
引数に受け取ったキーと値をMemtableに insert し、
Memtableに格納されているデータ数が、上限(limit)を超えた場合、flushを実行します。
impl LsmTree {
pub fn put(&mut self, key: &str, value: &str) -> Result<(), Error> {
self.memtable.insert(key.to_string(), value.to_string());
if self.limit < self.memtable.len() {
self.flush()?;
}
Ok(())
}
}
flushの実装は以下です。
SSTableを create して、Memtableを空にします。
impl LsmTree {
pub fn flush(&mut self) -> Result<(), Error> {
let timestamp: i64 = chrono::Local::now().timestamp();
match SSTable::create(&self.data_dir, &self.memtable, ×tamp.to_string()) {
Ok(sst) => self.sstable_list.push(sst),
Err(e) => return Err(e)
};
self.memtable.clear();
Ok(())
}
}
データの削除はdeleteで実行します。
deleteでは実際にデータは削除せず、削除されたことを示すデータをputします。
ここでは__delete flag__を削除されたことを示すデータとしています。
impl LsmTree {
pub fn delete(&mut self, key: &str) -> Result<(), Error> {
self.put(key, "__delete flag__")
}
}
読み取り処理 (get)
データの読み取り処理はgetで行います。
Memtable からデータを読み取り、Memtable にデータが存在しなければ SSTable からデータを読み取ります。
読み取ったデータが__delete flag__だった場合、そのデータは削除済みなので、Noneを返します。
impl LsmTree {
pub fn get(&mut self, key: &str) -> Result<Option<String>, Error> {
// Memtable から get する
if let Some(value) = self.memtable.get(key) {
if value == "__delete flag__" {
return Ok(None)
} else {
return Ok(Some(value.to_string()))
}
}
// SSTable から get する
for sstable in self.sstable_list.iter().rev() {
if let Some(value) = sstable.get(key)? {
if value == "__delete flag__" {
return Ok(None)
} else {
return Ok(Some(value))
}
}
}
Ok(None)
}
}
SSTable からのgetは以下の実装です。
SSTable の index からバイト列の開始位置を取得し、開始位置からデータを読み取ります。
impl SSTable {
pub fn get(&self, key: &str) -> Result<Option<String>, Error>{
// バイト列の開始位置を取得
let pointer: usize = match self.index.get(key) {
Some(p) => *p,
None => return Ok(None)
};
let mut reader: BufReader<File> = match File::open(&self.path) {
Ok(f) => BufReader::new(f),
Err(e) => return Err(e)
};
let (_, value) = read_key_value(&mut reader, pointer)?;
Ok(Some(value))
}
}
実際にファイルのデータを読むのはread_key_valueで行います。
keyの長さ -> key本体 -> valueの長さ -> value本体 の順でバイト列を読んでいきます。
pub fn read_key_value(buf_reader: &mut BufReader<File>, offset: usize) -> Result<(String, String), Error> {
buf_reader.seek(std::io::SeekFrom::Start(offset as u64))?;
// key の長さを read する
let mut bytes:[u8; 8] = [0; 8];
let key_length: usize = match buf_reader.read(&mut bytes) {
Ok(_) => usize::from_be_bytes(bytes),
Err(e) => return Err(e)
};
// key を read する
let mut key_bytes: Vec<u8> = vec![0; key_length];
let key: String = match buf_reader.read(&mut key_bytes) {
Ok(_) => String::from_utf8(key_bytes).unwrap(),
Err(e) => return Err(e)
};
// value の長さを read する
let mut bytes:[u8; 8] = [0; 8];
let value_length: usize = match buf_reader.read(&mut bytes) {
Ok(_) => usize::from_be_bytes(bytes),
Err(e) => return Err(e)
};
// value を read する
let mut value_bytes: Vec<u8> = vec![0; value_length];
let value: String = match buf_reader.read(&mut value_bytes) {
Ok(_) => String::from_utf8(value_bytes).unwrap(),
Err(e) => return Err(e)
};
Ok((key, value))
}
動作確認
put, get, delete
put, get, deleteの動作を確認します。
putでデータを書き込んだ後、getを実行し、
その後deleteでデータを削除し、再度getを実行します。
// limit=10 で LsmTree を生成
let mut lsm: LsmTree = match LsmTree::new(10) {
Ok(lsm) => lsm,
Err(e) => return eprintln!("{e}"),
};
let key: &str = "test_key";
let value: &str = "test_value";
// put 実行
if let Err(e) = lsm.put(key, value) {
eprintln!("{e}")
};
// get 実行 (put の結果を確認)
match lsm.get(key) {
Ok(opt) => println!("put後: {:?}", opt),
Err(e) => eprintln!("{e}")
};
// delete 実行
if let Err(e) = lsm.delete(key) {
eprintln!("{e}")
};
// get 実行 (delete の結果を確認)
match lsm.get(key) {
Ok(opt) => println!("delete後: {:?}", opt),
Err(e) => eprintln!("{e}")
}
実行結果は以下です。
putの後には、書き込んだデータをgetできており、
deleteの後には、データはgetできていないことが確認できます。
1回目のget(putの後) :Some("test_value")
2回目のget(deleteの後) :None
flush
flushの動作(= ディスクへの書き込みの動作)を確認します。
// limit=10 で LsmTree を生成
let mut lsm: LsmTree = match LsmTree::new(10) {
Ok(lsm) => lsm,
Err(e) => return eprintln!("{e}"),
};
let key: &str = "test_key";
let value: String = "test_value"
// put 実行
if let Err(e) = lsm.put(key, &value) {
eprintln!("{e}")
};
// flush 前後でファイルが増えることを確認する
let path: PathBuf = PathBuf::from("./data");
// flush 前、memtable にデータがあることを確認
println!("flush 前のmemtable\t: {:?}", lsm.memtable);
println!("--- flush 前のファイル一覧 --- ");
let files: fs::ReadDir = fs::read_dir(&path).unwrap();
for result in files {
println!("{:?}", result.unwrap());
}
// flush 実行
if let Err(e) = lsm.flush() {
eprintln!("{e}")
};
println!("--- flush 後のファイル一覧 ---");
let files: fs::ReadDir = fs::read_dir(&path).unwrap();
for result in files {
println!("{:?}", result.unwrap());
}
// flush 後、memtable が空になっていることを確認
println!("flush 後のmemtable\t: {:?}", lsm.memtable);
実行結果は以下です。
flush前はファイルが存在していませんが、flush後には1つファイルが存在しています。
またflushの後に、memtableが空になっていることが確認できます。
flush 前のmemtable : {"test_key": "test_value"}
--- flush 前のファイル一覧 ---
--- flush 後のファイル一覧 ---
DirEntry("./data/1729907176.dat")
flush 後のmemtable : {}
あとがき
ここまで読んでいただき、ありがとうございます。
今回の実装を通して、LSMツリーの概念を掴むことができたと思います。
また、Rustでのファイルの読み書きについても、学べたと思います。
今後はLSMツリーを使用した、キーバリューストアなどを作っていければよいかなと思っています。
Discussion