#[cfg(feature = "write-support")]
pub mod cql_to_mutation;
#[cfg(feature = "write-support")]
pub mod export;
#[cfg(feature = "write-support")]
pub mod memtable;
#[cfg(feature = "write-support")]
pub mod merge;
#[cfg(feature = "write-support")]
pub mod merge_policy;
#[cfg(feature = "write-support")]
pub mod mutation;
#[cfg(feature = "write-support")]
pub(crate) mod reconcile_rules;
#[cfg(feature = "write-support")]
pub mod wal;
#[cfg(feature = "write-support")]
pub use export::{ExportOptions, ExportReport};
#[cfg(feature = "write-support")]
pub use memtable::Memtable;
#[cfg(feature = "write-support")]
pub use merge::build_single_partition_merger;
#[cfg(feature = "write-support")]
pub use merge::build_single_partition_merger_from_readers;
#[cfg(feature = "write-support")]
pub use merge::KWayMerger;
#[cfg(feature = "write-support")]
pub use merge_policy::STCSPolicy;
#[cfg(feature = "write-support")]
pub use mutation::{
CellOperation, ClusteringBound, ClusteringKey, DecoratedKey, Mutation, PartitionKey,
PartitionTombstone, RangeTombstone, TableId,
};
#[cfg(feature = "write-support")]
pub use wal::{RecoveryReport, WriteAheadLog};
#[cfg(feature = "write-support")]
mod compaction;
#[cfg(feature = "write-support")]
pub(crate) mod durability;
#[cfg(feature = "write-support")]
mod maintenance;
#[cfg(feature = "write-support")]
mod stats;
#[cfg(feature = "write-support")]
mod sweep;
#[cfg(all(test, feature = "write-support"))]
mod test_support;
#[cfg(all(test, feature = "write-support"))]
mod admission_tests;
#[cfg(feature = "write-support")]
pub use maintenance::MaintenanceReport;
#[cfg(feature = "write-support")]
pub use stats::CompactionStats;
use crate::error::{Error, Result};
use crate::schema::{TableSchema, UdtRegistry};
use crate::storage::sstable::writer::SSTableInfo;
#[cfg(feature = "write-support")]
use maintenance::ActiveMerge;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;
#[cfg(feature = "write-support")]
pub trait MergePolicy: Send + std::fmt::Debug {
fn select_merge(&self, candidates: &[PathBuf]) -> Result<Vec<PathBuf>>;
}
#[cfg(feature = "write-support")]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Durability {
#[default]
SyncEachWrite,
Disabled,
}
#[cfg(feature = "write-support")]
#[derive(Debug, Clone)]
pub struct WriteEngineConfig {
pub data_dir: PathBuf,
pub wal_dir: PathBuf,
pub memtable_flush_threshold: usize,
pub memtable_hard_limit: usize,
pub schema: TableSchema,
pub durability: Durability,
pub udt_registry: Option<UdtRegistry>,
pub auto_compaction: bool,
pub compaction_min_threshold: usize,
pub compaction_max_threshold: usize,
}
#[cfg(feature = "write-support")]
impl WriteEngineConfig {
pub const DEFAULT_FLUSH_THRESHOLD: usize = 64 * 1024 * 1024;
pub const DEFAULT_HARD_LIMIT: usize = 256 * 1024 * 1024;
pub const DEFAULT_COMPACTION_MIN_THRESHOLD: usize = 4;
pub const DEFAULT_COMPACTION_MAX_THRESHOLD: usize = 32;
pub fn new(data_dir: PathBuf, wal_dir: PathBuf, schema: TableSchema) -> Self {
Self {
data_dir,
wal_dir,
memtable_flush_threshold: Self::DEFAULT_FLUSH_THRESHOLD,
memtable_hard_limit: Self::DEFAULT_HARD_LIMIT,
schema,
durability: Durability::default(),
udt_registry: None,
auto_compaction: true,
compaction_min_threshold: Self::DEFAULT_COMPACTION_MIN_THRESHOLD,
compaction_max_threshold: Self::DEFAULT_COMPACTION_MAX_THRESHOLD,
}
}
pub fn with_udt_registry(mut self, registry: UdtRegistry) -> Self {
self.udt_registry = Some(registry);
self
}
pub fn with_flush_threshold(mut self, threshold: usize) -> Self {
self.memtable_flush_threshold = threshold;
self
}
pub fn with_hard_limit(mut self, limit: usize) -> Self {
self.memtable_hard_limit = limit;
self
}
pub fn with_durability(mut self, durability: Durability) -> Self {
self.durability = durability;
self
}
pub fn with_compaction_config(mut self, compaction: &crate::config::CompactionConfig) -> Self {
self.auto_compaction = compaction.auto_compaction;
self
}
}
#[cfg(feature = "write-support")]
#[derive(Debug)]
pub struct WriteEngine {
config: WriteEngineConfig,
wal: WriteAheadLog,
memtable: Memtable,
wal_recovery: RecoveryReport,
generation: u64,
closed: AtomicBool,
active_merge: Option<ActiveMerge>,
merge_policy: Option<Box<dyn MergePolicy>>,
cumulative_stats: CompactionStats,
rows_written: u64,
l0_count: u64,
total_flushed_bytes: u64,
dir_lock: std::fs::File,
warned_over_threshold: bool,
}
#[cfg(feature = "write-support")]
fn reject_counter_cells(mutation: &Mutation) -> Result<()> {
for op in &mutation.operations {
match op {
CellOperation::Write { value, .. } | CellOperation::WriteWithTtl { value, .. } => {
if matches!(value, crate::types::Value::Counter(_)) {
return Err(Error::invalid_operation(
"counter writes are not supported via the standard mutation path; \
counter columns require server-side distributed increment semantics",
));
}
}
_ => {}
}
}
Ok(())
}
#[cfg(feature = "write-support")]
impl WriteEngine {
fn create_dir_all_durable(dir: &Path, label: &str) -> Result<()> {
std::fs::create_dir_all(dir).map_err(|e| {
Error::Storage(format!(
"Failed to create {} directory {:?}: {}",
label, dir, e
))
})?;
let mut next = dir.parent();
while let Some(cur) = next {
if cur.as_os_str().is_empty() {
wal::sync_directory(Path::new("."))?;
break;
}
wal::sync_directory(cur)?;
next = cur.parent();
}
Ok(())
}
pub fn new(config: WriteEngineConfig) -> Result<Self> {
Self::create_dir_all_durable(&config.data_dir, "data")?;
Self::create_dir_all_durable(&config.wal_dir, "WAL")?;
let lock_path = config.wal_dir.join(".lock");
let dir_lock = std::fs::OpenOptions::new()
.create(true)
.truncate(false)
.write(true)
.open(&lock_path)
.map_err(|e| {
Error::Storage(format!("Failed to create lock file {:?}: {}", lock_path, e))
})?;
fs2::FileExt::try_lock_exclusive(&dir_lock)
.map_err(|_| Error::write_dir_locked(config.wal_dir.to_string_lossy().into_owned()))?;
Self::sweep_startup_orphans(&config);
let wal_path = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
let mut wal = if wal_path.exists() {
WriteAheadLog::open_existing(&wal_path)?
} else {
WriteAheadLog::create(&config.wal_dir)?
};
let mut memtable = Memtable::new();
let mut recovered = 0usize;
let mut apply_error: Option<Error> = None;
let wal_recovery = wal.replay_each(|mutation| {
if apply_error.is_some() {
return Ok(());
}
match mutation.decorated_key(&config.schema) {
Ok(decorated_key) => match memtable.insert_with_key(decorated_key, mutation) {
Ok(()) => recovered += 1,
Err(e) => apply_error = Some(e),
},
Err(e) => apply_error = Some(e),
}
Ok(())
})?;
if !wal_recovery.is_clean() {
let preserved = Self::preserve_corrupt_wal(&wal_path)?;
let reset_to = wal.reset_to_valid_prefix()?;
tracing::error!(
"WAL recovery at {:?} was LOSSY: recovered {} mutation(s), {} corrupt entry \
(entries), stopped_early={}, {} byte(s) not recovered. Raw segment preserved at \
{:?}; live WAL reset to valid prefix ({:?}). Investigate before relying on this \
data.",
wal_path,
recovered,
wal_recovery.corrupt_entries,
wal_recovery.stopped_early,
wal_recovery.bytes_skipped,
preserved,
reset_to,
);
}
if let Some(e) = apply_error {
return Err(e);
}
if recovered > 0 {
tracing::info!(
"WAL replay complete: replayed {} mutation(s); {} rows in memtable, {} bytes",
recovered,
memtable.row_count(),
memtable.size_bytes()
);
}
let generation = Self::determine_next_generation(&config.data_dir)?;
let merge_policy: Option<Box<dyn MergePolicy>> = if config.auto_compaction {
Some(Box::new(STCSPolicy::new(
config.compaction_min_threshold,
config.compaction_max_threshold,
0.5,
1.5,
STCSPolicy::DEFAULT_MIN_SSTABLE_SIZE,
)?))
} else {
None
};
Ok(Self {
config,
wal,
memtable,
wal_recovery,
generation,
closed: AtomicBool::new(false),
active_merge: None,
merge_policy,
cumulative_stats: CompactionStats::default(),
rows_written: 0,
l0_count: 0,
total_flushed_bytes: 0,
dir_lock,
warned_over_threshold: false,
})
}
pub fn wal_recovery(&self) -> &RecoveryReport {
&self.wal_recovery
}
fn preserve_corrupt_wal(wal_path: &Path) -> Result<PathBuf> {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
let file_name = wal_path
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| WriteAheadLog::WAL_FILENAME.to_string());
let preserved = match wal_path.parent() {
Some(parent) => parent.join(format!("{}.corrupt.{}", file_name, nanos)),
None => PathBuf::from(format!("{}.corrupt.{}", file_name, nanos)),
};
std::fs::copy(wal_path, &preserved).map_err(|e| {
Error::Storage(format!(
"Failed to preserve corrupt WAL {:?} aside to {:?}: {}",
wal_path, preserved, e
))
})?;
{
let copy_file = std::fs::File::open(&preserved).map_err(|e| {
Error::Storage(format!(
"Failed to open preserved corrupt WAL {:?} for fsync: {}",
preserved, e
))
})?;
copy_file.sync_all().map_err(|e| {
Error::Storage(format!(
"Failed to fsync preserved corrupt WAL {:?}: {}",
preserved, e
))
})?;
}
if let Some(parent) = preserved.parent() {
wal::sync_directory(parent)?;
}
Ok(preserved)
}
#[tracing::instrument(name = "write.mutation", level = "debug", skip(self, mutation))]
pub fn write(&mut self, mutation: Mutation) -> Result<()> {
crate::observability::record_result("write", self.write_inner(mutation))
}
fn write_inner(&mut self, mutation: Mutation) -> Result<()> {
self.write_into_memtable(mutation)?;
if self
.memtable
.should_flush(self.config.memtable_flush_threshold)
{
if !self.warned_over_threshold {
tracing::warn!(
"Memtable size {} exceeds threshold {} - call flush() manually in async context",
self.memtable.size_bytes(),
self.config.memtable_flush_threshold
);
self.warned_over_threshold = true;
}
if tokio::runtime::Handle::try_current().is_err() {
tracing::info!("Triggering automatic flush");
self.flush_internal()?;
}
}
Ok(())
}
fn write_into_memtable(&mut self, mutation: Mutation) -> Result<()> {
if self.closed.load(Ordering::SeqCst) {
return Err(Error::InvalidInput(
"WriteEngine has been closed".to_string(),
));
}
reject_counter_cells(&mutation)?;
self.check_admission(&mutation)?;
if self.config.durability == Durability::SyncEachWrite {
self.wal.append(&mutation)?;
self.wal.sync()?;
}
let decorated_key = mutation.decorated_key(&self.config.schema)?;
self.memtable.insert_with_key(decorated_key, mutation)?;
self.rows_written += 1;
crate::observability::add_counter(crate::observability::catalog::WRITE_MUTATIONS, 1, &[]);
self.record_memtable_gauges();
Ok(())
}
fn check_admission(&self, mutation: &Mutation) -> Result<()> {
let hard_limit = self.config.memtable_hard_limit;
let incoming = self.memtable.estimate_mutation_size(mutation);
if incoming == usize::MAX {
return Err(Error::Storage(
"Write rejected: mutation size could not be bounded (estimator \
fail-closed sentinel); refusing admission"
.to_string(),
));
}
if incoming > hard_limit {
return Err(Error::Storage(format!(
"Write rejected: single mutation estimated at {incoming} bytes exceeds \
memtable hard limit {hard_limit} bytes (a single mutation may not exceed \
the hard limit)"
)));
}
let projected = self.memtable.size_bytes().saturating_add(incoming);
if projected > hard_limit {
return Err(Error::Storage(format!(
"Write rejected: memtable {} + mutation {incoming} would exceed hard limit {hard_limit}",
self.memtable.size_bytes()
)));
}
Ok(())
}
#[tracing::instrument(name = "write.mutation", level = "debug", skip(self, mutation))]
pub async fn write_async(&mut self, mutation: Mutation) -> Result<()> {
crate::observability::record_result("write", self.write_async_inner(mutation).await)
}
async fn write_async_inner(&mut self, mutation: Mutation) -> Result<()> {
if self.closed.load(Ordering::SeqCst) {
return Err(Error::InvalidInput(
"WriteEngine has been closed".to_string(),
));
}
reject_counter_cells(&mutation)?;
self.check_admission(&mutation)?;
if self.config.durability == Durability::SyncEachWrite {
self.wal.append(&mutation)?;
self.wal.sync()?;
}
let decorated_key = mutation.decorated_key(&self.config.schema)?;
self.memtable.insert_with_key(decorated_key, mutation)?;
self.rows_written += 1;
crate::observability::add_counter(crate::observability::catalog::WRITE_MUTATIONS, 1, &[]);
self.record_memtable_gauges();
if self
.memtable
.should_flush(self.config.memtable_flush_threshold)
{
tracing::info!(
"Memtable size {} exceeds threshold {}, triggering flush",
self.memtable.size_bytes(),
self.config.memtable_flush_threshold
);
self.flush_internal_async().await?;
}
Ok(())
}
fn record_memtable_gauges(&self) {
crate::observability::record_gauge(
crate::observability::catalog::MEMTABLE_SIZE_BYTES,
self.memtable.size_bytes() as i64,
&[],
);
crate::observability::record_gauge(
crate::observability::catalog::MEMTABLE_ROWS,
self.memtable.row_count() as i64,
&[],
);
}
#[tracing::instrument(name = "write.cql_execute", level = "debug", skip(self, statement))]
pub fn execute(&mut self, statement: &str) -> Result<()> {
crate::observability::record_result("write", self.execute_inner(statement))
}
fn execute_inner(&mut self, statement: &str) -> Result<()> {
self.execute_inner_counted(statement).map(|_| ())
}
fn execute_inner_counted(&mut self, statement: &str) -> Result<u64> {
if self.closed.load(Ordering::SeqCst) {
return Err(Error::InvalidInput(
"WriteEngine has been closed".to_string(),
));
}
let mutations = self.parse_statement_to_mutations(statement)?;
let n = mutations.len() as u64;
for mutation in mutations {
self.write_inner(mutation)?;
}
Ok(n)
}
fn parse_statement_to_mutations(&self, statement: &str) -> Result<Vec<Mutation>> {
let trimmed = statement.trim();
if trimmed.len() >= 5 && trimmed.as_bytes()[..5].eq_ignore_ascii_case(b"BEGIN") {
cql_to_mutation::convert_cql_to_mutations(trimmed, &self.config.schema)
} else {
Ok(vec![self.parse_cql_to_mutation(statement)?])
}
}
#[tracing::instrument(
name = "write.cql_execute_flushing",
level = "debug",
skip(self, statement)
)]
pub async fn execute_flushing(&mut self, statement: &str) -> Result<u64> {
crate::observability::record_result("write", self.execute_flushing_inner(statement).await)
}
async fn execute_flushing_inner(&mut self, statement: &str) -> Result<u64> {
if self.closed.load(Ordering::SeqCst) {
return Err(Error::InvalidInput(
"WriteEngine has been closed".to_string(),
));
}
let mutations = self.parse_statement_to_mutations(statement)?;
let n = mutations.len() as u64;
for mutation in mutations {
self.write_into_memtable(mutation)?;
if self
.memtable
.should_flush(self.config.memtable_flush_threshold)
{
self.flush_internal_async().await?;
}
}
Ok(n)
}
#[tracing::instrument(name = "flush.public", level = "debug", skip(self))]
pub async fn flush(&mut self) -> Result<Option<SSTableInfo>> {
crate::observability::record_result(
"write",
async {
if self.closed.load(Ordering::SeqCst) {
return Err(Error::InvalidInput(
"WriteEngine has been closed".to_string(),
));
}
self.flush_internal_async().await
}
.await,
)
}
fn flush_internal(&mut self) -> Result<()> {
merge::block_on_async(self.flush_internal_async())?;
Ok(())
}
#[tracing::instrument(name = "flush.memtable", level = "debug", skip(self))]
async fn flush_internal_async(&mut self) -> Result<Option<SSTableInfo>> {
if self.memtable.is_empty() {
return Ok(None);
}
let flush_start = Instant::now();
let rows_to_flush = self.memtable.row_count() as u64;
tracing::info!(
"Flushing memtable: {} partitions, {} rows, {} bytes",
self.memtable.iter().count(),
self.memtable.row_count(),
self.memtable.size_bytes()
);
let partition_count_hint = self.memtable.iter().count();
let mut writer =
crate::storage::sstable::writer::SSTableWriter::with_expected_partitions_and_registry(
self.config.data_dir.clone(),
self.generation,
&self.config.schema,
partition_count_hint,
self.config.udt_registry.as_ref(),
)?;
let mut baseline_min_ts = i64::MAX;
let mut baseline_min_ldt = i32::MAX;
let mut baseline_min_ttl = i32::MAX;
for (_, mutations) in self.memtable.iter() {
let (ts, ldt, ttl) =
crate::storage::sstable::writer::SSTableWriter::compute_mutations_baseline_stats(
mutations,
);
baseline_min_ts = baseline_min_ts.min(ts);
baseline_min_ldt = baseline_min_ldt.min(ldt);
baseline_min_ttl = baseline_min_ttl.min(ttl);
}
writer.pre_seed_encoding_baselines(baseline_min_ts, baseline_min_ldt, baseline_min_ttl);
for (decorated_key, mutations) in self.memtable.iter() {
writer.write_partition(decorated_key.clone(), mutations.to_vec())?;
}
let info = writer.finish().await?;
tracing::info!(
"SSTable flush complete: generation {}, {} partitions, {} bytes",
self.generation,
info.partition_count,
info.data_size
);
let durability_outcome = durability::finalize_flush_durability(
&durability::RealDurabilityBarrier,
&info.data_path,
&self.config.data_dir,
&mut self.wal,
)?;
self.memtable.clear();
self.warned_over_threshold = false;
self.l0_count += 1;
self.total_flushed_bytes = self.total_flushed_bytes.saturating_add(info.data_size);
self.generation += 1;
{
use crate::observability::{self as obs, catalog};
obs::record_histogram(
catalog::FLUSH_DURATION,
flush_start.elapsed().as_secs_f64(),
&[],
);
obs::add_counter(catalog::FLUSH_ROWS, rows_to_flush, &[]);
obs::add_counter(catalog::FLUSH_BYTES, info.data_size, &[]);
obs::add_counter(catalog::FLUSH_SSTABLES, 1, &[]);
obs::record_gauge(catalog::COMPACTION_LAG, self.l0_count as i64, &[]);
}
self.record_memtable_gauges();
match durability_outcome {
durability::FlushDurabilityOutcome::Durable => Ok(Some(info)),
durability::FlushDurabilityOutcome::WalTruncateFailedAfterCommit(err) => Err(err),
}
}
pub async fn close(&mut self) -> Result<()> {
if self.closed.swap(true, Ordering::SeqCst) {
return Ok(());
}
tracing::info!("Closing WriteEngine");
if !self.memtable.is_empty() {
tracing::info!("Flushing memtable before close");
match self.flush_internal_async().await {
Ok(_) => {
tracing::info!("Memtable flushed successfully");
}
Err(e) => {
tracing::error!("Failed to flush memtable during close: {}", e);
crate::observability::record_error(&e, "write");
self.closed.store(false, Ordering::SeqCst);
return Err(e);
}
}
}
if let Err(e) = self.wal.sync() {
tracing::warn!("Failed to sync WAL during close: {}", e);
}
if let Err(e) = fs2::FileExt::unlock(&self.dir_lock) {
tracing::warn!("Failed to release write_dir advisory lock: {}", e);
}
tracing::info!("WriteEngine closed");
Ok(())
}
fn parse_cql_to_mutation(&self, statement: &str) -> Result<Mutation> {
cql_to_mutation::convert_cql_to_mutation(statement, &self.config.schema)
}
fn determine_next_generation(data_dir: &Path) -> Result<u64> {
let mut max_generation = 0u64;
if !data_dir.exists() {
return Ok(1);
}
Self::scan_generations(
data_dir,
&mut max_generation,
crate::storage::sstable::MAX_SSTABLE_SCAN_DEPTH,
)?;
Ok(max_generation + 1)
}
fn scan_generations(dir: &Path, max_generation: &mut u64, depth: usize) -> Result<()> {
for entry in std::fs::read_dir(dir)
.map_err(|e| Error::Storage(format!("Failed to read data directory: {}", e)))?
{
let entry = entry
.map_err(|e| Error::Storage(format!("Failed to read directory entry: {}", e)))?;
let filename = entry.file_name();
let filename_str = filename.to_string_lossy();
if filename_str.starts_with("nb-") && filename_str.contains("-big-") {
if let Some(gen_str) = filename_str
.strip_prefix("nb-")
.and_then(|s| s.split('-').next())
{
if let Ok(gen) = gen_str.parse::<u64>() {
*max_generation = (*max_generation).max(gen);
}
}
} else if depth > 0 {
let path = entry.path();
if path.is_dir() {
Self::scan_generations(&path, max_generation, depth - 1)?;
}
}
}
Ok(())
}
}
#[cfg(feature = "write-support")]
impl Drop for WriteEngine {
fn drop(&mut self) {
if !self.memtable.is_empty() {
match self.config.durability {
Durability::SyncEachWrite => tracing::warn!(
"WriteEngine dropped without close(): {} row(s) in the memtable were NOT \
flushed to an SSTable and remain only in the WAL (durability now relies on \
WAL replay at next startup). Call `close().await` for a graceful shutdown.",
self.memtable.row_count()
),
Durability::Disabled => tracing::warn!(
"WriteEngine dropped without close(): {} row(s) in the memtable were NOT \
flushed to an SSTable and are LOST — durability is Disabled so these rows \
were never written to the WAL and cannot be recovered (they existed in \
memory only). Call `close().await` for a graceful shutdown.",
self.memtable.row_count()
),
}
}
if let Err(e) = fs2::FileExt::unlock(&self.dir_lock) {
tracing::debug!(
"WriteEngine drop: advisory lock release returned: {} \
(may have been released by close() already)",
e
);
}
}
}
#[cfg(all(test, feature = "write-support"))]
mod tests {
use super::*;
use crate::storage::write_engine::test_support::{create_test_mutation, create_test_schema};
use tempfile::TempDir;
#[test]
fn test_write_engine_config() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
assert_eq!(
config.memtable_flush_threshold,
WriteEngineConfig::DEFAULT_FLUSH_THRESHOLD
);
assert_eq!(
config.memtable_hard_limit,
WriteEngineConfig::DEFAULT_HARD_LIMIT
);
let config = config.with_flush_threshold(128 * 1024 * 1024);
assert_eq!(config.memtable_flush_threshold, 128 * 1024 * 1024);
let config = config.with_hard_limit(512 * 1024 * 1024);
assert_eq!(config.memtable_hard_limit, 512 * 1024 * 1024);
}
#[test]
fn test_write_engine_new() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let engine = WriteEngine::new(config).unwrap();
assert_eq!(engine.generation(), 1);
assert_eq!(engine.memtable_size(), 0);
assert_eq!(engine.memtable_row_count(), 0);
assert!(!engine.closed.load(std::sync::atomic::Ordering::Relaxed));
}
#[test]
fn test_create_dir_all_durable_fsyncs_parent() {
let temp_dir = TempDir::new().unwrap();
let root = temp_dir.path().join("root");
std::fs::create_dir_all(&root).unwrap();
let data = root.join("data");
assert!(!data.exists());
WriteEngine::create_dir_all_durable(&data, "data").unwrap();
assert!(data.is_dir(), "new root directory must exist");
WriteEngine::create_dir_all_durable(&data, "data").unwrap();
assert!(data.is_dir());
}
#[test]
fn test_create_dir_all_durable_fsyncs_full_chain_on_existing_tree() {
let temp_dir = TempDir::new().unwrap();
let base = temp_dir.path().join("base");
let a = base.join("a");
let b = a.join("b");
let data = b.join("data");
std::fs::create_dir_all(&data).unwrap();
assert!(data.is_dir(), "precondition: full tree pre-exists");
WriteEngine::create_dir_all_durable(&data, "data").unwrap();
for ancestor in [b.as_path(), a.as_path(), base.as_path()] {
assert!(
ancestor.is_dir(),
"ancestor {:?} must be present and fsyncable up the full chain",
ancestor
);
}
WriteEngine::create_dir_all_durable(&data, "data").unwrap();
assert!(data.is_dir());
}
#[test]
fn test_create_dir_all_durable_walk_terminates_at_root() {
let temp_dir = TempDir::new().unwrap();
let data = temp_dir.path().join("a").join("b").join("data");
WriteEngine::create_dir_all_durable(&data, "data").unwrap();
assert!(data.is_dir());
assert!(
std::path::Path::new("/").parent().is_none(),
"filesystem root must have no parent so the fsync walk terminates"
);
}
#[test]
fn test_write_engine_write_single_mutation() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
let mutation = create_test_mutation(1, "Alice", 1000000);
engine.write(mutation).unwrap();
assert_eq!(engine.memtable_row_count(), 1);
assert!(engine.memtable_size() > 0);
assert!(engine.wal_size() > 0);
}
#[test]
fn test_write_engine_write_multiple_mutations() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
for i in 0..10 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
engine.write(mutation).unwrap();
}
assert_eq!(engine.memtable_row_count(), 10);
assert!(engine.memtable_size() > 0);
}
#[tokio::test]
async fn test_write_engine_flush_empty() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
let result = engine.flush().await.unwrap();
assert!(result.is_none());
}
#[tokio::test]
async fn test_write_engine_flush_with_data() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
for i in 0..5 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
engine.write(mutation).unwrap();
}
let initial_generation = engine.generation();
let info = engine.flush().await.unwrap();
assert!(info.is_some());
let info = info.unwrap();
assert_eq!(info.partition_count, 5);
assert!(info.data_size > 0);
assert!(info.data_path.exists());
assert_eq!(engine.memtable_row_count(), 0);
assert_eq!(engine.memtable_size(), 0);
assert_eq!(engine.wal_size(), 0);
assert_eq!(engine.generation(), initial_generation + 1);
}
#[test]
fn test_write_engine_automatic_flush() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_flush_threshold(1024);
let mut engine = WriteEngine::new(config).unwrap();
for i in 0..100 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
engine.write(mutation).unwrap();
}
assert!(engine.generation() > 1 || engine.memtable_size() < 10000);
}
#[tokio::test]
async fn test_write_engine_close_with_data() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
for i in 0..5 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
engine.write(mutation).unwrap();
}
engine.close().await.unwrap();
let data_dir = temp_dir.path().join("data");
let entries: Vec<_> = std::fs::read_dir(&data_dir).unwrap().collect();
assert!(!entries.is_empty(), "SSTable files should exist");
}
#[tokio::test]
async fn test_write_engine_close_empty() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
engine.close().await.unwrap();
}
#[test]
fn test_write_engine_write_after_close() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
tokio::runtime::Runtime::new()
.unwrap()
.block_on(engine.close())
.unwrap();
let schema2 = create_test_schema();
let config2 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema2,
);
let mut engine2 = WriteEngine::new(config2).unwrap();
let mutation = create_test_mutation(1, "Alice", 1000000);
engine2.write(mutation).unwrap();
assert_eq!(engine2.memtable_row_count(), 1);
}
#[test]
fn test_write_engine_wal_recovery() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema.clone(),
);
{
let mut engine = WriteEngine::new(config.clone()).unwrap();
for i in 0..5 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
engine.write(mutation).unwrap();
}
}
let config2 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let engine = WriteEngine::new(config2).unwrap();
assert_eq!(engine.memtable_row_count(), 5);
assert!(engine.memtable_size() > 0);
}
#[test]
fn test_write_engine_recovers_across_torn_tail_crash() {
use std::io::Write as _;
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
{
let mut engine = WriteEngine::new(config.clone()).unwrap();
engine
.write(create_test_mutation(1, "Alice", 1_000_000))
.unwrap();
}
let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
{
let mut f = std::fs::OpenOptions::new()
.append(true)
.open(&wal_file)
.unwrap();
f.write_all(&100u32.to_le_bytes()).unwrap();
f.write_all(&0u32.to_le_bytes()).unwrap();
f.write_all(&[0xAB; 10]).unwrap();
f.sync_all().unwrap();
}
{
let mut engine = WriteEngine::new(config.clone()).unwrap();
assert_eq!(
engine.memtable_row_count(),
1,
"must recover A across the torn tail"
);
engine
.write(create_test_mutation(3, "Carol", 3_000_000))
.unwrap();
}
{
let engine = WriteEngine::new(config).unwrap();
assert_eq!(
engine.memtable_row_count(),
2,
"both A and C must survive the second recovery (C must not be lost)"
);
}
}
fn seed_wal_abc(wal_dir: &Path) -> u64 {
std::fs::create_dir_all(wal_dir).unwrap();
let mut wal = WriteAheadLog::create(wal_dir).unwrap();
wal.append(&create_test_mutation(1, "A", 1_000_000))
.unwrap();
wal.sync().unwrap();
let end_a = wal.size();
wal.append(&create_test_mutation(2, "B", 2_000_000))
.unwrap();
wal.sync().unwrap();
wal.append(&create_test_mutation(3, "C", 3_000_000))
.unwrap();
wal.sync().unwrap();
drop(wal);
end_a
}
fn count_corrupt_aside(wal_dir: &Path) -> usize {
std::fs::read_dir(wal_dir)
.unwrap()
.filter_map(|e| e.ok())
.filter(|e| {
e.file_name()
.to_string_lossy()
.contains("commitlog.wal.corrupt.")
})
.count()
}
#[tokio::test]
async fn test_write_engine_lossy_wal_preserved_before_flush_truncate() {
use std::io::{Read, Seek, SeekFrom, Write};
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let end_a = seed_wal_abc(&config.wal_dir);
let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
{
let mut f = std::fs::OpenOptions::new()
.read(true)
.write(true)
.open(&wal_file)
.unwrap();
f.seek(SeekFrom::Start(end_a + 8)).unwrap();
let mut byte = [0u8; 1];
f.read_exact(&mut byte).unwrap();
f.seek(SeekFrom::Start(end_a + 8)).unwrap();
f.write_all(&[byte[0] ^ 0x01]).unwrap();
f.sync_all().unwrap();
}
let mut engine = WriteEngine::new(config.clone()).unwrap();
assert!(
!engine.wal_recovery().is_clean(),
"engine must expose a non-clean RecoveryReport for a lossy WAL"
);
assert_eq!(engine.wal_recovery().corrupt_entries, 1);
assert!(engine.wal_recovery().stopped_early);
assert_eq!(
engine.memtable_row_count(),
1,
"only the valid prefix [A] is recovered"
);
assert_eq!(
count_corrupt_aside(&config.wal_dir),
1,
"the raw corrupt WAL segment must be preserved aside BEFORE any flush"
);
engine.flush().await.unwrap();
assert_eq!(
count_corrupt_aside(&config.wal_dir),
1,
"flush must NOT destroy the preserved corrupt WAL segment"
);
}
#[test]
fn test_write_engine_apply_error_still_preserves_and_resets_corrupt_tail() {
use crate::storage::write_engine::mutation::{CellOperation, PartitionKey, TableId};
use crate::types::Value;
use std::io::{Read, Seek, SeekFrom, Write};
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema(); let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
std::fs::create_dir_all(&config.wal_dir).unwrap();
let end_a = {
let mut wal = WriteAheadLog::create(&config.wal_dir).unwrap();
let bad = Mutation::new(
TableId::new("test_ks", "test_table"),
PartitionKey::single("id", Value::text("not-an-int".to_string())),
None,
vec![CellOperation::Write {
column: "name".to_string(),
value: Value::text("A".to_string()),
}],
1_000_000,
None,
);
wal.append(&bad).unwrap();
wal.sync().unwrap();
let end_a = wal.size();
wal.append(&create_test_mutation(2, "B", 2_000_000))
.unwrap();
wal.sync().unwrap();
end_a
};
let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
{
let mut f = std::fs::OpenOptions::new()
.read(true)
.write(true)
.open(&wal_file)
.unwrap();
f.seek(SeekFrom::Start(end_a + 8)).unwrap();
let mut byte = [0u8; 1];
f.read_exact(&mut byte).unwrap();
f.seek(SeekFrom::Start(end_a + 8)).unwrap();
f.write_all(&[byte[0] ^ 0x01]).unwrap();
f.sync_all().unwrap();
}
let result = WriteEngine::new(config.clone());
assert!(
result.is_err(),
"engine open must surface the deferred memtable-application error"
);
assert_eq!(
count_corrupt_aside(&config.wal_dir),
1,
"corrupt segment must be preserved aside even though an earlier apply failed"
);
assert_eq!(
std::fs::metadata(&wal_file).unwrap().len(),
end_a,
"live WAL must be reset to its valid prefix [A] before the apply error surfaces"
);
}
#[tokio::test]
async fn test_write_engine_clean_wal_recovery_truncates_normally() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
seed_wal_abc(&config.wal_dir);
let mut engine = WriteEngine::new(config.clone()).unwrap();
assert!(engine.wal_recovery().is_clean());
assert_eq!(engine.memtable_row_count(), 3, "A, B, C all recovered");
assert_eq!(
count_corrupt_aside(&config.wal_dir),
0,
"a clean recovery must not preserve any aside segment"
);
engine.flush().await.unwrap();
let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
assert_eq!(
std::fs::metadata(&wal_file).unwrap().len(),
0,
"clean recovery: the WAL is truncated normally after flush"
);
assert_eq!(count_corrupt_aside(&config.wal_dir), 0);
}
#[tokio::test]
async fn test_write_engine_reset_to_valid_prefix_keeps_post_recovery_write() {
use crate::types::Value;
use std::io::{Read, Seek, SeekFrom, Write};
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let end_a = seed_wal_abc(&config.wal_dir);
let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
{
let mut f = std::fs::OpenOptions::new()
.read(true)
.write(true)
.open(&wal_file)
.unwrap();
f.seek(SeekFrom::Start(end_a + 8)).unwrap();
let mut byte = [0u8; 1];
f.read_exact(&mut byte).unwrap();
f.seek(SeekFrom::Start(end_a + 8)).unwrap();
f.write_all(&[byte[0] ^ 0x01]).unwrap();
f.sync_all().unwrap();
}
{
let mut engine = WriteEngine::new(config.clone()).unwrap();
assert!(!engine.wal_recovery().is_clean());
assert_eq!(engine.memtable_row_count(), 1, "valid prefix [A] recovered");
assert_eq!(
count_corrupt_aside(&config.wal_dir),
1,
"corrupt segment must be preserved aside"
);
engine
.write(create_test_mutation(4, "D", 4_000_000))
.unwrap();
drop(engine);
}
let wal = WriteAheadLog::open_existing(&wal_file).unwrap();
let report = wal.replay().unwrap();
assert!(
report.is_clean(),
"the reset live WAL is a clean [A, D] prefix"
);
let names: Vec<&str> = report
.mutations
.iter()
.map(|m| match &m.operations[0] {
CellOperation::Write {
value: Value::Text(name),
..
} => std::str::from_utf8(name).unwrap_or_default(),
other => panic!("expected Write op, got {other:?}"),
})
.collect();
assert_eq!(
names,
vec!["A", "D"],
"replay must yield the valid prefix AND the post-recovery write D, \
never a set missing D (and never the corrupt B/C)"
);
assert!(
!names.contains(&"B"),
"the corrupt entry B must not resurface"
);
assert!(
!names.contains(&"C"),
"C (behind corruption) must not resurface"
);
assert_eq!(
count_corrupt_aside(&config.wal_dir),
1,
"the corrupt WAL evidence must remain preserved aside after the reset"
);
}
#[test]
fn test_write_engine_generation_tracking() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema.clone(),
);
{
let mut engine = WriteEngine::new(config.clone()).unwrap();
assert_eq!(engine.generation(), 1);
let mutation = create_test_mutation(1, "Alice", 1000000);
engine.write(mutation).unwrap();
tokio::runtime::Runtime::new()
.unwrap()
.block_on(engine.flush())
.unwrap();
assert_eq!(engine.generation(), 2);
}
let config2 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let engine = WriteEngine::new(config2).unwrap();
assert_eq!(engine.generation(), 2);
}
#[test]
fn test_write_engine_execute_table_mismatch() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
let result = engine.execute("INSERT INTO users (id, name) VALUES (1, 'Alice')");
let err_msg = result.unwrap_err().to_string();
assert!(
err_msg.contains("targets table 'users'")
&& err_msg.contains("schema is for 'test_table'"),
"Expected table mismatch error, got: {}",
err_msg
);
}
#[test]
fn test_write_engine_execute_insert_success() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
assert_eq!(engine.memtable_row_count(), 0);
let result = engine.execute("INSERT INTO test_table (id, name) VALUES (1, 'Alice')");
assert!(
result.is_ok(),
"execute() failed: {:?}",
result.unwrap_err()
);
assert_eq!(engine.memtable_row_count(), 1);
}
#[test]
fn test_determine_next_generation_empty_dir() {
let temp_dir = TempDir::new().unwrap();
let data_dir = temp_dir.path().join("data");
std::fs::create_dir_all(&data_dir).unwrap();
let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
assert_eq!(generation, 1);
}
#[test]
fn test_determine_next_generation_with_sstables() {
let temp_dir = TempDir::new().unwrap();
let data_dir = temp_dir.path().join("data");
std::fs::create_dir_all(&data_dir).unwrap();
std::fs::write(data_dir.join("nb-1-big-Data.db"), b"").unwrap();
std::fs::write(data_dir.join("nb-2-big-Data.db"), b"").unwrap();
std::fs::write(data_dir.join("nb-5-big-Data.db"), b"").unwrap();
let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
assert_eq!(generation, 6);
}
#[tokio::test]
async fn test_write_engine_close_idempotent() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
engine.close().await.unwrap();
assert!(engine.closed.load(Ordering::SeqCst));
engine.close().await.unwrap();
assert!(engine.closed.load(Ordering::SeqCst));
}
#[tokio::test]
async fn test_write_engine_close_syncs_wal() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
let mutation = create_test_mutation(1, "Alice", 1000000);
engine.write(mutation).unwrap();
engine.close().await.unwrap();
assert_eq!(engine.wal_size(), 0);
}
#[test]
fn test_write_engine_closed_flag_atomic() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let engine = WriteEngine::new(config).unwrap();
assert!(!engine.closed.load(Ordering::SeqCst));
engine.closed.store(true, Ordering::SeqCst);
assert!(engine.closed.load(Ordering::SeqCst));
let prev = engine.closed.swap(false, Ordering::SeqCst);
assert!(prev);
assert!(!engine.closed.load(Ordering::SeqCst));
}
#[tokio::test]
async fn test_write_engine_write_after_close_fails() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
engine.close().await.unwrap();
let mutation = create_test_mutation(1, "Alice", 1000000);
let result = engine.write(mutation);
assert!(result.is_err());
match result {
Err(Error::InvalidInput(_)) => {}
_ => panic!("Expected InvalidInput error"),
}
}
#[tokio::test]
async fn test_write_engine_flush_after_close_fails() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
engine.close().await.unwrap();
let result = engine.flush().await;
assert!(result.is_err());
match result {
Err(Error::InvalidInput(_)) => {}
_ => panic!("Expected InvalidInput error"),
}
}
#[test]
fn test_write_engine_hard_limit_enforcement() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_flush_threshold(10 * 1024) .with_hard_limit(2048);
let mut engine = WriteEngine::new(config).unwrap();
let mut write_count = 0;
for i in 0..1000 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
let result = engine.write(mutation);
match result {
Ok(()) => {
write_count += 1;
}
Err(Error::Storage(msg)) => {
assert!(msg.contains("hard limit"));
break;
}
Err(e) => panic!("Expected Storage error, got: {:?}", e),
}
}
assert!(
write_count < 1000,
"Should have hit hard limit before 1000 writes"
);
assert!(
write_count > 0,
"Should have accepted at least some writes before hitting limit"
);
}
#[tokio::test]
async fn test_write_engine_hard_limit_enforcement_async() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_flush_threshold(10 * 1024) .with_hard_limit(2048);
let mut engine = WriteEngine::new(config).unwrap();
let mut write_count = 0;
for i in 0..1000 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
let result = engine.write_async(mutation).await;
match result {
Ok(()) => {
write_count += 1;
}
Err(Error::Storage(msg)) => {
assert!(msg.contains("hard limit"));
break;
}
Err(e) => panic!("Expected Storage error, got: {:?}", e),
}
}
assert!(
write_count < 1000,
"Should have hit hard limit before 1000 writes"
);
assert!(
write_count > 0,
"Should have accepted at least some writes before hitting limit"
);
}
#[tokio::test]
async fn test_write_engine_hard_limit_recovery_after_flush() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_flush_threshold(1024)
.with_hard_limit(2048);
let mut engine = WriteEngine::new(config).unwrap();
let mut first_batch_count = 0;
for i in 0..1000 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
let result = engine.write(mutation);
if result.is_err() {
break;
}
first_batch_count += 1;
}
assert!(
first_batch_count > 0,
"Should have accepted some writes before limit"
);
engine.flush().await.unwrap();
let mutation = create_test_mutation(9999, "After flush", 2000000);
let result = engine.write(mutation);
assert!(result.is_ok(), "Should accept writes after flush");
assert_eq!(engine.memtable_row_count(), 1);
}
#[tokio::test]
async fn test_execute_flushing_auto_flushes_in_runtime() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_flush_threshold(4096)
.with_durability(Durability::Disabled);
let mut engine = WriteEngine::new(config).unwrap();
assert_eq!(engine.generation(), 1);
for i in 0..2000 {
let stmt = format!(
"INSERT INTO test_ks.test_table (id, name) VALUES ({}, 'User{}')",
i, i
);
let n = engine
.execute_flushing(&stmt)
.await
.expect("execute_flushing must not dead-end at the hard limit");
assert_eq!(n, 1, "single INSERT applies exactly one mutation");
}
assert!(
engine.generation() > 1,
"expected auto-flush to advance generation, got {} (would stay 1 on the old sync path in a runtime)",
engine.generation()
);
}
#[tokio::test]
async fn test_over_threshold_warn_fires_at_most_once_per_crossing() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_flush_threshold(256)
.with_hard_limit(64 * 1024)
.with_durability(Durability::Disabled);
let mut engine = WriteEngine::new(config).unwrap();
assert!(
!engine.warned_over_threshold,
"guard starts false before any crossing"
);
let mut crossed = false;
for i in 0..500 {
engine
.write(create_test_mutation(
i,
&format!("User{}", i),
1_000_000 + i as i64,
))
.unwrap();
if engine.warned_over_threshold {
crossed = true;
break;
}
}
assert!(crossed, "memtable should have crossed the flush threshold");
assert!(
engine.warned_over_threshold,
"guard set after first crossing"
);
for i in 500..600 {
engine
.write(create_test_mutation(
i,
&format!("User{}", i),
1_000_000 + i as i64,
))
.unwrap();
assert!(
engine.warned_over_threshold,
"guard must remain set across subsequent over-threshold writes"
);
}
engine.flush().await.unwrap();
assert!(
!engine.warned_over_threshold,
"guard resets to false after a successful flush"
);
}
#[test]
fn test_generation_counter_is_u64() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let engine = WriteEngine::new(config).unwrap();
let generation: u64 = engine.generation();
assert_eq!(generation, 1u64);
let _type_check: u64 = generation;
let large_generation: u64 = u32::MAX as u64 + 1000;
assert!(large_generation > u32::MAX as u64);
assert_eq!(large_generation, 4_294_968_295u64);
}
#[test]
fn test_determine_next_generation_large_numbers() {
let temp_dir = TempDir::new().unwrap();
let data_dir = temp_dir.path().join("data");
std::fs::create_dir_all(&data_dir).unwrap();
let large_gen: u64 = u32::MAX as u64 + 100;
std::fs::write(data_dir.join(format!("nb-{}-big-Data.db", large_gen)), b"").unwrap();
let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
assert_eq!(generation, large_gen + 1);
assert!(generation > u32::MAX as u64);
}
#[test]
fn test_durability_default_is_sync_each_write() {
assert_eq!(Durability::default(), Durability::SyncEachWrite);
}
#[test]
fn test_config_default_durability() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
assert_eq!(config.durability, Durability::SyncEachWrite);
}
#[test]
fn test_config_with_durability_builder() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_durability(Durability::Disabled);
assert_eq!(config.durability, Durability::Disabled);
}
#[test]
fn test_wal_on_produces_wal_growth() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_durability(Durability::SyncEachWrite);
let mut engine = WriteEngine::new(config).unwrap();
assert_eq!(engine.wal_size(), 0, "WAL must start empty");
let mutation = create_test_mutation(1, "Alice", 1_000_000);
engine.write(mutation).unwrap();
assert!(
engine.wal_size() > 0,
"WAL must grow after write with SyncEachWrite"
);
}
#[test]
fn test_wal_off_produces_no_wal_growth() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_durability(Durability::Disabled);
let mut engine = WriteEngine::new(config).unwrap();
assert_eq!(engine.wal_size(), 0, "WAL must start empty");
for i in 0..10 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
engine.write(mutation).unwrap();
}
assert_eq!(
engine.wal_size(),
0,
"WAL must remain empty with Durability::Disabled"
);
assert_eq!(
engine.memtable_row_count(),
10,
"Mutations must reach the memtable even without WAL"
);
}
#[tokio::test]
async fn test_wal_off_write_async_produces_no_wal_growth() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_durability(Durability::Disabled);
let mut engine = WriteEngine::new(config).unwrap();
for i in 0..5 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
engine.write_async(mutation).await.unwrap();
}
assert_eq!(
engine.wal_size(),
0,
"WAL must remain empty with Durability::Disabled (async path)"
);
assert_eq!(engine.memtable_row_count(), 5);
}
#[test]
fn test_wal_off_no_replay_on_restart() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
{
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema.clone(),
)
.with_durability(Durability::Disabled);
let mut engine = WriteEngine::new(config).unwrap();
for i in 0..5 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
engine.write(mutation).unwrap();
}
}
let config2 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let engine2 = WriteEngine::new(config2).unwrap();
assert_eq!(
engine2.memtable_row_count(),
0,
"No WAL entries were written with Disabled, so no replay is possible"
);
}
#[test]
fn test_wal_on_replays_on_restart() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
{
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema.clone(),
)
.with_durability(Durability::SyncEachWrite);
let mut engine = WriteEngine::new(config).unwrap();
for i in 0..5 {
let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
engine.write(mutation).unwrap();
}
}
let config2 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_durability(Durability::SyncEachWrite);
let engine2 = WriteEngine::new(config2).unwrap();
assert_eq!(
engine2.memtable_row_count(),
5,
"SyncEachWrite must replay mutations durably on restart"
);
}
#[test]
fn test_write_dir_lock_second_engine_fails_fast() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config1 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema.clone(),
);
let _engine1 = WriteEngine::new(config1).unwrap();
let config2 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let result = WriteEngine::new(config2);
assert!(
result.is_err(),
"A second WriteEngine on the same write_dir must fail"
);
let err = result.unwrap_err();
assert!(
matches!(err, Error::WriteDirLocked { .. }),
"Expected WriteDirLocked error, got: {:?}",
err
);
let msg = err.to_string();
assert!(
msg.contains("already locked"),
"Error message must mention the lock: {}",
msg
);
assert!(
msg.contains("Only one Database instance"),
"Error message must explain the constraint: {}",
msg
);
}
#[tokio::test]
async fn test_write_dir_lock_reacquired_after_close() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config1 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema.clone(),
);
let mut engine1 = WriteEngine::new(config1).unwrap();
engine1.close().await.unwrap();
let config2 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let engine2 = WriteEngine::new(config2);
assert!(
engine2.is_ok(),
"WriteEngine must acquire lock after the previous engine closed: {:?}",
engine2.err()
);
}
#[test]
fn test_write_dir_lock_reacquired_after_drop() {
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config1 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema.clone(),
);
{
let _engine1 = WriteEngine::new(config1).unwrap();
}
let config2 = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let engine2 = WriteEngine::new(config2);
assert!(
engine2.is_ok(),
"WriteEngine must acquire lock after the previous engine was dropped: {:?}",
engine2.err()
);
}
#[tokio::test]
async fn post_mutation_truncate_failure_commits_and_retry_writes_new_generation() {
use crate::platform::Platform;
use crate::storage::sstable::SSTableManager;
use crate::Config;
use std::sync::Arc;
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let data_dir = temp_dir.path().join("data");
let config = WriteEngineConfig::new(
data_dir.clone(),
temp_dir.path().join("wal"),
schema.clone(),
);
let mut engine = WriteEngine::new(config).unwrap();
engine
.write(create_test_mutation(1, "Alice", 1_000_000))
.unwrap();
let gen_before = engine.generation();
engine.wal.set_fail_sync_after_truncate(true);
let err = engine
.flush()
.await
.expect_err("a post-mutation WAL-truncate failure must surface as an error");
assert!(
matches!(err, Error::Storage(_)),
"post-mutation truncate failure must surface as a storage error, got {err:?}"
);
assert_eq!(
engine.memtable_row_count(),
0,
"memtable must be cleared once the SSTable is durably published"
);
assert_eq!(
engine.generation(),
gen_before + 1,
"generation must advance so a retry writes a NEW generation"
);
let sstable_dir = data_dir.join("test_ks").join("test_table");
let gen1_data = sstable_dir.join("nb-1-big-Data.db");
assert!(
gen1_data.exists(),
"the published gen-1 Data.db must exist on disk after the faulted flush"
);
let gen1_bytes = std::fs::read(&gen1_data).unwrap();
engine.wal.set_fail_sync_after_truncate(false);
engine
.write(create_test_mutation(2, "Bob", 2_000_000))
.unwrap();
engine
.flush()
.await
.expect("second flush must succeed")
.expect("second flush must produce an SSTable");
let gen2_data = sstable_dir.join("nb-2-big-Data.db");
assert!(
gen2_data.exists(),
"the retry must write a NEW generation (nb-2), not overwrite nb-1"
);
assert_eq!(
std::fs::read(&gen1_data).unwrap(),
gen1_bytes,
"the retry must NOT overwrite the published gen-1 SSTable"
);
let cqlite_config = Config::default();
let platform = Arc::new(Platform::new(&cqlite_config).await.unwrap());
let manager = SSTableManager::new(
&data_dir,
&cqlite_config,
platform,
#[cfg(feature = "state_machine")]
None,
)
.await
.expect("SSTableManager must load both published generations");
let table_id = crate::types::TableId::from("test_ks.test_table");
let rows = manager
.scan(&table_id, None, None, None, Some(&schema))
.await
.expect("scan across both generations must succeed");
assert_eq!(
rows.len(),
2,
"both mutations must be readable exactly once (no loss, no duplication)"
);
}
mod drop_warn_capture {
use std::cell::RefCell;
use std::fmt;
use tracing::field::{Field, Visit};
use tracing::subscriber::{set_default, DefaultGuard};
use tracing::{Event, Level, Subscriber};
use tracing_subscriber::layer::{Context, Layer, SubscriberExt};
use tracing_subscriber::Registry;
thread_local! {
static BUFFER: RefCell<Option<Vec<String>>> = const { RefCell::new(None) };
static GUARD: RefCell<Option<DefaultGuard>> = const { RefCell::new(None) };
}
#[derive(Default)]
struct MessageVisitor {
message: String,
}
impl Visit for MessageVisitor {
fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
if field.name() == "message" {
self.message = format!("{value:?}");
}
}
}
struct WarnCaptureLayer;
impl<S: Subscriber> Layer<S> for WarnCaptureLayer {
fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
if *event.metadata().level() != Level::WARN {
return;
}
let mut visitor = MessageVisitor::default();
event.record(&mut visitor);
BUFFER.with(|b| {
if let Some(buf) = b.borrow_mut().as_mut() {
buf.push(visitor.message);
}
});
}
}
pub(super) fn start() {
BUFFER.with(|b| *b.borrow_mut() = Some(Vec::new()));
let guard = set_default(Registry::default().with(WarnCaptureLayer));
GUARD.with(|g| *g.borrow_mut() = Some(guard));
}
pub(super) fn take_warnings() -> Vec<String> {
GUARD.with(|g| *g.borrow_mut() = None);
BUFFER.with(|b| b.borrow_mut().take().unwrap_or_default())
}
}
#[test]
fn test_drop_without_close_warns_on_nonempty_memtable() {
drop_warn_capture::start();
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
engine
.write(create_test_mutation(1, "Alice", 1_000_000))
.unwrap();
assert!(
engine.memtable_row_count() > 0,
"precondition: memtable must be non-empty before drop"
);
drop(engine);
let warnings = drop_warn_capture::take_warnings();
assert!(
warnings
.iter()
.any(|m| m.contains("dropped") && m.contains("without close")),
"expected an unflushed-drop warning, captured warnings: {warnings:?}"
);
assert!(
warnings.iter().any(|m| m.contains("WAL replay")),
"WAL-durable drop warning must mention WAL replay recovery, captured: {warnings:?}"
);
assert!(
!warnings.iter().any(|m| m.contains("LOST")),
"WAL-durable drop warning must NOT claim data loss, captured: {warnings:?}"
);
}
#[test]
fn test_drop_without_close_warns_data_loss_when_durability_disabled() {
drop_warn_capture::start();
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
)
.with_durability(Durability::Disabled);
let mut engine = WriteEngine::new(config).unwrap();
engine
.write(create_test_mutation(1, "Alice", 1_000_000))
.unwrap();
assert!(
engine.memtable_row_count() > 0,
"precondition: memtable must be non-empty before drop"
);
drop(engine);
let warnings = drop_warn_capture::take_warnings();
assert!(
warnings
.iter()
.any(|m| m.contains("dropped") && m.contains("without close")),
"expected an unflushed-drop warning, captured warnings: {warnings:?}"
);
assert!(
warnings.iter().any(|m| m.contains("LOST")),
"Durability::Disabled drop warning must signal data loss, captured: {warnings:?}"
);
assert!(
!warnings.iter().any(|m| m.contains("WAL replay")),
"Durability::Disabled drop warning must NOT promise WAL recovery, captured: {warnings:?}"
);
}
#[tokio::test]
async fn test_drop_after_close_does_not_warn() {
drop_warn_capture::start();
let temp_dir = TempDir::new().unwrap();
let schema = create_test_schema();
let config = WriteEngineConfig::new(
temp_dir.path().join("data"),
temp_dir.path().join("wal"),
schema,
);
let mut engine = WriteEngine::new(config).unwrap();
engine
.write(create_test_mutation(1, "Alice", 1_000_000))
.unwrap();
engine.close().await.unwrap();
drop(engine);
let warnings = drop_warn_capture::take_warnings();
assert!(
!warnings
.iter()
.any(|m| m.contains("dropped") && m.contains("without close")),
"close() flushed the memtable; Drop must not warn, captured: {warnings:?}"
);
}
}