use std::{
fs::{self, File},
io::{BufReader, BufWriter},
path::Path,
};
use crate::{
wal::{
entry::Lsn,
reader::{ReadResult, WalReader},
writer::WalWriter,
},
StorageError,
};
pub const DEFAULT_SAFETY_BUFFER: u64 = 100;
#[derive(Debug, Clone)]
pub struct TruncateResult {
pub entries_removed: u64,
pub entries_kept: u64,
pub bytes_saved: u64,
pub oldest_lsn: Option<Lsn>,
pub newest_lsn: Option<Lsn>,
}
pub fn truncate_wal(
wal_path: &Path,
checkpoint_lsn: Lsn,
safety_buffer: Option<u64>,
) -> Result<TruncateResult, StorageError> {
let safety_buffer = safety_buffer.unwrap_or(DEFAULT_SAFETY_BUFFER);
let cutoff_lsn = checkpoint_lsn.saturating_sub(safety_buffer);
let original_size = fs::metadata(wal_path)
.map_err(|e| StorageError::IoError(format!("Failed to stat WAL: {}", e)))?
.len();
let file = File::open(wal_path)
.map_err(|e| StorageError::IoError(format!("Failed to open WAL: {}", e)))?;
let reader = BufReader::new(file);
let mut wal_reader = WalReader::open(reader)?;
let temp_path = wal_path.with_extension("wal.tmp");
let mut entries_to_keep = Vec::new();
let mut entries_removed = 0u64;
loop {
match wal_reader.read_entry()? {
ReadResult::Entry(entry) => {
if entry.lsn > cutoff_lsn {
entries_to_keep.push(entry);
} else {
entries_removed += 1;
}
}
ReadResult::Eof => break,
ReadResult::Corruption { position } => {
log::warn!("WAL corruption detected at position {} during truncation", position);
break;
}
}
}
if entries_removed == 0 {
return Ok(TruncateResult {
entries_removed: 0,
entries_kept: entries_to_keep.len() as u64,
bytes_saved: 0,
oldest_lsn: entries_to_keep.first().map(|e| e.lsn),
newest_lsn: entries_to_keep.last().map(|e| e.lsn),
});
}
{
let temp_file = File::create(&temp_path)
.map_err(|e| StorageError::IoError(format!("Failed to create temp WAL: {}", e)))?;
let writer = BufWriter::new(temp_file);
let mut wal_writer = WalWriter::create(writer)?;
for entry in &entries_to_keep {
wal_writer.append(entry)?;
}
wal_writer.flush()?;
}
let new_size = fs::metadata(&temp_path)
.map_err(|e| StorageError::IoError(format!("Failed to stat temp WAL: {}", e)))?
.len();
fs::rename(&temp_path, wal_path)
.map_err(|e| StorageError::IoError(format!("Failed to replace WAL: {}", e)))?;
let bytes_saved = original_size.saturating_sub(new_size);
log::info!(
"WAL truncated: removed {} entries, kept {}, saved {} bytes",
entries_removed,
entries_to_keep.len(),
bytes_saved
);
Ok(TruncateResult {
entries_removed,
entries_kept: entries_to_keep.len() as u64,
bytes_saved,
oldest_lsn: entries_to_keep.first().map(|e| e.lsn),
newest_lsn: entries_to_keep.last().map(|e| e.lsn),
})
}
pub fn truncate_wal_inplace(
wal_path: &Path,
checkpoint_lsn: Lsn,
safety_buffer: Option<u64>,
) -> Result<TruncateResult, StorageError> {
truncate_wal(wal_path, checkpoint_lsn, safety_buffer)
}
pub fn preview_truncation(
wal_path: &Path,
checkpoint_lsn: Lsn,
safety_buffer: Option<u64>,
) -> Result<TruncateResult, StorageError> {
let safety_buffer = safety_buffer.unwrap_or(DEFAULT_SAFETY_BUFFER);
let cutoff_lsn = checkpoint_lsn.saturating_sub(safety_buffer);
let file = File::open(wal_path)
.map_err(|e| StorageError::IoError(format!("Failed to open WAL: {}", e)))?;
let reader = BufReader::new(file);
let mut wal_reader = WalReader::open(reader)?;
let mut entries_to_remove = 0u64;
let mut entries_to_keep = 0u64;
let mut oldest_kept_lsn: Option<Lsn> = None;
let mut newest_kept_lsn: Option<Lsn> = None;
let mut bytes_to_remove = 0u64;
loop {
let position_before = wal_reader.position();
match wal_reader.read_entry()? {
ReadResult::Entry(entry) => {
let entry_size = wal_reader.position() - position_before;
if entry.lsn > cutoff_lsn {
entries_to_keep += 1;
if oldest_kept_lsn.is_none() {
oldest_kept_lsn = Some(entry.lsn);
}
newest_kept_lsn = Some(entry.lsn);
} else {
entries_to_remove += 1;
bytes_to_remove += entry_size;
}
}
ReadResult::Eof => break,
ReadResult::Corruption { .. } => break,
}
}
Ok(TruncateResult {
entries_removed: entries_to_remove,
entries_kept: entries_to_keep,
bytes_saved: bytes_to_remove,
oldest_lsn: oldest_kept_lsn,
newest_lsn: newest_kept_lsn,
})
}
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use tempfile::TempDir;
use vibesql_types::SqlValue;
use super::*;
use crate::wal::entry::{WalEntry, WalOp};
fn create_test_wal(dir: &Path, num_entries: u64) -> PathBuf {
let wal_path = dir.join("test.wal");
let file = File::create(&wal_path).unwrap();
let writer = BufWriter::new(file);
let mut wal_writer = WalWriter::create(writer).unwrap();
for i in 1..=num_entries {
let entry = WalEntry::new(
i,
1234567890 + i,
WalOp::Insert { table_id: 1, row_id: i, values: vec![SqlValue::Integer(i as i64)] },
);
wal_writer.append(&entry).unwrap();
}
wal_writer.flush().unwrap();
wal_path
}
#[test]
fn test_truncate_wal_basic() {
let temp_dir = TempDir::new().unwrap();
let wal_path = create_test_wal(temp_dir.path(), 100);
let result = truncate_wal(&wal_path, 50, Some(10)).unwrap();
assert_eq!(result.entries_removed, 40);
assert_eq!(result.entries_kept, 60);
assert_eq!(result.oldest_lsn, Some(41));
assert_eq!(result.newest_lsn, Some(100));
let file = File::open(&wal_path).unwrap();
let mut reader = WalReader::open(BufReader::new(file)).unwrap();
let entries = reader.read_all().unwrap();
assert_eq!(entries.len(), 60);
assert_eq!(entries[0].lsn, 41);
}
#[test]
fn test_truncate_wal_no_removal() {
let temp_dir = TempDir::new().unwrap();
let wal_path = create_test_wal(temp_dir.path(), 100);
let result = truncate_wal(&wal_path, 0, Some(0)).unwrap();
assert_eq!(result.entries_removed, 0);
assert_eq!(result.entries_kept, 100);
}
#[test]
fn test_truncate_wal_remove_all() {
let temp_dir = TempDir::new().unwrap();
let wal_path = create_test_wal(temp_dir.path(), 100);
let result = truncate_wal(&wal_path, 100, Some(0)).unwrap();
assert_eq!(result.entries_removed, 100);
assert_eq!(result.entries_kept, 0);
let file = File::open(&wal_path).unwrap();
let mut reader = WalReader::open(BufReader::new(file)).unwrap();
let entries = reader.read_all().unwrap();
assert_eq!(entries.len(), 0);
}
#[test]
fn test_preview_truncation() {
let temp_dir = TempDir::new().unwrap();
let wal_path = create_test_wal(temp_dir.path(), 100);
let original_size = fs::metadata(&wal_path).unwrap().len();
let preview = preview_truncation(&wal_path, 50, Some(10)).unwrap();
assert_eq!(preview.entries_removed, 40);
assert_eq!(preview.entries_kept, 60);
let new_size = fs::metadata(&wal_path).unwrap().len();
assert_eq!(original_size, new_size);
let file = File::open(&wal_path).unwrap();
let mut reader = WalReader::open(BufReader::new(file)).unwrap();
let entries = reader.read_all().unwrap();
assert_eq!(entries.len(), 100);
}
#[test]
fn test_truncate_with_default_safety_buffer() {
let temp_dir = TempDir::new().unwrap();
let wal_path = create_test_wal(temp_dir.path(), 200);
let result = truncate_wal(&wal_path, 150, None).unwrap();
assert_eq!(result.entries_removed, 50);
assert_eq!(result.entries_kept, 150);
}
}