use crate::tombstone::{TombstoneManager, TombstoneStats};
use crate::types::Manifest;
use anyhow::{Context, Result};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompactionStats {
pub started_at: u64,
pub completed_at: u64,
pub duration_secs: u64,
pub tombstones_removed: usize,
pub segments_compacted: usize,
pub space_reclaimed_bytes: u64,
pub before_epoch: u64,
pub after_epoch: u64,
}
#[derive(Debug, Clone)]
pub struct CompactionConfig {
pub min_tombstone_age_epochs: u64,
pub min_tombstone_count: usize,
pub max_duration_secs: u64,
}
impl Default for CompactionConfig {
fn default() -> Self {
Self {
min_tombstone_age_epochs: 5, min_tombstone_count: 100, max_duration_secs: 300, }
}
}
pub struct CollectionCompactor {
collection_path: PathBuf,
config: CompactionConfig,
}
impl CollectionCompactor {
pub fn new(collection_path: &Path) -> Self {
Self {
collection_path: collection_path.to_path_buf(),
config: CompactionConfig::default(),
}
}
pub fn with_config(mut self, config: CompactionConfig) -> Self {
self.config = config;
self
}
pub fn needs_compaction(&self) -> Result<bool> {
let tombstone_manager = TombstoneManager::new(&self.collection_path);
let stats = tombstone_manager.get_stats()?;
let manifest_path = self.collection_path.join("MANIFEST.a");
let manifest = Manifest::read_from_file(&manifest_path)?;
let cutoff_epoch = manifest.epoch.saturating_sub(self.config.min_tombstone_age_epochs);
let old_tombstones = tombstone_manager.get_compaction_candidates(cutoff_epoch)?;
Ok(stats.total_count >= self.config.min_tombstone_count &&
old_tombstones.len() >= self.config.min_tombstone_count)
}
pub fn compact(&self) -> Result<CompactionStats> {
let started_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
println!("🧹 Starting collection compaction...");
let manifest_path = self.collection_path.join("MANIFEST.a");
let manifest = Manifest::read_from_file(&manifest_path)?;
let before_epoch = manifest.epoch;
let tombstone_manager = TombstoneManager::new(&self.collection_path);
let before_stats = tombstone_manager.get_stats()?;
println!(" Current epoch: {}", before_epoch);
println!(" Tombstones before: {}", before_stats.total_count);
let cutoff_epoch = before_epoch.saturating_sub(self.config.min_tombstone_age_epochs);
let candidates = tombstone_manager.get_compaction_candidates(cutoff_epoch)?;
if candidates.is_empty() {
println!(" No tombstones eligible for compaction");
let completed_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
return Ok(CompactionStats {
started_at,
completed_at,
duration_secs: completed_at - started_at,
tombstones_removed: 0,
segments_compacted: 0,
space_reclaimed_bytes: 0,
before_epoch,
after_epoch: before_epoch,
});
}
println!(" Compacting {} tombstones older than epoch {}", candidates.len(), cutoff_epoch);
let mut segments_to_compact: HashMap<u32, Vec<_>> = HashMap::new();
for tombstone in &candidates {
segments_to_compact
.entry(tombstone.segment_id)
.or_insert_with(Vec::new)
.push(tombstone);
}
println!(" Segments affected: {}", segments_to_compact.len());
let tombstones_removed = tombstone_manager.compact_tombstones(cutoff_epoch)?;
let space_reclaimed_bytes = (tombstones_removed * 256) as u64;
let new_epoch = before_epoch + 1;
let mut new_manifest = manifest;
new_manifest.epoch = new_epoch;
new_manifest.write_to_file(&manifest_path)?;
let completed_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
let stats = CompactionStats {
started_at,
completed_at,
duration_secs: completed_at - started_at,
tombstones_removed,
segments_compacted: segments_to_compact.len(),
space_reclaimed_bytes,
before_epoch,
after_epoch: new_epoch,
};
println!("✅ Compaction completed:");
println!(" Duration: {}s", stats.duration_secs);
println!(" Tombstones removed: {}", stats.tombstones_removed);
println!(" Segments compacted: {}", stats.segments_compacted);
println!(" Space reclaimed: ~{} bytes", stats.space_reclaimed_bytes);
println!(" New epoch: {}", stats.after_epoch);
self.save_compaction_stats(&stats)?;
Ok(stats)
}
pub fn get_compaction_history(&self) -> Result<Vec<CompactionStats>> {
let stats_dir = self.collection_path.join("gc");
let mut history = Vec::new();
if !stats_dir.exists() {
return Ok(history);
}
for entry in fs::read_dir(&stats_dir)? {
let entry = entry?;
let path = entry.path();
if let Some(filename) = path.file_name() {
if let Some(filename_str) = filename.to_str() {
if filename_str.starts_with("compaction-") && filename_str.ends_with(".json") {
if let Ok(content) = fs::read_to_string(&path) {
if let Ok(stats) = serde_json::from_str::<CompactionStats>(&content) {
history.push(stats);
}
}
}
}
}
}
history.sort_by_key(|s| s.completed_at);
Ok(history)
}
pub fn auto_compact_if_needed(&self) -> Result<bool> {
if self.needs_compaction()? {
println!("🔄 Auto-compaction triggered");
self.compact()?;
Ok(true)
} else {
Ok(false)
}
}
fn save_compaction_stats(&self, stats: &CompactionStats) -> Result<()> {
let stats_dir = self.collection_path.join("gc");
fs::create_dir_all(&stats_dir)?;
let stats_file = stats_dir.join(format!("compaction-{:010}.json", stats.completed_at));
let json = serde_json::to_string_pretty(stats)
.context("Failed to serialize compaction stats")?;
fs::write(&stats_file, json)
.context("Failed to write compaction stats")?;
Ok(())
}
}
pub struct CompactionManager {
collection_path: PathBuf,
}
impl CompactionManager {
pub fn new(collection_path: &Path) -> Self {
Self {
collection_path: collection_path.to_path_buf(),
}
}
pub fn compact_with_config(&self, config: CompactionConfig) -> Result<CompactionStats> {
let compactor = CollectionCompactor::new(&self.collection_path)
.with_config(config);
compactor.compact()
}
pub fn status(&self) -> Result<CompactionStatus> {
let tombstone_manager = TombstoneManager::new(&self.collection_path);
let tombstone_stats = tombstone_manager.get_stats()?;
let compactor = CollectionCompactor::new(&self.collection_path);
let needs_compaction = compactor.needs_compaction()?;
let history = compactor.get_compaction_history()?;
let last_compaction = history.last().map(|s| s.completed_at);
Ok(CompactionStatus {
needs_compaction,
total_tombstones: tombstone_stats.total_count,
oldest_tombstone_epoch: tombstone_stats.oldest_epoch,
newest_tombstone_epoch: tombstone_stats.newest_epoch,
last_compaction_at: last_compaction,
compaction_count: history.len(),
})
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct CompactionStatus {
pub needs_compaction: bool,
pub total_tombstones: usize,
pub oldest_tombstone_epoch: u64,
pub newest_tombstone_epoch: u64,
pub last_compaction_at: Option<u64>,
pub compaction_count: usize,
}