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