use std::fs::File as SysFile;
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock, RwLockWriteGuard};
use crate::bplustree::tree::DiskBPlusTree;
use crate::compaction::{CompactionChoice, CompactionInput, CompactionStrategy};
use crate::error::{BackgroundErrorHandler, Result};
use crate::iter::{BoxedLSMIterator, CompactionIterator};
use crate::levels::{write_manifest_to_disk, LevelManifest, ManifestChangeSet};
use crate::lsm::{cleanup_vlog_and_index, CoreInner};
use crate::memtable::ImmutableMemtables;
use crate::snapshot::SnapshotTracker;
use crate::sstable::table::{Table, TableWriter};
use crate::vfs::File;
use crate::vlog::VLog;
use crate::{Comparator, Options as LSMOptions};
struct HiddenTablesGuard {
level_manifest: Arc<RwLock<LevelManifest>>,
table_ids: Vec<u64>,
committed: bool,
}
impl HiddenTablesGuard {
fn new(level_manifest: Arc<RwLock<LevelManifest>>, table_ids: &[u64]) -> Self {
Self {
level_manifest,
table_ids: table_ids.to_vec(),
committed: false,
}
}
fn commit(&mut self) {
self.committed = true;
}
}
impl Drop for HiddenTablesGuard {
fn drop(&mut self) {
if !self.committed {
if let Ok(mut levels) = self.level_manifest.write() {
levels.unhide_tables(&self.table_ids);
}
}
}
}
pub(crate) struct CompactionOptions {
pub(crate) lopts: Arc<LSMOptions>,
pub(crate) level_manifest: Arc<RwLock<LevelManifest>>,
pub(crate) immutable_memtables: Arc<RwLock<ImmutableMemtables>>,
pub(crate) vlog: Option<Arc<VLog>>,
pub(crate) error_handler: Arc<BackgroundErrorHandler>,
pub(crate) snapshot_tracker: SnapshotTracker,
pub(crate) versioned_index: Option<Arc<parking_lot::RwLock<DiskBPlusTree>>>,
}
impl CompactionOptions {
pub(crate) fn from(tree: &CoreInner) -> Self {
Self {
lopts: Arc::clone(&tree.opts),
level_manifest: Arc::clone(&tree.level_manifest),
immutable_memtables: Arc::clone(&tree.immutable_memtables),
vlog: tree.vlog.clone(),
error_handler: Arc::clone(&tree.error_handler),
snapshot_tracker: tree.snapshot_tracker.clone(),
versioned_index: tree.versioned_index.clone(),
}
}
}
pub(crate) struct Compactor {
pub(crate) options: CompactionOptions,
pub(crate) strategy: Arc<dyn CompactionStrategy>,
}
impl Compactor {
pub(crate) fn new(options: CompactionOptions, strategy: Arc<dyn CompactionStrategy>) -> Self {
Self {
options,
strategy,
}
}
pub(crate) fn compact(&self) -> Result<()> {
let levels_guard = self.options.level_manifest.write()?;
let choice = self.strategy.pick_levels(&levels_guard)?;
match choice {
CompactionChoice::Merge(input) => self.merge_tables(levels_guard, &input),
CompactionChoice::Skip => Ok(()),
}
}
fn merge_tables(
&self,
mut levels: RwLockWriteGuard<'_, LevelManifest>,
input: &CompactionInput,
) -> Result<()> {
levels.hide_tables(&input.tables_to_merge);
let mut guard = HiddenTablesGuard::new(
Arc::clone(&self.options.level_manifest),
&input.tables_to_merge,
);
let tables = levels.get_all_tables();
let to_merge: Vec<_> =
input.tables_to_merge.iter().filter_map(|&id| tables.get(&id).cloned()).collect();
let iterators: Vec<BoxedLSMIterator<'_>> = to_merge
.iter()
.filter_map(|table| table.iter(None).ok())
.map(|iter| Box::new(iter) as BoxedLSMIterator<'_>)
.collect();
drop(levels);
let new_table_id = self.options.level_manifest.read().unwrap().next_table_id();
let new_table_path = self.get_table_path(new_table_id);
let table_created =
match self.write_merged_table(&new_table_path, new_table_id, iterators, input) {
Ok(result) => result,
Err(e) => {
return Err(e);
}
};
let new_table = if table_created {
match self.open_table(new_table_id, &new_table_path) {
Ok(table) => Some(table),
Err(e) => {
return Err(e);
}
}
} else {
None
};
self.update_manifest(input, new_table, &mut guard)?;
self.cleanup_old_tables(input);
Ok(())
}
fn write_merged_table(
&self,
path: &Path,
table_id: u64,
merge_iter: Vec<BoxedLSMIterator<'_>>,
input: &CompactionInput,
) -> Result<bool> {
let file = SysFile::create(path)?;
let mut writer =
TableWriter::new(file, table_id, Arc::clone(&self.options.lopts), input.target_level);
let snapshots = self.options.snapshot_tracker.get_all_snapshots();
let max_level = self.options.lopts.level_count - 1;
let is_bottom_level = input.target_level >= max_level;
let mut comp_iter = CompactionIterator::new(
merge_iter,
Arc::clone(&self.options.lopts.internal_comparator) as Arc<dyn Comparator>,
is_bottom_level,
self.options.lopts.enable_versioning,
self.options.lopts.versioned_history_retention_ns,
Arc::clone(&self.options.lopts.clock),
snapshots,
);
let mut entries = 0;
for item in &mut comp_iter {
let (key, value) = item?;
writer.add(key, &value)?;
entries += 1;
}
if entries == 0 {
drop(writer);
let _ = std::fs::remove_file(path);
return Ok(false);
}
writer.finish()?;
Ok(true)
}
fn update_manifest(
&self,
input: &CompactionInput,
new_table: Option<Arc<Table>>,
guard: &mut HiddenTablesGuard,
) -> Result<()> {
let mut manifest = self.options.level_manifest.write()?;
let _imm_guard = self.options.immutable_memtables.write();
if let Some(ref table) = new_table {
if input.tables_to_merge.contains(&table.id) {
return Err(crate::error::Error::TableIDCollision(table.id));
}
}
let mut changeset = ManifestChangeSet::default();
for (level_idx, level) in manifest.levels.get_levels().iter().enumerate() {
for &table_id in &input.tables_to_merge {
if level.tables.iter().any(|t| t.id == table_id) {
changeset.deleted_tables.insert((level_idx as u8, table_id));
}
}
}
if let Some(table) = new_table {
changeset.new_tables.push((input.target_level, table));
}
let rollback = manifest.apply_changeset(&changeset)?;
if let Err(e) = write_manifest_to_disk(&manifest) {
manifest.revert_changeset(rollback);
self.options
.error_handler
.set_error(e.clone(), crate::error::BackgroundErrorReason::ManifestWrite);
return Err(e);
}
manifest.unhide_tables(&input.tables_to_merge);
guard.commit();
let min_oldest_vlog = manifest.min_oldest_vlog_file_id();
cleanup_vlog_and_index(
&self.options.vlog,
&self.options.versioned_index,
min_oldest_vlog,
"compaction",
);
Ok(())
}
fn get_table_path(&self, table_id: u64) -> PathBuf {
self.options.lopts.sstable_file_path(table_id)
}
fn cleanup_old_tables(&self, input: &CompactionInput) {
for &table_id in &input.tables_to_merge {
let path = self.options.lopts.sstable_file_path(table_id);
if let Err(e) = std::fs::remove_file(path) {
log::warn!("Failed to remove old table file: {e}");
}
}
}
fn open_table(&self, table_id: u64, table_path: &Path) -> Result<Arc<Table>> {
let file = SysFile::open(table_path)?;
let file: Arc<dyn File> = Arc::new(file);
let file_size = file.size()?;
Ok(Arc::new(Table::new(table_id, Arc::clone(&self.options.lopts), file, file_size)?))
}
}