use std::{
collections::HashMap,
fs::{self, File},
io::{BufRead, BufReader, Write},
path::{Path, PathBuf},
sync::{Arc, Mutex},
};
use crate::lsm_tree::{LSMTree, LSMTreeOptions};
use crate::lsm_tree::sst::BlockCache;
pub struct LsmDB {
root: PathBuf,
default_name: String,
cfs: HashMap<String, Arc<Mutex<LSMTree>>>,
default_opts: LSMTreeOptions,
shared_cache: Arc<Mutex<BlockCache>>,
}
impl LsmDB {
pub fn open(root: impl AsRef<Path>) -> std::io::Result<Self> {
Self::open_with_config(root, LSMTreeOptions::default())
}
pub fn open_with_config(
root: impl AsRef<Path>,
opts: LSMTreeOptions,
) -> std::io::Result<Self> {
let root = root.as_ref().to_path_buf();
fs::create_dir_all(&root)?;
let current = root.join("CURRENT");
let mut names = vec!["default".to_string()];
if current.exists() {
names = load_current(¤t)?;
} else {
write_current(¤t, &names)?;
}
let shared_cache = Arc::new(Mutex::new(BlockCache::new(opts.block_cache_blocks)));
let mut cfs = HashMap::new();
for name in &names {
let dir = cf_dir(&root, name);
let mut tree = LSMTree::open_with_config(&dir, opts.clone())?;
tree.install_shared_block_cache(shared_cache.clone());
cfs.insert(name.clone(), Arc::new(Mutex::new(tree)));
}
Ok(Self {
root,
default_name: "default".to_string(),
cfs,
default_opts: opts,
shared_cache,
})
}
pub fn root(&self) -> &Path {
&self.root
}
pub fn list_cfs(&self) -> Vec<String> {
let mut v: Vec<_> = self.cfs.keys().cloned().collect();
v.sort();
v
}
pub fn default_cf(&self) -> Arc<Mutex<LSMTree>> {
self.cf(&self.default_name).expect("default")
}
pub fn cf(&self, name: &str) -> Option<Arc<Mutex<LSMTree>>> {
self.cfs.get(name).cloned()
}
pub fn create_cf(&mut self, name: &str) -> std::io::Result<Arc<Mutex<LSMTree>>> {
if let Some(c) = self.cfs.get(name) {
return Ok(c.clone());
}
if name.is_empty() || name.contains('/') || name.contains('\\') {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"非法 CF 名",
));
}
let dir = cf_dir(&self.root, name);
let mut tree = LSMTree::open_with_config(&dir, self.default_opts.clone())?;
tree.install_shared_block_cache(self.shared_cache.clone());
let arc = Arc::new(Mutex::new(tree));
self.cfs.insert(name.to_string(), arc.clone());
self.persist_current()?;
Ok(arc)
}
pub fn drop_cf(&mut self, name: &str) -> std::io::Result<()> {
if name == "default" {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"不能删除 default CF",
));
}
if self.cfs.remove(name).is_none() {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
"CF 不存在",
));
}
let dir = cf_dir(&self.root, name);
let _ = fs::remove_dir_all(dir);
self.persist_current()?;
Ok(())
}
pub fn flush_all(&self) -> std::io::Result<()> {
for name in self.list_cfs() {
if let Some(cf) = self.cf(&name) {
cf.lock().unwrap().flush()?;
}
}
Ok(())
}
pub fn set_options_all(&self, opts: LSMTreeOptions) {
for name in self.list_cfs() {
if let Some(cf) = self.cf(&name) {
cf.lock().unwrap().set_options(opts.clone());
}
}
}
fn persist_current(&self) -> std::io::Result<()> {
let mut names: Vec<_> = self.cfs.keys().cloned().collect();
names.sort();
names.retain(|n| n != "default");
names.insert(0, "default".to_string());
write_current(&self.root.join("CURRENT"), &names)
}
}
fn cf_dir(root: &Path, name: &str) -> PathBuf {
if name == "default" {
root.join("default")
} else {
root.join("cf").join(name)
}
}
fn write_current(path: &Path, names: &[String]) -> std::io::Result<()> {
let mut f = File::create(path)?;
writeln!(f, "LSMCF01")?;
for n in names {
writeln!(f, "{n}")?;
}
f.sync_all()?;
Ok(())
}
fn load_current(path: &Path) -> std::io::Result<Vec<String>> {
let f = File::open(path)?;
let mut lines = BufReader::new(f).lines();
let magic = lines
.next()
.ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "空 CURRENT"))??;
if magic.trim() != "LSMCF01" {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"非法 CURRENT",
));
}
let mut names = Vec::new();
for line in lines {
let line = line?;
let line = line.trim();
if !line.is_empty() {
names.push(line.to_string());
}
}
if names.is_empty() {
names.push("default".to_string());
}
Ok(names)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_multi_cf() {
let dir = std::env::temp_dir().join(format!(
"lsmdb_{}",
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
let _ = fs::remove_dir_all(&dir);
{
let mut db = LsmDB::open(&dir).unwrap();
{
let def = db.default_cf();
let mut t = def.lock().unwrap();
t.put(b"k".to_vec(), b"d".to_vec());
t.flush().unwrap();
}
let users = db.create_cf("users").unwrap();
{
let mut t = users.lock().unwrap();
t.put(b"k".to_vec(), b"u".to_vec());
t.flush().unwrap();
}
assert_eq!(
db.default_cf().lock().unwrap().get(b"k"),
Some(b"d".to_vec())
);
assert_eq!(
db.cf("users").unwrap().lock().unwrap().get(b"k"),
Some(b"u".to_vec())
);
let list = db.list_cfs();
assert!(list.contains(&"default".to_string()));
assert!(list.contains(&"users".to_string()));
}
{
let db = LsmDB::open(&dir).unwrap();
assert_eq!(
db.cf("users").unwrap().lock().unwrap().get(b"k"),
Some(b"u".to_vec())
);
}
let _ = fs::remove_dir_all(&dir);
}
}