use super::error::{StorageError, StorageResult};
use super::merge::merge_iter;
use super::sstable::SSTable;
use crate::log::{JudgeReal, LogManager};
use crate::storage::SyncDatabase;
use crate::MemDatabase;
use agilulf_protocol::Slice;
use crossbeam::sync::ShardedLock;
use futures::channel::mpsc::{unbounded, UnboundedSender};
use futures::executor::LocalPool;
use futures::stream::StreamExt;
use futures::task::LocalSpawnExt;
use std::collections::{BTreeMap, VecDeque};
use std::path::Path;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
#[repr(packed)]
#[derive(Clone)]
pub struct RawManifestLogEntry {
pub real_flag: u8,
pub add_flag: u8,
pub level: u8,
pub id: u8,
}
impl JudgeReal for RawManifestLogEntry {
fn is_real(&self) -> bool {
self.real_flag == 1
}
fn set_real(&mut self, real: bool) {
if real {
self.real_flag = 1;
} else {
self.real_flag = 0;
}
}
}
pub struct ManifestManager {
base_dir: String,
log_manager: Arc<LogManager<RawManifestLogEntry>>,
frozen_databases: Arc<ShardedLock<VecDeque<Arc<MemDatabase>>>>,
sstables: Arc<[ShardedLock<BTreeMap<usize, SSTable>>; 6]>, level_counter: Arc<[AtomicUsize; 6]>,
}
impl ManifestManager {
pub fn create_new(
base_dir: &str,
frozen_databases: Arc<ShardedLock<VecDeque<Arc<MemDatabase>>>>,
) -> StorageResult<ManifestManager> {
let base_path = Path::new(base_dir);
let manifest_path = base_path.join("MANIFEST");
let manifest_path = match manifest_path.to_str() {
Some(str) => str,
None => {
log::error!("Manifest path is not UTF-8: {:#?}", manifest_path);
return Err(StorageError::UnicodeError);
}
};
Ok(ManifestManager {
base_dir: base_dir.to_string(),
log_manager: Arc::new(LogManager::create_new(manifest_path, 4 * 1024)?),
frozen_databases,
sstables: Arc::new([
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
]),
level_counter: Arc::new([
AtomicUsize::new(0),
AtomicUsize::new(0),
AtomicUsize::new(0),
AtomicUsize::new(0),
AtomicUsize::new(0),
AtomicUsize::new(0),
]), })
}
pub fn open(
base_dir: &str,
frozen_databases: Arc<ShardedLock<VecDeque<Arc<MemDatabase>>>>,
) -> StorageResult<ManifestManager> {
let base_path = Path::new(base_dir);
let manifest_path = base_path.join("MANIFEST");
let manifest_path = match manifest_path.to_str() {
Some(str) => str,
None => {
log::error!("Manifest path is not UTF-8: {:#?}", manifest_path);
return Err(StorageError::UnicodeError);
}
};
let log_manager: Arc<LogManager<RawManifestLogEntry>> =
Arc::new(LogManager::open(manifest_path, 4 * 1024)?);
let sstables = Arc::new([
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
ShardedLock::new(BTreeMap::new()),
]);
let level_counter = Arc::new([
AtomicUsize::new(0),
AtomicUsize::new(0),
AtomicUsize::new(0),
AtomicUsize::new(0),
AtomicUsize::new(0),
AtomicUsize::new(0),
]);
for log in log_manager.iter() {
match log.add_flag {
1 => {
if log.level >= 6 {
return Err(StorageError::ManifestLogFormatError);
}
let table_path = base_path.join(format!("sstable_{}_{}", log.level, log.id));
log::info!("Restoring sstable from {:#?}", table_path);
let sstable_file = std::fs::OpenOptions::new()
.read(true)
.write(true)
.open(table_path)?;
let sstable = SSTable::open(sstable_file)?;
sstables
.get(log.level as usize)
.unwrap() .write()
.unwrap()
.insert(log.id as usize, sstable);
level_counter
.get(log.level as usize)
.unwrap() .fetch_max(log.id as usize, Ordering::SeqCst);
}
0 => {
if log.level >= 6 {
return Err(StorageError::ManifestLogFormatError);
}
sstables
.get(log.level as usize)
.unwrap() .write()
.unwrap()
.remove(&(log.id as usize));
}
_ => unreachable!(),
}
}
Ok(ManifestManager {
base_dir: base_dir.to_string(),
log_manager,
frozen_databases,
sstables,
level_counter,
})
}
pub fn background_work(&self) -> StorageResult<UnboundedSender<usize>> {
let frozen_databases = self.frozen_databases.clone();
let (freeze_sender, freeze_receiver) = unbounded::<usize>();
let mut freeze_receiver = freeze_receiver.fuse();
let base_dir = self.base_dir.clone();
let level_counter = self.level_counter.clone();
let sstables = self.sstables.clone();
let log_manager = self.log_manager.clone();
std::thread::Builder::new()
.name("background_worker".to_string())
.spawn(move || {
let mut local_pool = LocalPool::new();
let spawn_result = local_pool.spawner().spawn_local(async move {
let base_path = Path::new(&base_dir);
loop {
let newest_log_id = match freeze_receiver.next().await {
Some(id) => id,
None => {
break;
}
};
let new_log_path = base_path.join(format!("log.{}", newest_log_id));
let db_guard = frozen_databases.read().unwrap();
match db_guard.back() {
Some(db) => {
let sstable = SSTable::from(db.clone());
drop(db_guard);
let id = level_counter[0].fetch_add(1, Ordering::SeqCst);
let table_path = base_path.join(format! {"sstable_0_{}", id});
let table_path = match table_path.to_str() {
Some(str) => str,
None => {
log::error!("Table path is not UTF-8: {:#?}", table_path);
continue;
}
};
match sstable.save(table_path).await {
Ok(()) => {}
Err(err) => {
log::error!("Error while storing SSTable: {}", err);
continue;
}
}
log_manager.add_entry(RawManifestLogEntry {
real_flag: 1,
add_flag: 1,
level: 0,
id: id as u8,
});
sstables[0].write().unwrap().insert(id, sstable);
frozen_databases.write().unwrap().pop_back();
match std::fs::remove_file(new_log_path) {
Ok(()) => {}
Err(err) => {
log::error!("Error while remove log: {}", err);
continue;
}
};
}
None => {}
}
}
});
match spawn_result {
Ok(()) => {}
Err(err) => log::error!("Error while spawning: {}", err),
}
local_pool.run();
})?;
Ok(freeze_sender)
}
pub fn find_key(&self, key: Slice) -> Option<Slice> {
for level in 0..6 {
let level = self.sstables[level].read().unwrap();
for (_id, table) in level.iter() {
match table.get_sync(key.clone()) {
Ok(value) => return Some(value),
Err(_) => {}
}
}
}
None
}
pub fn scan(&self, start: Slice, end: Slice) -> impl Iterator<Item = (Slice, Slice)> {
let mut merge_vec = Vec::new();
for level in 0..6 {
let level = self.sstables[level].read().unwrap();
for (_id, table) in level.iter() {
merge_vec.push(table.scan_sync(start.clone(), end.clone()).into_iter())
}
}
merge_iter(merge_vec)
}
}