cqlite_core/storage/write_engine/mod.rs
1//! Write engine for SSTable generation and persistence (M5)
2//!
3//! This module provides the write path for CQLite, implementing WAL-backed
4//! memtable flushing and K-way merge for producing valid Cassandra 5.0 SSTables.
5//!
6//! ## Architecture
7//!
8//! The WriteEngine is the public API that coordinates:
9//! 1. WAL (Write-Ahead Log) - Durability
10//! 2. Memtable - In-memory buffer
11//! 3. SSTableWriter - On-disk persistence
12//!
13//! ## Write Flow
14//!
15//! 1. User calls `write(mutation)` or `execute(cql_statement)`
16//! 2. WriteEngine appends to WAL (durability)
17//! 3. WriteEngine inserts into Memtable
18//! 4. When Memtable exceeds threshold → flush to SSTable
19//! 5. After successful flush → truncate WAL
20//!
21//! ## Recovery
22//!
23//! On startup, the WriteEngine replays WAL entries into the memtable.
24
25#[cfg(feature = "write-support")]
26pub mod cql_to_mutation;
27#[cfg(feature = "write-support")]
28pub mod export;
29#[cfg(feature = "write-support")]
30pub mod memtable;
31#[cfg(feature = "write-support")]
32pub mod merge;
33#[cfg(feature = "write-support")]
34pub mod merge_policy;
35#[cfg(feature = "write-support")]
36pub mod mutation;
37#[cfg(feature = "write-support")]
38pub(crate) mod reconcile_rules;
39#[cfg(feature = "write-support")]
40pub mod wal;
41
42#[cfg(feature = "write-support")]
43pub use export::{ExportOptions, ExportReport};
44#[cfg(feature = "write-support")]
45pub use memtable::Memtable;
46#[cfg(feature = "write-support")]
47// Issue #2346: the reader-based analogue of `build_single_partition_merger`,
48// re-exported alongside it so both surfaces are reachable the same way.
49pub use merge::build_single_partition_merger_from_readers;
50#[cfg(feature = "write-support")]
51pub use merge::KWayMerger;
52#[cfg(feature = "write-support")]
53pub use merge::{build_single_partition_merger, build_single_partition_merger_with_registry};
54#[cfg(feature = "write-support")]
55pub use merge_policy::STCSPolicy;
56
57/// Read-time TTL "now" clock (epoch seconds), the authoritative reconciliation
58/// instant the core read path uses (issue #1741/#1853). Re-exported (issue
59/// #2789) so an out-of-crate merge consumer — the Flight producer — threads the
60/// SAME `now` into its k-way merger's TTL expiry as a `SELECT`, honoring the
61/// debug-only `CQLITE_TTL_NOW_OVERRIDE_SECS` pin in tests and wall-clock in
62/// production, rather than re-deriving it (no divergence, no new env read).
63#[doc(hidden)]
64pub use crate::storage::sstable::reader::parsing::row_decoder::now_clock::now_epoch_secs as read_time_now_secs;
65#[cfg(feature = "write-support")]
66pub use mutation::{
67 CellOperation, ClusteringBound, ClusteringKey, DecoratedKey, Mutation, PartitionKey,
68 PartitionTombstone, RangeTombstone, TableId,
69};
70#[cfg(feature = "write-support")]
71pub use wal::{RecoveryReport, WriteAheadLog};
72
73#[cfg(feature = "write-support")]
74mod compaction;
75#[cfg(feature = "write-support")]
76pub(crate) mod durability;
77#[cfg(feature = "write-support")]
78mod maintenance;
79#[cfg(feature = "write-support")]
80mod stats;
81#[cfg(feature = "write-support")]
82mod sweep;
83#[cfg(all(test, feature = "write-support"))]
84mod test_support;
85// Issue #1625: honest memtable hard-limit admission tests live in a sibling
86// module to avoid growing the already-oversized `mod.rs` (epic #1116/#1135).
87#[cfg(all(test, feature = "write-support"))]
88mod admission_tests;
89
90#[cfg(feature = "write-support")]
91pub use maintenance::MaintenanceReport;
92#[cfg(feature = "write-support")]
93pub use stats::CompactionStats;
94
95use crate::error::{Error, Result};
96use crate::schema::{TableSchema, UdtRegistry};
97use crate::storage::sstable::writer::SSTableInfo;
98#[cfg(feature = "write-support")]
99use maintenance::ActiveMerge;
100use std::path::{Path, PathBuf};
101use std::sync::atomic::{AtomicBool, Ordering};
102use std::time::Instant;
103
104/// Trait for merge policy implementations (M5.2, Issue #383)
105///
106/// A merge policy decides which SSTables should be compacted together.
107/// This trait allows different compaction strategies (STCS, LCS, TWCS, etc.)
108/// to be plugged into the WriteEngine.
109#[cfg(feature = "write-support")]
110pub trait MergePolicy: Send + std::fmt::Debug {
111 /// Select SSTables for the next compaction
112 ///
113 /// # Arguments
114 ///
115 /// * `candidates` - Available SSTable paths in the data directory
116 ///
117 /// # Returns
118 ///
119 /// Paths to SSTables that should be merged, ordered newest to oldest.
120 /// Returns empty Vec if no compaction is needed.
121 fn select_merge(&self, candidates: &[PathBuf]) -> Result<Vec<PathBuf>>;
122}
123
124/// WAL durability mode for the write engine.
125///
126/// Controls whether `write` and `write_async` append to and fsync the
127/// write-ahead log on every call. The default (`SyncEachWrite`) matches the
128/// pre-existing behavior and is the **only safe choice for production
129/// workloads** — a process crash between a successful `write` and a later
130/// `flush` will lose mutations written with `Disabled`.
131///
132/// ## When to use `Disabled`
133///
134/// - **Bulk-load / import pipelines** where the source data is replayable and
135/// you are willing to re-run the load on failure.
136/// - **Benchmarking** where you want to isolate CPU-bound write throughput from
137/// fsync latency. The companion `write/ingest_wal_off` Criterion bench uses
138/// this variant (see `cqlite-core/benches/write.rs`).
139///
140/// In both cases, call [`WriteEngine::flush`] (and, optionally,
141/// [`WriteEngine::close`]) when the load is finished so the data is durably
142/// persisted to SSTables.
143///
144/// ## WAL replay on restart
145///
146/// When `Disabled`, no WAL entries are written. Reopening the engine on the
147/// same `wal_dir` after a crash will replay **zero** mutations, even if
148/// `flush` was never called. If you need crash-safe recovery, use
149/// `SyncEachWrite`.
150///
151/// # Example
152///
153/// ```rust,ignore
154/// use cqlite_core::storage::write_engine::{Durability, WriteEngineConfig};
155///
156/// // Production (default)
157/// let config = WriteEngineConfig::new(data, wal, schema);
158///
159/// // Bulk-load / benchmarking
160/// let config = WriteEngineConfig::new(data, wal, schema)
161/// .with_durability(Durability::Disabled);
162/// ```
163#[cfg(feature = "write-support")]
164#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
165pub enum Durability {
166 /// Append to the WAL and call `fsync` on every `write` / `write_async`
167 /// call. A successful return guarantees the mutation is durable on disk.
168 ///
169 /// This is the **default** and the safe choice for all production
170 /// workloads.
171 #[default]
172 SyncEachWrite,
173
174 /// Skip WAL append **and** fsync on every `write` / `write_async` call.
175 /// Mutations are buffered in the memtable only. Data is durable only
176 /// after a successful [`WriteEngine::flush`].
177 ///
178 /// **Use only for bulk-load pipelines and benchmarks where durability can
179 /// be traded for throughput.**
180 Disabled,
181}
182
183/// Write engine configuration
184#[cfg(feature = "write-support")]
185#[derive(Debug, Clone)]
186pub struct WriteEngineConfig {
187 /// Directory for SSTable data files
188 pub data_dir: PathBuf,
189 /// Directory for WAL files
190 pub wal_dir: PathBuf,
191 /// Memtable flush threshold in bytes (default: 64MB)
192 pub memtable_flush_threshold: usize,
193 /// Memtable hard limit in bytes (default: 256MB)
194 /// When this limit is reached, writes will fail with an error
195 pub memtable_hard_limit: usize,
196 /// Table schema for column metadata
197 pub schema: TableSchema,
198 /// WAL durability mode (default: [`Durability::SyncEachWrite`])
199 pub durability: Durability,
200 /// Optional UDT registry for resolving bare CQL UDT column types to their
201 /// `UserType(...)` marshal form at flush time (issue #929). When `None`
202 /// (the default), a column whose `data_type` is a bare UDT name is written
203 /// as a single simple cell (documented fallback).
204 pub udt_registry: Option<UdtRegistry>,
205 /// Whether the engine installs a default STCS compaction policy so that
206 /// [`WriteEngine::maintenance_step`] performs size-tiered compaction
207 /// (issue #1619). Defaults to `true` (compaction on). Set to `false` to
208 /// disable compaction entirely — `maintenance_step` then becomes a no-op.
209 pub auto_compaction: bool,
210 /// Minimum number of SSTables in a size bucket required to trigger a
211 /// compaction (STCS `min_threshold`). Defaults to `4`. Ignored when
212 /// [`WriteEngineConfig::auto_compaction`] is `false`.
213 pub compaction_min_threshold: usize,
214 /// Maximum number of SSTables compacted together in one step (STCS
215 /// `max_threshold`). Defaults to `32`. Ignored when
216 /// [`WriteEngineConfig::auto_compaction`] is `false`.
217 pub compaction_max_threshold: usize,
218}
219
220#[cfg(feature = "write-support")]
221impl WriteEngineConfig {
222 /// Default flush threshold (64 MB)
223 pub const DEFAULT_FLUSH_THRESHOLD: usize = 64 * 1024 * 1024;
224 /// Default hard limit (256 MB)
225 pub const DEFAULT_HARD_LIMIT: usize = 256 * 1024 * 1024;
226 /// Default STCS `min_threshold` (issue #1619)
227 pub const DEFAULT_COMPACTION_MIN_THRESHOLD: usize = 4;
228 /// Default STCS `max_threshold` (issue #1619)
229 pub const DEFAULT_COMPACTION_MAX_THRESHOLD: usize = 32;
230
231 /// Create a new configuration with default flush threshold
232 pub fn new(data_dir: PathBuf, wal_dir: PathBuf, schema: TableSchema) -> Self {
233 Self {
234 data_dir,
235 wal_dir,
236 memtable_flush_threshold: Self::DEFAULT_FLUSH_THRESHOLD,
237 memtable_hard_limit: Self::DEFAULT_HARD_LIMIT,
238 schema,
239 durability: Durability::default(),
240 udt_registry: None,
241 auto_compaction: true,
242 compaction_min_threshold: Self::DEFAULT_COMPACTION_MIN_THRESHOLD,
243 compaction_max_threshold: Self::DEFAULT_COMPACTION_MAX_THRESHOLD,
244 }
245 }
246
247 /// Attach a [`UdtRegistry`] used to resolve bare CQL UDT column types at
248 /// flush time (issue #929). See [`WriteEngineConfig::udt_registry`].
249 pub fn with_udt_registry(mut self, registry: UdtRegistry) -> Self {
250 self.udt_registry = Some(registry);
251 self
252 }
253
254 /// Set a custom flush threshold
255 pub fn with_flush_threshold(mut self, threshold: usize) -> Self {
256 self.memtable_flush_threshold = threshold;
257 self
258 }
259
260 /// Set a custom hard limit
261 pub fn with_hard_limit(mut self, limit: usize) -> Self {
262 self.memtable_hard_limit = limit;
263 self
264 }
265
266 /// Set the WAL durability mode.
267 ///
268 /// Mirrors `with_flush_threshold` in style. See [`Durability`] for the
269 /// trade-offs between [`Durability::SyncEachWrite`] (default, production)
270 /// and [`Durability::Disabled`] (bulk-load / benchmarking).
271 ///
272 /// # Example
273 ///
274 /// ```rust,ignore
275 /// use cqlite_core::storage::write_engine::{Durability, WriteEngineConfig};
276 ///
277 /// let config = WriteEngineConfig::new(data, wal, schema)
278 /// .with_durability(Durability::Disabled);
279 /// ```
280 pub fn with_durability(mut self, durability: Durability) -> Self {
281 self.durability = durability;
282 self
283 }
284
285 /// Apply a [`crate::config::CompactionConfig`] onto this write-engine
286 /// configuration (issue #1619). This makes `Config.storage.compaction`
287 /// authoritative for the write path: `auto_compaction = false` disables the
288 /// default STCS policy so [`WriteEngine::maintenance_step`] becomes a no-op.
289 ///
290 /// The compaction thresholds (`compaction_min_threshold` /
291 /// `compaction_max_threshold`) keep their defaults; only the on/off switch
292 /// is surfaced through the public `Config` today.
293 pub fn with_compaction_config(mut self, compaction: &crate::config::CompactionConfig) -> Self {
294 self.auto_compaction = compaction.auto_compaction;
295 self
296 }
297}
298
299/// Write engine coordinator
300///
301/// Orchestrates WAL, memtable, and SSTable flushing for write operations.
302/// This is the primary public API for all write operations in CQLite.
303///
304/// ## Durability contract: you MUST call [`close`](WriteEngine::close)
305///
306/// **`Drop` is not a flush.** Rows written with [`write`](WriteEngine::write) /
307/// [`execute`](WriteEngine::execute) live in the in-memory memtable (and the
308/// WAL) until a flush turns them into an SSTable. Only
309/// [`close`](WriteEngine::close) (or an explicit flush) guarantees the memtable
310/// is persisted to a Data.db. Because Tokio has no async drop, `Drop` CANNOT
311/// flush — doing so would require a `block_on` inside `drop`, which is
312/// forbidden (issue #1693/AG3). An engine dropped with a non-empty memtable
313/// logs a `warn!` and leaves those rows recoverable only via WAL replay on the
314/// next startup.
315///
316/// Embedders (and every long-lived writer) MUST therefore call
317/// `engine.close().await` for a graceful shutdown — e.g. from a `SIGINT`
318/// handler — before the process exits.
319///
320/// ## Thread Safety
321///
322/// WriteEngine follows a single-writer model. It is NOT thread-safe and
323/// should be used from a single thread or protected by external locking.
324/// The `closed` flag uses atomic operations for safe concurrent access checking.
325///
326/// ## Example
327///
328/// ```rust,ignore
329/// use cqlite_core::storage::write_engine::{WriteEngine, WriteEngineConfig, Mutation};
330/// use std::path::PathBuf;
331///
332/// // Create configuration
333/// let config = WriteEngineConfig::new(
334/// PathBuf::from("data"),
335/// PathBuf::from("wal"),
336/// schema
337/// );
338///
339/// // Create engine
340/// let mut engine = WriteEngine::new(config)?;
341///
342/// // Write a mutation
343/// engine.write(mutation)?;
344///
345/// // Execute CQL statement
346/// engine.execute("INSERT INTO users (id, name) VALUES (1, 'Alice')")?;
347///
348/// // Flush to SSTable
349/// engine.flush()?;
350///
351/// // Close cleanly
352/// engine.close()?;
353/// ```
354#[cfg(feature = "write-support")]
355#[derive(Debug)]
356pub struct WriteEngine {
357 /// Configuration
358 config: WriteEngineConfig,
359 /// Write-ahead log for durability
360 wal: WriteAheadLog,
361 /// In-memory write buffer
362 memtable: Memtable,
363 /// Summary of the WAL crash-recovery replay performed in `new()` (issue
364 /// #1391). Its `mutations` are drained into the memtable, leaving only the
365 /// lossiness metadata (`corrupt_entries`, `stopped_early`, `bytes_skipped`).
366 /// Exposed via [`WriteEngine::wal_recovery`] so a caller can detect that a
367 /// recovery was lossy BEFORE the next flush truncates the WAL. A non-clean
368 /// recovery also preserves the raw WAL segment aside (see `new()`).
369 wal_recovery: RecoveryReport,
370 /// SSTable generation counter (increments on each flush)
371 generation: u64,
372 /// Whether the engine has been closed (atomic for thread safety)
373 closed: AtomicBool,
374 /// Active merge state for incremental compaction (M5.2)
375 active_merge: Option<ActiveMerge>,
376 /// Merge policy for compaction decisions (M5.2)
377 merge_policy: Option<Box<dyn MergePolicy>>,
378 /// Cumulative compaction statistics (M5.2, Issue #474)
379 cumulative_stats: CompactionStats,
380 /// Cumulative count of rows written since the engine was opened (Issue #486).
381 ///
382 /// Incremented for each row that enters the memtable via `write()` or
383 /// `write_async()`. Unlike `memtable.row_count()`, this counter is NOT
384 /// reset when the memtable is flushed, so it reflects the total number of
385 /// rows written across the lifetime of the session.
386 rows_written: u64,
387 /// Number of L0 SSTable files successfully flushed since the engine was
388 /// opened (Issue #486).
389 ///
390 /// Incremented once per successful `flush_internal_async()` call that
391 /// produces a non-empty SSTable. This is an in-process counter; it is
392 /// reset to zero when the engine is re-opened (a directory scan could
393 /// also be used for persistence, but the counter is more robust and avoids
394 /// scanning the filesystem on every stats query).
395 l0_count: u64,
396 /// Cumulative bytes written to flushed L0 SSTables (Data.db + all sibling
397 /// components) since the engine was opened (issue #1620). Incremented once
398 /// per successful `flush_internal_async()` by the flushed SSTable's
399 /// `data_size`. In-process only; reset to zero on re-open. Read via
400 /// [`WriteEngine::total_flushed_bytes`] so binding stats stay accurate for
401 /// automatic flushes (not just explicit `flush()` calls).
402 total_flushed_bytes: u64,
403 /// Advisory exclusive lock on `write_dir` (`.lock` file).
404 ///
405 /// Held for the lifetime of the `WriteEngine`; released on `close()` or
406 /// `Drop`. The lock is acquired in `new()` via
407 /// `fs2::FileExt::try_lock_exclusive`, which returns an error immediately
408 /// if another process already holds it (fail-fast, no blocking).
409 ///
410 /// The lock prevents two `Database` / `WriteEngine` instances from sharing
411 /// the same `write_dir`, which would corrupt WAL files and SSTables
412 /// (Issue #485).
413 dir_lock: std::fs::File,
414 /// Whether the "memtable over flush threshold" warning has already been
415 /// emitted for the current threshold crossing (issue #1620).
416 ///
417 /// The sync `write()` path does NOT auto-flush when a Tokio runtime is
418 /// present, so once the memtable crosses the flush threshold it stays over
419 /// it until an explicit (or async) flush. Without this guard the
420 /// over-threshold `tracing::warn!` fired on EVERY subsequent write (log spam).
421 /// It is set to `true` the first time the warning is emitted and reset to
422 /// `false` on the next successful flush, so the warning fires at most once
423 /// per threshold crossing.
424 warned_over_threshold: bool,
425}
426
427/// Reject any mutation that contains a counter cell write.
428///
429/// Counter columns require server-side distributed increment semantics and
430/// cannot be expressed as a last-write-wins mutation. Both the sync
431/// `write()` and the async `write_async()` paths call this guard immediately
432/// after the closed-check.
433#[cfg(feature = "write-support")]
434fn reject_counter_cells(mutation: &Mutation) -> Result<()> {
435 for op in &mutation.operations {
436 match op {
437 CellOperation::Write { value, .. } | CellOperation::WriteWithTtl { value, .. } => {
438 if matches!(value, crate::types::Value::Counter(_)) {
439 return Err(Error::invalid_operation(
440 "counter writes are not supported via the standard mutation path; \
441 counter columns require server-side distributed increment semantics",
442 ));
443 }
444 }
445 _ => {}
446 }
447 }
448 Ok(())
449}
450
451#[cfg(feature = "write-support")]
452impl WriteEngine {
453 /// Create `dir` (and any missing ancestors), then fsync the FULL parent
454 /// chain from `dir`'s parent walking UP to and including the filesystem
455 /// root — unconditionally, on every call (issue #1392).
456 ///
457 /// We deliberately do NOT track "what this invocation created." Across a
458 /// crash that is unknowable: a prior startup may have created the whole
459 /// nested `data_dir`/`wal_dir` tree but crashed partway through the
460 /// parent-fsync sequence, leaving higher ancestors' dirents unpersisted. A
461 /// retry that inspected "already exists" and re-fsynced only the immediate
462 /// parent would leave those ancestors un-durable forever — the "which
463 /// ancestor is durable on retry" hole.
464 ///
465 /// Instead we fsync EVERY ancestor unconditionally. Fsyncing an
466 /// already-durable directory is idempotent and cheap, and the walk
467 /// terminates at the filesystem root (`Path::parent()` returns `None`), so
468 /// no ancestor level is ever left unsynced and there is no per-attempt
469 /// state to track. Every ancestor of a successfully-created directory
470 /// exists, so each fsync targets a real, present directory. Together with
471 /// the flush-path barrier (which fsyncs leaf → data_root on every publish),
472 /// this makes the full durability chain persistent on first flush AND on
473 /// any retry after a partial crash.
474 fn create_dir_all_durable(dir: &Path, label: &str) -> Result<()> {
475 std::fs::create_dir_all(dir).map_err(|e| {
476 Error::Storage(format!(
477 "Failed to create {} directory {:?}: {}",
478 label, dir, e
479 ))
480 })?;
481
482 // Ascend the parent chain, fsyncing each ancestor's dirent. The walk
483 // strictly shrinks (one component per step) and ends when `parent()`
484 // yields `None` at the filesystem root, so it always terminates.
485 let mut next = dir.parent();
486 while let Some(cur) = next {
487 if cur.as_os_str().is_empty() {
488 // A relative single-component path (e.g. "data") has an empty
489 // parent whose real directory is the current working
490 // directory. Fsync "." and stop: it has no ancestor to ascend.
491 wal::sync_directory(Path::new("."))?;
492 break;
493 }
494 wal::sync_directory(cur)?;
495 next = cur.parent();
496 }
497
498 Ok(())
499 }
500
501 /// Create a new write engine
502 ///
503 /// This initializes the WAL and memtable. If a WAL exists in the
504 /// wal_dir, it will be replayed to recover in-flight writes.
505 ///
506 /// # Arguments
507 ///
508 /// * `config` - Write engine configuration
509 ///
510 /// # Returns
511 ///
512 /// A new WriteEngine ready to accept writes.
513 ///
514 /// # Errors
515 ///
516 /// Returns an error if:
517 /// - WAL directory doesn't exist
518 /// - Data directory doesn't exist
519 /// - WAL replay fails
520 pub fn new(config: WriteEngineConfig) -> Result<Self> {
521 // Ensure directories exist AND that EVERY ancestor's own dirent is
522 // durable. `create_dir_all` alone persists the entries *inside* a
523 // freshly-created root once we fsync the root, but NOT the dirents of
524 // the newly-created directories themselves in their parents — and it
525 // may create several missing ancestors at once. Fsyncing only the
526 // immediate parent would leave intermediate dirents unpersisted, so a
527 // crash after WAL truncation could still lose the whole data/WAL root.
528 // `create_dir_all_durable` fsyncs the FULL parent chain up to and
529 // including the filesystem root, unconditionally on every call — this
530 // closes the "which ancestor is durable on a partial-crash retry" hole
531 // and completes the durability chain (SSTable component contents →
532 // leaf → ancestors → data_root → every ancestor up to root). #1392.
533 Self::create_dir_all_durable(&config.data_dir, "data")?;
534 Self::create_dir_all_durable(&config.wal_dir, "WAL")?;
535
536 // Acquire an exclusive advisory lock on the WAL directory to prevent
537 // two WriteEngine / Database instances from sharing the same write_dir
538 // and silently corrupting WAL files or SSTables (Issue #485).
539 //
540 // We use `try_lock_exclusive` (non-blocking) so that callers get an
541 // immediate, actionable error rather than hanging forever.
542 let lock_path = config.wal_dir.join(".lock");
543 let dir_lock = std::fs::OpenOptions::new()
544 .create(true)
545 .truncate(false)
546 .write(true)
547 .open(&lock_path)
548 .map_err(|e| {
549 Error::Storage(format!("Failed to create lock file {:?}: {}", lock_path, e))
550 })?;
551 fs2::FileExt::try_lock_exclusive(&dir_lock)
552 .map_err(|_| Error::write_dir_locked(config.wal_dir.to_string_lossy().into_owned()))?;
553
554 // Startup sweep: remove orphaned compaction artifacts left by a previous crash.
555 //
556 // Two kinds of orphans can be left if the process crashes mid-rename in
557 // `finalize_merge_async`:
558 //
559 // (a) A `.compaction-tmp-{gen}/` directory under `data_dir` with partial
560 // component files.
561 //
562 // (b) A partial set of renamed components in `data_dir/{keyspace}/{table}/`
563 // — specifically one or more `nb-{gen}-big-*.db` files without a
564 // matching `TOC.txt`. Because `scan_data_files` discovers SSTables by
565 // `nb-*-big-Data.db` glob, an orphaned Data.db without TOC.txt will be
566 // picked up by the merge policy and fed to `KWayMerger`, which may
567 // produce garbled output.
568 //
569 // Both sweeps are best-effort: individual failures are logged as warnings but
570 // do not abort engine startup.
571 Self::sweep_startup_orphans(&config);
572
573 // Initialize WAL
574 let wal_path = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
575 let mut wal = if wal_path.exists() {
576 // Recover from existing WAL
577 WriteAheadLog::open_existing(&wal_path)?
578 } else {
579 // Create new WAL
580 WriteAheadLog::create(&config.wal_dir)?
581 };
582
583 // Replay WAL into memtable. The report distinguishes a clean recovery
584 // from a lossy one (issue #1391); a bare Vec cannot.
585 //
586 // Stream each recovered mutation straight into the memtable via
587 // `replay_each` (issue #1661) instead of materialising a whole-log
588 // `Vec<Mutation>` and then copying it in: peak memory is bounded by the
589 // memtable itself, not ~2× it. `recovered` is an authoritative count of
590 // mutations actually recovered (used only for the lossy-recovery log,
591 // preserving its prior `mutations.len()` meaning); the completion log
592 // counts from the memtable. The report still carries all corruption
593 // metadata — mutations stream via the closure, so its `mutations` vec
594 // stays empty.
595 let mut memtable = Memtable::new();
596 let mut recovered = 0usize;
597 // DEFER the first memtable-application error instead of aborting the scan
598 // (issue #1661, roborev). The pre-streaming flow scanned the ENTIRE WAL
599 // first (`replay()`), ran the lossy-recovery preserve/reset, and only THEN
600 // applied mutations — so a schema/key-conversion or insert failure on an
601 // earlier valid mutation could never prevent a LATER corrupt tail from
602 // being detected, preserved aside, and reset. If the streaming callback
603 // returned that error via `?`, replay would stop early and skip
604 // preserve/reset, silently changing the #1390/#1391 recovery semantics
605 // this refactor must preserve. So we record the FIRST apply error, stop
606 // applying (the engine will not be constructed), but keep scanning for the
607 // full RecoveryReport; the error is surfaced only after preserve/reset
608 // below. The streaming memory win is unaffected — no whole-log Vec.
609 let mut apply_error: Option<Error> = None;
610 let wal_recovery = wal.replay_each(|mutation| {
611 if apply_error.is_some() {
612 // An earlier mutation already failed; keep scanning to reach any
613 // later corruption but skip further application.
614 return Ok(());
615 }
616 match mutation.decorated_key(&config.schema) {
617 Ok(decorated_key) => match memtable.insert_with_key(decorated_key, mutation) {
618 Ok(()) => recovered += 1,
619 Err(e) => apply_error = Some(e),
620 },
621 Err(e) => apply_error = Some(e),
622 }
623 Ok(())
624 })?;
625
626 if !wal_recovery.is_clean() {
627 // Preserve the raw WAL segment aside BEFORE anything (a later flush,
628 // or the reset below) can truncate it, so the corruption evidence
629 // survives for manual recovery. The report is also retained on the
630 // engine and exposed via `wal_recovery()` so a caller sees the loss
631 // pre-truncate.
632 let preserved = Self::preserve_corrupt_wal(&wal_path)?;
633
634 // With the evidence safely aside, trim the LIVE WAL back to its last
635 // CRC-valid prefix (issue #1391). Otherwise the engine would remain
636 // writable with the corrupt tail still on disk, and the FIRST synced
637 // write after this lossy recovery would be appended AFTER the corrupt
638 // entry — where the next replay stops — so that acknowledged write
639 // would be silently lost. Resetting to the valid prefix makes
640 // post-recovery appends land at a replayable position. (A torn tail
641 // was already trimmed in `open_existing`, so this only fires for
642 // mid-stream corruption.)
643 let reset_to = wal.reset_to_valid_prefix()?;
644
645 tracing::error!(
646 "WAL recovery at {:?} was LOSSY: recovered {} mutation(s), {} corrupt entry \
647 (entries), stopped_early={}, {} byte(s) not recovered. Raw segment preserved at \
648 {:?}; live WAL reset to valid prefix ({:?}). Investigate before relying on this \
649 data.",
650 wal_path,
651 recovered,
652 wal_recovery.corrupt_entries,
653 wal_recovery.stopped_early,
654 wal_recovery.bytes_skipped,
655 preserved,
656 reset_to,
657 );
658 }
659
660 // Surface any deferred memtable-application error now — AFTER the
661 // lossy-recovery preserve/reset above has run (issue #1661, roborev).
662 // This restores the pre-streaming ordering: on-disk WAL corruption is
663 // preserved aside and the live WAL is reset to its valid prefix before
664 // engine construction fails on the apply error.
665 if let Some(e) = apply_error {
666 return Err(e);
667 }
668
669 if recovered > 0 {
670 tracing::info!(
671 "WAL replay complete: replayed {} mutation(s); {} rows in memtable, {} bytes",
672 recovered,
673 memtable.row_count(),
674 memtable.size_bytes()
675 );
676 }
677
678 // Determine next generation number by scanning data directory
679 let generation = Self::determine_next_generation(&config.data_dir)?;
680
681 // Install the default STCS compaction policy so `maintenance_step`
682 // performs size-tiered compaction out of the box (issue #1619). The
683 // off-switch is `WriteEngineConfig::auto_compaction = false`, which
684 // leaves the policy unset and makes `maintenance_step` a no-op. Read the
685 // config values into locals BEFORE the struct-init moves `config`.
686 let merge_policy: Option<Box<dyn MergePolicy>> = if config.auto_compaction {
687 Some(Box::new(STCSPolicy::new(
688 config.compaction_min_threshold,
689 config.compaction_max_threshold,
690 0.5,
691 1.5,
692 STCSPolicy::DEFAULT_MIN_SSTABLE_SIZE,
693 )?))
694 } else {
695 None
696 };
697
698 Ok(Self {
699 config,
700 wal,
701 memtable,
702 wal_recovery,
703 generation,
704 closed: AtomicBool::new(false),
705 active_merge: None,
706 merge_policy,
707 cumulative_stats: CompactionStats::default(),
708 rows_written: 0,
709 l0_count: 0,
710 total_flushed_bytes: 0,
711 dir_lock,
712 warned_over_threshold: false,
713 })
714 }
715
716 /// Summary of the WAL crash-recovery replay performed when this engine was
717 /// opened (issue #1391).
718 ///
719 /// The `mutations` field has been drained into the memtable, so only the
720 /// lossiness metadata remains. Use [`RecoveryReport::is_clean`] to detect a
721 /// lossy recovery. A non-clean report means the raw WAL segment was
722 /// preserved aside (as `commitlog.wal.corrupt.<nanos>`) so a subsequent
723 /// flush cannot destroy the evidence.
724 pub fn wal_recovery(&self) -> &RecoveryReport {
725 &self.wal_recovery
726 }
727
728 /// Copy the raw WAL file aside before any truncation can destroy it, so a
729 /// lossy recovery leaves forensic evidence (issue #1391). Returns the path
730 /// of the preserved copy.
731 ///
732 /// The copy is a sibling `commitlog.wal.corrupt.<unix_nanos>`; the original
733 /// is left untouched (a later flush truncates it, but the aside copy — and
734 /// the retained [`RecoveryReport`] — survive).
735 fn preserve_corrupt_wal(wal_path: &Path) -> Result<PathBuf> {
736 let nanos = std::time::SystemTime::now()
737 .duration_since(std::time::UNIX_EPOCH)
738 .map(|d| d.as_nanos())
739 .unwrap_or(0);
740 let file_name = wal_path
741 .file_name()
742 .map(|n| n.to_string_lossy().into_owned())
743 .unwrap_or_else(|| WriteAheadLog::WAL_FILENAME.to_string());
744 let preserved = match wal_path.parent() {
745 Some(parent) => parent.join(format!("{}.corrupt.{}", file_name, nanos)),
746 None => PathBuf::from(format!("{}.corrupt.{}", file_name, nanos)),
747 };
748 std::fs::copy(wal_path, &preserved).map_err(|e| {
749 Error::Storage(format!(
750 "Failed to preserve corrupt WAL {:?} aside to {:?}: {}",
751 wal_path, preserved, e
752 ))
753 })?;
754
755 // Make the forensic copy durable BEFORE the caller resets/truncates the
756 // live WAL (issue #1391). `std::fs::copy` does not fsync; without this a
757 // crash after the copy but before it reaches disk could leave the live
758 // WAL already trimmed while the preserved copy is missing/incomplete —
759 // destroying the very evidence we set aside. Order: copy → fsync copy →
760 // fsync parent dir → (caller) reset live WAL.
761 {
762 let copy_file = std::fs::File::open(&preserved).map_err(|e| {
763 Error::Storage(format!(
764 "Failed to open preserved corrupt WAL {:?} for fsync: {}",
765 preserved, e
766 ))
767 })?;
768 copy_file.sync_all().map_err(|e| {
769 Error::Storage(format!(
770 "Failed to fsync preserved corrupt WAL {:?}: {}",
771 preserved, e
772 ))
773 })?;
774 }
775 if let Some(parent) = preserved.parent() {
776 wal::sync_directory(parent)?;
777 }
778
779 Ok(preserved)
780 }
781
782 /// Write a mutation to the write engine
783 ///
784 /// This appends the mutation to the WAL for durability, then inserts it
785 /// into the memtable. If the memtable exceeds the flush threshold,
786 /// an automatic flush is triggered.
787 ///
788 /// **Note**: Automatic flush is disabled when called from an async context.
789 /// Use `write_async()` for async contexts with automatic flush support.
790 ///
791 /// # Arguments
792 ///
793 /// * `mutation` - The mutation to write
794 ///
795 /// # Returns
796 ///
797 /// Ok(()) on success, or an error if the write fails.
798 ///
799 /// # Errors
800 ///
801 /// Returns an error if:
802 /// - Engine has been closed
803 /// - WAL append fails
804 /// - Memtable insert fails
805 /// - Automatic flush fails (sync context only)
806 #[tracing::instrument(name = "write.mutation", level = "debug", skip(self, mutation))]
807 pub fn write(&mut self, mutation: Mutation) -> Result<()> {
808 crate::observability::record_result("write", self.write_inner(mutation))
809 }
810
811 fn write_inner(&mut self, mutation: Mutation) -> Result<()> {
812 // Insert into WAL + memtable, then handle sync-context auto-flush.
813 self.write_into_memtable(mutation)?;
814
815 // Check if memtable should be flushed (sync callers only). The binding
816 // write path runs inside a Tokio runtime and does NOT reach here — it
817 // uses `execute_flushing`, which owns its own real async flush and
818 // deliberately skips the "call flush() manually" warning below.
819 if self
820 .memtable
821 .should_flush(self.config.memtable_flush_threshold)
822 {
823 // Rate-limit the over-threshold warning (issue #1620). A sync caller
824 // that lives inside a runtime never auto-flushes below, so the
825 // memtable stays over threshold across many writes; without this
826 // guard the warning fired on EVERY such write. Emit it at most once
827 // per crossing; it is reset on the next flush.
828 if !self.warned_over_threshold {
829 tracing::warn!(
830 "Memtable size {} exceeds threshold {} - call flush() manually in async context",
831 self.memtable.size_bytes(),
832 self.config.memtable_flush_threshold
833 );
834 self.warned_over_threshold = true;
835 }
836
837 // Try to flush synchronously only if we're not in an async context.
838 if tokio::runtime::Handle::try_current().is_err() {
839 tracing::info!("Triggering automatic flush");
840 self.flush_internal()?;
841 }
842 }
843
844 Ok(())
845 }
846
847 /// Insert a mutation into the WAL + memtable WITHOUT any auto-flush handling.
848 ///
849 /// Shared by the sync [`write_inner`] path (which then performs the
850 /// sync-context warn + flush) and the async [`execute_flushing`] path (which
851 /// owns its own real async flush and must NOT emit the sync "call flush()
852 /// manually in async context" warning) (issue #1620, roborev job 2879).
853 fn write_into_memtable(&mut self, mutation: Mutation) -> Result<()> {
854 if self.closed.load(Ordering::SeqCst) {
855 return Err(Error::InvalidInput(
856 "WriteEngine has been closed".to_string(),
857 ));
858 }
859
860 reject_counter_cells(&mutation)?;
861
862 // Honest hard-limit admission: gate on current size + the INCOMING
863 // mutation's estimated size, and reject a lone jumbo write (issue #1625).
864 self.check_admission(&mutation)?;
865
866 // 1. Append to WAL (durability) — skipped when Durability::Disabled
867 if self.config.durability == Durability::SyncEachWrite {
868 self.wal.append(&mutation)?;
869 self.wal.sync()?;
870 }
871
872 // 2. Compute decorated key from partition key
873 let decorated_key = mutation.decorated_key(&self.config.schema)?;
874
875 // 3. Insert into memtable
876 self.memtable.insert_with_key(decorated_key, mutation)?;
877
878 // 4. Increment the cumulative rows-written counter (Issue #486).
879 // Done after a successful insert so that failed writes are not counted.
880 self.rows_written += 1;
881 crate::observability::add_counter(crate::observability::catalog::WRITE_MUTATIONS, 1, &[]);
882 self.record_memtable_gauges();
883
884 Ok(())
885 }
886
887 /// Honest memtable hard-limit admission (issue #1625).
888 ///
889 /// The pre-#1625 check compared only the *current* memtable size against the
890 /// hard limit — it never measured the incoming mutation, so (a) a single
891 /// jumbo mutation could be admitted into an empty memtable and blow straight
892 /// past the limit, and (b) a mutation could push a nearly-full memtable over
893 /// the limit before the next write was rejected. This gate measures the
894 /// incoming mutation with the SAME estimator the memtable uses for
895 /// accounting, then:
896 ///
897 /// 1. Rejects any single mutation whose estimate alone exceeds the hard
898 /// limit (the ceiling for one mutation is the hard limit itself) with a
899 /// distinct error, so a lone jumbo write can never be admitted.
900 /// 2. Rejects the write when `current_size + incoming` would exceed the hard
901 /// limit (`saturating_add`, so the sum can never overflow `usize`).
902 fn check_admission(&self, mutation: &Mutation) -> Result<()> {
903 let hard_limit = self.config.memtable_hard_limit;
904 let incoming = self.memtable.estimate_mutation_size(mutation);
905
906 // (0) Fail-closed sentinel: the estimator returns `usize::MAX` when a
907 // mutation is pathological/unmeasurable (node-cap hit). Reject it
908 // EXPLICITLY and unconditionally — a `>` comparison against a
909 // configurable `hard_limit == usize::MAX` would otherwise be false and
910 // admit the very mutation the sentinel is meant to fence off.
911 if incoming == usize::MAX {
912 return Err(Error::Storage(
913 "Write rejected: mutation size could not be bounded (estimator \
914 fail-closed sentinel); refusing admission"
915 .to_string(),
916 ));
917 }
918
919 // (1) Single-mutation ceiling: one mutation may not exceed the hard
920 // limit on its own, regardless of how empty the memtable is.
921 if incoming > hard_limit {
922 return Err(Error::Storage(format!(
923 "Write rejected: single mutation estimated at {incoming} bytes exceeds \
924 memtable hard limit {hard_limit} bytes (a single mutation may not exceed \
925 the hard limit)"
926 )));
927 }
928
929 // (2) Projected sum: never admit a write that would push the memtable
930 // over the hard limit.
931 let projected = self.memtable.size_bytes().saturating_add(incoming);
932 if projected > hard_limit {
933 return Err(Error::Storage(format!(
934 "Write rejected: memtable {} + mutation {incoming} would exceed hard limit {hard_limit}",
935 self.memtable.size_bytes()
936 )));
937 }
938
939 Ok(())
940 }
941
942 /// Write a mutation with async automatic flush support
943 ///
944 /// This is the async version of `write()` that supports automatic flushing
945 /// in async contexts. Use this method when calling from async code.
946 ///
947 /// # Arguments
948 ///
949 /// * `mutation` - The mutation to write
950 ///
951 /// # Returns
952 ///
953 /// Ok(()) on success, or an error if the write fails.
954 ///
955 /// # Errors
956 ///
957 /// Returns an error if:
958 /// - Engine has been closed
959 /// - WAL append fails
960 /// - Memtable insert fails
961 /// - Automatic flush fails
962 #[tracing::instrument(name = "write.mutation", level = "debug", skip(self, mutation))]
963 pub async fn write_async(&mut self, mutation: Mutation) -> Result<()> {
964 crate::observability::record_result("write", self.write_async_inner(mutation).await)
965 }
966
967 async fn write_async_inner(&mut self, mutation: Mutation) -> Result<()> {
968 if self.closed.load(Ordering::SeqCst) {
969 return Err(Error::InvalidInput(
970 "WriteEngine has been closed".to_string(),
971 ));
972 }
973
974 reject_counter_cells(&mutation)?;
975
976 // Honest hard-limit admission: gate on current size + the INCOMING
977 // mutation's estimated size, and reject a lone jumbo write (issue #1625).
978 self.check_admission(&mutation)?;
979
980 // 1. Append to WAL (durability) — skipped when Durability::Disabled
981 if self.config.durability == Durability::SyncEachWrite {
982 self.wal.append(&mutation)?;
983 self.wal.sync()?;
984 }
985
986 // 2. Compute decorated key from partition key
987 let decorated_key = mutation.decorated_key(&self.config.schema)?;
988
989 // 3. Insert into memtable
990 self.memtable.insert_with_key(decorated_key, mutation)?;
991
992 // 4. Increment the cumulative rows-written counter (Issue #486).
993 // Done after a successful insert so that failed writes are not counted.
994 self.rows_written += 1;
995 crate::observability::add_counter(crate::observability::catalog::WRITE_MUTATIONS, 1, &[]);
996 self.record_memtable_gauges();
997
998 // 5. Check if memtable should be flushed
999 if self
1000 .memtable
1001 .should_flush(self.config.memtable_flush_threshold)
1002 {
1003 tracing::info!(
1004 "Memtable size {} exceeds threshold {}, triggering flush",
1005 self.memtable.size_bytes(),
1006 self.config.memtable_flush_threshold
1007 );
1008 self.flush_internal_async().await?;
1009 }
1010
1011 Ok(())
1012 }
1013
1014 /// Emit the current memtable size/row gauges (issue #1036). No-op when the
1015 /// `observability` feature is off.
1016 fn record_memtable_gauges(&self) {
1017 crate::observability::record_gauge(
1018 crate::observability::catalog::MEMTABLE_SIZE_BYTES,
1019 self.memtable.size_bytes() as i64,
1020 &[],
1021 );
1022 crate::observability::record_gauge(
1023 crate::observability::catalog::MEMTABLE_ROWS,
1024 self.memtable.row_count() as i64,
1025 &[],
1026 );
1027 }
1028
1029 /// Execute a CQL statement (INSERT, UPDATE, DELETE)
1030 ///
1031 /// This parses the CQL statement and converts it to a mutation,
1032 /// then writes it using the `write()` method.
1033 ///
1034 /// # Arguments
1035 ///
1036 /// * `statement` - CQL statement string
1037 ///
1038 /// # Returns
1039 ///
1040 /// Ok(()) on success, or an error if parsing or writing fails.
1041 ///
1042 /// # Errors
1043 ///
1044 /// Returns an error if:
1045 /// - CQL parsing fails
1046 /// - Statement is not a mutation (INSERT/UPDATE/DELETE)
1047 /// - Mutation conversion fails
1048 /// - Write fails
1049 ///
1050 /// # Example
1051 ///
1052 /// ```rust,ignore
1053 /// engine.execute("INSERT INTO users (id, name) VALUES (1, 'Alice')")?;
1054 /// engine.execute("UPDATE users SET name = 'Bob' WHERE id = 1")?;
1055 /// engine.execute("DELETE FROM users WHERE id = 1")?;
1056 /// ```
1057 #[tracing::instrument(name = "write.cql_execute", level = "debug", skip(self, statement))]
1058 pub fn execute(&mut self, statement: &str) -> Result<()> {
1059 crate::observability::record_result("write", self.execute_inner(statement))
1060 }
1061
1062 fn execute_inner(&mut self, statement: &str) -> Result<()> {
1063 // The public sync `execute` signature is unchanged: delegate to the
1064 // counted implementation and discard the mutation count (issue #1620).
1065 self.execute_inner_counted(statement).map(|_| ())
1066 }
1067
1068 /// Parse + write a CQL statement into the memtable, returning the number of
1069 /// mutations applied (N for a BATCH, else 1).
1070 ///
1071 /// This is the shared core of the sync `execute` path and the async
1072 /// `execute_flushing` path (issue #1620). It is *unrecorded* (no
1073 /// `record_result`); callers record at the public boundary.
1074 fn execute_inner_counted(&mut self, statement: &str) -> Result<u64> {
1075 if self.closed.load(Ordering::SeqCst) {
1076 return Err(Error::InvalidInput(
1077 "WriteEngine has been closed".to_string(),
1078 ));
1079 }
1080
1081 // Single-boundary error recording (issue #1036): call the *unrecorded*
1082 // `write_inner` here rather than the public `write`. `execute` already
1083 // wraps this method in `record_result`, so the escaping error is counted
1084 // exactly once at the public boundary instead of twice.
1085 let mutations = self.parse_statement_to_mutations(statement)?;
1086 let n = mutations.len() as u64;
1087 for mutation in mutations {
1088 self.write_inner(mutation)?;
1089 }
1090 Ok(n)
1091 }
1092
1093 /// Parse a CQL DML statement into its mutation(s). A `BEGIN BATCH` yields one
1094 /// mutation per statement in the batch; any other DML yields exactly one.
1095 ///
1096 /// Shared by the sync `execute` path ([`execute_inner_counted`]) and the
1097 /// async [`execute_flushing`] path so both agree on batch semantics
1098 /// (issue #1620).
1099 fn parse_statement_to_mutations(&self, statement: &str) -> Result<Vec<Mutation>> {
1100 let trimmed = statement.trim();
1101 if trimmed.len() >= 5 && trimmed.as_bytes()[..5].eq_ignore_ascii_case(b"BEGIN") {
1102 cql_to_mutation::convert_cql_to_mutations(trimmed, &self.config.schema)
1103 } else {
1104 Ok(vec![self.parse_cql_to_mutation(statement)?])
1105 }
1106 }
1107
1108 /// Execute a DML statement and, if the memtable has crossed the flush
1109 /// threshold, await a REAL async flush (issue #1620, DECIDED: write_async).
1110 ///
1111 /// This is the entry point for the Node/Python binding write path, which runs
1112 /// inside a Tokio runtime where the sync auto-flush in `write()` is
1113 /// intentionally skipped. It restores auto-flush there WITHOUT the surprise
1114 /// inline-flush latency the plain sync `write()`/`execute()` path avoids.
1115 /// Returns the number of mutations applied (N for BATCH, else 1).
1116 #[tracing::instrument(
1117 name = "write.cql_execute_flushing",
1118 level = "debug",
1119 skip(self, statement)
1120 )]
1121 pub async fn execute_flushing(&mut self, statement: &str) -> Result<u64> {
1122 // Single-boundary error recording (issue #1036): `execute_flushing` is
1123 // the public Node/Python DML entry point, so — like `execute`,
1124 // `write_async`, and `flush` — it records the escaping result exactly
1125 // once here. The inner helper and everything it calls (`write_inner`,
1126 // `flush_internal_async`) are *unrecorded* to avoid double counting.
1127 crate::observability::record_result("write", self.execute_flushing_inner(statement).await)
1128 }
1129
1130 async fn execute_flushing_inner(&mut self, statement: &str) -> Result<u64> {
1131 if self.closed.load(Ordering::SeqCst) {
1132 return Err(Error::InvalidInput(
1133 "WriteEngine has been closed".to_string(),
1134 ));
1135 }
1136
1137 // Parse first, then write each mutation and flush as soon as the memtable
1138 // crosses the threshold. Checking after EVERY mutation (not just once at
1139 // the end) means a large `BEGIN BATCH` cannot accumulate all the way to
1140 // the hard limit before flushing and dead-ending (roborev job 2854,
1141 // issue #1620).
1142 let mutations = self.parse_statement_to_mutations(statement)?;
1143 let n = mutations.len() as u64;
1144 for mutation in mutations {
1145 // Use the warn-free memtable insert: this path owns the real async
1146 // flush below, so it must not emit the sync "call flush() manually"
1147 // warning (roborev job 2879, issue #1620).
1148 self.write_into_memtable(mutation)?;
1149 if self
1150 .memtable
1151 .should_flush(self.config.memtable_flush_threshold)
1152 {
1153 self.flush_internal_async().await?;
1154 }
1155 }
1156 Ok(n)
1157 }
1158
1159 /// Force a flush of the memtable to SSTable
1160 ///
1161 /// This writes all data in the memtable to a new SSTable generation,
1162 /// then truncates the WAL. The memtable is cleared after a successful flush.
1163 ///
1164 /// # Returns
1165 ///
1166 /// Returns `Some(SSTableInfo)` if data was flushed, or `None` if the
1167 /// memtable was empty.
1168 ///
1169 /// # Errors
1170 ///
1171 /// Returns an error if:
1172 /// - Engine has been closed
1173 /// - SSTable write fails
1174 /// - WAL truncate fails
1175 #[tracing::instrument(name = "flush.public", level = "debug", skip(self))]
1176 pub async fn flush(&mut self) -> Result<Option<SSTableInfo>> {
1177 // Single-boundary error recording (issue #1036): `flush` is a public API
1178 // entry point, so it records here. The nested `flush_internal_async`
1179 // helper is intentionally *unrecorded* to avoid double-counting when the
1180 // write/close paths trigger a flush internally.
1181 crate::observability::record_result(
1182 "write",
1183 async {
1184 if self.closed.load(Ordering::SeqCst) {
1185 return Err(Error::InvalidInput(
1186 "WriteEngine has been closed".to_string(),
1187 ));
1188 }
1189
1190 self.flush_internal_async().await
1191 }
1192 .await,
1193 )
1194 }
1195
1196 /// Internal synchronous flush helper.
1197 ///
1198 /// Bridges to the async flush via [`merge::block_on_async`], which is safe to
1199 /// call whether or not a Tokio runtime is already running on this thread
1200 /// (Issue #587).
1201 fn flush_internal(&mut self) -> Result<()> {
1202 merge::block_on_async(self.flush_internal_async())?;
1203 Ok(())
1204 }
1205
1206 /// Internal async flush implementation.
1207 ///
1208 /// This is an *unrecorded* inner helper (issue #1036): it never calls
1209 /// `record_error`/`record_result` itself. Error counting happens exactly once
1210 /// at the public boundary that invoked it (`flush`, `close`, or the
1211 /// `write`/`write_async` auto-flush path, all of which wrap their work in
1212 /// `record_result`).
1213 #[tracing::instrument(name = "flush.memtable", level = "debug", skip(self))]
1214 async fn flush_internal_async(&mut self) -> Result<Option<SSTableInfo>> {
1215 // Check if memtable is empty
1216 if self.memtable.is_empty() {
1217 return Ok(None);
1218 }
1219
1220 let flush_start = Instant::now();
1221 let rows_to_flush = self.memtable.row_count() as u64;
1222
1223 tracing::info!(
1224 "Flushing memtable: {} partitions, {} rows, {} bytes",
1225 self.memtable.iter().count(),
1226 self.memtable.row_count(),
1227 self.memtable.size_bytes()
1228 );
1229
1230 // Create SSTable writer with hint for Bloom filter sizing
1231 let partition_count_hint = self.memtable.iter().count();
1232 let mut writer =
1233 crate::storage::sstable::writer::SSTableWriter::with_expected_partitions_and_registry(
1234 self.config.data_dir.clone(),
1235 self.generation,
1236 &self.config.schema,
1237 partition_count_hint,
1238 self.config.udt_registry.as_ref(),
1239 )?;
1240
1241 // Two-pass flush (issue #729): compute FINAL encoding baselines from all
1242 // partitions before writing any data. Delta encoding in Data.db must use
1243 // the same baselines that will end up in Statistics.db's EncodingStats.
1244 // Without this, an early partition encoded against a higher baseline becomes
1245 // silently corrupted when a later partition lowers the minimum.
1246 let mut baseline_min_ts = i64::MAX;
1247 let mut baseline_min_ldt = i32::MAX;
1248 let mut baseline_min_ttl = i32::MAX;
1249 for (_, mutations) in self.memtable.iter() {
1250 let (ts, ldt, ttl) =
1251 crate::storage::sstable::writer::SSTableWriter::compute_mutations_baseline_stats(
1252 mutations,
1253 );
1254 baseline_min_ts = baseline_min_ts.min(ts);
1255 baseline_min_ldt = baseline_min_ldt.min(ldt);
1256 baseline_min_ttl = baseline_min_ttl.min(ttl);
1257 }
1258 writer.pre_seed_encoding_baselines(baseline_min_ts, baseline_min_ldt, baseline_min_ttl);
1259
1260 // Write all partitions from memtable (already in token order)
1261 for (decorated_key, mutations) in self.memtable.iter() {
1262 writer.write_partition(decorated_key.clone(), mutations.to_vec())?;
1263 }
1264
1265 // Finalize SSTable
1266 let info = writer.finish().await?;
1267
1268 tracing::info!(
1269 "SSTable flush complete: generation {}, {} partitions, {} bytes",
1270 self.generation,
1271 info.partition_count,
1272 info.data_size
1273 );
1274
1275 // Durability handoff (issue #1392): fsync the SSTable's parent directory
1276 // BEFORE truncating the WAL. See `durability::finalize_flush_durability`.
1277 //
1278 // A `?`-propagated error here (e.g. a directory fsync failure) happens
1279 // BEFORE the WAL is truncated, so the WAL is still intact and replayable:
1280 // returning early WITHOUT advancing the generation is correct, because a
1281 // retry re-flushes to the same generation and the untouched WAL still
1282 // recovers the data on a crash mid-retry.
1283 let durability_outcome = durability::finalize_flush_durability(
1284 &durability::RealDurabilityBarrier,
1285 &info.data_path,
1286 &self.config.data_dir,
1287 &mut self.wal,
1288 )?;
1289
1290 // The SSTable is now durably published. Commit flush state — clear the
1291 // memtable and advance the generation — for BOTH outcomes, including
1292 // `WalTruncateFailedAfterCommit`. In that case the WAL has already been
1293 // zeroed, so the published SSTable is the ONLY durable copy of these
1294 // mutations; committing state guarantees a retry writes a NEW generation
1295 // and never lets `File::create` overwrite the sole durable copy (which
1296 // would lose the data on a crash mid-retry). See issue #1392.
1297
1298 // Clear memtable
1299 self.memtable.clear();
1300
1301 // Reset the over-threshold warn guard (issue #1620): the memtable is now
1302 // empty, so the next threshold crossing is a fresh event that should warn
1303 // again (at most once).
1304 self.warned_over_threshold = false;
1305
1306 // Increment the L0 SSTable counter (Issue #486).
1307 self.l0_count += 1;
1308
1309 // Accumulate cumulative flushed bytes (issue #1620) so binding write
1310 // stats reflect automatic flushes, not only explicit `flush()` calls.
1311 self.total_flushed_bytes = self.total_flushed_bytes.saturating_add(info.data_size);
1312
1313 // Increment generation for next flush
1314 self.generation += 1;
1315
1316 // Flush metrics (issue #1036): latency in seconds, rows/bytes flushed,
1317 // and one L0 SSTable created. Emit the post-clear memtable gauges (now
1318 // zero) and the current compaction lag (L0 pending) gauge.
1319 {
1320 use crate::observability::{self as obs, catalog};
1321 obs::record_histogram(
1322 catalog::FLUSH_DURATION,
1323 flush_start.elapsed().as_secs_f64(),
1324 &[],
1325 );
1326 obs::add_counter(catalog::FLUSH_ROWS, rows_to_flush, &[]);
1327 obs::add_counter(catalog::FLUSH_BYTES, info.data_size, &[]);
1328 obs::add_counter(catalog::FLUSH_SSTABLES, 1, &[]);
1329 obs::record_gauge(catalog::COMPACTION_LAG, self.l0_count as i64, &[]);
1330 }
1331 self.record_memtable_gauges();
1332
1333 // Surface a post-mutation WAL-truncate failure AFTER state has been
1334 // committed above (issue #1392). The data is durable in the published
1335 // SSTable and the generation has advanced, so a retry cannot overwrite
1336 // it; the error informs the caller that the WAL is no longer a valid
1337 // replay marker.
1338 match durability_outcome {
1339 durability::FlushDurabilityOutcome::Durable => Ok(Some(info)),
1340 durability::FlushDurabilityOutcome::WalTruncateFailedAfterCommit(err) => Err(err),
1341 }
1342 }
1343
1344 /// Close the write engine
1345 ///
1346 /// This flushes any remaining data in the memtable to SSTable,
1347 /// syncs the WAL, then marks the engine as closed. After calling close(),
1348 /// the engine cannot be used for further writes.
1349 ///
1350 /// **This is the durability boundary.** `Drop` does not (and cannot — no
1351 /// async drop in Tokio) flush; callers MUST `close().await` before exit to
1352 /// guarantee written rows reach an SSTable rather than relying on WAL replay
1353 /// (issue #1693). See the type-level docs on [`WriteEngine`].
1354 ///
1355 /// This method is idempotent - calling it multiple times is safe.
1356 ///
1357 /// # Returns
1358 ///
1359 /// Ok(()) on success.
1360 ///
1361 /// # Errors
1362 ///
1363 /// Returns an error if the final flush fails.
1364 ///
1365 /// WAL-truncate handling during that flush is phase-aware (issue #1392):
1366 ///
1367 /// * A truncate failure that leaves the WAL intact (it faulted *before*
1368 /// mutating the WAL) is logged and swallowed — the WAL stays a valid,
1369 /// idempotent replay marker, so no error is surfaced.
1370 /// * A truncate failure *after* `set_len(0)` has already zeroed the WAL
1371 /// (`WalTruncateFailedAfterCommit`) is **propagated**. By then flush state
1372 /// has already been committed — the SSTable is durable and the generation
1373 /// has advanced — so the data is safe, but the error is surfaced so the
1374 /// caller knows the WAL is no longer a replay marker.
1375 ///
1376 /// When the WAL is already empty (e.g. `Durability::Disabled`) the truncate
1377 /// phase is skipped, so no truncate-phase error can arise.
1378 pub async fn close(&mut self) -> Result<()> {
1379 // Check if already closed (idempotent)
1380 if self.closed.swap(true, Ordering::SeqCst) {
1381 return Ok(());
1382 }
1383
1384 tracing::info!("Closing WriteEngine");
1385
1386 // Flush any remaining data
1387 if !self.memtable.is_empty() {
1388 tracing::info!("Flushing memtable before close");
1389
1390 // Attempt to flush to SSTable
1391 match self.flush_internal_async().await {
1392 Ok(_) => {
1393 tracing::info!("Memtable flushed successfully");
1394 }
1395 Err(e) => {
1396 // If flush fails, log error and return it
1397 tracing::error!("Failed to flush memtable during close: {}", e);
1398 // Single-boundary error recording (issue #1036): `close` is a
1399 // public entry point and `flush_internal_async` is unrecorded,
1400 // so record the escaping flush failure here, exactly once.
1401 crate::observability::record_error(&e, "write");
1402 // Reset closed flag since we failed to close cleanly
1403 self.closed.store(false, Ordering::SeqCst);
1404 return Err(e);
1405 }
1406 }
1407 }
1408
1409 // Sync WAL before closing
1410 if let Err(e) = self.wal.sync() {
1411 tracing::warn!("Failed to sync WAL during close: {}", e);
1412 // Don't fail close if sync fails - data is already persisted to SSTable
1413 }
1414
1415 // Release the exclusive advisory lock on write_dir so a subsequent
1416 // WriteEngine on the same directory can acquire it immediately.
1417 // On Drop the OS would release it anyway, but explicit unlock is more
1418 // deterministic in async / multi-phase shutdown sequences.
1419 if let Err(e) = fs2::FileExt::unlock(&self.dir_lock) {
1420 tracing::warn!("Failed to release write_dir advisory lock: {}", e);
1421 }
1422
1423 tracing::info!("WriteEngine closed");
1424
1425 Ok(())
1426 }
1427
1428 /// Parse a CQL statement to a Mutation
1429 ///
1430 /// Supports INSERT, UPDATE, and DELETE statements.
1431 fn parse_cql_to_mutation(&self, statement: &str) -> Result<Mutation> {
1432 cql_to_mutation::convert_cql_to_mutation(statement, &self.config.schema)
1433 }
1434
1435 /// Determine the next SSTable generation number
1436 ///
1437 /// Scans the data directory for existing SSTable files and returns
1438 /// the next generation number.
1439 fn determine_next_generation(data_dir: &Path) -> Result<u64> {
1440 let mut max_generation = 0u64;
1441
1442 if !data_dir.exists() {
1443 return Ok(1);
1444 }
1445
1446 // Recursively scan for SSTable files (writer places them in keyspace/table/ subdirs)
1447 Self::scan_generations(
1448 data_dir,
1449 &mut max_generation,
1450 crate::storage::sstable::MAX_SSTABLE_SCAN_DEPTH,
1451 )?;
1452
1453 Ok(max_generation + 1)
1454 }
1455
1456 /// Recursively scan directory for SSTable generation numbers
1457 fn scan_generations(dir: &Path, max_generation: &mut u64, depth: usize) -> Result<()> {
1458 for entry in std::fs::read_dir(dir)
1459 .map_err(|e| Error::Storage(format!("Failed to read data directory: {}", e)))?
1460 {
1461 let entry = entry
1462 .map_err(|e| Error::Storage(format!("Failed to read directory entry: {}", e)))?;
1463
1464 let filename = entry.file_name();
1465 let filename_str = filename.to_string_lossy();
1466
1467 // Parse generation from filename: nb-{generation}-big-{Component}.db
1468 if filename_str.starts_with("nb-") && filename_str.contains("-big-") {
1469 if let Some(gen_str) = filename_str
1470 .strip_prefix("nb-")
1471 .and_then(|s| s.split('-').next())
1472 {
1473 if let Ok(gen) = gen_str.parse::<u64>() {
1474 *max_generation = (*max_generation).max(gen);
1475 }
1476 }
1477 } else if depth > 0 {
1478 let path = entry.path();
1479 if path.is_dir() {
1480 Self::scan_generations(&path, max_generation, depth - 1)?;
1481 }
1482 }
1483 }
1484 Ok(())
1485 }
1486}
1487
1488/// Safety-net `Drop` implementation for `WriteEngine`.
1489///
1490/// When a `WriteEngine` is dropped without calling `close()` (e.g. due to an
1491/// early return or a panic), the OS would release the advisory lock anyway once
1492/// the file descriptor is closed. This explicit `Drop` makes the release
1493/// deterministic and logs a warning so callers can distinguish a normal shutdown
1494/// from an ungraceful one.
1495#[cfg(feature = "write-support")]
1496impl Drop for WriteEngine {
1497 fn drop(&mut self) {
1498 // Issue #1693 (AG4): `Drop` is NOT a flush — there is no async drop in
1499 // Tokio, and re-introducing a `block_on` here is the AG3 defect. So the
1500 // best we can do on an ungraceful drop is make the silent data-loss mode
1501 // visible: if the memtable still holds rows that `close()` never flushed,
1502 // warn. The `is_empty()` guard is a cheap length check; the format
1503 // arguments (a single `usize`) allocate only when a warn-level logger is
1504 // actually enabled.
1505 if !self.memtable.is_empty() {
1506 // The recovery guidance depends on the durability mode: with WAL
1507 // durability the rows survive in the WAL and replay on next open,
1508 // but with `Durability::Disabled` the WAL was skipped entirely, so
1509 // an ungraceful drop loses the un-flushed rows permanently.
1510 match self.config.durability {
1511 Durability::SyncEachWrite => tracing::warn!(
1512 "WriteEngine dropped without close(): {} row(s) in the memtable were NOT \
1513 flushed to an SSTable and remain only in the WAL (durability now relies on \
1514 WAL replay at next startup). Call `close().await` for a graceful shutdown.",
1515 self.memtable.row_count()
1516 ),
1517 Durability::Disabled => tracing::warn!(
1518 "WriteEngine dropped without close(): {} row(s) in the memtable were NOT \
1519 flushed to an SSTable and are LOST — durability is Disabled so these rows \
1520 were never written to the WAL and cannot be recovered (they existed in \
1521 memory only). Call `close().await` for a graceful shutdown.",
1522 self.memtable.row_count()
1523 ),
1524 }
1525 }
1526
1527 // If close() was already called the lock was already released; a second
1528 // unlock is a no-op on most platforms (Linux returns ENOLCK which we
1529 // ignore here). The important invariant is that the lock is always
1530 // released before the file descriptor is closed by the OS on drop.
1531 if let Err(e) = fs2::FileExt::unlock(&self.dir_lock) {
1532 tracing::debug!(
1533 "WriteEngine drop: advisory lock release returned: {} \
1534 (may have been released by close() already)",
1535 e
1536 );
1537 }
1538 }
1539}
1540
1541#[cfg(all(test, feature = "write-support"))]
1542mod tests {
1543 use super::*;
1544 use crate::storage::write_engine::test_support::{create_test_mutation, create_test_schema};
1545 use tempfile::TempDir;
1546
1547 #[test]
1548 fn test_write_engine_config() {
1549 let temp_dir = TempDir::new().unwrap();
1550 let schema = create_test_schema();
1551
1552 let config = WriteEngineConfig::new(
1553 temp_dir.path().join("data"),
1554 temp_dir.path().join("wal"),
1555 schema,
1556 );
1557
1558 assert_eq!(
1559 config.memtable_flush_threshold,
1560 WriteEngineConfig::DEFAULT_FLUSH_THRESHOLD
1561 );
1562 assert_eq!(
1563 config.memtable_hard_limit,
1564 WriteEngineConfig::DEFAULT_HARD_LIMIT
1565 );
1566
1567 let config = config.with_flush_threshold(128 * 1024 * 1024);
1568 assert_eq!(config.memtable_flush_threshold, 128 * 1024 * 1024);
1569
1570 let config = config.with_hard_limit(512 * 1024 * 1024);
1571 assert_eq!(config.memtable_hard_limit, 512 * 1024 * 1024);
1572 }
1573
1574 #[test]
1575 fn test_write_engine_new() {
1576 let temp_dir = TempDir::new().unwrap();
1577 let schema = create_test_schema();
1578
1579 let config = WriteEngineConfig::new(
1580 temp_dir.path().join("data"),
1581 temp_dir.path().join("wal"),
1582 schema,
1583 );
1584
1585 let engine = WriteEngine::new(config).unwrap();
1586
1587 assert_eq!(engine.generation(), 1);
1588 assert_eq!(engine.memtable_size(), 0);
1589 assert_eq!(engine.memtable_row_count(), 0);
1590 assert!(!engine.closed.load(std::sync::atomic::Ordering::Relaxed));
1591 }
1592
1593 /// Issue #1392: `create_dir_all_durable` is the init-time link in the
1594 /// durability chain — it must create a fresh (possibly nested) root and
1595 /// fsync its parent so the new root's own dirent is persisted. There is no
1596 /// fs-op recorder seam on the init path, so this exercises the helper's
1597 /// call path directly (create + parent fsync succeeds and is idempotent);
1598 /// a crash-injection observation of the fsync would require a syscall seam
1599 /// that does not exist here.
1600 #[test]
1601 fn test_create_dir_all_durable_fsyncs_parent() {
1602 let temp_dir = TempDir::new().unwrap();
1603
1604 // Fresh nested root: parent (`root`) exists, `data` is newly created.
1605 let root = temp_dir.path().join("root");
1606 std::fs::create_dir_all(&root).unwrap();
1607 let data = root.join("data");
1608 assert!(!data.exists());
1609
1610 WriteEngine::create_dir_all_durable(&data, "data").unwrap();
1611 assert!(data.is_dir(), "new root directory must exist");
1612
1613 // Idempotent: re-running on an already-existing dir still succeeds and
1614 // re-fsyncs the parent without error.
1615 WriteEngine::create_dir_all_durable(&data, "data").unwrap();
1616 assert!(data.is_dir());
1617 }
1618
1619 /// Issue #1392 (roborev HIGH): the "track what this attempt created"
1620 /// approach had a retry hole — if a PRIOR startup created the nested tree
1621 /// but crashed partway through its parent-fsync sequence, a retry sees the
1622 /// tree already exists, records an empty newly-created set, and re-fsyncs
1623 /// only the immediate parent, leaving higher ancestors un-durable forever.
1624 ///
1625 /// The definitive fix is to stop tracking anything and fsync the FULL
1626 /// parent chain (up to and including the filesystem root) on EVERY call.
1627 /// This test simulates that partial-crash retry: it pre-creates the entire
1628 /// `base/a/b/data` tree first, so nothing is "newly created" on the call,
1629 /// then asserts every ancestor is still openable/fsyncable up to the root.
1630 ///
1631 /// Red-before: the tracking helper would fsync only `data`'s immediate
1632 /// parent (`b`) on the already-exists path, leaving `a`/`base` unsynced.
1633 /// Green-after: the unconditional full-chain walk fsyncs `b`, `a`, `base`,
1634 /// and every ancestor up to the filesystem root.
1635 #[test]
1636 fn test_create_dir_all_durable_fsyncs_full_chain_on_existing_tree() {
1637 let temp_dir = TempDir::new().unwrap();
1638
1639 // Simulate a prior partial-crash: the ENTIRE nested tree already exists
1640 // before the call, so no directory is "newly created" this invocation.
1641 let base = temp_dir.path().join("base");
1642 let a = base.join("a");
1643 let b = a.join("b");
1644 let data = b.join("data");
1645 std::fs::create_dir_all(&data).unwrap();
1646 assert!(data.is_dir(), "precondition: full tree pre-exists");
1647
1648 // The unconditional full-chain fsync must succeed on the already-exists
1649 // path — it walks `data`'s parent up to the filesystem root, fsyncing
1650 // every ancestor regardless of what (if anything) this call created.
1651 WriteEngine::create_dir_all_durable(&data, "data").unwrap();
1652
1653 // Every ancestor from the leaf's parent up to the root is a real,
1654 // present directory that the walk fsynced (open-for-sync would fail on
1655 // a missing/non-durable dirent). Assert the whole chain is present.
1656 for ancestor in [b.as_path(), a.as_path(), base.as_path()] {
1657 assert!(
1658 ancestor.is_dir(),
1659 "ancestor {:?} must be present and fsyncable up the full chain",
1660 ancestor
1661 );
1662 }
1663
1664 // Idempotent: a second unconditional pass still succeeds.
1665 WriteEngine::create_dir_all_durable(&data, "data").unwrap();
1666 assert!(data.is_dir());
1667 }
1668
1669 /// Issue #1392: the full-chain walk must TERMINATE at the filesystem root
1670 /// (a path whose `parent()` is `None`) rather than looping. Exercising the
1671 /// helper against an absolute nested path implicitly walks to `/`; this
1672 /// test guards the termination contract by confirming the call returns
1673 /// (does not hang) and the root itself has no parent to ascend into.
1674 #[test]
1675 fn test_create_dir_all_durable_walk_terminates_at_root() {
1676 let temp_dir = TempDir::new().unwrap();
1677 let data = temp_dir.path().join("a").join("b").join("data");
1678
1679 // Absolute path: the parent chain ascends through `b`, `a`, the temp
1680 // dir, ... up to `/`. `Path::new("/").parent()` is `None`, so the walk
1681 // stops there — no infinite loop. The call simply returns.
1682 WriteEngine::create_dir_all_durable(&data, "data").unwrap();
1683 assert!(data.is_dir());
1684
1685 // Contract the walk relies on: the filesystem root terminates the walk.
1686 assert!(
1687 std::path::Path::new("/").parent().is_none(),
1688 "filesystem root must have no parent so the fsync walk terminates"
1689 );
1690 }
1691
1692 #[test]
1693 fn test_write_engine_write_single_mutation() {
1694 let temp_dir = TempDir::new().unwrap();
1695 let schema = create_test_schema();
1696
1697 let config = WriteEngineConfig::new(
1698 temp_dir.path().join("data"),
1699 temp_dir.path().join("wal"),
1700 schema,
1701 );
1702
1703 let mut engine = WriteEngine::new(config).unwrap();
1704
1705 let mutation = create_test_mutation(1, "Alice", 1000000);
1706 engine.write(mutation).unwrap();
1707
1708 assert_eq!(engine.memtable_row_count(), 1);
1709 assert!(engine.memtable_size() > 0);
1710 assert!(engine.wal_size() > 0);
1711 }
1712
1713 #[test]
1714 fn test_write_engine_write_multiple_mutations() {
1715 let temp_dir = TempDir::new().unwrap();
1716 let schema = create_test_schema();
1717
1718 let config = WriteEngineConfig::new(
1719 temp_dir.path().join("data"),
1720 temp_dir.path().join("wal"),
1721 schema,
1722 );
1723
1724 let mut engine = WriteEngine::new(config).unwrap();
1725
1726 // Write 10 mutations
1727 for i in 0..10 {
1728 let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
1729 engine.write(mutation).unwrap();
1730 }
1731
1732 assert_eq!(engine.memtable_row_count(), 10);
1733 assert!(engine.memtable_size() > 0);
1734 }
1735
1736 #[tokio::test]
1737 async fn test_write_engine_flush_empty() {
1738 let temp_dir = TempDir::new().unwrap();
1739 let schema = create_test_schema();
1740
1741 let config = WriteEngineConfig::new(
1742 temp_dir.path().join("data"),
1743 temp_dir.path().join("wal"),
1744 schema,
1745 );
1746
1747 let mut engine = WriteEngine::new(config).unwrap();
1748
1749 // Flush empty memtable
1750 let result = engine.flush().await.unwrap();
1751 assert!(result.is_none());
1752 }
1753
1754 #[tokio::test]
1755 async fn test_write_engine_flush_with_data() {
1756 let temp_dir = TempDir::new().unwrap();
1757 let schema = create_test_schema();
1758
1759 let config = WriteEngineConfig::new(
1760 temp_dir.path().join("data"),
1761 temp_dir.path().join("wal"),
1762 schema,
1763 );
1764
1765 let mut engine = WriteEngine::new(config).unwrap();
1766
1767 // Write mutations
1768 for i in 0..5 {
1769 let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
1770 engine.write(mutation).unwrap();
1771 }
1772
1773 let initial_generation = engine.generation();
1774
1775 // Flush
1776 let info = engine.flush().await.unwrap();
1777 assert!(info.is_some());
1778
1779 let info = info.unwrap();
1780 assert_eq!(info.partition_count, 5);
1781 assert!(info.data_size > 0);
1782 assert!(info.data_path.exists());
1783
1784 // Memtable should be empty after flush
1785 assert_eq!(engine.memtable_row_count(), 0);
1786 assert_eq!(engine.memtable_size(), 0);
1787
1788 // WAL should be truncated
1789 assert_eq!(engine.wal_size(), 0);
1790
1791 // Generation should increment
1792 assert_eq!(engine.generation(), initial_generation + 1);
1793 }
1794
1795 #[test]
1796 fn test_write_engine_automatic_flush() {
1797 let temp_dir = TempDir::new().unwrap();
1798 let schema = create_test_schema();
1799
1800 // Set very low flush threshold (1KB)
1801 let config = WriteEngineConfig::new(
1802 temp_dir.path().join("data"),
1803 temp_dir.path().join("wal"),
1804 schema,
1805 )
1806 .with_flush_threshold(1024);
1807
1808 let mut engine = WriteEngine::new(config).unwrap();
1809
1810 // Write enough mutations to trigger automatic flush
1811 for i in 0..100 {
1812 let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
1813 engine.write(mutation).unwrap();
1814 }
1815
1816 // Should have automatically flushed
1817 // Memtable may have some data if writes continued after flush
1818 // But generation should have incremented
1819 assert!(engine.generation() > 1 || engine.memtable_size() < 10000);
1820 }
1821
1822 #[tokio::test]
1823 async fn test_write_engine_close_with_data() {
1824 let temp_dir = TempDir::new().unwrap();
1825 let schema = create_test_schema();
1826
1827 let config = WriteEngineConfig::new(
1828 temp_dir.path().join("data"),
1829 temp_dir.path().join("wal"),
1830 schema,
1831 );
1832
1833 let mut engine = WriteEngine::new(config).unwrap();
1834
1835 // Write mutations
1836 for i in 0..5 {
1837 let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
1838 engine.write(mutation).unwrap();
1839 }
1840
1841 // Close should flush
1842 engine.close().await.unwrap();
1843
1844 // Verify SSTable was created
1845 let data_dir = temp_dir.path().join("data");
1846 let entries: Vec<_> = std::fs::read_dir(&data_dir).unwrap().collect();
1847 assert!(!entries.is_empty(), "SSTable files should exist");
1848 }
1849
1850 #[tokio::test]
1851 async fn test_write_engine_close_empty() {
1852 let temp_dir = TempDir::new().unwrap();
1853 let schema = create_test_schema();
1854
1855 let config = WriteEngineConfig::new(
1856 temp_dir.path().join("data"),
1857 temp_dir.path().join("wal"),
1858 schema,
1859 );
1860
1861 let mut engine = WriteEngine::new(config).unwrap();
1862
1863 // Close empty engine
1864 engine.close().await.unwrap();
1865 }
1866
1867 #[test]
1868 fn test_write_engine_write_after_close() {
1869 let temp_dir = TempDir::new().unwrap();
1870 let schema = create_test_schema();
1871
1872 let config = WriteEngineConfig::new(
1873 temp_dir.path().join("data"),
1874 temp_dir.path().join("wal"),
1875 schema,
1876 );
1877
1878 let mut engine = WriteEngine::new(config).unwrap();
1879
1880 // Close
1881 tokio::runtime::Runtime::new()
1882 .unwrap()
1883 .block_on(engine.close())
1884 .unwrap();
1885
1886 // Create new engine with same config (simulating restart)
1887 let schema2 = create_test_schema();
1888 let config2 = WriteEngineConfig::new(
1889 temp_dir.path().join("data"),
1890 temp_dir.path().join("wal"),
1891 schema2,
1892 );
1893
1894 let mut engine2 = WriteEngine::new(config2).unwrap();
1895
1896 // Write should fail on closed engine (if we still had reference)
1897 // But new engine should work
1898 let mutation = create_test_mutation(1, "Alice", 1000000);
1899 engine2.write(mutation).unwrap();
1900 assert_eq!(engine2.memtable_row_count(), 1);
1901 }
1902
1903 #[test]
1904 fn test_write_engine_wal_recovery() {
1905 let temp_dir = TempDir::new().unwrap();
1906 let schema = create_test_schema();
1907
1908 let config = WriteEngineConfig::new(
1909 temp_dir.path().join("data"),
1910 temp_dir.path().join("wal"),
1911 schema.clone(),
1912 );
1913
1914 // Write mutations and close without flushing
1915 {
1916 let mut engine = WriteEngine::new(config.clone()).unwrap();
1917
1918 for i in 0..5 {
1919 let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
1920 engine.write(mutation).unwrap();
1921 }
1922
1923 // Don't flush - just drop engine (simulating crash)
1924 }
1925
1926 // Create new engine - should recover from WAL
1927 let config2 = WriteEngineConfig::new(
1928 temp_dir.path().join("data"),
1929 temp_dir.path().join("wal"),
1930 schema,
1931 );
1932
1933 let engine = WriteEngine::new(config2).unwrap();
1934
1935 // Should have recovered 5 mutations
1936 assert_eq!(engine.memtable_row_count(), 5);
1937 assert!(engine.memtable_size() > 0);
1938 }
1939
1940 /// Issue #1390 criterion 5: engine-level crash simulation through the
1941 /// WriteEngine (not just the WAL unit). Write A (fsync-acknowledged), crash
1942 /// (drop without flush), inject a torn tail into the commitlog, recover in a
1943 /// new engine, write C (fsync-acknowledged), crash again, then recover once
1944 /// more and assert every acknowledged mutation is present exactly once.
1945 ///
1946 /// Before the fix, C lands AFTER the retained torn bytes and is silently
1947 /// lost on the second recovery.
1948 #[test]
1949 fn test_write_engine_recovers_across_torn_tail_crash() {
1950 use std::io::Write as _;
1951
1952 let temp_dir = TempDir::new().unwrap();
1953 let schema = create_test_schema();
1954 let config = WriteEngineConfig::new(
1955 temp_dir.path().join("data"),
1956 temp_dir.path().join("wal"),
1957 schema,
1958 );
1959
1960 // 1. Write A and "crash" (drop without flush). Default durability is
1961 // SyncEachWrite, so A is fsync'd to the WAL.
1962 {
1963 let mut engine = WriteEngine::new(config.clone()).unwrap();
1964 engine
1965 .write(create_test_mutation(1, "Alice", 1_000_000))
1966 .unwrap();
1967 }
1968
1969 // 2. Inject a torn tail: a complete 8-byte header declaring a 100-byte
1970 // payload, followed by only 10 payload bytes (interrupted append).
1971 let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
1972 {
1973 let mut f = std::fs::OpenOptions::new()
1974 .append(true)
1975 .open(&wal_file)
1976 .unwrap();
1977 f.write_all(&100u32.to_le_bytes()).unwrap();
1978 f.write_all(&0u32.to_le_bytes()).unwrap();
1979 f.write_all(&[0xAB; 10]).unwrap();
1980 f.sync_all().unwrap();
1981 }
1982
1983 // 3. Recover: the torn tail must be trimmed. Then write C and crash.
1984 {
1985 let mut engine = WriteEngine::new(config.clone()).unwrap();
1986 assert_eq!(
1987 engine.memtable_row_count(),
1988 1,
1989 "must recover A across the torn tail"
1990 );
1991 engine
1992 .write(create_test_mutation(3, "Carol", 3_000_000))
1993 .unwrap();
1994 }
1995
1996 // 4. Recover again: both acknowledged mutations must be present exactly
1997 // once (2 distinct partition keys => 2 memtable rows).
1998 {
1999 let engine = WriteEngine::new(config).unwrap();
2000 assert_eq!(
2001 engine.memtable_row_count(),
2002 2,
2003 "both A and C must survive the second recovery (C must not be lost)"
2004 );
2005 }
2006 }
2007
2008 // ---- Issue #1391: engine-level flush guard over a lossy WAL replay ----
2009 //
2010 // A lossy WAL recovery must be surfaced to the caller (via the retained
2011 // RecoveryReport) AND its raw segment preserved aside BEFORE the next flush
2012 // truncates the WAL. Otherwise the loss becomes permanent and invisible.
2013
2014 /// Build a 3-entry (A, B, C) WAL in `wal_dir` and return the byte offset at
2015 /// the end of entry A (so B's payload, which starts at `end_a + 8`, can be
2016 /// corrupted). All entries are fsync'd.
2017 fn seed_wal_abc(wal_dir: &Path) -> u64 {
2018 std::fs::create_dir_all(wal_dir).unwrap();
2019 let mut wal = WriteAheadLog::create(wal_dir).unwrap();
2020 wal.append(&create_test_mutation(1, "A", 1_000_000))
2021 .unwrap();
2022 wal.sync().unwrap();
2023 let end_a = wal.size();
2024 wal.append(&create_test_mutation(2, "B", 2_000_000))
2025 .unwrap();
2026 wal.sync().unwrap();
2027 wal.append(&create_test_mutation(3, "C", 3_000_000))
2028 .unwrap();
2029 wal.sync().unwrap();
2030 drop(wal);
2031 end_a
2032 }
2033
2034 fn count_corrupt_aside(wal_dir: &Path) -> usize {
2035 std::fs::read_dir(wal_dir)
2036 .unwrap()
2037 .filter_map(|e| e.ok())
2038 .filter(|e| {
2039 e.file_name()
2040 .to_string_lossy()
2041 .contains("commitlog.wal.corrupt.")
2042 })
2043 .count()
2044 }
2045
2046 /// Criterion 5: a lossy WAL must not be silently truncated by the first
2047 /// flush. Opening the engine over an A,B(corrupt),C WAL must expose a
2048 /// non-clean RecoveryReport and preserve the raw segment aside; a subsequent
2049 /// flush must NOT destroy that evidence.
2050 #[tokio::test]
2051 async fn test_write_engine_lossy_wal_preserved_before_flush_truncate() {
2052 use std::io::{Read, Seek, SeekFrom, Write};
2053
2054 let temp_dir = TempDir::new().unwrap();
2055 let schema = create_test_schema();
2056 let config = WriteEngineConfig::new(
2057 temp_dir.path().join("data"),
2058 temp_dir.path().join("wal"),
2059 schema,
2060 );
2061
2062 let end_a = seed_wal_abc(&config.wal_dir);
2063
2064 // Bit-flip the first payload byte of B (just past its 8-byte header).
2065 let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
2066 {
2067 let mut f = std::fs::OpenOptions::new()
2068 .read(true)
2069 .write(true)
2070 .open(&wal_file)
2071 .unwrap();
2072 f.seek(SeekFrom::Start(end_a + 8)).unwrap();
2073 let mut byte = [0u8; 1];
2074 f.read_exact(&mut byte).unwrap();
2075 f.seek(SeekFrom::Start(end_a + 8)).unwrap();
2076 f.write_all(&[byte[0] ^ 0x01]).unwrap();
2077 f.sync_all().unwrap();
2078 }
2079
2080 // Open the engine: recovery must surface as lossy and preserve evidence.
2081 let mut engine = WriteEngine::new(config.clone()).unwrap();
2082 assert!(
2083 !engine.wal_recovery().is_clean(),
2084 "engine must expose a non-clean RecoveryReport for a lossy WAL"
2085 );
2086 assert_eq!(engine.wal_recovery().corrupt_entries, 1);
2087 assert!(engine.wal_recovery().stopped_early);
2088 assert_eq!(
2089 engine.memtable_row_count(),
2090 1,
2091 "only the valid prefix [A] is recovered"
2092 );
2093 assert_eq!(
2094 count_corrupt_aside(&config.wal_dir),
2095 1,
2096 "the raw corrupt WAL segment must be preserved aside BEFORE any flush"
2097 );
2098
2099 // Flush: this truncates the live WAL, but the preserved evidence must
2100 // survive so the loss is not destroyed.
2101 engine.flush().await.unwrap();
2102 assert_eq!(
2103 count_corrupt_aside(&config.wal_dir),
2104 1,
2105 "flush must NOT destroy the preserved corrupt WAL segment"
2106 );
2107 }
2108
2109 /// Issue #1661 (roborev): a memtable-application error on an EARLIER valid
2110 /// mutation MUST NOT abort the WAL scan before a LATER corrupt tail can be
2111 /// detected, preserved aside, and reset. The streaming refactor applies each
2112 /// mutation as it is scanned; if that apply error propagated immediately, the
2113 /// #1390/#1391 lossy-recovery preserve/reset would be skipped — a forbidden
2114 /// behavior change. Build a WAL of [A(applies-fail), B(corrupt)]: A is a
2115 /// CRC-valid, decodable mutation whose partition key value (`Text`) cannot be
2116 /// encoded against the `int` schema, so `decorated_key` fails at apply time;
2117 /// B is bit-flipped so replay reports the recovery as lossy. `WriteEngine::new`
2118 /// must return the deferred apply error, but ONLY after the corrupt segment was
2119 /// preserved aside and the live WAL reset to its valid prefix [A].
2120 #[test]
2121 fn test_write_engine_apply_error_still_preserves_and_resets_corrupt_tail() {
2122 use crate::storage::write_engine::mutation::{CellOperation, PartitionKey, TableId};
2123 use crate::types::Value;
2124 use std::io::{Read, Seek, SeekFrom, Write};
2125
2126 let temp_dir = TempDir::new().unwrap();
2127 let schema = create_test_schema(); // partition key `id` is `int`
2128 let config = WriteEngineConfig::new(
2129 temp_dir.path().join("data"),
2130 temp_dir.path().join("wal"),
2131 schema,
2132 );
2133
2134 // Seed A = a mutation whose partition key is Text (mismatches the int
2135 // schema) so it serializes/CRCs/deserializes fine but fails at apply
2136 // (`decorated_key`); then B = a valid mutation we will corrupt.
2137 std::fs::create_dir_all(&config.wal_dir).unwrap();
2138 let end_a = {
2139 let mut wal = WriteAheadLog::create(&config.wal_dir).unwrap();
2140 let bad = Mutation::new(
2141 TableId::new("test_ks", "test_table"),
2142 PartitionKey::single("id", Value::text("not-an-int".to_string())),
2143 None,
2144 vec![CellOperation::Write {
2145 column: "name".to_string(),
2146 value: Value::text("A".to_string()),
2147 }],
2148 1_000_000,
2149 None,
2150 );
2151 wal.append(&bad).unwrap();
2152 wal.sync().unwrap();
2153 let end_a = wal.size();
2154 wal.append(&create_test_mutation(2, "B", 2_000_000))
2155 .unwrap();
2156 wal.sync().unwrap();
2157 end_a
2158 };
2159
2160 // Bit-flip B's first payload byte (just past its 8-byte header) → CRC
2161 // mismatch → mid-stream corruption.
2162 let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
2163 {
2164 let mut f = std::fs::OpenOptions::new()
2165 .read(true)
2166 .write(true)
2167 .open(&wal_file)
2168 .unwrap();
2169 f.seek(SeekFrom::Start(end_a + 8)).unwrap();
2170 let mut byte = [0u8; 1];
2171 f.read_exact(&mut byte).unwrap();
2172 f.seek(SeekFrom::Start(end_a + 8)).unwrap();
2173 f.write_all(&[byte[0] ^ 0x01]).unwrap();
2174 f.sync_all().unwrap();
2175 }
2176
2177 // Opening must FAIL on the deferred apply error...
2178 let result = WriteEngine::new(config.clone());
2179 assert!(
2180 result.is_err(),
2181 "engine open must surface the deferred memtable-application error"
2182 );
2183
2184 // ...but ONLY after the lossy-recovery preserve/reset ran. Pre-fix, the
2185 // apply error aborted the scan before B's corruption was seen, so neither
2186 // of these held.
2187 assert_eq!(
2188 count_corrupt_aside(&config.wal_dir),
2189 1,
2190 "corrupt segment must be preserved aside even though an earlier apply failed"
2191 );
2192 assert_eq!(
2193 std::fs::metadata(&wal_file).unwrap().len(),
2194 end_a,
2195 "live WAL must be reset to its valid prefix [A] before the apply error surfaces"
2196 );
2197 }
2198
2199 /// Regression guard: a CLEAN recovery is unchanged — no aside file, and the
2200 /// flush truncates the WAL normally.
2201 #[tokio::test]
2202 async fn test_write_engine_clean_wal_recovery_truncates_normally() {
2203 let temp_dir = TempDir::new().unwrap();
2204 let schema = create_test_schema();
2205 let config = WriteEngineConfig::new(
2206 temp_dir.path().join("data"),
2207 temp_dir.path().join("wal"),
2208 schema,
2209 );
2210
2211 seed_wal_abc(&config.wal_dir);
2212
2213 let mut engine = WriteEngine::new(config.clone()).unwrap();
2214 assert!(engine.wal_recovery().is_clean());
2215 assert_eq!(engine.memtable_row_count(), 3, "A, B, C all recovered");
2216 assert_eq!(
2217 count_corrupt_aside(&config.wal_dir),
2218 0,
2219 "a clean recovery must not preserve any aside segment"
2220 );
2221
2222 engine.flush().await.unwrap();
2223 let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
2224 assert_eq!(
2225 std::fs::metadata(&wal_file).unwrap().len(),
2226 0,
2227 "clean recovery: the WAL is truncated normally after flush"
2228 );
2229 assert_eq!(count_corrupt_aside(&config.wal_dir), 0);
2230 }
2231
2232 /// Issue #1391 (Finding B): recovery-write-crash. After a LOSSY recovery the
2233 /// engine stays writable, so a synced write issued right after recovery MUST
2234 /// survive the next replay. Before the fix the live WAL still carried the
2235 /// corrupt tail, so the acknowledged write landed AFTER the corrupt entry —
2236 /// where replay stops — and was silently lost. With the valid-prefix reset,
2237 /// the write lands at a replayable position: reopening the live WAL and
2238 /// replaying must yield the valid prefix PLUS the post-recovery write (never
2239 /// a set missing it), while the corrupt evidence is still preserved aside.
2240 #[tokio::test]
2241 async fn test_write_engine_reset_to_valid_prefix_keeps_post_recovery_write() {
2242 use crate::types::Value;
2243 use std::io::{Read, Seek, SeekFrom, Write};
2244
2245 let temp_dir = TempDir::new().unwrap();
2246 let schema = create_test_schema();
2247 let config = WriteEngineConfig::new(
2248 temp_dir.path().join("data"),
2249 temp_dir.path().join("wal"),
2250 schema,
2251 );
2252
2253 // Seed A, B, C then bit-flip B's payload → mid-stream corruption.
2254 let end_a = seed_wal_abc(&config.wal_dir);
2255 let wal_file = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
2256 {
2257 let mut f = std::fs::OpenOptions::new()
2258 .read(true)
2259 .write(true)
2260 .open(&wal_file)
2261 .unwrap();
2262 f.seek(SeekFrom::Start(end_a + 8)).unwrap();
2263 let mut byte = [0u8; 1];
2264 f.read_exact(&mut byte).unwrap();
2265 f.seek(SeekFrom::Start(end_a + 8)).unwrap();
2266 f.write_all(&[byte[0] ^ 0x01]).unwrap();
2267 f.sync_all().unwrap();
2268 }
2269
2270 // Open the engine (recovery is lossy → evidence preserved, live WAL reset
2271 // to the valid prefix [A]), then issue a synced post-recovery write D.
2272 {
2273 let mut engine = WriteEngine::new(config.clone()).unwrap();
2274 assert!(!engine.wal_recovery().is_clean());
2275 assert_eq!(engine.memtable_row_count(), 1, "valid prefix [A] recovered");
2276 assert_eq!(
2277 count_corrupt_aside(&config.wal_dir),
2278 1,
2279 "corrupt segment must be preserved aside"
2280 );
2281 // Default durability is SyncEachWrite, so D is fsync'd to the WAL.
2282 engine
2283 .write(create_test_mutation(4, "D", 4_000_000))
2284 .unwrap();
2285 drop(engine);
2286 }
2287
2288 // Reopen the LIVE WAL directly and replay: it must recover [A, D] — the
2289 // acknowledged post-recovery write is present, NOT lost behind the old
2290 // corrupt tail.
2291 let wal = WriteAheadLog::open_existing(&wal_file).unwrap();
2292 let report = wal.replay().unwrap();
2293 assert!(
2294 report.is_clean(),
2295 "the reset live WAL is a clean [A, D] prefix"
2296 );
2297 let names: Vec<&str> = report
2298 .mutations
2299 .iter()
2300 .map(|m| match &m.operations[0] {
2301 CellOperation::Write {
2302 value: Value::Text(name),
2303 ..
2304 } => std::str::from_utf8(name).unwrap_or_default(),
2305 other => panic!("expected Write op, got {other:?}"),
2306 })
2307 .collect();
2308 assert_eq!(
2309 names,
2310 vec!["A", "D"],
2311 "replay must yield the valid prefix AND the post-recovery write D, \
2312 never a set missing D (and never the corrupt B/C)"
2313 );
2314 assert!(
2315 !names.contains(&"B"),
2316 "the corrupt entry B must not resurface"
2317 );
2318 assert!(
2319 !names.contains(&"C"),
2320 "C (behind corruption) must not resurface"
2321 );
2322
2323 // The forensic evidence must still be on disk.
2324 assert_eq!(
2325 count_corrupt_aside(&config.wal_dir),
2326 1,
2327 "the corrupt WAL evidence must remain preserved aside after the reset"
2328 );
2329 }
2330
2331 #[test]
2332 fn test_write_engine_generation_tracking() {
2333 let temp_dir = TempDir::new().unwrap();
2334 let schema = create_test_schema();
2335
2336 let config = WriteEngineConfig::new(
2337 temp_dir.path().join("data"),
2338 temp_dir.path().join("wal"),
2339 schema.clone(),
2340 );
2341
2342 // First engine
2343 {
2344 let mut engine = WriteEngine::new(config.clone()).unwrap();
2345 assert_eq!(engine.generation(), 1);
2346
2347 // Write and flush
2348 let mutation = create_test_mutation(1, "Alice", 1000000);
2349 engine.write(mutation).unwrap();
2350
2351 tokio::runtime::Runtime::new()
2352 .unwrap()
2353 .block_on(engine.flush())
2354 .unwrap();
2355
2356 assert_eq!(engine.generation(), 2);
2357 }
2358
2359 // Second engine - should detect existing generation
2360 let config2 = WriteEngineConfig::new(
2361 temp_dir.path().join("data"),
2362 temp_dir.path().join("wal"),
2363 schema,
2364 );
2365
2366 let engine = WriteEngine::new(config2).unwrap();
2367 assert_eq!(engine.generation(), 2);
2368 }
2369
2370 #[test]
2371 fn test_write_engine_execute_table_mismatch() {
2372 let temp_dir = TempDir::new().unwrap();
2373 let schema = create_test_schema();
2374
2375 let config = WriteEngineConfig::new(
2376 temp_dir.path().join("data"),
2377 temp_dir.path().join("wal"),
2378 schema,
2379 );
2380
2381 let mut engine = WriteEngine::new(config).unwrap();
2382
2383 // Schema defines test_table, but statement targets users → table mismatch
2384 let result = engine.execute("INSERT INTO users (id, name) VALUES (1, 'Alice')");
2385 let err_msg = result.unwrap_err().to_string();
2386 assert!(
2387 err_msg.contains("targets table 'users'")
2388 && err_msg.contains("schema is for 'test_table'"),
2389 "Expected table mismatch error, got: {}",
2390 err_msg
2391 );
2392 }
2393
2394 #[test]
2395 fn test_write_engine_execute_insert_success() {
2396 let temp_dir = TempDir::new().unwrap();
2397 let schema = create_test_schema();
2398
2399 let config = WriteEngineConfig::new(
2400 temp_dir.path().join("data"),
2401 temp_dir.path().join("wal"),
2402 schema,
2403 );
2404
2405 let mut engine = WriteEngine::new(config).unwrap();
2406
2407 assert_eq!(engine.memtable_row_count(), 0);
2408
2409 // INSERT matching the test schema: test_ks.test_table(id int PK, name text)
2410 let result = engine.execute("INSERT INTO test_table (id, name) VALUES (1, 'Alice')");
2411 assert!(
2412 result.is_ok(),
2413 "execute() failed: {:?}",
2414 result.unwrap_err()
2415 );
2416
2417 assert_eq!(engine.memtable_row_count(), 1);
2418 }
2419
2420 #[test]
2421 fn test_determine_next_generation_empty_dir() {
2422 let temp_dir = TempDir::new().unwrap();
2423 let data_dir = temp_dir.path().join("data");
2424 std::fs::create_dir_all(&data_dir).unwrap();
2425
2426 let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
2427 assert_eq!(generation, 1);
2428 }
2429
2430 #[test]
2431 fn test_determine_next_generation_with_sstables() {
2432 let temp_dir = TempDir::new().unwrap();
2433 let data_dir = temp_dir.path().join("data");
2434 std::fs::create_dir_all(&data_dir).unwrap();
2435
2436 // Create dummy SSTable files
2437 std::fs::write(data_dir.join("nb-1-big-Data.db"), b"").unwrap();
2438 std::fs::write(data_dir.join("nb-2-big-Data.db"), b"").unwrap();
2439 std::fs::write(data_dir.join("nb-5-big-Data.db"), b"").unwrap();
2440
2441 let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
2442 assert_eq!(generation, 6);
2443 }
2444
2445 #[tokio::test]
2446 async fn test_write_engine_close_idempotent() {
2447 let temp_dir = TempDir::new().unwrap();
2448 let schema = create_test_schema();
2449
2450 let config = WriteEngineConfig::new(
2451 temp_dir.path().join("data"),
2452 temp_dir.path().join("wal"),
2453 schema,
2454 );
2455
2456 let mut engine = WriteEngine::new(config).unwrap();
2457
2458 // Close once
2459 engine.close().await.unwrap();
2460 assert!(engine.closed.load(Ordering::SeqCst));
2461
2462 // Close again - should be idempotent
2463 engine.close().await.unwrap();
2464 assert!(engine.closed.load(Ordering::SeqCst));
2465 }
2466
2467 #[tokio::test]
2468 async fn test_write_engine_close_syncs_wal() {
2469 let temp_dir = TempDir::new().unwrap();
2470 let schema = create_test_schema();
2471
2472 let config = WriteEngineConfig::new(
2473 temp_dir.path().join("data"),
2474 temp_dir.path().join("wal"),
2475 schema,
2476 );
2477
2478 let mut engine = WriteEngine::new(config).unwrap();
2479
2480 // Write a mutation
2481 let mutation = create_test_mutation(1, "Alice", 1000000);
2482 engine.write(mutation).unwrap();
2483
2484 // Close should sync WAL before completing
2485 engine.close().await.unwrap();
2486
2487 // Verify WAL was truncated (because data was flushed to SSTable)
2488 assert_eq!(engine.wal_size(), 0);
2489 }
2490
2491 #[test]
2492 fn test_write_engine_closed_flag_atomic() {
2493 let temp_dir = TempDir::new().unwrap();
2494 let schema = create_test_schema();
2495
2496 let config = WriteEngineConfig::new(
2497 temp_dir.path().join("data"),
2498 temp_dir.path().join("wal"),
2499 schema,
2500 );
2501
2502 let engine = WriteEngine::new(config).unwrap();
2503
2504 // Verify closed flag is atomic
2505 assert!(!engine.closed.load(Ordering::SeqCst));
2506
2507 // Store true
2508 engine.closed.store(true, Ordering::SeqCst);
2509 assert!(engine.closed.load(Ordering::SeqCst));
2510
2511 // Swap back to false
2512 let prev = engine.closed.swap(false, Ordering::SeqCst);
2513 assert!(prev);
2514 assert!(!engine.closed.load(Ordering::SeqCst));
2515 }
2516
2517 #[tokio::test]
2518 async fn test_write_engine_write_after_close_fails() {
2519 let temp_dir = TempDir::new().unwrap();
2520 let schema = create_test_schema();
2521
2522 let config = WriteEngineConfig::new(
2523 temp_dir.path().join("data"),
2524 temp_dir.path().join("wal"),
2525 schema,
2526 );
2527
2528 let mut engine = WriteEngine::new(config).unwrap();
2529
2530 // Close the engine
2531 engine.close().await.unwrap();
2532
2533 // Try to write - should fail
2534 let mutation = create_test_mutation(1, "Alice", 1000000);
2535 let result = engine.write(mutation);
2536
2537 assert!(result.is_err());
2538 match result {
2539 Err(Error::InvalidInput(_)) => {}
2540 _ => panic!("Expected InvalidInput error"),
2541 }
2542 }
2543
2544 #[tokio::test]
2545 async fn test_write_engine_flush_after_close_fails() {
2546 let temp_dir = TempDir::new().unwrap();
2547 let schema = create_test_schema();
2548
2549 let config = WriteEngineConfig::new(
2550 temp_dir.path().join("data"),
2551 temp_dir.path().join("wal"),
2552 schema,
2553 );
2554
2555 let mut engine = WriteEngine::new(config).unwrap();
2556
2557 // Close the engine
2558 engine.close().await.unwrap();
2559
2560 // Try to flush - should fail
2561 let result = engine.flush().await;
2562
2563 assert!(result.is_err());
2564 match result {
2565 Err(Error::InvalidInput(_)) => {}
2566 _ => panic!("Expected InvalidInput error"),
2567 }
2568 }
2569
2570 #[test]
2571 fn test_write_engine_hard_limit_enforcement() {
2572 let temp_dir = TempDir::new().unwrap();
2573 let schema = create_test_schema();
2574
2575 // Set very low hard limit (2KB) with flush threshold higher (to prevent auto-flush)
2576 let config = WriteEngineConfig::new(
2577 temp_dir.path().join("data"),
2578 temp_dir.path().join("wal"),
2579 schema,
2580 )
2581 .with_flush_threshold(10 * 1024) // 10KB flush threshold (higher than hard limit for test)
2582 .with_hard_limit(2048); // 2KB hard limit
2583
2584 let mut engine = WriteEngine::new(config).unwrap();
2585
2586 // Write mutations until we hit the hard limit
2587 let mut write_count = 0;
2588 for i in 0..1000 {
2589 let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
2590 let result = engine.write(mutation);
2591
2592 match result {
2593 Ok(()) => {
2594 write_count += 1;
2595 }
2596 Err(Error::Storage(msg)) => {
2597 assert!(msg.contains("hard limit"));
2598 break;
2599 }
2600 Err(e) => panic!("Expected Storage error, got: {:?}", e),
2601 }
2602 }
2603
2604 // Should have stopped before 1000 writes due to hard limit
2605 assert!(
2606 write_count < 1000,
2607 "Should have hit hard limit before 1000 writes"
2608 );
2609 assert!(
2610 write_count > 0,
2611 "Should have accepted at least some writes before hitting limit"
2612 );
2613 }
2614
2615 #[tokio::test]
2616 async fn test_write_engine_hard_limit_enforcement_async() {
2617 let temp_dir = TempDir::new().unwrap();
2618 let schema = create_test_schema();
2619
2620 // Set very low hard limit (2KB) with flush threshold higher (to prevent auto-flush)
2621 let config = WriteEngineConfig::new(
2622 temp_dir.path().join("data"),
2623 temp_dir.path().join("wal"),
2624 schema,
2625 )
2626 .with_flush_threshold(10 * 1024) // 10KB flush threshold (higher than hard limit for test)
2627 .with_hard_limit(2048); // 2KB hard limit
2628
2629 let mut engine = WriteEngine::new(config).unwrap();
2630
2631 // Write mutations until we hit the hard limit
2632 let mut write_count = 0;
2633 for i in 0..1000 {
2634 let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
2635 let result = engine.write_async(mutation).await;
2636
2637 match result {
2638 Ok(()) => {
2639 write_count += 1;
2640 }
2641 Err(Error::Storage(msg)) => {
2642 assert!(msg.contains("hard limit"));
2643 break;
2644 }
2645 Err(e) => panic!("Expected Storage error, got: {:?}", e),
2646 }
2647 }
2648
2649 // Should have stopped before 1000 writes due to hard limit
2650 assert!(
2651 write_count < 1000,
2652 "Should have hit hard limit before 1000 writes"
2653 );
2654 assert!(
2655 write_count > 0,
2656 "Should have accepted at least some writes before hitting limit"
2657 );
2658 }
2659
2660 #[tokio::test]
2661 async fn test_write_engine_hard_limit_recovery_after_flush() {
2662 let temp_dir = TempDir::new().unwrap();
2663 let schema = create_test_schema();
2664
2665 // Set low hard limit
2666 let config = WriteEngineConfig::new(
2667 temp_dir.path().join("data"),
2668 temp_dir.path().join("wal"),
2669 schema,
2670 )
2671 .with_flush_threshold(1024)
2672 .with_hard_limit(2048);
2673
2674 let mut engine = WriteEngine::new(config).unwrap();
2675
2676 // Write until hard limit
2677 let mut first_batch_count = 0;
2678 for i in 0..1000 {
2679 let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
2680 let result = engine.write(mutation);
2681
2682 if result.is_err() {
2683 break;
2684 }
2685
2686 first_batch_count += 1;
2687 }
2688
2689 assert!(
2690 first_batch_count > 0,
2691 "Should have accepted some writes before limit"
2692 );
2693
2694 // Flush to clear memtable
2695 engine.flush().await.unwrap();
2696
2697 // Should be able to write again after flush
2698 let mutation = create_test_mutation(9999, "After flush", 2000000);
2699 let result = engine.write(mutation);
2700 assert!(result.is_ok(), "Should accept writes after flush");
2701
2702 assert_eq!(engine.memtable_row_count(), 1);
2703 }
2704
2705 /// Issue #1620 (N2 — auto-flush cliff): the binding write path runs INSIDE a
2706 /// Tokio runtime, where the sync `write()`/`execute()` auto-flush is
2707 /// intentionally skipped (`Handle::try_current().is_err()` is false). On main
2708 /// that meant the memtable grew unbounded until the hard limit, then every
2709 /// write dead-ended. `execute_flushing` restores auto-flush via a REAL async
2710 /// flush.
2711 ///
2712 /// Negative control: with the OLD sync `execute()` path in this same
2713 /// runtime-present topology, no flush would ever fire, so `generation()`
2714 /// would stay `1` and eventually a write would hit the hard limit. Here we
2715 /// assert a real flush advanced the generation and no write dead-ended.
2716 #[tokio::test]
2717 async fn test_execute_flushing_auto_flushes_in_runtime() {
2718 let temp_dir = TempDir::new().unwrap();
2719 let schema = create_test_schema();
2720
2721 // Tiny flush threshold (4KB), default 256MB hard limit.
2722 let config = WriteEngineConfig::new(
2723 temp_dir.path().join("data"),
2724 temp_dir.path().join("wal"),
2725 schema,
2726 )
2727 .with_flush_threshold(4096)
2728 .with_durability(Durability::Disabled);
2729
2730 let mut engine = WriteEngine::new(config).unwrap();
2731 assert_eq!(engine.generation(), 1);
2732
2733 // We are inside a runtime (this is #[tokio::test]) — the exact topology
2734 // the bindings hit. Drive enough inserts to cross the tiny threshold
2735 // many times.
2736 for i in 0..2000 {
2737 let stmt = format!(
2738 "INSERT INTO test_ks.test_table (id, name) VALUES ({}, 'User{}')",
2739 i, i
2740 );
2741 let n = engine
2742 .execute_flushing(&stmt)
2743 .await
2744 .expect("execute_flushing must not dead-end at the hard limit");
2745 assert_eq!(n, 1, "single INSERT applies exactly one mutation");
2746 }
2747
2748 // A real async flush happened: generation advanced past 1.
2749 assert!(
2750 engine.generation() > 1,
2751 "expected auto-flush to advance generation, got {} (would stay 1 on the old sync path in a runtime)",
2752 engine.generation()
2753 );
2754 }
2755
2756 /// Issue #1620: the over-threshold `tracing::warn!` must fire at most once per
2757 /// threshold crossing. With a runtime present the sync `write()` path does
2758 /// NOT auto-flush, so the memtable stays over threshold across many writes;
2759 /// the `warned_over_threshold` guard prevents re-warning until the next
2760 /// flush resets it. This asserts the private flag directly (same module) so
2761 /// it is deterministic without capturing log output.
2762 #[tokio::test]
2763 async fn test_over_threshold_warn_fires_at_most_once_per_crossing() {
2764 let temp_dir = TempDir::new().unwrap();
2765 let schema = create_test_schema();
2766
2767 // Tiny flush threshold so a couple of writes cross it; generous hard
2768 // limit so the sync path keeps accepting writes over threshold.
2769 let config = WriteEngineConfig::new(
2770 temp_dir.path().join("data"),
2771 temp_dir.path().join("wal"),
2772 schema,
2773 )
2774 .with_flush_threshold(256)
2775 .with_hard_limit(64 * 1024)
2776 .with_durability(Durability::Disabled);
2777
2778 let mut engine = WriteEngine::new(config).unwrap();
2779 assert!(
2780 !engine.warned_over_threshold,
2781 "guard starts false before any crossing"
2782 );
2783
2784 // Sync path inside a runtime: no auto-flush, so the memtable climbs and
2785 // stays over threshold. First over-threshold write flips the guard.
2786 let mut crossed = false;
2787 for i in 0..500 {
2788 engine
2789 .write(create_test_mutation(
2790 i,
2791 &format!("User{}", i),
2792 1_000_000 + i as i64,
2793 ))
2794 .unwrap();
2795 if engine.warned_over_threshold {
2796 crossed = true;
2797 break;
2798 }
2799 }
2800 assert!(crossed, "memtable should have crossed the flush threshold");
2801 assert!(
2802 engine.warned_over_threshold,
2803 "guard set after first crossing"
2804 );
2805
2806 // Subsequent over-threshold writes must NOT re-arm/clear the guard.
2807 for i in 500..600 {
2808 engine
2809 .write(create_test_mutation(
2810 i,
2811 &format!("User{}", i),
2812 1_000_000 + i as i64,
2813 ))
2814 .unwrap();
2815 assert!(
2816 engine.warned_over_threshold,
2817 "guard must remain set across subsequent over-threshold writes"
2818 );
2819 }
2820
2821 // An explicit flush clears the memtable and resets the guard, so the
2822 // next crossing warns again.
2823 engine.flush().await.unwrap();
2824 assert!(
2825 !engine.warned_over_threshold,
2826 "guard resets to false after a successful flush"
2827 );
2828 }
2829
2830 #[test]
2831 fn test_generation_counter_is_u64() {
2832 // Verify that generation counter is u64 to prevent overflow (Issue #410)
2833 let temp_dir = TempDir::new().unwrap();
2834 let schema = create_test_schema();
2835
2836 let config = WriteEngineConfig::new(
2837 temp_dir.path().join("data"),
2838 temp_dir.path().join("wal"),
2839 schema,
2840 );
2841
2842 let engine = WriteEngine::new(config).unwrap();
2843
2844 // Verify type by checking the value is within u64 range
2845 let generation: u64 = engine.generation();
2846 assert_eq!(generation, 1u64);
2847
2848 // This compile-time check ensures generation() returns u64
2849 // If it returned u32, this assignment would be a no-op but still compile
2850 let _type_check: u64 = generation;
2851
2852 // Verify that u64 can handle generations beyond u32::MAX
2853 // This would overflow with u32 (max value: 4,294,967,295)
2854 let large_generation: u64 = u32::MAX as u64 + 1000;
2855 assert!(large_generation > u32::MAX as u64);
2856 assert_eq!(large_generation, 4_294_968_295u64);
2857 }
2858
2859 #[test]
2860 fn test_determine_next_generation_large_numbers() {
2861 // Verify that generation parsing handles large u64 values (Issue #410)
2862 let temp_dir = TempDir::new().unwrap();
2863 let data_dir = temp_dir.path().join("data");
2864 std::fs::create_dir_all(&data_dir).unwrap();
2865
2866 // Create dummy SSTable files with large generation numbers
2867 // These would overflow if we used u32 (max: 4,294,967,295)
2868 let large_gen: u64 = u32::MAX as u64 + 100;
2869 std::fs::write(data_dir.join(format!("nb-{}-big-Data.db", large_gen)), b"").unwrap();
2870
2871 let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
2872 assert_eq!(generation, large_gen + 1);
2873 assert!(generation > u32::MAX as u64);
2874 }
2875
2876 // ============================================================================
2877 // Issue #547: WAL durability toggle tests
2878 // ============================================================================
2879
2880 /// `Durability::SyncEachWrite` is the default variant.
2881 #[test]
2882 fn test_durability_default_is_sync_each_write() {
2883 assert_eq!(Durability::default(), Durability::SyncEachWrite);
2884 }
2885
2886 /// `WriteEngineConfig` defaults to `Durability::SyncEachWrite`.
2887 #[test]
2888 fn test_config_default_durability() {
2889 let temp_dir = TempDir::new().unwrap();
2890 let schema = create_test_schema();
2891 let config = WriteEngineConfig::new(
2892 temp_dir.path().join("data"),
2893 temp_dir.path().join("wal"),
2894 schema,
2895 );
2896 assert_eq!(config.durability, Durability::SyncEachWrite);
2897 }
2898
2899 /// `with_durability` builder sets the field and returns `Self`.
2900 #[test]
2901 fn test_config_with_durability_builder() {
2902 let temp_dir = TempDir::new().unwrap();
2903 let schema = create_test_schema();
2904 let config = WriteEngineConfig::new(
2905 temp_dir.path().join("data"),
2906 temp_dir.path().join("wal"),
2907 schema,
2908 )
2909 .with_durability(Durability::Disabled);
2910 assert_eq!(config.durability, Durability::Disabled);
2911 }
2912
2913 /// With `Durability::SyncEachWrite`, the WAL grows after each `write`.
2914 #[test]
2915 fn test_wal_on_produces_wal_growth() {
2916 let temp_dir = TempDir::new().unwrap();
2917 let schema = create_test_schema();
2918 let config = WriteEngineConfig::new(
2919 temp_dir.path().join("data"),
2920 temp_dir.path().join("wal"),
2921 schema,
2922 )
2923 .with_durability(Durability::SyncEachWrite);
2924
2925 let mut engine = WriteEngine::new(config).unwrap();
2926 assert_eq!(engine.wal_size(), 0, "WAL must start empty");
2927
2928 let mutation = create_test_mutation(1, "Alice", 1_000_000);
2929 engine.write(mutation).unwrap();
2930
2931 assert!(
2932 engine.wal_size() > 0,
2933 "WAL must grow after write with SyncEachWrite"
2934 );
2935 }
2936
2937 /// With `Durability::Disabled`, the WAL is never written — `wal_size()` stays 0.
2938 #[test]
2939 fn test_wal_off_produces_no_wal_growth() {
2940 let temp_dir = TempDir::new().unwrap();
2941 let schema = create_test_schema();
2942 let config = WriteEngineConfig::new(
2943 temp_dir.path().join("data"),
2944 temp_dir.path().join("wal"),
2945 schema,
2946 )
2947 .with_durability(Durability::Disabled);
2948
2949 let mut engine = WriteEngine::new(config).unwrap();
2950 assert_eq!(engine.wal_size(), 0, "WAL must start empty");
2951
2952 // Write several mutations — none should touch the WAL.
2953 for i in 0..10 {
2954 let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
2955 engine.write(mutation).unwrap();
2956 }
2957
2958 assert_eq!(
2959 engine.wal_size(),
2960 0,
2961 "WAL must remain empty with Durability::Disabled"
2962 );
2963 assert_eq!(
2964 engine.memtable_row_count(),
2965 10,
2966 "Mutations must reach the memtable even without WAL"
2967 );
2968 }
2969
2970 /// With `Durability::Disabled`, async writes also skip the WAL.
2971 #[tokio::test]
2972 async fn test_wal_off_write_async_produces_no_wal_growth() {
2973 let temp_dir = TempDir::new().unwrap();
2974 let schema = create_test_schema();
2975 let config = WriteEngineConfig::new(
2976 temp_dir.path().join("data"),
2977 temp_dir.path().join("wal"),
2978 schema,
2979 )
2980 .with_durability(Durability::Disabled);
2981
2982 let mut engine = WriteEngine::new(config).unwrap();
2983
2984 for i in 0..5 {
2985 let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
2986 engine.write_async(mutation).await.unwrap();
2987 }
2988
2989 assert_eq!(
2990 engine.wal_size(),
2991 0,
2992 "WAL must remain empty with Durability::Disabled (async path)"
2993 );
2994 assert_eq!(engine.memtable_row_count(), 5);
2995 }
2996
2997 /// With `Durability::Disabled`, data that was never WAL'd is NOT replayed on
2998 /// restart — confirming the documented durability trade-off.
2999 #[test]
3000 fn test_wal_off_no_replay_on_restart() {
3001 let temp_dir = TempDir::new().unwrap();
3002 let schema = create_test_schema();
3003
3004 {
3005 let config = WriteEngineConfig::new(
3006 temp_dir.path().join("data"),
3007 temp_dir.path().join("wal"),
3008 schema.clone(),
3009 )
3010 .with_durability(Durability::Disabled);
3011
3012 let mut engine = WriteEngine::new(config).unwrap();
3013
3014 for i in 0..5 {
3015 let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
3016 engine.write(mutation).unwrap();
3017 }
3018
3019 // Drop without flushing — simulating crash.
3020 }
3021
3022 // Reopen with default durability. Because the WAL was never written, the
3023 // memtable must be empty.
3024 let config2 = WriteEngineConfig::new(
3025 temp_dir.path().join("data"),
3026 temp_dir.path().join("wal"),
3027 schema,
3028 );
3029 let engine2 = WriteEngine::new(config2).unwrap();
3030
3031 assert_eq!(
3032 engine2.memtable_row_count(),
3033 0,
3034 "No WAL entries were written with Disabled, so no replay is possible"
3035 );
3036 }
3037
3038 /// With `Durability::SyncEachWrite`, mutations ARE replayed after a simulated crash.
3039 #[test]
3040 fn test_wal_on_replays_on_restart() {
3041 let temp_dir = TempDir::new().unwrap();
3042 let schema = create_test_schema();
3043
3044 {
3045 let config = WriteEngineConfig::new(
3046 temp_dir.path().join("data"),
3047 temp_dir.path().join("wal"),
3048 schema.clone(),
3049 )
3050 .with_durability(Durability::SyncEachWrite);
3051
3052 let mut engine = WriteEngine::new(config).unwrap();
3053
3054 for i in 0..5 {
3055 let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
3056 engine.write(mutation).unwrap();
3057 }
3058
3059 // Drop without flushing — WAL entries remain on disk.
3060 }
3061
3062 // Reopen — WAL replay must restore the 5 mutations.
3063 let config2 = WriteEngineConfig::new(
3064 temp_dir.path().join("data"),
3065 temp_dir.path().join("wal"),
3066 schema,
3067 )
3068 .with_durability(Durability::SyncEachWrite);
3069
3070 let engine2 = WriteEngine::new(config2).unwrap();
3071
3072 assert_eq!(
3073 engine2.memtable_row_count(),
3074 5,
3075 "SyncEachWrite must replay mutations durably on restart"
3076 );
3077 }
3078
3079 // ============================================================================
3080 // Issue #485: write_dir file lock tests
3081 // ============================================================================
3082
3083 /// A second WriteEngine on the same write_dir must fail fast with a clear error
3084 /// while the first engine is still open.
3085 #[test]
3086 fn test_write_dir_lock_second_engine_fails_fast() {
3087 let temp_dir = TempDir::new().unwrap();
3088 let schema = create_test_schema();
3089
3090 let config1 = WriteEngineConfig::new(
3091 temp_dir.path().join("data"),
3092 temp_dir.path().join("wal"),
3093 schema.clone(),
3094 );
3095
3096 // First engine acquires the lock.
3097 let _engine1 = WriteEngine::new(config1).unwrap();
3098
3099 // Second engine on the same wal_dir must fail immediately.
3100 let config2 = WriteEngineConfig::new(
3101 temp_dir.path().join("data"),
3102 temp_dir.path().join("wal"),
3103 schema,
3104 );
3105
3106 let result = WriteEngine::new(config2);
3107 assert!(
3108 result.is_err(),
3109 "A second WriteEngine on the same write_dir must fail"
3110 );
3111
3112 let err = result.unwrap_err();
3113 assert!(
3114 matches!(err, Error::WriteDirLocked { .. }),
3115 "Expected WriteDirLocked error, got: {:?}",
3116 err
3117 );
3118
3119 // Error message must contain the path and actionable advice.
3120 let msg = err.to_string();
3121 assert!(
3122 msg.contains("already locked"),
3123 "Error message must mention the lock: {}",
3124 msg
3125 );
3126 assert!(
3127 msg.contains("Only one Database instance"),
3128 "Error message must explain the constraint: {}",
3129 msg
3130 );
3131 }
3132
3133 /// After the first engine is closed, a new engine must successfully acquire the lock.
3134 #[tokio::test]
3135 async fn test_write_dir_lock_reacquired_after_close() {
3136 let temp_dir = TempDir::new().unwrap();
3137 let schema = create_test_schema();
3138
3139 let config1 = WriteEngineConfig::new(
3140 temp_dir.path().join("data"),
3141 temp_dir.path().join("wal"),
3142 schema.clone(),
3143 );
3144
3145 // First engine acquires the lock.
3146 let mut engine1 = WriteEngine::new(config1).unwrap();
3147
3148 // Close releases the lock.
3149 engine1.close().await.unwrap();
3150
3151 // A new engine on the same directory must now succeed.
3152 let config2 = WriteEngineConfig::new(
3153 temp_dir.path().join("data"),
3154 temp_dir.path().join("wal"),
3155 schema,
3156 );
3157
3158 let engine2 = WriteEngine::new(config2);
3159 assert!(
3160 engine2.is_ok(),
3161 "WriteEngine must acquire lock after the previous engine closed: {:?}",
3162 engine2.err()
3163 );
3164 }
3165
3166 /// After the first engine is dropped (without calling close()), a new engine
3167 /// must also successfully acquire the lock.
3168 #[test]
3169 fn test_write_dir_lock_reacquired_after_drop() {
3170 let temp_dir = TempDir::new().unwrap();
3171 let schema = create_test_schema();
3172
3173 let config1 = WriteEngineConfig::new(
3174 temp_dir.path().join("data"),
3175 temp_dir.path().join("wal"),
3176 schema.clone(),
3177 );
3178
3179 // First engine acquires the lock; drop releases it.
3180 {
3181 let _engine1 = WriteEngine::new(config1).unwrap();
3182 // engine1 dropped here, lock released
3183 }
3184
3185 // A new engine on the same directory must now succeed.
3186 let config2 = WriteEngineConfig::new(
3187 temp_dir.path().join("data"),
3188 temp_dir.path().join("wal"),
3189 schema,
3190 );
3191
3192 let engine2 = WriteEngine::new(config2);
3193 assert!(
3194 engine2.is_ok(),
3195 "WriteEngine must acquire lock after the previous engine was dropped: {:?}",
3196 engine2.err()
3197 );
3198 }
3199
3200 // Issue #1392 (roborev r4): when the WAL truncate fails AFTER `set_len(0)`
3201 // has already zeroed the WAL, the just-published SSTable is the ONLY durable
3202 // copy of the flushed mutations. The flush must therefore treat the SSTable
3203 // handoff as COMMITTED for state purposes — clear the memtable and advance
3204 // the generation BEFORE surfacing the error — so that a retry writes a NEW
3205 // generation instead of `File::create`-overwriting the published SSTable
3206 // (which would lose the data on a crash mid-retry).
3207 //
3208 // This drives the real `RealDurabilityBarrier` path via the WAL's test-only
3209 // post-`set_len(0)` sync fault, then proves: (1) the flush surfaces the
3210 // error, (2) state is committed (memtable cleared, generation advanced),
3211 // (3) the published gen-1 SSTable survives byte-for-byte, (4) a subsequent
3212 // flush writes a NEW generation, and (5) a full read returns each mutation
3213 // exactly once — crash-safe, no loss, no duplication.
3214 #[tokio::test]
3215 async fn post_mutation_truncate_failure_commits_and_retry_writes_new_generation() {
3216 use crate::platform::Platform;
3217 use crate::storage::sstable::SSTableManager;
3218 use crate::Config;
3219 use std::sync::Arc;
3220
3221 let temp_dir = TempDir::new().unwrap();
3222 let schema = create_test_schema();
3223 let data_dir = temp_dir.path().join("data");
3224 let config = WriteEngineConfig::new(
3225 data_dir.clone(),
3226 temp_dir.path().join("wal"),
3227 schema.clone(),
3228 );
3229 let mut engine = WriteEngine::new(config).unwrap();
3230
3231 // Write the first mutation, then arm a post-set_len(0) WAL-truncate fault
3232 // so the flush finalize hits the `AfterMutation` path.
3233 engine
3234 .write(create_test_mutation(1, "Alice", 1_000_000))
3235 .unwrap();
3236 let gen_before = engine.generation();
3237 engine.wal.set_fail_sync_after_truncate(true);
3238
3239 let err = engine
3240 .flush()
3241 .await
3242 .expect_err("a post-mutation WAL-truncate failure must surface as an error");
3243 assert!(
3244 matches!(err, Error::Storage(_)),
3245 "post-mutation truncate failure must surface as a storage error, got {err:?}"
3246 );
3247
3248 // State is COMMITTED despite the error: the memtable is cleared and the
3249 // generation has advanced, so a retry cannot reuse the published gen.
3250 assert_eq!(
3251 engine.memtable_row_count(),
3252 0,
3253 "memtable must be cleared once the SSTable is durably published"
3254 );
3255 assert_eq!(
3256 engine.generation(),
3257 gen_before + 1,
3258 "generation must advance so a retry writes a NEW generation"
3259 );
3260
3261 // The published gen-1 SSTable exists on disk. Snapshot its bytes to prove
3262 // a later flush does NOT overwrite it.
3263 let sstable_dir = data_dir.join("test_ks").join("test_table");
3264 let gen1_data = sstable_dir.join("nb-1-big-Data.db");
3265 assert!(
3266 gen1_data.exists(),
3267 "the published gen-1 Data.db must exist on disk after the faulted flush"
3268 );
3269 let gen1_bytes = std::fs::read(&gen1_data).unwrap();
3270
3271 // Disarm the fault and flush a SECOND mutation: it must write a NEW
3272 // generation and leave the first SSTable byte-for-byte intact.
3273 engine.wal.set_fail_sync_after_truncate(false);
3274 engine
3275 .write(create_test_mutation(2, "Bob", 2_000_000))
3276 .unwrap();
3277 engine
3278 .flush()
3279 .await
3280 .expect("second flush must succeed")
3281 .expect("second flush must produce an SSTable");
3282
3283 let gen2_data = sstable_dir.join("nb-2-big-Data.db");
3284 assert!(
3285 gen2_data.exists(),
3286 "the retry must write a NEW generation (nb-2), not overwrite nb-1"
3287 );
3288 assert_eq!(
3289 std::fs::read(&gen1_data).unwrap(),
3290 gen1_bytes,
3291 "the retry must NOT overwrite the published gen-1 SSTable"
3292 );
3293
3294 // Full read across both generations: each mutation appears exactly once.
3295 let cqlite_config = Config::default();
3296 let platform = Arc::new(Platform::new(&cqlite_config).await.unwrap());
3297 let manager = SSTableManager::new(
3298 &data_dir,
3299 &cqlite_config,
3300 platform,
3301 #[cfg(feature = "state_machine")]
3302 None,
3303 )
3304 .await
3305 .expect("SSTableManager must load both published generations");
3306
3307 let table_id = crate::types::TableId::from("test_ks.test_table");
3308 let rows = manager
3309 .scan(&table_id, None, None, None, Some(&schema))
3310 .await
3311 .expect("scan across both generations must succeed");
3312 assert_eq!(
3313 rows.len(),
3314 2,
3315 "both mutations must be readable exactly once (no loss, no duplication)"
3316 );
3317 }
3318
3319 /// Thread-local `tracing` WARN capturer for the Drop-warn tests (#1693).
3320 /// `start()` installs a thread-local `tracing` subscriber with NO
3321 /// `tracing-log` bridge (the tracing-only path an embedder wires) that
3322 /// records WARN messages per-thread, isolating concurrent tests.
3323 mod drop_warn_capture {
3324 use std::cell::RefCell;
3325 use std::fmt;
3326 use tracing::field::{Field, Visit};
3327 use tracing::subscriber::{set_default, DefaultGuard};
3328 use tracing::{Event, Level, Subscriber};
3329 use tracing_subscriber::layer::{Context, Layer, SubscriberExt};
3330 use tracing_subscriber::Registry;
3331
3332 thread_local! {
3333 static BUFFER: RefCell<Option<Vec<String>>> = const { RefCell::new(None) };
3334 static GUARD: RefCell<Option<DefaultGuard>> = const { RefCell::new(None) };
3335 }
3336 #[derive(Default)]
3337 struct MessageVisitor {
3338 message: String,
3339 }
3340 impl Visit for MessageVisitor {
3341 fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
3342 if field.name() == "message" {
3343 self.message = format!("{value:?}");
3344 }
3345 }
3346 }
3347 struct WarnCaptureLayer;
3348 impl<S: Subscriber> Layer<S> for WarnCaptureLayer {
3349 fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
3350 if *event.metadata().level() != Level::WARN {
3351 return;
3352 }
3353 let mut visitor = MessageVisitor::default();
3354 event.record(&mut visitor);
3355 BUFFER.with(|b| {
3356 if let Some(buf) = b.borrow_mut().as_mut() {
3357 buf.push(visitor.message);
3358 }
3359 });
3360 }
3361 }
3362
3363 pub(super) fn start() {
3364 BUFFER.with(|b| *b.borrow_mut() = Some(Vec::new()));
3365 let guard = set_default(Registry::default().with(WarnCaptureLayer));
3366 GUARD.with(|g| *g.borrow_mut() = Some(guard));
3367 }
3368
3369 pub(super) fn take_warnings() -> Vec<String> {
3370 GUARD.with(|g| *g.borrow_mut() = None);
3371 BUFFER.with(|b| b.borrow_mut().take().unwrap_or_default())
3372 }
3373 }
3374
3375 /// Issue #1693 (AG4): dropping a `WriteEngine` whose memtable still holds
3376 /// un-flushed rows (i.e. `close()` was never called) must emit a `warn!` so
3377 /// the silent data-loss mode is visible in logs.
3378 #[test]
3379 fn test_drop_without_close_warns_on_nonempty_memtable() {
3380 drop_warn_capture::start();
3381
3382 let temp_dir = TempDir::new().unwrap();
3383 let schema = create_test_schema();
3384 let config = WriteEngineConfig::new(
3385 temp_dir.path().join("data"),
3386 temp_dir.path().join("wal"),
3387 schema,
3388 );
3389 let mut engine = WriteEngine::new(config).unwrap();
3390 engine
3391 .write(create_test_mutation(1, "Alice", 1_000_000))
3392 .unwrap();
3393 assert!(
3394 engine.memtable_row_count() > 0,
3395 "precondition: memtable must be non-empty before drop"
3396 );
3397
3398 // Drop WITHOUT close() — the durability-loss path.
3399 drop(engine);
3400
3401 let warnings = drop_warn_capture::take_warnings();
3402 assert!(
3403 warnings
3404 .iter()
3405 .any(|m| m.contains("dropped") && m.contains("without close")),
3406 "expected an unflushed-drop warning, captured warnings: {warnings:?}"
3407 );
3408 // WAL-durable engine: the guidance must point at WAL recovery, NOT loss.
3409 assert!(
3410 warnings.iter().any(|m| m.contains("WAL replay")),
3411 "WAL-durable drop warning must mention WAL replay recovery, captured: {warnings:?}"
3412 );
3413 assert!(
3414 !warnings.iter().any(|m| m.contains("LOST")),
3415 "WAL-durable drop warning must NOT claim data loss, captured: {warnings:?}"
3416 );
3417 }
3418
3419 /// Issue #1693 (roborev): with `Durability::Disabled` the WAL is skipped, so
3420 /// an ungraceful drop loses the un-flushed rows. The warning must say the
3421 /// rows are LOST rather than giving false "recoverable from the WAL" guidance.
3422 #[test]
3423 fn test_drop_without_close_warns_data_loss_when_durability_disabled() {
3424 drop_warn_capture::start();
3425
3426 let temp_dir = TempDir::new().unwrap();
3427 let schema = create_test_schema();
3428 let config = WriteEngineConfig::new(
3429 temp_dir.path().join("data"),
3430 temp_dir.path().join("wal"),
3431 schema,
3432 )
3433 .with_durability(Durability::Disabled);
3434 let mut engine = WriteEngine::new(config).unwrap();
3435 engine
3436 .write(create_test_mutation(1, "Alice", 1_000_000))
3437 .unwrap();
3438 assert!(
3439 engine.memtable_row_count() > 0,
3440 "precondition: memtable must be non-empty before drop"
3441 );
3442
3443 drop(engine);
3444
3445 let warnings = drop_warn_capture::take_warnings();
3446 assert!(
3447 warnings
3448 .iter()
3449 .any(|m| m.contains("dropped") && m.contains("without close")),
3450 "expected an unflushed-drop warning, captured warnings: {warnings:?}"
3451 );
3452 // No-WAL engine: the warning must signal data loss, NOT WAL recovery.
3453 assert!(
3454 warnings.iter().any(|m| m.contains("LOST")),
3455 "Durability::Disabled drop warning must signal data loss, captured: {warnings:?}"
3456 );
3457 assert!(
3458 !warnings.iter().any(|m| m.contains("WAL replay")),
3459 "Durability::Disabled drop warning must NOT promise WAL recovery, captured: {warnings:?}"
3460 );
3461 }
3462
3463 /// Control for issue #1693: after a graceful `close()` the memtable is empty,
3464 /// so `Drop` must NOT emit the unflushed-drop warning.
3465 #[tokio::test]
3466 async fn test_drop_after_close_does_not_warn() {
3467 drop_warn_capture::start();
3468
3469 let temp_dir = TempDir::new().unwrap();
3470 let schema = create_test_schema();
3471 let config = WriteEngineConfig::new(
3472 temp_dir.path().join("data"),
3473 temp_dir.path().join("wal"),
3474 schema,
3475 );
3476 let mut engine = WriteEngine::new(config).unwrap();
3477 engine
3478 .write(create_test_mutation(1, "Alice", 1_000_000))
3479 .unwrap();
3480 engine.close().await.unwrap();
3481
3482 drop(engine);
3483
3484 let warnings = drop_warn_capture::take_warnings();
3485 assert!(
3486 !warnings
3487 .iter()
3488 .any(|m| m.contains("dropped") && m.contains("without close")),
3489 "close() flushed the memtable; Drop must not warn, captured: {warnings:?}"
3490 );
3491 }
3492}