85 lines
2.2 KiB
Rust
85 lines
2.2 KiB
Rust
|
use std::collections::HashSet;
|
||
|
|
use std::sync::Arc;
|
||
|
|
|
||
|
|
use clap::Parser;
|
||
|
|
use mini_lsm::compact::TieredCompactionController;
|
||
|
|
use mini_lsm::lsm_storage::LsmStorageInner;
|
||
|
|
use mini_lsm::mem_table::MemTable;
|
||
|
|
|
||
|
|
#[derive(Parser, Debug)]
|
||
|
|
#[command(author, version, about, long_about = None)]
|
||
|
|
enum Args {
|
||
|
|
Tiered {},
|
||
|
|
Leveled {},
|
||
|
|
}
|
||
|
|
|
||
|
|
pub struct MockStorage {
|
||
|
|
snapshot: LsmStorageInner,
|
||
|
|
next_sst_id: usize,
|
||
|
|
file_list: HashSet<usize>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl MockStorage {
|
||
|
|
pub fn new() -> Self {
|
||
|
|
let snapshot = LsmStorageInner {
|
||
|
|
memtable: Arc::new(MemTable::create()),
|
||
|
|
imm_memtables: Vec::new(),
|
||
|
|
l0_sstables: Vec::new(),
|
||
|
|
levels: Vec::new(),
|
||
|
|
sstables: Default::default(),
|
||
|
|
};
|
||
|
|
Self {
|
||
|
|
snapshot,
|
||
|
|
next_sst_id: 0,
|
||
|
|
file_list: Default::default(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn generate_sst_id(&mut self) -> usize {
|
||
|
|
let id = self.next_sst_id;
|
||
|
|
self.next_sst_id += 1;
|
||
|
|
id
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn flush_sst(&mut self) {
|
||
|
|
let id = self.generate_sst_id();
|
||
|
|
self.snapshot.l0_sstables.push(id);
|
||
|
|
self.file_list.insert(id);
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn remove(&mut self, files_to_remove: &[usize]) {
|
||
|
|
for file_id in files_to_remove {
|
||
|
|
self.file_list.remove(file_id);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn dump(&self) {
|
||
|
|
print!("L0: {:?}", self.snapshot.l0_sstables);
|
||
|
|
for (level, files) in &self.snapshot.levels {
|
||
|
|
print!("L{level}: {:?}", files);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn main() {
|
||
|
|
let args = Args::parse();
|
||
|
|
match args {
|
||
|
|
Args::Tiered {} => {
|
||
|
|
let controller = TieredCompactionController {};
|
||
|
|
let mut storage = MockStorage::new();
|
||
|
|
for i in 0..500 {
|
||
|
|
println!("Iteration {i}");
|
||
|
|
storage.flush_sst();
|
||
|
|
let task = controller.generate_compaction_task(&storage.snapshot);
|
||
|
|
let sst_id = storage.generate_sst_id();
|
||
|
|
let (snapshot, del) =
|
||
|
|
controller.apply_compaction_result(&storage.snapshot, &task, &[sst_id]);
|
||
|
|
storage.snapshot = snapshot;
|
||
|
|
storage.remove(&del);
|
||
|
|
storage.dump();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Args::Leveled {} => {}
|
||
|
|
}
|
||
|
|
}
|