From 6cac58fd4f88e1c6f6972a6b76a6e8c7e1e38a9d Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 11:32:23 +0700 Subject: [PATCH 01/15] fix(disk-offload): block instead of dropping spill completions (data loss) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A dropped SpillCompletion is permanent data loss, not a benign event as the old comment claimed. By the time a completion is produced the key has already been evicted from RAM, and the manifest add_file + cold_index insert happen ONLY when the event loop consumes the completion (apply_spill_completions). So try_send(Full) -> drop left the .mpf file orphaned on disk (never recorded in the manifest) and lost its keys permanently after restart. CodeRabbit flagged this on PR #136 (review 4420584384); validated against the merged code. Fixes (issue #137): - send_one_completion now blocks (send_timeout loop) until the event loop drains a slot, applying backpressure instead of dropping. Deadlock-free: the dedicated spill thread is the only blocker and the eviction path enqueues SpillRequests with try_send (and refuses to free RAM unless the enqueue succeeds), so there is no cyclic wait. stop_flag keeps shutdown from wedging the join. - shutdown() now drains and returns the thread's final-flush completions so the caller (drain_and_shutdown_spill) can still apply them — closing the shutdown-time loss window. - The run loop drains request_rx on stop_flag before its final flush, so spills queued at shutdown are processed rather than abandoned in the channel. - persistence_tick: apply_spill_completions + the shutdown path share a new apply_completion_vec helper. Validated in OrbStack VM (cgroup-capped), red->green TDD: - full_completion_channel_blocks_instead_of_dropping (RED: 1/2 delivered -> GREEN: 2/2, none dropped) - shutdown_drains_and_returns_unapplied_completions (RED: 0 -> GREEN: 1) - full spill_thread module 10/10 green; fmt + clippy clean under default and --no-default-features --features runtime-tokio,jemalloc. author: Tin Dang --- src/shard/persistence_tick.rs | 18 ++- src/storage/tiered/spill_thread.rs | 190 ++++++++++++++++++++++++----- 2 files changed, 179 insertions(+), 29 deletions(-) diff --git a/src/shard/persistence_tick.rs b/src/shard/persistence_tick.rs index 00441de6..4745a734 100644 --- a/src/shard/persistence_tick.rs +++ b/src/shard/persistence_tick.rs @@ -369,7 +369,11 @@ pub(crate) fn drain_and_shutdown_spill( apply_spill_completions(spill_t, shard_manifest, shard_databases, shard_id); } if let Some(st) = spill_thread.take() { - st.shutdown(); + // shutdown() returns any completions from the thread's final buffer + // flush that the drain above did not see; apply them so those cold keys + // are not lost (file on disk but never recorded in the manifest). + let leftover = st.shutdown(); + apply_completion_vec(leftover, shard_manifest, shard_databases, shard_id); tracing::info!("Shard {}: spill background thread shut down", shard_id); } } @@ -387,6 +391,18 @@ pub(crate) fn apply_spill_completions( shard_id: usize, ) { let completions = spill_thread.drain_completions(); + apply_completion_vec(completions, shard_manifest, shard_databases, shard_id); +} + +/// Apply a batch of spill completions: ONE manifest `add_file`+commit per file, +/// one `cold_index` insert per KV entry within it. Shared by the live drain +/// (`apply_spill_completions`) and the shutdown final-flush drain. +fn apply_completion_vec( + completions: Vec, + shard_manifest: &mut Option, + shard_databases: &std::sync::Arc, + shard_id: usize, +) { if completions.is_empty() { return; } diff --git a/src/storage/tiered/spill_thread.rs b/src/storage/tiered/spill_thread.rs index 0c6f5282..5c38f13c 100644 --- a/src/storage/tiered/spill_thread.rs +++ b/src/storage/tiered/spill_thread.rs @@ -357,10 +357,15 @@ impl SpillThread { let mut buffer: Vec = Vec::with_capacity(FLUSH_ENTRY_CAP); loop { - // Check stop flag — but flush the buffer before exiting. + // Check stop flag — drain any still-queued requests and flush them + // before exiting, so spills queued at shutdown are not lost (their + // keys may already be evicted from RAM). if stop_flag.load(Ordering::Acquire) { + while let Ok(req) = request_rx.try_recv() { + buffer.push(req); + } if !buffer.is_empty() { - Self::send_completions(&completion_tx, flush_buffer(&mut buffer)); + Self::send_completions(&completion_tx, flush_buffer(&mut buffer), &stop_flag); } break; } @@ -369,19 +374,31 @@ impl SpillThread { Ok(req) => { buffer.push(req); if buffer.len() >= FLUSH_ENTRY_CAP { - Self::send_completions(&completion_tx, flush_buffer(&mut buffer)); + Self::send_completions( + &completion_tx, + flush_buffer(&mut buffer), + &stop_flag, + ); } } Err(flume::RecvTimeoutError::Timeout) => { // Latency guard: flush non-empty buffer on tick. if !buffer.is_empty() { - Self::send_completions(&completion_tx, flush_buffer(&mut buffer)); + Self::send_completions( + &completion_tx, + flush_buffer(&mut buffer), + &stop_flag, + ); } } Err(flume::RecvTimeoutError::Disconnected) => { // Drain guard: flush remaining entries then exit. if !buffer.is_empty() { - Self::send_completions(&completion_tx, flush_buffer(&mut buffer)); + Self::send_completions( + &completion_tx, + flush_buffer(&mut buffer), + &stop_flag, + ); } break; } @@ -389,32 +406,63 @@ impl SpillThread { } } - /// Send multiple completions, dropping on full channel and bumping the counter. + /// Send multiple completions, applying backpressure on a full channel. fn send_completions( completion_tx: &flume::Sender, completions: Vec, + stop_flag: &Arc, ) { for completion in completions { - Self::send_one_completion(completion_tx, completion); + Self::send_one_completion(completion_tx, completion, stop_flag); } } - /// Send a single completion, dropping on full channel and bumping the counter. + /// Send a single completion to the event loop. + /// + /// Dropping a completion is **data loss**: by the time it is produced the + /// key has already been evicted from RAM, and the manifest `add_file` + + /// `cold_index` insert happen only when the event loop consumes the + /// completion (`apply_spill_completions`). A dropped completion therefore + /// orphans the on-disk `.mpf` file (never recorded in the manifest) and + /// loses its keys permanently after restart. + /// + /// So we block until the event loop drains a slot rather than dropping. + /// This is safe: the dedicated spill thread is the only blocker, and the + /// eviction path enqueues `SpillRequest`s with `try_send` (never blocking + /// on us), so there is no cyclic wait. We honour `stop_flag` so shutdown + /// cannot wedge the join — and because `shutdown()` drains the channel + /// after the join, the final-flush completions are still applied. fn send_one_completion( completion_tx: &flume::Sender, completion: SpillCompletion, + stop_flag: &Arc, ) { - match completion_tx.try_send(completion) { - Ok(()) => {} - Err(flume::TrySendError::Full(_)) => { - SPILL_COMPLETION_DROPPED.fetch_add(1, Ordering::Relaxed); - warn!( - "spill_thread: completion channel full, dropping completion (total dropped: {})", - SPILL_COMPLETION_DROPPED.load(Ordering::Relaxed) - ); - } - Err(flume::TrySendError::Disconnected(_)) => { - // Event loop dropped its receiver -- shutting down; ignore. + let mut pending = completion; + loop { + match completion_tx.send_timeout(pending, std::time::Duration::from_millis(100)) { + Ok(()) => return, + Err(flume::SendTimeoutError::Timeout(c)) => { + if stop_flag.load(Ordering::Acquire) { + // Shutting down and the event loop is no longer draining; + // we cannot persist the manifest entry ourselves. Count + + // warn. In practice unreachable: shutdown() drains the + // channel before the final flush, and that flush is bounded + // by FLUSH_ENTRY_CAP (< channel capacity), so it always fits. + SPILL_COMPLETION_DROPPED.fetch_add(1, Ordering::Relaxed); + warn!( + "spill_thread: completion channel full at shutdown; dropping completion (total dropped: {})", + SPILL_COMPLETION_DROPPED.load(Ordering::Relaxed) + ); + return; + } + // Backpressure: the event loop will drain a slot. Retry rather + // than drop — a dropped completion is permanent key loss. + pending = c; + } + Err(flume::SendTimeoutError::Disconnected(_)) => { + // Event loop dropped its receiver -- shutting down; ignore. + return; + } } } } @@ -444,11 +492,18 @@ impl SpillThread { /// still alive: the background thread polls the flag every 100 ms and /// exits without waiting for channel close. This avoids the deadlock where /// connection futures held cloned senders past shutdown. - pub fn shutdown(mut self) { + /// Returns any completions produced by the thread's final buffer flush that + /// the event loop never drained, so the caller can still apply them + /// (manifest `add_file` + `cold_index` insert) instead of losing those keys. + #[must_use] + pub fn shutdown(mut self) -> Vec { self.stop_flag.store(true, Ordering::Release); if let Some(handle) = self.join_handle.take() { let _ = handle.join(); } + // Thread has exited; surface any completions from its final buffer flush + // that the event loop never drained, so the caller can still apply them. + self.completion_rx.try_iter().collect() } } @@ -486,7 +541,7 @@ mod tests { let st = SpillThread::new(0); assert!(!st.request_tx.is_disconnected()); assert!(!st.completion_rx.is_disconnected()); - st.shutdown(); + let _ = st.shutdown(); } /// Single request produces a successful per-FILE completion with one entry. @@ -551,7 +606,86 @@ mod tests { _ => panic!("expected String"), } - st.shutdown(); + let _ = st.shutdown(); + } + + /// A full completion channel must apply backpressure (block until the event + /// loop drains a slot), never drop — a dropped completion permanently loses + /// the already-evicted keys (orphaned `.mpf`, never recorded in the manifest). + #[test] + fn full_completion_channel_blocks_instead_of_dropping() { + use std::sync::atomic::AtomicUsize; + let (tx, rx) = flume::bounded::(1); + let stop = Arc::new(AtomicBool::new(false)); + let dummy = || SpillCompletion { + file_entry: make_file_entry(7, 1, PAGE_4K as u64), + entries: Vec::new(), + success: true, + }; + + // Saturate the single slot so the next send must wait for a free slot. + tx.try_send(dummy()).unwrap(); + + let received = Arc::new(AtomicUsize::new(0)); + let r2 = received.clone(); + let drainer = std::thread::spawn(move || { + // Delay draining so send_one_completion is forced to block first. + std::thread::sleep(std::time::Duration::from_millis(50)); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2); + while std::time::Instant::now() < deadline { + if rx + .recv_timeout(std::time::Duration::from_millis(50)) + .is_ok() + { + r2.fetch_add(1, Ordering::Relaxed); + } else if r2.load(Ordering::Relaxed) >= 2 { + break; + } + } + }); + + // Must block until a slot frees, then deliver — never drop on Full. + SpillThread::send_one_completion(&tx, dummy(), &stop); + drop(tx); + drainer.join().unwrap(); + + assert_eq!( + received.load(Ordering::Relaxed), + 2, + "both completions must be delivered; a Full channel must block, not drop" + ); + } + + /// `shutdown()` must surface the thread's final-flush completions so the + /// caller can still apply them — they would otherwise be silently lost + /// (file on disk, never added to the manifest). + #[test] + fn shutdown_drains_and_returns_unapplied_completions() { + let tmp = tempfile::tempdir().unwrap(); + let st = SpillThread::new(9); + let sender = st.sender(); + sender + .send(SpillRequest { + key: Bytes::from_static(b"shutdown_key"), + db_index: 0, + value_bytes: Bytes::from_static(b"v"), + value_type: ValueType::String, + flags: 0, + ttl_ms: None, + file_id: 1, + shard_dir: tmp.path().to_path_buf(), + }) + .unwrap(); + drop(sender); + + // Deliberately do NOT drain via the event loop. The completion is + // produced by the thread's final flush; shutdown() must return it. + let leftover = st.shutdown(); + let total: usize = leftover.iter().map(|c| c.entries.len()).sum(); + assert_eq!( + total, 1, + "shutdown must return the unapplied completion, not drop it" + ); } #[test] @@ -584,7 +718,7 @@ mod tests { assert!(c.success); assert_eq!(c.file_entry.file_type, PageType::KvLeaf as u8); - st.shutdown(); + let _ = st.shutdown(); } #[test] @@ -592,7 +726,7 @@ mod tests { let st = SpillThread::new(3); let sender = st.sender(); drop(sender); - st.shutdown(); + let _ = st.shutdown(); // Reaching here without hang = clean exit. } @@ -645,7 +779,7 @@ mod tests { } } - st.shutdown(); + let _ = st.shutdown(); } /// Full pipeline: 5 requests, verify round-trip via cold_read. @@ -680,7 +814,7 @@ mod tests { assert!(c.success); } - st.shutdown(); + let _ = st.shutdown(); } #[test] @@ -717,7 +851,7 @@ mod tests { assert_eq!(received, sent, "should receive all sent entries"); drop(sender); - st.shutdown(); + let _ = st.shutdown(); } #[test] @@ -742,7 +876,7 @@ mod tests { drop(sender); let start = std::time::Instant::now(); - st.shutdown(); + let _ = st.shutdown(); let elapsed = start.elapsed(); assert!( From deae60d6c94ee3af90a22f056c8503c10dc04e0a Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 11:36:54 +0700 Subject: [PATCH 02/15] fix(disk-offload): salvage inline-batch spill failures per-entry, not wholesale MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An entry that passes the `value_bytes.len() <= INLINE_MAX_VALUE_BYTES` (3500) pre-screen can still fail `build_kv_spill_batch` with InvalidData when it does not fit a fresh 4KB inline leaf (e.g. a large key plus an incompressible ~3.5KB value). The old code turned that single bad candidate into one wholesale `success: false` completion for the ENTIRE inline buffer — and those keys were already evicted from RAM, so the whole flush was lost. CodeRabbit flagged this on PR #136 (review 4420584384); validated against the merged code. Fix (issue #137): on inline-batch build OR write failure, fall back to writing each inline entry as its own single-page (+overflow) file via the existing `build_kv_spill_pages` path — the same path oversized entries already use, which succeeds where the batch fails because the value goes to overflow pages. Only an entry that cannot be written even with overflow (key too large for a leaf) now fails, and it fails in isolation. Extracted the per-entry logic into a shared `spill_single_entry` helper used by both the oversized loop and the fallback. Also corrected the stale `SPILL_COMPLETION_DROPPED` doc comment. Validated in OrbStack VM (cgroup-capped), red->green TDD: - inline_batch_failure_falls_back_to_per_entry_spill: a 800B key + incompressible 3500B value (RED: 0 salvaged -> GREEN: 1 salvaged, file readable on disk). - full spill_thread module 11/11 green; fmt + clippy clean under default and --no-default-features --features runtime-tokio,jemalloc. author: Tin Dang --- src/storage/tiered/spill_thread.rs | 188 +++++++++++++++++++---------- 1 file changed, 127 insertions(+), 61 deletions(-) diff --git a/src/storage/tiered/spill_thread.rs b/src/storage/tiered/spill_thread.rs index 5c38f13c..73bed015 100644 --- a/src/storage/tiered/spill_thread.rs +++ b/src/storage/tiered/spill_thread.rs @@ -32,10 +32,12 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; /// reasonable memory budget and keeps file sizes manageable. const FLUSH_ENTRY_CAP: usize = 256; -/// Cumulative count of `SpillCompletion`s dropped because the event-loop-side -/// completion channel was full. Each drop means the data is on disk but the -/// in-memory `cold_index` slot was not refreshed; the next checkpoint repairs -/// it from the manifest. +/// Cumulative count of `SpillCompletion`s dropped. Dropping is **data loss**: +/// the `.mpf` file is on disk but its manifest `add_file` + `cold_index` insert +/// happen only when the event loop consumes the completion, and the keys are +/// already evicted from RAM. The sender therefore blocks on a full channel +/// instead of dropping, so this only increments in the rare +/// shutdown-with-full-channel edge. Exposed for INFO / metrics scraping. static SPILL_COMPLETION_DROPPED: AtomicU64 = AtomicU64::new(0); /// Returns the cumulative number of dropped spill completions across all @@ -183,7 +185,7 @@ fn flush_buffer(buffer: &mut Vec) -> Vec { }) .collect(); - let completion = match build_kv_spill_batch(&spill_entries, file_id) { + match build_kv_spill_batch(&spill_entries, file_id) { Ok(batch) => { let total_pages = batch.leaves.len() as u32; // overflow is always empty (inline-only) match write_kv_spill_batch(&shard_dir, file_id, &batch) { @@ -198,23 +200,23 @@ fn flush_buffer(buffer: &mut Vec) -> Vec { slot_idx, }) .collect(); - SpillCompletion { + completions.push(SpillCompletion { file_entry: make_file_entry(file_id, total_pages, byte_size), entries, success: true, - } + }); } Err(e) => { warn!( file_id, error = %e, count = inline_indices.len(), - "spill_thread: inline batch write failed" + "spill_thread: inline batch write failed; falling back to per-entry spill" ); - SpillCompletion { - file_entry: make_file_entry(file_id, 0, 0), - entries: Vec::new(), - success: false, + // Salvage each entry individually rather than dropping the + // whole already-evicted flush. + for &i in &inline_indices { + completions.push(spill_single_entry(&buffer[i], buffer[i].file_id)); } } } @@ -224,63 +226,61 @@ fn flush_buffer(buffer: &mut Vec) -> Vec { file_id, error = %e, count = inline_indices.len(), - "spill_thread: inline batch build failed" + "spill_thread: inline batch build failed; falling back to per-entry spill" ); - SpillCompletion { - file_entry: make_file_entry(file_id, 0, 0), - entries: Vec::new(), - success: false, + // One entry that does not fit a fresh inline leaf must not drop the + // whole batch — write each via the overflow path instead. + for &i in &inline_indices { + completions.push(spill_single_entry(&buffer[i], buffer[i].file_id)); } } - }; - completions.push(completion); + } } // ── Write ONE single-page file per oversized entry ──────────────────────── for &i in &oversized_indices { - let req = &buffer[i]; - let file_id = req.file_id; - let shard_dir = req.shard_dir.clone(); - - let completion = match build_kv_spill_pages( - &req.key, - &req.value_bytes, - req.value_type, - req.flags, - req.ttl_ms, - file_id, - ) { - Ok(pages) => match write_kv_spill_pages(&shard_dir, file_id, &pages) { - Ok(byte_size) => SpillCompletion { - file_entry: make_file_entry(file_id, pages.total_pages, byte_size), - entries: vec![SpillCompletionEntry { - key: req.key.clone(), - db_index: req.db_index, - page_idx: 0, - slot_idx: 0, - }], - success: true, - }, - Err(e) => { - warn!( - file_id, - error = %e, - key_len = req.key.len(), - "spill_thread: oversized single-file write failed" - ); - SpillCompletion { - file_entry: make_file_entry(file_id, 0, 0), - entries: Vec::new(), - success: false, - } - } + completions.push(spill_single_entry(&buffer[i], buffer[i].file_id)); + } + + buffer.clear(); + completions +} + +/// Spill ONE entry as its own single-page (+overflow) file via +/// `build_kv_spill_pages`. +/// +/// Used both for oversized entries and as the inline-batch fallback: when an +/// inline batch fails (an entry that passed the `value_bytes.len()` pre-screen +/// still does not fit a fresh leaf), salvaging each entry here prevents one bad +/// candidate from dropping a whole flush whose keys are already out of RAM. +/// Returns a failed completion only if this single entry cannot be written even +/// with overflow pages (key too large for a leaf). +fn spill_single_entry(req: &SpillRequest, file_id: u64) -> SpillCompletion { + match build_kv_spill_pages( + &req.key, + &req.value_bytes, + req.value_type, + req.flags, + req.ttl_ms, + file_id, + ) { + Ok(pages) => match write_kv_spill_pages(&req.shard_dir, file_id, &pages) { + Ok(byte_size) => SpillCompletion { + file_entry: make_file_entry(file_id, pages.total_pages, byte_size), + entries: vec![SpillCompletionEntry { + key: req.key.clone(), + db_index: req.db_index, + page_idx: 0, + slot_idx: 0, + }], + success: true, }, Err(e) => { warn!( file_id, error = %e, key_len = req.key.len(), - "spill_thread: oversized single-file build failed (key too large)" + "spill_thread: single-file write failed" ); SpillCompletion { file_entry: make_file_entry(file_id, 0, 0), @@ -288,12 +288,21 @@ fn flush_buffer(buffer: &mut Vec) -> Vec { success: false, } } - }; - completions.push(completion); + }, + Err(e) => { + warn!( + file_id, + error = %e, + key_len = req.key.len(), + "spill_thread: single-file build failed (key too large)" + ); + SpillCompletion { + file_entry: make_file_entry(file_id, 0, 0), + entries: Vec::new(), + success: false, + } + } } - - buffer.clear(); - completions } /// Background thread that performs pwrite for evicted KV entries. @@ -688,6 +697,63 @@ mod tests { ); } + /// An entry that passes the `INLINE_MAX_VALUE_BYTES` pre-screen but does not + /// fit a fresh inline leaf (large key + incompressible value) makes + /// `build_kv_spill_batch` fail. That must NOT fail the whole inline flush — + /// the offender is salvaged via the per-entry (overflow) path instead, since + /// its key is already evicted from RAM. + #[test] + fn inline_batch_failure_falls_back_to_per_entry_spill() { + let tmp = tempfile::tempdir().unwrap(); + + // High-entropy (LZ4-incompressible) value at the inline threshold. + let mut value = Vec::with_capacity(INLINE_MAX_VALUE_BYTES); + let mut s: u32 = 0x1234_5678; + for _ in 0..INLINE_MAX_VALUE_BYTES { + s = s.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + value.push((s >> 24) as u8); + } + // Sizable key so key + value overflows a 4KB leaf (forces batch failure), + // yet the key alone fits a leaf with an overflow pointer (pages succeed). + let key = vec![b'k'; 800]; + + let mut buffer = vec![SpillRequest { + key: Bytes::from(key), + db_index: 0, + value_bytes: Bytes::from(value), + value_type: ValueType::String, + flags: 0, + ttl_ms: None, + file_id: 1, + shard_dir: tmp.path().to_path_buf(), + }]; + + let completions = flush_buffer(&mut buffer); + let succeeded: usize = completions + .iter() + .filter(|c| c.success) + .map(|c| c.entries.len()) + .sum(); + assert_eq!( + succeeded, 1, + "inline entry that overflows a leaf must be salvaged via per-entry fallback, not dropped" + ); + + // And the salvaged entry must be readable back from its on-disk file. + let c = completions + .iter() + .find(|c| c.success && !c.entries.is_empty()) + .expect("expected a successful fallback completion"); + let file_path = tmp + .path() + .join("data") + .join(format!("heap-{:06}.mpf", c.file_entry.file_id)); + assert!( + file_path.exists(), + "fallback spill file should exist on disk" + ); + } + #[test] fn test_spill_request_with_ttl() { let tmp = tempfile::tempdir().unwrap(); From cf908002dc4c4ebeebfc8989e935395ab9114d82 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 11:41:29 +0700 Subject: [PATCH 03/15] fix(disk-offload): preserve spill context across tokio connection migration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit spawn_tokio_connection, spawn_monoio_connection and spawn_migrated_monoio_ connection all thread the disk-offload spill context (spill_sender / spill_file_id / disk_offload_dir) into the ConnectionContext, but spawn_migrated_tokio_connection built it with None / Cell(0) / None. After an affinity migration the tokio connection therefore fell back to the non-spilling eviction path even when disk-offload was enabled, so cold-tier durability silently depended on whether the connection had migrated. CodeRabbit flagged this on PR #136 (review 4420584384); validated against the merged code (issue #141). Fix: thread spill_sender / spill_file_id / disk_offload_dir through spawn_migrated_tokio_connection and pass them to ConnectionContext::new, mirroring the three sibling spawn paths. Updated both call sites in event_loop.rs. Validation: this is mechanical parity with three already-correct sibling functions in the same file. A genuine red/green unit test is disproportionate here — ConnectionContext::new takes 26 args and is consumed inside a spawned task, so asserting the migrated context's spill fields would require stubbing the entire context or simulating a live cross-shard connection migration. Verified instead by: clean build + clippy under both default (monoio) and --no-default-features --features runtime-tokio,jemalloc, and structural parity (all four spawn paths now pass the spill context). The behavioural guarantee is covered by the existing disk-offload read-through integration suite. author: Tin Dang --- src/shard/conn_accept.rs | 16 +++++++++++++--- src/shard/event_loop.rs | 2 ++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/src/shard/conn_accept.rs b/src/shard/conn_accept.rs index f3352a59..351ad40f 100644 --- a/src/shard/conn_accept.rs +++ b/src/shard/conn_accept.rs @@ -301,6 +301,9 @@ pub(crate) fn spawn_migrated_tokio_connection( shard_id: usize, num_shards: usize, config_port: u16, + spill_sender: &Option>, + spill_file_id: &Rc>, + disk_offload_dir: &Option, ) { use std::os::unix::io::FromRawFd; @@ -365,6 +368,13 @@ pub(crate) fn spawn_migrated_tokio_connection( // Pool is built by the spawn site and threaded through here. let pool_for_ctx = aof_pool.as_ref().map(Arc::clone); + // Preserve the disk-offload spill context across migration — without + // this a migrated connection silently falls back to the non-spilling + // eviction path even when disk-offload is enabled, so cold-tier + // durability would depend on whether the connection had migrated. + let spill_tx = spill_sender.clone(); + let spill_fid = spill_file_id.clone(); + let do_dir = disk_offload_dir.clone(); let conn_ctx = crate::server::conn::ConnectionContext::new( sdbs, shard_id, @@ -390,9 +400,9 @@ pub(crate) fn spawn_migrated_tokio_connection( all_regs, all_rsm, aff, - None, // spill_sender - Rc::new(std::cell::Cell::new(0)), // spill_file_id - None, // disk_offload_dir + spill_tx, // spill_sender (preserved across migration) + spill_fid, // spill_file_id + do_dir, // disk_offload_dir ); // State restoration happens directly via migrated_state parameter — diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index f2bffa41..34d9f239 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -1244,6 +1244,7 @@ impl super::Shard { &cached_clock, &remote_sub_map_arc, &all_pubsub_registries, &all_remote_sub_maps, &affinity_tracker, shard_id, num_shards, config_port, + &spill_sender, &spill_file_id, &disk_offload_dir, ); } #[cfg(feature = "runtime-monoio")] @@ -1343,6 +1344,7 @@ impl super::Shard { &cached_clock, &remote_sub_map_arc, &all_pubsub_registries, &all_remote_sub_maps, &affinity_tracker, shard_id, num_shards, config_port, + &spill_sender, &spill_file_id, &disk_offload_dir, ); } #[cfg(feature = "runtime-monoio")] From 4828b3d723dcdcbdbb73248d6bced675624cffd5 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 11:46:48 +0700 Subject: [PATCH 04/15] fix(embedded): recover cold tier under appendonly=no + disable per-shard accept MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit run_embedded() diverged from the production main.rs startup path on two functional points (CodeRabbit PR #136 review 4420584384, issue #142): 1. Recovery only ran when `persistence_dir` was set, so `--appendonly no` + disk-offload (persistence_dir = None, disk_offload_base = Some) skipped `Shard::restore_from_persistence` entirely — cold-tier keys were unrecoverable after an embedded restart. Now gated on `should_run_recovery(persistence_dir, disk_offload_enabled)` = persistence OR disk-offload, recovering from `config.dir` when persistence_dir is None (mirrors main.rs:702-704). 2. Embedded forced `per_shard_accept = cfg!(target_os = "linux")`, re-enabling the tokio/Linux per-shard accept bind-order race that PR #136 disabled in main.rs. Now `false` (central accept), matching main.rs. Validated in OrbStack VM (cgroup-capped), red->green TDD under --no-default-features --features runtime-tokio,jemalloc (embedded is runtime-tokio-only): recovery_runs_for_disk_offload_without_persistence (RED with the disk-offload term dropped -> GREEN). fmt + clippy clean under both default (monoio) and tokio,jemalloc. Note: a full embedded+disk-offload+restart integration test was not added — there is no run_embedded test harness to extend and standing it up for a 2-line parity fix is disproportionate; the behavioural recovery path is the same one covered by tests/crash_recovery_disk_offload_no_aof.rs on main.rs. The extracted should_run_recovery predicate unit-tests the exact gate that regressed. author: Tin Dang --- src/server/embedded.rs | 49 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 45 insertions(+), 4 deletions(-) diff --git a/src/server/embedded.rs b/src/server/embedded.rs index 01ca43be..7cd1c750 100644 --- a/src/server/embedded.rs +++ b/src/server/embedded.rs @@ -57,6 +57,18 @@ use crate::shard::Shard; use crate::shard::mesh::{CHANNEL_BUFFER_SIZE, ChannelMesh}; use crate::shard::shared_databases::ShardDatabases; +/// Whether pre-loop recovery (`Shard::restore_from_persistence`) must run. +/// +/// Recovery is required when EITHER durable persistence (`persistence_dir`) OR +/// disk-offload is configured. The disk-offload case is the subtle one: under +/// `--appendonly no` + disk-offload, `persistence_dir` is `None`, yet the cold +/// tier still has on-disk state whose baseline must be restored on restart — +/// skipping it leaves cold keys unrecoverable. Mirrors the `main.rs` gate. +#[inline] +fn should_run_recovery(persistence_dir: Option<&str>, disk_offload_enabled: bool) -> bool { + persistence_dir.is_some() || disk_offload_enabled +} + /// Run an embedded sharded Moon server until `cancel` is fired. /// /// Behaves like the production `main.rs` startup path but with cluster, @@ -218,8 +230,13 @@ pub async fn run_embedded( config.initial_keyspace_hint, config.to_runtime_config(), ); - if let Some(ref dir) = persistence_dir { - shard.restore_from_persistence(dir, disk_offload_base.as_deref()); + // Recover whenever persistence OR disk-offload is configured. Under + // `--appendonly no` + disk-offload, persistence_dir is None but the + // cold tier still needs its baseline restored on restart, else cold + // keys are unrecoverable (mirrors main.rs). + if should_run_recovery(persistence_dir.as_deref(), disk_offload_base.is_some()) { + let recover_dir = persistence_dir.as_deref().unwrap_or(config.dir.as_str()); + shard.restore_from_persistence(recover_dir, disk_offload_base.as_deref()); } if let Some(ref offload_base) = disk_offload_base { let shard_dir = offload_base.join(format!("shard-{}", id)); @@ -355,8 +372,10 @@ pub async fn run_embedded( } let _ = snap_tx; // keep the watch sender alive for the shard's snap_rx clones - // Run the sharded listener until cancelled. - let per_shard_accept = cfg!(target_os = "linux"); + // Run the sharded listener until cancelled. Use the central accept path + // (per_shard_accept = false) like main.rs: the per-shard tokio/Linux accept + // loop has a bind-order race that can hang serving (fixed in PR #136). + let per_shard_accept = false; let listener_result = server::listener::run_sharded( config, conn_txs, @@ -463,3 +482,25 @@ fn panic_message<'a>(payload: &'a Box) -> &' "" } } + +#[cfg(test)] +mod tests { + use super::should_run_recovery; + + #[test] + fn recovery_runs_for_disk_offload_without_persistence() { + // The regression: `--appendonly no` + disk-offload means persistence_dir + // is None but the cold tier must still be recovered on restart. + assert!( + should_run_recovery(None, true), + "disk-offload-only (appendonly=no) must still run cold-tier recovery" + ); + // Other combinations behave as before. + assert!(should_run_recovery(Some("/var/lib/moon"), false)); + assert!(should_run_recovery(Some("/var/lib/moon"), true)); + assert!( + !should_run_recovery(None, false), + "no persistence and no disk-offload: nothing to recover" + ); + } +} From fa9afe909d2a2bbeb03fab0eb9fb17dd58ec883b Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 12:00:02 +0700 Subject: [PATCH 05/15] fix(persistence): clear rewrite flag when per-shard fan-out fails partway MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit try_send_rewrite_per_shard fans the RewritePerShard message out to every shard writer with a blocking send. If a writer channel is disconnected the send fails and the function returned Err immediately — but the writers that ALREADY received the message decrement the PerShardRewriteCoord countdown when they fold, while the failed writer and the unsent tail never do. The countdown therefore never reaches zero, no shard runs the terminal branch, and AOF_REWRITE_IN_PROGRESS (set by bgrewriteaof_start_sharded before dispatch) is wedged true forever — every future BGREWRITEAOF is silently rejected for the lifetime of the process. Same wedged-flag family CodeRabbit flagged on PR #136 review 4420584384 (related to #138). Fix: on send failure, call coord.mark_failed() then coord.shard_done() once for each shard from the failing index through the tail (the failed shard plus every unsent one). The already-dispatched writers still decrement on fold, so the countdown reaches zero exactly once; mark_failed() is published before the tail decrements, so whichever decrement is terminal takes the abort path — keeps old_seq authoritative (no manifest commit) and clears AOF_REWRITE_IN_PROGRESS. The error log now names the offending writer index. Validated in OrbStack VM (cgroup-capped), red/green TDD: rewrite_fan_out_partial_failure_clears_in_progress_flag disconnects shard-0's writer, asserts the fan-out returns Err AND the in-progress flag is cleared (RED before the accounting fix — flag stayed true; GREEN after). fmt + clippy clean under both default (monoio) and tokio,jemalloc. author: Tin Dang --- src/persistence/aof.rs | 58 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 55 insertions(+), 3 deletions(-) diff --git a/src/persistence/aof.rs b/src/persistence/aof.rs index 77b8d09b..6fd3660a 100644 --- a/src/persistence/aof.rs +++ b/src/persistence/aof.rs @@ -791,7 +791,7 @@ impl AofWriterPool { let n_shards = self.senders.len(); let shared_manifest = Arc::new(parking_lot::Mutex::new(manifest)); let coord = PerShardRewriteCoord::new(shared_manifest, current_seq, n_shards); - for s in &self.senders { + for (idx, s) in self.senders.iter().enumerate() { // Blocking send for guaranteed delivery — see the doc comment. if s.send(AofMessage::RewritePerShard { shard_dbs: shard_dbs.clone(), @@ -800,10 +800,21 @@ impl AofWriterPool { .is_err() { error!( - "F6 per-shard rewrite: a writer channel is disconnected; \ + "F6 per-shard rewrite: writer {} channel disconnected; \ rewrite aborted (no manifest commit, old generation remains \ - authoritative). Inspect AOF writer threads." + authoritative). Inspect AOF writer threads.", + idx ); + // Account for the shards that did NOT receive the message (this + // one plus the unsent tail). The writers that DID receive it will + // fold and call `shard_done`; decrementing for the rest here lets + // the countdown still reach zero, so the final `shard_done` runs + // the abort path (keeps `old_seq`, clears AOF_REWRITE_IN_PROGRESS) + // instead of wedging the rewrite flag forever. + coord.mark_failed(); + for _ in idx..n_shards { + coord.shard_done(); + } return Err(AofPoolSendError::SendFailed); } } @@ -846,6 +857,47 @@ mod pool_tests { use crate::persistence::aof_manifest::AofLayout; use crate::runtime::channel; + /// A per-shard BGREWRITEAOF fan-out that fails partway (a writer channel is + /// disconnected) must NOT leave `AOF_REWRITE_IN_PROGRESS` stuck true. The + /// failed/unsent shards are accounted for (`mark_failed` + `shard_done`) so + /// the countdown still reaches zero and the abort path clears the flag, + /// leaving the old generation authoritative. + #[test] + fn rewrite_fan_out_partial_failure_clears_in_progress_flag() { + use crate::command::persistence::AOF_REWRITE_IN_PROGRESS; + use std::sync::atomic::Ordering; + + let tmp = tempfile::tempdir().unwrap(); + // PerShard manifest on disk so the rewrite reaches the fan-out loop. + crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); + + // Disconnect shard-0's writer so the FIRST send fails (before any + // successful send), making the abort accounting deterministic. + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + drop(rx0); + let pool = AofWriterPool::per_shard_with_base_dir( + vec![tx0, tx1], + FsyncPolicy::EverySec, + DEFAULT_AOF_FSYNC_TIMEOUT, + tmp.path().to_path_buf(), + ); + + let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![ + vec![crate::storage::Database::new()], + vec![crate::storage::Database::new()], + ]); + + // The command handler sets this before dispatching the rewrite. + AOF_REWRITE_IN_PROGRESS.store(true, Ordering::SeqCst); + let res = pool.try_send_rewrite_per_shard(shard_dbs); + assert!(res.is_err(), "disconnected writer must fail the fan-out"); + assert!( + !AOF_REWRITE_IN_PROGRESS.load(Ordering::SeqCst), + "partial fan-out failure must clear AOF_REWRITE_IN_PROGRESS, not wedge it" + ); + } + #[test] fn top_level_pool_routes_all_shards_to_writer_zero() { let (tx, rx) = channel::mpsc_bounded::(8); From c919ff5d9291d3bf64f5cccb7e37b63f0b590b0b Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 12:18:27 +0700 Subject: [PATCH 06/15] fix(persistence): roll per-shard rewrite writers back to committed gen on abort MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A per-shard BGREWRITEAOF reopens each shard's append file onto the NEW generation in phase 6 (do_rewrite_per_shard) BEFORE the coordinator decides whether the rewrite commits. If the rewrite then aborts — another shard failed to fold, or the final manifest commit errored — the manifest keeps old_seq and the abort path prunes the new-seq files, but the already-folded writers keep appending into the now-discarded new-seq incr. Recovery resolves base/incr by the committed seq (old_seq), so every post-abort write is silently lost. The old code acknowledged this with a "RESTART recommended" log; a server that is not restarted loses data. CodeRabbit flagged this on PR #136 review 4420584384. Fix — barrier-before-resume. PerShardRewriteCoord now publishes the committed generation exactly once from the terminal shard_done (new_seq on success, old_seq on abort/commit-failure) via a parking_lot Mutex> + Condvar. After phase 7 each folded writer blocks in await_outcome and, if the committed seq != new_seq, reopens its append file onto the committed seq's incr (shard_incr_path_seq) — so it resumes appending into the authoritative generation with no restart and no loss. On the happy path committed == new_seq and the barrier is a no-op beyond unblocking. The barrier turns "every shard calls shard_done exactly once" into a LIVENESS requirement: a writer that decrements then blocks hangs forever if another shard never decrements. The one path that skips shard_done is a panic mid-fold (save_snapshot_to_bytes OOM-unwinding — issue #138), which under the default panic=unwind would now hang every healthy shard's AOF writer thread. Closed with a ShardDoneGuard: the fold creates it on entry and calls complete() on success (clean shard_done); on any ?-error or panic-unwind its Drop fires mark_failed + shard_done, so the countdown always closes, the terminal writer always publishes, and every waiter wakes and rolls back. The guard now owns the single decrement for ALL exits, so both writer-task error arms (monoio + tokio) no longer decrement after do_rewrite_per_shard (was a double-decrement risk). This incidentally closes the #138 panic-wedge. Each per-shard writer runs on its own dedicated OS thread (main.rs spawns aof-writer-{sid} → block_on_local), so blocking one at the barrier never starves the thread that must publish — no deadlock on either runtime. Validated in OrbStack VM (cgroup-capped), red/green TDD: - rewrite_abort_reopens_writer_onto_committed_old_generation: drives one shard's real fold with a second shard pre-failed; asserts the new-gen incr is pruned AND post-abort appends through *file land in the committed old-gen incr (RED before phase 8 — writes went to the pruned inode; GREEN after). - rewrite_abort_wakes_waiter_cross_thread: real condvar wait/notify across threads observes old_seq. - rewrite_fold_panic_releases_barrier_for_other_writers: a ShardDoneGuard dropped during catch_unwind releases a blocked waiter with old_seq (hangs forever without the guard). No regression: per-shard AOF SIGKILL recovery (crash_matrix_per_shard_aof, 3) and per-shard BGREWRITEAOF crash recovery (crash_matrix_per_shard_bgrewriteaof, 2) stay green. fmt + clippy clean under both default (monoio) and tokio,jemalloc. author: Tin Dang --- src/persistence/aof.rs | 331 ++++++++++++++++++++++++++++++++++++++--- 1 file changed, 312 insertions(+), 19 deletions(-) diff --git a/src/persistence/aof.rs b/src/persistence/aof.rs index 6fd3660a..b06bb9e6 100644 --- a/src/persistence/aof.rs +++ b/src/persistence/aof.rs @@ -234,6 +234,15 @@ pub struct PerShardRewriteCoord { /// refuse to start. On abort the old generation (`old_seq`) stays the /// committed state for all shards (crash-safe). failed: std::sync::atomic::AtomicBool, + /// The committed generation, published exactly once by the terminal writer + /// (`new_seq` on success, `old_seq` on abort/commit-failure). Every folded + /// writer blocks on this in `await_outcome` after `shard_done` so it can + /// reopen its append file onto the COMMITTED generation before resuming + /// appends — without this barrier a writer that reopened onto `new_seq` in + /// phase 6 keeps appending there after an abort, into a generation recovery + /// ignores (silent data loss; the old "RESTART recommended" hazard). + outcome: parking_lot::Mutex>, + outcome_cv: parking_lot::Condvar, } impl PerShardRewriteCoord { @@ -251,9 +260,41 @@ impl PerShardRewriteCoord { old_seq: current_seq, n_shards, failed: std::sync::atomic::AtomicBool::new(false), + outcome: parking_lot::Mutex::new(None), + outcome_cv: parking_lot::Condvar::new(), }) } + /// Publish the committed generation to every writer blocked in + /// `await_outcome`. Called exactly once, by the terminal `shard_done`. + #[inline] + fn publish_outcome(&self, committed_seq: u64) { + let mut slot = self.outcome.lock(); + *slot = Some(committed_seq); + self.outcome_cv.notify_all(); + } + + /// Block until the terminal writer publishes the committed generation, then + /// return it (`new_seq` on success, `old_seq` on abort/commit-failure). Each + /// folded writer calls this AFTER `shard_done` so it can reopen its append + /// file onto the committed generation. Safe against missed wake-ups: the + /// outcome is set under the same lock the condvar waits on, so a writer that + /// arrives after the publish observes `Some` and skips the wait. + /// + /// Deadlock-free: all `n_shards` writers call `shard_done` exactly once + /// (success via phase 7, fold-error/send-failure via the caller), so the + /// countdown always reaches zero and the terminal branch always publishes. + /// Each per-shard writer runs on its own dedicated OS thread, so blocking + /// one here never starves the thread that must publish. + #[must_use] + fn await_outcome(&self) -> u64 { + let mut slot = self.outcome.lock(); + while slot.is_none() { + self.outcome_cv.wait(&mut slot); + } + slot.expect("outcome is Some after the wait loop") + } + /// The generation writers advance to. #[inline] pub fn new_seq(&self) -> u64 { @@ -293,11 +334,14 @@ impl PerShardRewriteCoord { drop(m); error!( "F6 per-shard rewrite ABORTED: a shard failed to fold; seq stays {}. \ - Old generation remains authoritative (crash-safe). A RESTART is \ - recommended so successful shards' writers stop appending to the \ - discarded new generation.", + Old generation remains authoritative (crash-safe). Folded writers \ + roll their append files back to the old generation at the \ + barrier — no restart required.", self.old_seq ); + // Publish BEFORE clearing the in-progress flag so every blocked + // writer wakes and reopens onto the (still-authoritative) old gen. + self.publish_outcome(self.old_seq); crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); return; } @@ -311,6 +355,10 @@ impl PerShardRewriteCoord { ); // Do NOT prune — old generation is still the committed state. drop(m); + // Commit failed: old_seq is still authoritative, so folded + // writers must roll back to it (their phase-6 reopen targeted + // new_seq, which never committed). + self.publish_outcome(self.old_seq); crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); return; } @@ -322,11 +370,63 @@ impl PerShardRewriteCoord { "F6 per-shard rewrite complete: committed seq {} across {} shards, pruned seq {}", self.new_seq, self.n_shards, self.old_seq ); + // Success: new_seq is committed. Folded writers already reopened onto + // new_seq in phase 6, so the barrier is a no-op for them — but it must + // still unblock them. + self.publish_outcome(self.new_seq); crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); } } } +/// Panic-safety guard for a per-shard fold's `shard_done` obligation. +/// +/// The `await_outcome` barrier in `do_rewrite_per_shard` turns "every shard +/// calls `shard_done` exactly once" from a tidiness rule into a LIVENESS +/// requirement: any folded writer that decrements then blocks in `await_outcome` +/// hangs forever if some other shard never decrements. The one path that would +/// skip `shard_done` is a panic mid-fold (e.g. `save_snapshot_to_bytes` +/// OOM-unwinding — issue #138), which under `panic = "unwind"` would otherwise +/// hang every healthy shard's AOF writer thread. +/// +/// This guard makes the decrement unconditional. The fold creates it on entry +/// and calls [`complete`](Self::complete) on the success path (clean +/// `shard_done`). On ANY other exit — `?` error or panic unwind — `Drop` fires +/// `mark_failed` + `shard_done`, so the countdown still reaches zero, the +/// terminal writer publishes `old_seq`, and every barrier waiter wakes and rolls +/// back. Because the guard owns the decrement for ALL exits, callers MUST NOT +/// call `shard_done` again after invoking `do_rewrite_per_shard`. +struct ShardDoneGuard<'a> { + coord: &'a PerShardRewriteCoord, + done: bool, +} + +impl<'a> ShardDoneGuard<'a> { + #[inline] + fn new(coord: &'a PerShardRewriteCoord) -> Self { + Self { coord, done: false } + } + + /// Success-path completion: a single clean `shard_done` (no `mark_failed`). + /// Consumes the guard so the subsequent `Drop` is a no-op. + #[inline] + fn complete(mut self) { + self.done = true; + self.coord.shard_done(); + } +} + +impl Drop for ShardDoneGuard<'_> { + fn drop(&mut self) { + if !self.done { + // Unwind or early-error before completion: abort the rewrite (keep + // old_seq) and still decrement so the barrier releases every waiter. + self.coord.mark_failed(); + self.coord.shard_done(); + } + } +} + /// Reasons a pool send may be refused without queueing. #[derive(Debug, Clone, PartialEq, Eq)] pub enum AofPoolSendError { @@ -898,6 +998,161 @@ mod pool_tests { ); } + /// A per-shard rewrite that ABORTS (another shard failed to fold) must roll + /// THIS shard's append file back onto the committed old generation. Phase 6 + /// reopens `*file` onto the new-seq incr; if the rewrite then aborts, the + /// manifest keeps `old_seq` and prunes the new-seq files — so without a + /// barrier-before-resume the writer keeps appending into a discarded incr + /// that recovery ignores (silent data loss). This test drives one shard's + /// real fold to completion with a second shard pre-marked failed, then proves + /// post-abort appends land in the COMMITTED old-gen incr. + #[test] + fn rewrite_abort_reopens_writer_onto_committed_old_generation() { + use crate::command::persistence::AOF_REWRITE_IN_PROGRESS; + use std::io::Write; + use std::sync::atomic::Ordering; + + let tmp = tempfile::tempdir().unwrap(); + // 2-shard manifest at seq=1 on disk; share it through the coord's Arc. + let manifest = + crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); + let old_seq = manifest.seq; + let old_incr_s0 = manifest.shard_incr_path_seq(0, old_seq); + assert!(old_incr_s0.exists(), "old-gen incr must exist pre-rewrite"); + let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); + + let coord = PerShardRewriteCoord::new(manifest.clone(), old_seq, 2); + let new_seq = coord.new_seq(); + assert_ne!(new_seq, old_seq); + + // Shard 1 "fails to fold": mark_failed + shard_done (countdown 2 -> 1). + // Shard 0's shard_done (inside do_rewrite_per_shard) will then be the + // terminal decrement and take the abort path. + coord.mark_failed(); + coord.shard_done(); + + // Shard 0's append file starts on the OLD incr (where the live writer + // appends before a rewrite). Empty databases keep the fold trivial; the + // reopen behaviour is independent of key count. + let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![ + vec![crate::storage::Database::new()], + vec![crate::storage::Database::new()], + ]); + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&old_incr_s0) + .unwrap(); + let (_tx, rx) = channel::mpsc_bounded::(4); + + AOF_REWRITE_IN_PROGRESS.store(true, Ordering::SeqCst); + // The fold itself succeeds; aborting is a coordinator decision, so this + // returns Ok even though the rewrite is discarded. + do_rewrite_per_shard(0, &shard_dbs, &mut file, &rx, &coord).unwrap(); + + // Abort kept the old generation committed and pruned the new-gen incr. + assert_eq!(manifest.lock().seq, old_seq, "abort must keep old_seq"); + let new_incr_s0 = manifest.lock().shard_incr_path_seq(0, new_seq); + assert!( + !new_incr_s0.exists(), + "aborted new-gen incr must be pruned (the dangling target)" + ); + + // The writer must have been rolled back onto the old-gen incr: appends + // through `*file` after the rewrite must land in the committed file, not + // the pruned/unlinked new-gen inode. + let marker = b"*1\r\n$4\r\nPING\r\n"; + file.write_all(marker).unwrap(); + file.sync_data().unwrap(); + drop(file); + let committed = std::fs::read(&old_incr_s0).unwrap(); + assert!( + committed.windows(marker.len()).any(|w| w == marker), + "post-abort appends must land in the committed old-gen incr (no silent loss)" + ); + } + + /// Cross-thread barrier wakeup: a folded writer that decrements then blocks + /// in `await_outcome` on one thread must be woken with the committed + /// generation by the terminal `shard_done` running on ANOTHER thread. The + /// single-threaded behavioural test above never actually blocks (its + /// `await_outcome` returns immediately), so this exercises the real condvar + /// wait/notify path across threads. + #[test] + fn rewrite_abort_wakes_waiter_cross_thread() { + let tmp = tempfile::tempdir().unwrap(); + let manifest = + crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); + let old_seq = manifest.seq; + let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); + let coord = PerShardRewriteCoord::new(manifest, old_seq, 2); + + // Shard A: decrement (countdown 2 -> 1, non-terminal) then block. + let c_a = coord.clone(); + let waiter = std::thread::spawn(move || { + c_a.shard_done(); + c_a.await_outcome() + }); + + // Shard B fails: mark_failed + the terminal decrement (-> 0). Whichever + // of the two `shard_done` calls is terminal sees `failed` set (B sets it + // first) and publishes old_seq. + coord.mark_failed(); + coord.shard_done(); + + let observed = waiter.join().expect("waiter must wake, not hang or panic"); + assert_eq!( + observed, old_seq, + "aborted rewrite must publish old_seq to the cross-thread waiter" + ); + } + + /// Panic-safety / liveness: a fold that PANICS mid-flight (the OOM-unwind + /// hazard of issue #138) must not hang the other shards' writers at the + /// barrier. The `ShardDoneGuard`'s `Drop` must fire `mark_failed` + + /// `shard_done` on unwind so the countdown still closes and every waiter + /// wakes. Without the guard this test hangs forever (the panicking shard + /// never decrements). + #[test] + fn rewrite_fold_panic_releases_barrier_for_other_writers() { + use std::sync::atomic::Ordering; + + let tmp = tempfile::tempdir().unwrap(); + let manifest = + crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); + let old_seq = manifest.seq; + let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); + let coord = PerShardRewriteCoord::new(manifest, old_seq, 2); + + // Shard A folds, decrements, then blocks on the barrier in a thread. + let c_a = coord.clone(); + let waiter = std::thread::spawn(move || { + c_a.shard_done(); + c_a.await_outcome() + }); + + // Shard B "panics mid-fold": a ShardDoneGuard dropped during unwind. The + // guard's Drop must abort + decrement so A's barrier releases. + let c_b = coord.clone(); + let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _guard = ShardDoneGuard::new(&c_b); + panic!("simulated OOM unwind during fold"); + })); + assert!(panicked.is_err(), "the simulated fold must have panicked"); + + let observed = waiter + .join() + .expect("waiter must wake after the panicking shard's guard fires"); + assert_eq!( + observed, old_seq, + "panic-aborted rewrite must publish old_seq, not hang" + ); + assert!( + coord.failed.load(Ordering::Acquire), + "the dropped guard must have marked the rewrite failed" + ); + } + #[test] fn top_level_pool_routes_all_shards_to_writer_zero() { let (tx, rx) = channel::mpsc_bounded::(8); @@ -2201,20 +2456,22 @@ pub async fn per_shard_aof_writer_task( let res = do_rewrite_per_shard( shard_id, &shard_dbs, &mut sf, &rx, &coord, ); - // On success `sf` points at the NEW incr (the fold - // reopened it + already called `shard_done()`); on - // error it is still the OLD incr (pre-reopen). Wrap - // it back either way so the writer stays valid. + // `sf` is left pointing at the committed generation + // by the fold's internal barrier: NEW incr on + // success, OLD incr on abort (phase-8 rollback) or + // on a pre-reopen error. The fold's ShardDoneGuard + // already performed `shard_done` for every exit, so + // we MUST NOT decrement again here. Wrap `sf` back so + // the writer stays valid either way. writer = tokio::io::BufWriter::new(tokio::fs::File::from_std(sf)); if let Err(e) = res { error!( "F6 tokio per-shard rewrite: shard {} fold failed: {}. \ - Aborting commit; old generation stays authoritative.", + Rewrite aborted by the fold guard; old generation \ + stays authoritative.", shard_id, e ); - coord.mark_failed(); - coord.shard_done(); } } } @@ -2479,17 +2736,17 @@ pub async fn per_shard_aof_writer_task( if let Err(e) = do_rewrite_per_shard(shard_id, &shard_dbs, &mut file, &rx, &coord) { + // The fold's ShardDoneGuard already marked the rewrite + // failed and decremented on this error exit (committing + // new_seq with a shard missing its new base would break + // recovery), so do NOT decrement again here. `file` is left + // on the OLD incr (error exits are pre-reopen). error!( "F6 per-shard rewrite: shard {} fold failed: {}. \ - Aborting rewrite; old generation stays authoritative.", + Rewrite aborted by the fold guard; old generation \ + stays authoritative.", shard_id, e ); - // Mark the whole rewrite failed so the final writer - // aborts the commit (committing new_seq with a shard - // missing its new base would break recovery), then - // decrement so the countdown can still complete. - coord.mark_failed(); - coord.shard_done(); } } Ok(AofMessage::Shutdown) | Err(_) => { @@ -3106,6 +3363,13 @@ fn do_rewrite_per_shard( rx: &channel::MpscReceiver, coord: &PerShardRewriteCoord, ) -> Result<(), MoonError> { + // Panic/early-error safety: guarantees `shard_done` runs on EVERY exit + // (success via `complete()`, `?`-error or panic-unwind via `Drop`). The + // phase-8 `await_outcome` barrier makes that a liveness requirement, so + // callers MUST NOT call `shard_done` after invoking this function — the + // guard owns the single decrement for all exits. See `ShardDoneGuard`. + let guard = ShardDoneGuard::new(coord); + let sidx = shard_id as usize; let all_shards = shard_dbs.all_shard_dbs(); if sidx >= all_shards.len() { @@ -3188,8 +3452,37 @@ fn do_rewrite_per_shard( ); } - // Phase 7: signal completion; the last writer commits + prunes. - coord.shard_done(); + // Phase 7: signal completion; the last writer commits + prunes. `complete()` + // performs the single clean `shard_done` and disarms the guard's Drop. + guard.complete(); + + // Phase 8 (barrier-before-resume): block until the terminal writer publishes + // the committed generation, then make sure THIS writer's append file points + // at it. On the happy path committed == new_seq and *file already points at + // new_incr (phase 6) — nothing to do. On an abort/commit-failure the manifest + // kept old_seq and pruned our new_seq incr, so reopen *file onto old_seq's + // incr; otherwise we keep appending into a discarded generation that recovery + // ignores — silent data loss. Replaces the old "RESTART recommended" hazard. + let committed_seq = coord.await_outcome(); + if committed_seq != coord.new_seq { + let committed_incr = coord + .manifest + .lock() + .shard_incr_path_seq(shard_id, committed_seq); + *file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&committed_incr) + .map_err(|e| AofError::Io { + path: committed_incr, + source: e, + })?; + warn!( + "F6 per-shard rewrite ABORTED: shard {} rolled its append file back to \ + committed seq {} (no restart needed)", + shard_id, committed_seq + ); + } Ok(()) } From 73d95141c18072506af40b0dded70e6b9ae78cb8 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 12:26:58 +0700 Subject: [PATCH 07/15] fix(persistence): ack drained AppendSync only after the boundary fsync MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit During a BGREWRITEAOF drain, both drain_pending_appends (legacy/TopLevel, default --shards 1 path) and drain_pending_appends_framed (per-shard) wrote the appended bytes and immediately acked AofAck::Synced — BEFORE the post-drain boundary sync_data(). Under appendfsync=always a client received +OK (durable) for an entry that a crash in the ack→fsync window would lose. CodeRabbit flagged this on PR #136 review 4420584384 (issue #140); the always-fsync durability contract requires the bytes to be on disk before the client is told Synced. Fix: park the AppendSync ack senders in DrainOutcome.pending_acks instead of acking inside the drain. New sync_and_fulfill_drain() performs the boundary fsync and only then resolves the parked acks — Synced on success, or FsyncFailed plus a propagated IO error on failure (never a false Synced). All six drain sites across do_rewrite_per_shard, do_rewrite_single and do_rewrite_sharded route through it. A write error inside the drain still drops the parked acks with the outcome (RecvError → hard failure), unchanged. Validated in OrbStack VM (cgroup-capped), red/green TDD: - drain_framed_parks_appendsync_ack_until_boundary_fsync: asserts the drain PARKS the ack (pending_acks.len()==1, not sent) and that it resolves Synced only after sync_data() (RED before: no pending_acks field / acked in drain). - drain_single_fulfills_fsync_failure_as_fsync_failed: the DEFAULT --shards 1 non-framed path must resolve FsyncFailed (never Synced) when the boundary fsync fails. No regression: full aof lib suite (66) green; bgrewriteaof TTL emission tests green; per-shard AOF SIGKILL + BGREWRITEAOF crash-recovery suites green. fmt + clippy clean under both default (monoio) and tokio,jemalloc. Note: tests/aof_fsync_err_subscribe_ordering single-shard variant (an #[ignore]d normal-append-path test, multi-shard twin passes) fails identically on HEAD without this change — pre-existing and unrelated to the rewrite-drain path touched here. author: Tin Dang --- src/persistence/aof.rs | 212 ++++++++++++++++++++++++++++++++--------- 1 file changed, 167 insertions(+), 45 deletions(-) diff --git a/src/persistence/aof.rs b/src/persistence/aof.rs index b06bb9e6..5ed1b67c 100644 --- a/src/persistence/aof.rs +++ b/src/persistence/aof.rs @@ -1352,6 +1352,82 @@ mod pool_tests { assert_eq!(result, AofAck::FsyncFailed); } + /// Issue #140: an AppendSync drained during a BGREWRITEAOF must NOT be acked + /// `Synced` inside the drain — that reports the write durable before the + /// post-drain boundary fsync, so a crash in the window loses an entry the + /// client was told was safe. The drain must PARK the ack and only fulfil it + /// after the boundary `sync_data()`. (Framed / per-shard drain.) + #[test] + fn drain_framed_parks_appendsync_ack_until_boundary_fsync() { + let tmp = tempfile::tempdir().unwrap(); + let incr = tmp.path().join("incr.aof"); + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let recv = pool.try_send_append_sync(0, 99, Bytes::from_static(b"SET k v")); + + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr) + .unwrap(); + let mut outcome = drain_pending_appends_framed(&rx0, &mut file).unwrap(); + + // CONTRACT: drained + parked, NOT yet acked. + assert_eq!(outcome.drained, 1, "the AppendSync was drained"); + assert_eq!( + outcome.pending_acks.len(), + 1, + "the ack must be parked until the boundary fsync, not sent during drain" + ); + + // Boundary fsync succeeds → fulfil Synced; only NOW is the client told durable. + file.sync_data().unwrap(); + outcome.fulfill_acks(true); + assert_eq!( + recv.recv_blocking().expect("ack resolves"), + AofAck::Synced, + "post-fsync the parked ack must resolve Synced" + ); + } + + /// Issue #140 failure path: if the rewrite-boundary fsync FAILS, a drained + /// AppendSync must resolve `FsyncFailed`, never `Synced`. Exercises the + /// non-framed `drain_pending_appends` — the DEFAULT `--shards 1` rewrite + /// path (`do_rewrite_single`), which is reachable under appendfsync=always. + #[cfg(feature = "runtime-monoio")] + #[test] + fn drain_single_fulfills_fsync_failure_as_fsync_failed() { + let tmp = tempfile::tempdir().unwrap(); + let incr = tmp.path().join("incr.aof"); + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let recv = pool.try_send_append_sync(0, 7, Bytes::from_static(b"SET a b")); + + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr) + .unwrap(); + let mut outcome = drain_pending_appends(&rx0, &mut file).unwrap(); + assert_eq!( + outcome.pending_acks.len(), + 1, + "ack parked, not sent in drain" + ); + + // Simulate a failed boundary fsync: the parked ack must report FsyncFailed. + outcome.fulfill_acks(false); + assert_eq!( + recv.recv_blocking().expect("ack resolves"), + AofAck::FsyncFailed, + "a failed boundary fsync must NOT be reported Synced to the client" + ); + } + // F2 (design-for-failure): `appendfsync=always` must bound its fsync-ack // await. A stalled writer must surface a hard error within the budget, // never park the connection forever. Tokio-gated because it drives the @@ -3178,6 +3254,62 @@ fn snapshot_and_generate(db: &SharedDatabases) -> BytesMut { struct DrainOutcome { drained: usize, shutdown_requested: bool, + /// AppendSync ack senders for entries drained during a rewrite. Under + /// `appendfsync=always` the client must NOT be told `Synced` until the + /// post-drain boundary `sync_data()` makes those bytes durable, so the acks + /// are parked here and resolved by [`fulfill_acks`](Self::fulfill_acks) AFTER + /// the caller's fsync — `Synced` on success, `FsyncFailed` on failure. Acking + /// inside the drain (the old behaviour) reports a write durable before the + /// boundary fsync, so a crash in that window loses an entry the client was + /// told was safe. (Issue #140.) + pending_acks: Vec>, +} + +#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] +impl DrainOutcome { + /// Resolve every parked AppendSync ack after the rewrite-boundary fsync. + /// `synced=true` → `Synced`; `false` → `FsyncFailed`. A fresh `AofAck` is + /// built per sender so `AofAck` needs no `Copy`/`Clone`. Drains the vec so a + /// second call is a no-op. + fn fulfill_acks(&mut self, synced: bool) { + for tx in std::mem::take(&mut self.pending_acks) { + let _ = tx.send(if synced { + AofAck::Synced + } else { + AofAck::FsyncFailed + }); + } + } +} + +/// Fsync `file` at a rewrite drain boundary, then resolve the drained batch's +/// parked AppendSync acks against the result: `Synced` on success, or +/// `FsyncFailed` + propagate the IO error on failure. Centralizes the issue-#140 +/// durability ordering (ack strictly AFTER the bytes are durable) for every +/// `do_rewrite_*` drain site. +#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] +fn sync_and_fulfill_drain( + outcome: &mut DrainOutcome, + file: &mut std::fs::File, + incr_path: PathBuf, +) -> Result<(), MoonError> { + match file.sync_data() { + Ok(()) => { + outcome.fulfill_acks(true); + Ok(()) + } + Err(e) => { + // Boundary fsync failed: the drained writes are NOT durable. Tell the + // waiting clients FsyncFailed (never Synced) and propagate the error + // so the rewrite aborts. + outcome.fulfill_acks(false); + Err(AofError::Io { + path: incr_path, + source: e, + } + .into()) + } + } } #[cfg(feature = "runtime-monoio")] @@ -3201,11 +3333,12 @@ fn drain_pending_appends( })?; outcome.drained += 1; } - // AppendSync during a rewrite drain: bytes are written and counted; - // the post-drain fsync at the rewrite boundary covers durability, - // so we ack `Synced`. If the write itself fails the error is - // already propagated upward by the `?` and the ack is dropped — - // the caller observes `RecvError`, which it treats as failure. + // AppendSync during a rewrite drain: bytes are written and counted, + // but the ack is PARKED until the caller's post-drain boundary fsync + // (issue #140) — acking `Synced` here would report durability before + // the bytes are fsynced. If the write itself fails the `?` propagates + // the error and the parked ack is dropped with the outcome — the + // caller observes `RecvError`, which it treats as failure. AofMessage::AppendSync { bytes: data, lsn: _, @@ -3216,7 +3349,7 @@ fn drain_pending_appends( source: e, })?; outcome.drained += 1; - let _ = ack.send(AofAck::Synced); + outcome.pending_acks.push(ack); } AofMessage::Shutdown => { outcome.shutdown_requested = true; @@ -3272,9 +3405,10 @@ fn drain_pending_appends_framed( source: e, })?; outcome.drained += 1; - // Durability for these is covered by the post-drain fsync at - // the rewrite boundary (mirrors drain_pending_appends). - let _ = ack.send(AofAck::Synced); + // Park the ack until the caller's post-drain boundary fsync + // (issue #140); resolved Synced/FsyncFailed by + // `sync_and_fulfill_drain`. Mirrors `drain_pending_appends`. + outcome.pending_acks.push(ack); } AofMessage::Shutdown => { outcome.shutdown_requested = true; @@ -3383,23 +3517,19 @@ fn do_rewrite_per_shard( .into()); } - // Phase 1: drain pre-rewrite queued appends into old incr (framed). - let pre_drain = drain_pending_appends_framed(rx, file)?; - file.sync_data().map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; + // Phase 1: drain pre-rewrite queued appends into old incr (framed), fsync, + // then resolve their parked AppendSync acks (issue #140). + let mut pre_drain = drain_pending_appends_framed(rx, file)?; + sync_and_fulfill_drain(&mut pre_drain, file, PathBuf::from(""))?; // Phase 2: acquire write locks on this shard's db(s) (db_idx ascending). let shard_locks = &all_shards[sidx]; let guards: Vec<_> = shard_locks.iter().map(|lock| lock.write()).collect(); - // Phase 3: drain appends that completed between phase 1 and phase 2. - let mid_drain = drain_pending_appends_framed(rx, file)?; - file.sync_data().map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; + // Phase 3: drain appends that completed between phase 1 and phase 2, fsync, + // then resolve their parked AppendSync acks (issue #140). + let mut mid_drain = drain_pending_appends_framed(rx, file)?; + sync_and_fulfill_drain(&mut mid_drain, file, PathBuf::from(""))?; // Phase 4: snapshot this shard's databases under the locks. let now_ms = current_time_ms(); @@ -3513,12 +3643,10 @@ fn do_rewrite_single( file: &mut std::fs::File, rx: &channel::MpscReceiver, ) -> Result<(), MoonError> { - // Phase 1: drain pre-rewrite queued appends into old incr, fsync. - let pre_drain = drain_pending_appends(rx, file)?; - file.sync_data().map_err(|e| AofError::Io { - path: manifest.incr_path(), - source: e, - })?; + // Phase 1: drain pre-rewrite queued appends into old incr, fsync, then + // resolve their parked AppendSync acks (issue #140). + let mut pre_drain = drain_pending_appends(rx, file)?; + sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; // Phase 2: acquire write locks on every database in the shard. // Order is consistent (index-ascending) so concurrent callers would @@ -3526,12 +3654,10 @@ fn do_rewrite_single( // acquires multi-db locks. let guards: Vec<_> = db.iter().map(|lock| lock.write()).collect(); - // Phase 3: drain any appends the handlers sent between phase 1 and phase 2. - let mid_drain = drain_pending_appends(rx, file)?; - file.sync_data().map_err(|e| AofError::Io { - path: manifest.incr_path(), - source: e, - })?; + // Phase 3: drain any appends the handlers sent between phase 1 and phase 2, + // fsync, then resolve their parked AppendSync acks (issue #140). + let mut mid_drain = drain_pending_appends(rx, file)?; + sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; // Phase 4: snapshot under the write locks. No mutation is possible. let now_ms = current_time_ms(); @@ -3596,12 +3722,10 @@ fn do_rewrite_sharded( file: &mut std::fs::File, rx: &channel::MpscReceiver, ) -> Result<(), MoonError> { - // Phase 1: drain pre-rewrite queued appends into old incr. - let pre_drain = drain_pending_appends(rx, file)?; - file.sync_data().map_err(|e| AofError::Io { - path: manifest.incr_path(), - source: e, - })?; + // Phase 1: drain pre-rewrite queued appends into old incr, fsync, then + // resolve their parked AppendSync acks (issue #140). + let mut pre_drain = drain_pending_appends(rx, file)?; + sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; // Phase 2: acquire write locks on ALL (shard, db) pairs simultaneously. // Lock order is (shard_idx, db_idx) ascending — must match anywhere else @@ -3618,12 +3742,10 @@ fn do_rewrite_sharded( guards.push(shard_guards); } - // Phase 3: drain appends completed between phase 1 and phase 2. - let mid_drain = drain_pending_appends(rx, file)?; - file.sync_data().map_err(|e| AofError::Io { - path: manifest.incr_path(), - source: e, - })?; + // Phase 3: drain appends completed between phase 1 and phase 2, fsync, then + // resolve their parked AppendSync acks (issue #140). + let mut mid_drain = drain_pending_appends(rx, file)?; + sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; // Phase 4: snapshot under locks. let db_count = shard_dbs.db_count(); From 14c572b1b9be6bceda20468132e2048dbfb377f8 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 12:41:14 +0700 Subject: [PATCH 08/15] refactor(persistence): split aof_manifest.rs into submodules under 1500-line cap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit src/persistence/aof_manifest.rs had grown to 3058 lines — past the 1500-line cap in CLAUDE.md. CodeRabbit flagged this on PR #136 review 4420584384 (issue #143). Converted the file into a directory module and relocated cohesive chunks, preserving every public path via re-export. Layout (all three files now under the cap): - aof_manifest/mod.rs (1395): AofManifest state + manifest-level I/O (paths, initialize/load, parse_v1/v2, cleanup_orphans, write_manifest, advance, migrate, verify, global_max_lsn) + manifest-core unit tests. - aof_manifest/shard_replay.rs (1159): the multi-part / per-shard replay free functions (replay_multi_part, replay_incr_resp, replay_incr_framed, replay_per_shard, replay_ordered_merge) + OrderedEntry + the replay unit tests. - aof_manifest/shard_rewrite.rs (577): the per-shard rewrite inherent methods (initialize_multi, try_initialize_multi, advance_shard, prune_shard_files) as `impl AofManifest` + the shard-rewrite unit tests. Mechanics: - `use super::*` in each submodule pulls the parent's types, std/tracing imports, consts, and private helpers down (a child module sees the parent's private items), so no import duplication beyond the test modules. - Relocated public replay API is re-exported from mod.rs (`pub use shard_replay::{OrderedEntry, replay_multi_part, replay_ordered_merge, replay_per_shard}`) so external callers (`crate::persistence::aof_manifest::replay_*`) keep resolving unchanged. - Per CLAUDE.md the cap and the "tests stay in mod.rs" convention can't both hold for a 3058-line file (tests alone are ~1100 lines); tests were moved next to the code they cover, each submodule getting its own temp_dir() copy with a distinct directory prefix (replay-/rewrite-) so the independent static counters never collide. No behaviour change. Validated in OrbStack VM: 30 aof_manifest unit tests green; per-shard AOF SIGKILL recovery (3) and per-shard BGREWRITEAOF crash recovery (2) green (exercise the relocated replay_per_shard / advance_shard / initialize_multi via recovery). fmt + clippy clean under both default (monoio) and tokio,jemalloc. author: Tin Dang --- src/persistence/aof_manifest.rs | 3058 ----------------- src/persistence/aof_manifest/mod.rs | 1388 ++++++++ src/persistence/aof_manifest/shard_replay.rs | 1157 +++++++ src/persistence/aof_manifest/shard_rewrite.rs | 577 ++++ 4 files changed, 3122 insertions(+), 3058 deletions(-) delete mode 100644 src/persistence/aof_manifest.rs create mode 100644 src/persistence/aof_manifest/mod.rs create mode 100644 src/persistence/aof_manifest/shard_replay.rs create mode 100644 src/persistence/aof_manifest/shard_rewrite.rs diff --git a/src/persistence/aof_manifest.rs b/src/persistence/aof_manifest.rs deleted file mode 100644 index 572e16e8..00000000 --- a/src/persistence/aof_manifest.rs +++ /dev/null @@ -1,3058 +0,0 @@ -//! Multi-part AOF manifest: tracks base (RDB) and incremental (RESP) files. -//! -//! Part of the **storage format v1** umbrella commitment — see -//! [`docs/STORAGE-FORMAT-V1.md`](../../../docs/STORAGE-FORMAT-V1.md). The -//! manifest framing is the canonical on-disk marker; the human-readable -//! "v1" umbrella also covers WAL v3 and RDB v2 sub-formats. -//! -//! Two on-disk layouts coexist (selected at manifest creation time, never mixed -//! within one directory): -//! -//! **TopLevel (manifest v1, single-shard / legacy):** -//! ```text -//! appendonlydir/ -//! moon.aof.1.base.rdb # RDB snapshot base -//! moon.aof.1.incr.aof # Incremental RESP since base -//! moon.aof.manifest # v1 text format -//! ``` -//! -//! **PerShard (manifest v2, multi-shard durability):** -//! ```text -//! appendonlydir/ -//! moon.aof.manifest # v2 text format (carries shard count + max_lsn) -//! shard-0/ -//! moon.aof.1.base.rdb -//! moon.aof.1.incr.aof -//! shard-1/ -//! moon.aof.1.base.rdb -//! moon.aof.1.incr.aof -//! … -//! ``` -//! -//! The manifest text format is line-prefix based. v1 manifests have no -//! `version` line; v2 manifests begin with `version 2`. On BGREWRITEAOF the -//! sequence increments, a new base + incr pair is created per shard (PerShard) -//! or at top level (TopLevel), and old files are deleted. - -use std::io::Write; -use std::path::{Path, PathBuf}; - -use tracing::{error, info, warn}; - -use crate::persistence::fsync::fsync_directory; - -const MANIFEST_NAME: &str = "moon.aof.manifest"; -const AOF_DIR_NAME: &str = "appendonlydir"; - -/// Fsync the parent directory of `path` (best-effort). -/// -/// POSIX guarantees atomicity of `rename()` but does NOT guarantee that the -/// directory entry update is durable after a crash. On ext4 and XFS without -/// `data=ordered`, a crash between the rename and a directory fsync can leave -/// the old file name visible on the next boot even though the rename completed -/// in memory. Calling this after every manifest-visible rename closes that gap. -/// -/// Best-effort: logs on failure but does not propagate the error. A failed -/// dir fsync means the rename may not survive a crash — the worst case is -/// that recovery falls back to the previous manifest state, which is still -/// consistent (the atomic rename guarantees the file is either fully old or -/// fully new). Call sites that CAN propagate (i.e., are in a fallible fn that -/// returns `std::io::Result`) should call `fsync_directory(parent)?` directly. -fn fsync_parent_best_effort(path: &Path) { - let parent = match path.parent() { - Some(p) if !p.as_os_str().is_empty() => p, - _ => return, // root or no parent — nothing to fsync - }; - if let Err(e) = fsync_directory(parent) { - warn!( - "fsync_parent_best_effort: failed to fsync dir {} after rename of {}: {}", - parent.display(), - path.display(), - e - ); - } -} - -/// On-disk layout discriminator. -/// -/// `TopLevel` is the legacy single-shard layout from manifest v1. `PerShard` -/// is the multi-shard layout introduced with manifest v2 — used whenever -/// `num_shards >= 2`. A `--shards 1` deployment with an existing v1 manifest -/// stays TopLevel until explicitly migrated. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AofLayout { - /// Legacy single-shard layout: `appendonlydir/moon.aof.{seq}.{base|incr}.*`. - TopLevel, - /// Per-shard layout: `appendonlydir/shard-{N}/moon.aof.{seq}.{base|incr}.*`. - PerShard, -} - -/// Per-shard manifest entry. One per shard in `PerShard` layout. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ShardManifest { - /// Shard ID (0..num_shards). - pub shard_id: u16, - /// Max LSN persisted to this shard's incr file so far. Semantics defined - /// by step 3 (LSN tagging) of the per-shard AOF RFC — until then this is - /// 0 and recovery does not use it. Once step 3 ships, recovery seeds - /// `master_repl_offset = max(shards[*].max_lsn)` before accepting writes. - pub max_lsn: u64, -} - -/// Active AOF file set tracked by the manifest. -#[derive(Debug, Clone)] -pub struct AofManifest { - /// Base directory (parent of `appendonlydir/`) - pub dir: PathBuf, - /// Current sequence number (incremented on each rewrite). - pub seq: u64, - /// On-disk layout. Determines path computation for base/incr files. - pub layout: AofLayout, - /// Per-shard metadata. Length is 1 for `TopLevel`, `num_shards` for - /// `PerShard`. Indexed by `shard_id`. - pub shards: Vec, -} - -impl AofManifest { - /// Path to the `appendonlydir/` directory. - pub fn aof_dir(&self) -> PathBuf { - self.dir.join(AOF_DIR_NAME) - } - - /// Path to the manifest file. - pub fn manifest_path(&self) -> PathBuf { - self.aof_dir().join(MANIFEST_NAME) - } - - /// Path to the base RDB file for the current sequence. - /// - /// Layout-aware: TopLevel returns `appendonlydir/moon.aof.{seq}.base.rdb`; - /// PerShard routes to `appendonlydir/shard-0/moon.aof.{seq}.base.rdb`. - /// This single-file helper is meaningful only when there is one shard - /// (post-migration `--shards 1`); a multi-shard PerShard manifest has N - /// base files and the caller must use [`Self::shard_base_path`] instead. - /// In debug builds, calling this on a multi-shard PerShard manifest - /// asserts; in release it returns the shard-0 path so production stays - /// recoverable rather than panicking on a stale call site. - pub fn base_path(&self) -> PathBuf { - match self.layout { - AofLayout::TopLevel => self - .aof_dir() - .join(format!("moon.aof.{}.base.rdb", self.seq)), - AofLayout::PerShard => { - debug_assert!( - self.shards.len() == 1, - "base_path() called on multi-shard PerShard manifest; use shard_base_path(shard_id)", - ); - self.shard_base_path_seq(0, self.seq) - } - } - } - - /// Path to the incremental RESP file for the current sequence. - /// - /// Layout-aware — see [`Self::base_path`] for the same routing rules. - pub fn incr_path(&self) -> PathBuf { - match self.layout { - AofLayout::TopLevel => self - .aof_dir() - .join(format!("moon.aof.{}.incr.aof", self.seq)), - AofLayout::PerShard => { - debug_assert!( - self.shards.len() == 1, - "incr_path() called on multi-shard PerShard manifest; use shard_incr_path(shard_id)", - ); - self.shard_incr_path_seq(0, self.seq) - } - } - } - - /// Path to the base RDB file for a given sequence. Layout-aware — see - /// [`Self::base_path`]. - pub fn base_path_seq(&self, seq: u64) -> PathBuf { - match self.layout { - AofLayout::TopLevel => self.aof_dir().join(format!("moon.aof.{}.base.rdb", seq)), - AofLayout::PerShard => { - debug_assert!( - self.shards.len() == 1, - "base_path_seq() called on multi-shard PerShard manifest; use shard_base_path_seq(shard_id, seq)", - ); - self.shard_base_path_seq(0, seq) - } - } - } - - /// Path to the incremental RESP file for a given sequence. Layout-aware — - /// see [`Self::base_path`]. - pub fn incr_path_seq(&self, seq: u64) -> PathBuf { - match self.layout { - AofLayout::TopLevel => self.aof_dir().join(format!("moon.aof.{}.incr.aof", seq)), - AofLayout::PerShard => { - debug_assert!( - self.shards.len() == 1, - "incr_path_seq() called on multi-shard PerShard manifest; use shard_incr_path_seq(shard_id, seq)", - ); - self.shard_incr_path_seq(0, seq) - } - } - } - - /// Create the `appendonlydir/` and write the initial manifest. - /// - /// Prefer [`Self::initialize_with_base`] when the in-memory databases - /// already contain state (e.g. first upgrade from legacy single-file AOF - /// or per-shard WAL) — otherwise subsequent boots cannot reconstruct that - /// state because there is no base RDB for `replay_multi_part` to load. - /// - /// B4 fix: even on fresh install (no prior state), materialize an EMPTY - /// base RDB so the `(base + incr)` invariant always holds. Without this, - /// the recovery path refuses to replay incr-only state and the server - /// fails to restart after a graceful shutdown that only wrote incr. - pub fn initialize(dir: &Path) -> std::io::Result { - let manifest = Self { - dir: dir.to_path_buf(), - seq: 1, - layout: AofLayout::TopLevel, - shards: vec![ShardManifest { - shard_id: 0, - max_lsn: 0, - }], - }; - std::fs::create_dir_all(manifest.aof_dir())?; - - // Serialize an empty database vector to an empty base RDB so the - // (base + incr) invariant holds from the first boot. - let empty_dbs: [crate::storage::Database; 0] = []; - let empty_rdb = crate::persistence::rdb::save_to_bytes(&empty_dbs) - .map_err(|e| std::io::Error::other(format!("empty RDB serialize: {e}")))?; - let base_path = manifest.base_path(); - let tmp_path = base_path.with_extension("rdb.tmp"); - { - let mut f = std::fs::File::create(&tmp_path)?; - f.write_all(&empty_rdb)?; - f.sync_data()?; - } - std::fs::rename(&tmp_path, &base_path)?; - fsync_parent_best_effort(&base_path); - - // Create the empty incr file so the writer has a target. - std::fs::File::create(manifest.incr_path())?; - - manifest.write_manifest()?; - Ok(manifest) - } - - /// Create the `appendonlydir/` and write an initial manifest with a base RDB - /// capturing the current in-memory state. - /// - /// Used on first upgrade from legacy persistence formats: after - /// `restore_from_persistence` has loaded state from the per-shard WAL or - /// `appendonly.aof`, this call materializes that state as the seq 1 base - /// RDB. Without a base, on the next boot the multi-part replay path would - /// clear the databases and then fail (missing base with non-empty incr) - /// or silently restart from empty state. - pub fn initialize_with_base(dir: &Path, rdb_bytes: &[u8]) -> std::io::Result { - let manifest = Self { - dir: dir.to_path_buf(), - seq: 1, - layout: AofLayout::TopLevel, - shards: vec![ShardManifest { - shard_id: 0, - max_lsn: 0, - }], - }; - std::fs::create_dir_all(manifest.aof_dir())?; - - // Write base RDB atomically: tmp file + fsync + rename. - let base_path = manifest.base_path(); - let tmp_path = base_path.with_extension("rdb.tmp"); - { - let mut f = std::fs::File::create(&tmp_path)?; - f.write_all(rdb_bytes)?; - f.sync_data()?; - } - std::fs::rename(&tmp_path, &base_path)?; - fsync_parent_best_effort(&base_path); - - // Create empty incr file so the writer has something to append to. - std::fs::File::create(manifest.incr_path())?; - - manifest.write_manifest()?; - Ok(manifest) - } - - /// Load manifest from disk. - /// - /// Returns: - /// - `Ok(None)` — manifest file does not exist (fresh install or legacy single-file AOF) - /// - `Ok(Some(manifest))` — manifest loaded successfully - /// - `Err(_)` — manifest file exists but is unreadable or corrupt. - /// Callers MUST treat this as fatal: overwriting a corrupt manifest with a - /// fresh one silently destroys the reference to the real base RDB and loses data. - pub fn load(dir: &Path) -> std::io::Result> { - let aof_dir = dir.join(AOF_DIR_NAME); - let manifest_path = aof_dir.join(MANIFEST_NAME); - - if !manifest_path.exists() { - return Ok(None); - } - - let content = std::fs::read_to_string(&manifest_path)?; - - // Detect format version. v1 manifests have no `version` line and use - // line prefixes `seq`/`base`/`incr`. v2 manifests start with `version 2` - // and carry per-shard records. - let mut format_version: u8 = 1; - for line in content.lines() { - let line = line.trim(); - if let Some(val) = line.strip_prefix("version ") { - if let Ok(v) = val.parse::() { - format_version = v; - } - break; - } - if !line.is_empty() { - // First non-blank line is not a version header → v1. - break; - } - } - - let manifest = match format_version { - 1 => Self::parse_v1(&content, dir, &manifest_path)?, - 2 => Self::parse_v2(&content, dir, &manifest_path)?, - other => { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} has unsupported format version {} (max supported: 2)", - manifest_path.display(), - other, - ), - )); - } - }; - - // Best-effort orphan cleanup: delete stray base/incr files from aborted - // rewrites. A crash between advance() steps 1-3 leaves a new base RDB on - // disk that the active manifest never references. Without this sweep, - // repeated crashes during rewrite can fill the disk with zombie files. - // - // Safe to call here: parse_* verified the manifest has all required - // records, so cleanup_orphans won't delete the active files. - manifest.cleanup_orphans(); - - Ok(Some(manifest)) - } - - /// Parse a v1 (TopLevel, single-shard) manifest. - fn parse_v1(content: &str, dir: &Path, manifest_path: &Path) -> std::io::Result { - let mut seq = 0u64; - let mut has_base_record = false; - let mut has_incr_record = false; - for line in content.lines() { - let line = line.trim(); - if let Some(val) = line.strip_prefix("seq ") { - if let Ok(n) = val.parse::() { - seq = n; - } - } else if line.starts_with("base ") { - has_base_record = true; - } else if line.starts_with("incr ") { - has_incr_record = true; - } - } - - if seq == 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} has no valid sequence number", - manifest_path.display() - ), - )); - } - - if !has_base_record || !has_incr_record { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} is truncated: seq={} base={} incr={}", - manifest_path.display(), - seq, - has_base_record, - has_incr_record, - ), - )); - } - - Ok(Self { - dir: dir.to_path_buf(), - seq, - layout: AofLayout::TopLevel, - shards: vec![ShardManifest { - shard_id: 0, - max_lsn: 0, - }], - }) - } - - /// Parse a v2 (PerShard, multi-shard) manifest. - /// - /// Expected line format: - /// ```text - /// version 2 - /// seq N - /// shards K - /// shard 0 max_lsn LSN0 - /// shard 1 max_lsn LSN1 - /// ... - /// ``` - /// - /// Per-shard `base`/`incr` paths are derived from `shard-{N}/moon.aof.{seq}.*` - /// rather than stored explicitly — the layout is canonical, so storing - /// paths invites drift between the stored value and the computed one. - fn parse_v2(content: &str, dir: &Path, manifest_path: &Path) -> std::io::Result { - let mut seq = 0u64; - let mut num_shards: Option = None; - let mut shards: Vec = Vec::new(); - - for line in content.lines() { - let line = line.trim(); - if line.is_empty() || line.starts_with('#') { - continue; - } - if line == "version 2" { - continue; - } else if let Some(val) = line.strip_prefix("seq ") { - seq = val.parse::().map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} has invalid seq line `{}`: {}", - manifest_path.display(), - line, - e, - ), - ) - })?; - } else if let Some(val) = line.strip_prefix("shards ") { - num_shards = Some(val.parse::().map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} has invalid shards line `{}`: {}", - manifest_path.display(), - line, - e, - ), - ) - })?); - } else if let Some(rest) = line.strip_prefix("shard ") { - // Format: `shard max_lsn ` - let mut it = rest.split_whitespace(); - let id_str = it.next().ok_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} has shard line missing id: `{}`", - manifest_path.display(), - line, - ), - ) - })?; - let id: u16 = id_str.parse().map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} has shard line invalid id `{}`: {}", - manifest_path.display(), - id_str, - e, - ), - ) - })?; - // Expect `max_lsn `. - let label = it.next().unwrap_or(""); - let val_str = it.next().unwrap_or("0"); - if label != "max_lsn" { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} shard {} expected `max_lsn`, got `{}`", - manifest_path.display(), - id, - label, - ), - )); - } - let max_lsn: u64 = val_str.parse().map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} shard {} invalid max_lsn `{}`: {}", - manifest_path.display(), - id, - val_str, - e, - ), - ) - })?; - shards.push(ShardManifest { - shard_id: id, - max_lsn, - }); - } - // Unknown lines are tolerated (forward-compat). Strict parsers can - // be added at v3 if needed. - } - - if seq == 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} has no valid sequence number", - manifest_path.display() - ), - )); - } - - let expected = num_shards.ok_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} is missing required `shards N` line", - manifest_path.display() - ), - ) - })?; - - if shards.len() != expected as usize { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} declares shards={} but has {} shard records", - manifest_path.display(), - expected, - shards.len(), - ), - )); - } - - // Sort by shard_id and verify contiguous range [0, expected). - shards.sort_by_key(|s| s.shard_id); - for (i, s) in shards.iter().enumerate() { - if s.shard_id as usize != i { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!( - "AOF manifest at {} has non-contiguous shard ids (expected {} at position {}, got {})", - manifest_path.display(), - i, - i, - s.shard_id, - ), - )); - } - } - - Ok(Self { - dir: dir.to_path_buf(), - seq, - layout: AofLayout::PerShard, - shards, - }) - } - - /// Delete any base/incr files in `appendonlydir/` that do not match the - /// current sequence. Best-effort — logs but does not propagate errors. - /// - /// For `PerShard` layout, also recurses into every `shard-N/` subdirectory - /// and removes stale/tmp files there. Aborted BGREWRITEAOF runs leave - /// `.rdb.tmp` files in the shard subdirs that otherwise accumulate forever. - fn cleanup_orphans(&self) { - match self.layout { - AofLayout::TopLevel => { - self.cleanup_orphans_dir(&self.aof_dir(), self.seq); - } - AofLayout::PerShard => { - // Top-level appendonlydir/ holds only the manifest — no data files - // to clean up there. All data lives in shard-N/ subdirs. - for shard in &self.shards { - self.cleanup_orphans_shard(shard.shard_id); - } - } - } - } - - /// Scan a single shard's directory for orphan base/incr/tmp files that do - /// not correspond to the current manifest sequence. Best-effort. - fn cleanup_orphans_shard(&self, shard_id: u16) { - self.cleanup_orphans_dir(&self.shard_dir(shard_id), self.seq); - } - - /// Core orphan sweep: scan `dir` and remove any `moon.aof.*` files whose - /// sequence is not `keep_seq`. Skips the manifest file itself. - fn cleanup_orphans_dir(&self, dir: &Path, keep_seq: u64) { - let entries = match std::fs::read_dir(dir) { - Ok(e) => e, - Err(_) => return, - }; - let current_base = format!("moon.aof.{}.base.rdb", keep_seq); - let current_incr = format!("moon.aof.{}.incr.aof", keep_seq); - for entry in entries.flatten() { - let name = entry.file_name(); - let name_str = match name.to_str() { - Some(s) => s, - None => continue, - }; - // Keep manifest, current base, current incr. Delete any other moon.aof.*. - if name_str == MANIFEST_NAME || name_str == current_base || name_str == current_incr { - continue; - } - let is_moon_aof = name_str.starts_with("moon.aof.") - && (name_str.ends_with(".base.rdb") - || name_str.ends_with(".incr.aof") - || name_str.ends_with(".rdb.tmp") - || name_str.ends_with(".tmp")); - if !is_moon_aof { - continue; - } - let path = entry.path(); - match std::fs::remove_file(&path) { - Ok(()) => info!("AOF orphan cleanup: removed {}", path.display()), - Err(e) => warn!( - "AOF orphan cleanup: failed to remove {}: {}", - path.display(), - e - ), - } - } - } - - /// Write the manifest file atomically (write tmp + rename). - /// - /// Emits v1 format for `TopLevel` and v2 for `PerShard`. The format is - /// selected by `self.layout`, never by callers — preserving the invariant - /// that one directory holds one layout. - pub fn write_manifest(&self) -> std::io::Result<()> { - let manifest_path = self.manifest_path(); - let tmp_path = manifest_path.with_extension("tmp"); - - let content = match self.layout { - AofLayout::TopLevel => format!( - "seq {}\nbase moon.aof.{}.base.rdb\nincr moon.aof.{}.incr.aof\n", - self.seq, self.seq, self.seq - ), - AofLayout::PerShard => { - let mut s = String::with_capacity(64 + self.shards.len() * 40); - s.push_str("version 2\n"); - s.push_str(&format!("seq {}\n", self.seq)); - s.push_str(&format!("shards {}\n", self.shards.len())); - for shard in &self.shards { - s.push_str(&format!( - "shard {} max_lsn {}\n", - shard.shard_id, shard.max_lsn - )); - } - s - } - }; - - let mut f = std::fs::File::create(&tmp_path)?; - f.write_all(content.as_bytes())?; - f.sync_data()?; - std::fs::rename(&tmp_path, &manifest_path)?; - fsync_parent_best_effort(&manifest_path); - Ok(()) - } - - // ------------------------------------------------------------------ - // Per-shard layout helpers - // ------------------------------------------------------------------ - - /// Directory holding a shard's AOF files. - /// - /// - `TopLevel`: `appendonlydir/` (the shard_id argument is asserted to be 0). - /// - `PerShard`: `appendonlydir/shard-{shard_id}/`. - pub fn shard_dir(&self, shard_id: u16) -> PathBuf { - match self.layout { - AofLayout::TopLevel => { - debug_assert_eq!(shard_id, 0, "TopLevel layout only has shard 0"); - self.aof_dir() - } - AofLayout::PerShard => self.aof_dir().join(format!("shard-{}", shard_id)), - } - } - - /// Path to a shard's base RDB file for the current sequence. - pub fn shard_base_path(&self, shard_id: u16) -> PathBuf { - self.shard_dir(shard_id) - .join(format!("moon.aof.{}.base.rdb", self.seq)) - } - - /// Path to a shard's incremental RESP file for the current sequence. - pub fn shard_incr_path(&self, shard_id: u16) -> PathBuf { - self.shard_dir(shard_id) - .join(format!("moon.aof.{}.incr.aof", self.seq)) - } - - /// Path to a shard's base RDB file for a given sequence. - pub fn shard_base_path_seq(&self, shard_id: u16, seq: u64) -> PathBuf { - self.shard_dir(shard_id) - .join(format!("moon.aof.{}.base.rdb", seq)) - } - - /// Path to a shard's incremental RESP file for a given sequence. - pub fn shard_incr_path_seq(&self, shard_id: u16, seq: u64) -> PathBuf { - self.shard_dir(shard_id) - .join(format!("moon.aof.{}.incr.aof", seq)) - } - - /// Maximum LSN persisted across all shards. - /// - /// Computed (not stored) so the stored value can never drift from - /// the per-shard records. Returns 0 if `shards` is empty (defensive; - /// constructors guarantee at least one shard). - pub fn global_max_lsn(&self) -> u64 { - self.shards.iter().map(|s| s.max_lsn).max().unwrap_or(0) - } - - /// Verify that the manifest matches the runtime shard count. - /// - /// Returns the verbatim error from RFC § 3 if the shard count differs, - /// for operator-facing consistency. Callers (typically `main.rs` boot) - /// should treat this as fatal: continuing with a mismatched shard count - /// silently drops data from shards that no longer exist or replays a - /// shard's data into the wrong DashTable. - pub fn verify_shard_count(&self, expected: u16) -> Result<(), String> { - let actual = self.shards.len() as u16; - if actual != expected { - return Err(format!( - "ERR shard count changed (manifest={}, config={}); refusing to start to avoid data loss. See docs/runbooks/shard-count-change.md", - actual, expected - )); - } - Ok(()) - } - - /// Returns true if the on-disk layout under `appendonlydir/` matches the - /// legacy TopLevel format (files at top level, no `shard-N/` subdirs). - /// - /// Used by callers to detect when a v1 single-shard deployment is being - /// upgraded to v2 multi-shard and needs explicit migration. Does NOT - /// migrate — separate from `migrate_top_level_to_per_shard` so the side - /// effect is opt-in, not hidden in a load path. - pub fn is_legacy_top_level_layout(dir: &Path) -> bool { - let aof_dir = dir.join(AOF_DIR_NAME); - if !aof_dir.exists() { - return false; - } - - // Check manifest version first. If a valid v2 (PerShard) manifest exists, - // return false regardless of stray top-level files. Operators occasionally - // leave old base.rdb / incr.aof files at the top level during debugging - // or failed upgrades; scanning filenames without reading the manifest would - // produce a misleading "legacy detected" result and trigger unwanted - // migration on an already-upgraded deployment. - if let Ok(Some(m)) = Self::load(dir) { - if m.layout == AofLayout::PerShard { - return false; - } - } - - let entries = match std::fs::read_dir(&aof_dir) { - Ok(e) => e, - Err(_) => return false, - }; - for entry in entries.flatten() { - let name = entry.file_name(); - let Some(name_str) = name.to_str() else { - continue; - }; - if name_str.starts_with("moon.aof.") - && (name_str.ends_with(".base.rdb") || name_str.ends_with(".incr.aof")) - { - return true; - } - } - false - } - - /// Migrate a single-shard TopLevel layout in place to a single-shard - /// PerShard layout. - /// - /// Moves `appendonlydir/moon.aof.{seq}.{base.rdb,incr.aof}` into - /// `appendonlydir/shard-0/`, then rewrites the manifest as v2 with - /// `shards 1`. Idempotent: a second call on an already-PerShard manifest - /// returns Ok with no filesystem changes. - /// - /// This is the RFC § 5 case 1 migration — zero data movement (rename only), - /// safe to run on first boot after upgrading from v0.1.x. Multi-shard - /// migrations from legacy AOF (case 2) use the `moon migrate-aof` - /// subcommand and are NOT handled here. - pub fn migrate_top_level_to_per_shard(&mut self) -> std::io::Result<()> { - if self.layout == AofLayout::PerShard { - return Ok(()); - } - if self.shards.len() != 1 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - format!( - "migrate_top_level_to_per_shard called with {} shards; \ - only single-shard TopLevel can be migrated in place", - self.shards.len() - ), - )); - } - - // Compute paths up front. shard_dir/shard_*_path_seq for a single- - // shard target are pure path computations and do NOT depend on - // self.layout, so it is safe to derive them while layout is still - // TopLevel. - let old_base = self - .aof_dir() - .join(format!("moon.aof.{}.base.rdb", self.seq)); - let old_incr = self - .aof_dir() - .join(format!("moon.aof.{}.incr.aof", self.seq)); - let new_dir = self.aof_dir().join("shard-0"); - let new_base = new_dir.join(format!("moon.aof.{}.base.rdb", self.seq)); - let new_incr = new_dir.join(format!("moon.aof.{}.incr.aof", self.seq)); - - if !old_base.exists() { - // Pre-flight check: nothing moved yet, no rollback needed. - return Err(std::io::Error::new( - std::io::ErrorKind::NotFound, - format!( - "TopLevel→PerShard migration: source base {} not found", - old_base.display() - ), - )); - } - std::fs::create_dir_all(&new_dir)?; - - // Move base. If the rename itself fails, no on-disk mutation has - // happened yet — bail without rollback. Layout stays TopLevel until - // commit at the bottom. - std::fs::rename(&old_base, &new_base)?; - - // Fsync the target directory so the base rename is durable before we - // proceed. A crash after rename but before dir-fsync could leave the - // old filename visible on the next boot. - // - // NOTE: if this fsync fails, old_base has already moved to new_base — - // rollback the rename before returning so the manifest stays consistent. - if let Err(e) = fsync_directory(&new_dir) { - if let Err(re) = std::fs::rename(&new_base, &old_base) { - error!( - "Migration rollback: failed to restore base {} → {} after fsync_directory failure: {}", - new_base.display(), - old_base.display(), - re - ); - } - return Err(e); - } - - // Base is now durably in shard-0/. Any subsequent error must restore it. - let moved_incr: bool; - let created_incr: bool; - if old_incr.exists() { - if let Err(e) = std::fs::rename(&old_incr, &new_incr) { - if let Err(re) = std::fs::rename(&new_base, &old_base) { - error!( - "Migration rollback: failed to restore base {} → {}: {}", - new_base.display(), - old_base.display(), - re - ); - } - return Err(e); - } - // Fsync the shard directory to make the incr rename durable. - // If this fails, roll back both incr and base renames. - if let Err(e) = fsync_directory(&new_dir) { - if let Err(re) = std::fs::rename(&new_incr, &old_incr) { - error!( - "Migration rollback: failed to restore incr {} → {} after fsync_directory failure: {}", - new_incr.display(), - old_incr.display(), - re - ); - } - if let Err(re) = std::fs::rename(&new_base, &old_base) { - error!( - "Migration rollback: failed to restore base {} → {} after fsync_directory failure: {}", - new_base.display(), - old_base.display(), - re - ); - } - return Err(e); - } - moved_incr = true; - created_incr = false; - } else { - match std::fs::File::create(&new_incr) { - Ok(_) => { - moved_incr = false; - created_incr = true; - } - Err(e) => { - if let Err(re) = std::fs::rename(&new_base, &old_base) { - error!( - "Migration rollback: failed to restore base {} → {}: {}", - new_base.display(), - old_base.display(), - re - ); - } - return Err(e); - } - } - } - - // Commit: flip layout, persist as v2. If write_manifest fails, undo - // every filesystem mutation and restore layout so the next boot still - // sees a valid v1 TopLevel deployment. - self.layout = AofLayout::PerShard; - if let Err(e) = self.write_manifest() { - self.layout = AofLayout::TopLevel; - if moved_incr { - if let Err(re) = std::fs::rename(&new_incr, &old_incr) { - error!( - "Migration rollback: failed to restore incr {} → {}: {}", - new_incr.display(), - old_incr.display(), - re - ); - } - } else if created_incr { - if let Err(re) = std::fs::remove_file(&new_incr) { - warn!( - "Migration rollback: failed to remove freshly created incr {}: {}", - new_incr.display(), - re - ); - } - } - if let Err(re) = std::fs::rename(&new_base, &old_base) { - error!( - "Migration rollback: failed to restore base {} → {}: {}. \ - Manifest dir {} may be in an inconsistent state.", - new_base.display(), - old_base.display(), - re, - self.dir.display() - ); - } - return Err(e); - } - - info!( - "AOF migrated: TopLevel → PerShard (single shard) at {}", - self.aof_dir().display() - ); - Ok(()) - } - - /// Create the `appendonlydir/` and write an initial v2 manifest for the - /// given shard count. - /// - /// Each shard gets its own `shard-{N}/` subdirectory with an empty base - /// RDB and an empty incr file. Mirrors `initialize()` semantics: the - /// `(base + incr)` invariant holds from the first boot, so recovery can - /// replay incr-only state without complaint. - /// - /// **Idempotency pre-flight:** if `appendonlydir/moon.aof.manifest` already - /// exists, returns `Err(AlreadyExists)` without modifying any files. A - /// mid-loop crash followed by a retry would otherwise overwrite the already- - /// written shard-0 base RDB with an empty RDB, losing state. Callers that - /// want resume-or-skip semantics should use [`Self::try_initialize_multi`]. - /// - /// **Rollback on partial failure:** if the per-shard loop fails mid-way (e.g. - /// shard-1 write fails after shard-0 succeeded), all already-created shard - /// base RDB files are deleted before returning the error. - pub fn initialize_multi(dir: &Path, num_shards: u16) -> std::io::Result { - if num_shards == 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "initialize_multi requires num_shards >= 1", - )); - } - let manifest = Self { - dir: dir.to_path_buf(), - seq: 1, - layout: AofLayout::PerShard, - shards: (0..num_shards) - .map(|id| ShardManifest { - shard_id: id, - max_lsn: 0, - }) - .collect(), - }; - std::fs::create_dir_all(manifest.aof_dir())?; - - // Pre-flight: refuse if manifest already exists to avoid overwriting - // already-written shard base RDB files (idempotency guard). - let manifest_path = manifest.manifest_path(); - if manifest_path.exists() { - return Err(std::io::Error::new( - std::io::ErrorKind::AlreadyExists, - format!( - "initialize_multi: manifest already exists at {}; \ - use try_initialize_multi() for idempotent initialization", - manifest_path.display() - ), - )); - } - - // Per-shard empty RDB. Single Database::default() inside a 1-element - // slice matches `initialize()`'s empty-RDB shape for each shard. - let empty_dbs: [crate::storage::Database; 0] = []; - let empty_rdb = crate::persistence::rdb::save_to_bytes(&empty_dbs) - .map_err(|e| std::io::Error::other(format!("empty RDB serialize: {e}")))?; - - // Track which shard directories were successfully created so we can - // roll them back on partial failure. - let mut created_shards: Vec = Vec::with_capacity(num_shards as usize); - - let loop_result = (|| -> std::io::Result<()> { - for shard_id in 0..num_shards { - let shard_dir = manifest.shard_dir(shard_id); - std::fs::create_dir_all(&shard_dir)?; - - let base_path = manifest.shard_base_path(shard_id); - let tmp_path = base_path.with_extension("rdb.tmp"); - { - let mut f = std::fs::File::create(&tmp_path)?; - f.write_all(&empty_rdb)?; - f.sync_data()?; - } - std::fs::rename(&tmp_path, &base_path)?; - fsync_parent_best_effort(&base_path); - std::fs::File::create(manifest.shard_incr_path(shard_id))?; - created_shards.push(shard_id); - } - Ok(()) - })(); - - if let Err(e) = loop_result { - // Rollback: remove base RDB files for all successfully-created shards. - for sid in created_shards { - let base = manifest.shard_base_path(sid); - if let Err(re) = std::fs::remove_file(&base) { - warn!( - "initialize_multi rollback: failed to remove {}: {}", - base.display(), - re - ); - } - } - return Err(e); - } - - manifest.write_manifest()?; - Ok(manifest) - } - - /// Initialize a v2 multi-shard manifest only if one does not already exist. - /// - /// Returns `Ok(Some(manifest))` on successful creation, or `Ok(None)` if the - /// manifest file already existed (already initialized — no files modified). - /// Returns `Err(_)` only on actual I/O failures. - pub fn try_initialize_multi(dir: &Path, num_shards: u16) -> std::io::Result> { - match Self::initialize_multi(dir, num_shards) { - Ok(m) => Ok(Some(m)), - Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(None), - Err(e) => Err(e), - } - } - - /// Advance to the next sequence: write new base RDB, create new incr file, - /// update manifest, delete old files. - /// - /// Returns the path to the new incremental file (caller should switch writing to it). - pub fn advance(&mut self, rdb_bytes: &[u8]) -> Result { - let old_seq = self.seq; - let new_seq = old_seq + 1; - - let aof_dir = self.aof_dir(); - std::fs::create_dir_all(&aof_dir).map_err(|e| crate::error::AofError::Io { - path: aof_dir.clone(), - source: e, - })?; - - // 1. Write new base RDB (atomic: tmp + fsync + rename). - // Must fsync the data BEFORE renaming — a rename without prior fsync - // can publish a file whose contents aren't durable, so a crash leaves - // the manifest pointing at an empty/partial base RDB. - let new_base = self.base_path_seq(new_seq); - let tmp_base = new_base.with_extension("rdb.tmp"); - { - let mut f = - std::fs::File::create(&tmp_base).map_err(|e| crate::error::AofError::Io { - path: tmp_base.clone(), - source: e, - })?; - f.write_all(rdb_bytes) - .map_err(|e| crate::error::AofError::Io { - path: tmp_base.clone(), - source: e, - })?; - f.sync_data().map_err(|e| crate::error::AofError::Io { - path: tmp_base.clone(), - source: e, - })?; - } - std::fs::rename(&tmp_base, &new_base).map_err(|e| { - crate::error::AofError::RewriteFailed { - detail: format!("rename base: {}", e), - } - })?; - fsync_parent_best_effort(&new_base); - - // 2. Create empty new incremental file - let new_incr = self.incr_path_seq(new_seq); - std::fs::File::create(&new_incr).map_err(|e| crate::error::AofError::Io { - path: new_incr.clone(), - source: e, - })?; - - // 3. Update manifest (atomic) - self.seq = new_seq; - self.write_manifest() - .map_err(|e| crate::error::AofError::Io { - path: self.manifest_path(), - source: e, - })?; - - // 4. Delete old files (best-effort) - let old_base = self.base_path_seq(old_seq); - let old_incr = self.incr_path_seq(old_seq); - if old_base.exists() { - if let Err(e) = std::fs::remove_file(&old_base) { - warn!("Failed to delete old base {}: {}", old_base.display(), e); - } - } - if old_incr.exists() { - if let Err(e) = std::fs::remove_file(&old_incr) { - warn!("Failed to delete old incr {}: {}", old_incr.display(), e); - } - } - - info!( - "AOF advanced to seq {}: base={} bytes, incr={}", - new_seq, - rdb_bytes.len(), - new_incr.display() - ); - - Ok(new_incr) - } - - /// Advance a single shard to a new sequence: write the shard's new base RDB, - /// create a new empty incr file, then update the shard's `max_lsn` in the - /// in-memory manifest. - /// - /// **Does NOT delete the old generation's files.** Deletion is deferred to - /// the coordinator via [`prune_shard_files`](Self::prune_shard_files), - /// called only AFTER `write_manifest()` durably commits the new seq. - /// Deleting before the commit would leave a crash window where the - /// persisted (old) seq points at files that are already gone — recovery - /// resolves base/incr by `self.seq`, so a crash mid-fan-out would lose data - /// for any shard that had already advanced. This matches the post-commit - /// deletion ordering in [`advance`](Self::advance) (TopLevel layout). - /// - /// **Caller MUST call `write_manifest()` after all shards have been advanced** - /// (and set `self.seq` to the new seq) to persist the updated manifest - /// atomically — this is the single durable commit point for the rewrite. - /// Advancing shards one at a time and writing the manifest per-shard would - /// leave the manifest in an inconsistent state between calls. - /// - /// For `TopLevel` layout, `shard_id` must be 0 and this delegates to - /// `advance()` (which deletes post-commit internally). For `PerShard` - /// layout, files are written to `shard_dir(shard_id)/`. - /// - /// Returns the path to the new incremental file for this shard. - pub fn advance_shard( - &mut self, - shard_id: u16, - new_seq: u64, - rdb_bytes: &[u8], - ) -> Result { - if self.layout == AofLayout::TopLevel { - debug_assert_eq!(shard_id, 0, "TopLevel layout only has shard 0"); - return self.advance(rdb_bytes); - } - - // Validate shard_id is known in this manifest. - let shard_idx = self - .shards - .iter() - .position(|s| s.shard_id == shard_id) - .ok_or_else(|| crate::error::AofError::RewriteFailed { - detail: format!( - "advance_shard: shard_id {} not in manifest (shards: {})", - shard_id, - self.shards.len() - ), - })?; - - let shard_dir = self.shard_dir(shard_id); - std::fs::create_dir_all(&shard_dir).map_err(|e| crate::error::AofError::Io { - path: shard_dir.clone(), - source: e, - })?; - - // 1. Write new base RDB atomically: tmp + fsync + rename. - let new_base = self.shard_base_path_seq(shard_id, new_seq); - let tmp_base = new_base.with_extension("rdb.tmp"); - { - let mut f = - std::fs::File::create(&tmp_base).map_err(|e| crate::error::AofError::Io { - path: tmp_base.clone(), - source: e, - })?; - f.write_all(rdb_bytes) - .map_err(|e| crate::error::AofError::Io { - path: tmp_base.clone(), - source: e, - })?; - f.sync_data().map_err(|e| crate::error::AofError::Io { - path: tmp_base.clone(), - source: e, - })?; - } - std::fs::rename(&tmp_base, &new_base).map_err(|e| { - crate::error::AofError::RewriteFailed { - detail: format!( - "advance_shard {}: rename base {}: {}", - shard_id, - tmp_base.display(), - e - ), - } - })?; - fsync_parent_best_effort(&new_base); - - // 2. Create empty new incremental file. - let new_incr = self.shard_incr_path_seq(shard_id, new_seq); - std::fs::File::create(&new_incr).map_err(|e| crate::error::AofError::Io { - path: new_incr.clone(), - source: e, - })?; - - // 3. Update per-shard LSN in-memory (manifest write is the caller's job). - // Old-generation files are intentionally NOT deleted here — the - // coordinator prunes them via `prune_shard_files` only after - // `write_manifest()` durably commits the new seq (see the fn doc; - // delete-before-commit would lose data on a mid-fan-out crash). - self.shards[shard_idx].max_lsn = self.shards[shard_idx].max_lsn.max(new_seq); - - info!( - "AOF shard {} advanced to seq {}: base={} bytes, incr={}", - shard_id, - new_seq, - rdb_bytes.len(), - new_incr.display() - ); - - Ok(new_incr) - } - - /// Delete a shard's base + incr files for a specific `seq`. Best-effort. - /// - /// **Crash-safety contract:** the rewrite coordinator MUST call this only - /// AFTER `write_manifest()` has durably committed the new seq. Deleting an - /// old generation's files before the manifest flips would orphan the - /// persisted (old) seq whose files are already gone — recovery resolves - /// base/incr by `self.seq`, so it would read a missing base and lose data - /// for any shard that completed before the crash. This mirrors the - /// post-commit deletion ordering in `advance()` (TopLevel layout). - pub fn prune_shard_files(&self, shard_id: u16, seq: u64) { - let base = self.shard_base_path_seq(shard_id, seq); - let incr = self.shard_incr_path_seq(shard_id, seq); - if base.exists() { - if let Err(e) = std::fs::remove_file(&base) { - warn!( - "prune_shard_files {}: failed to delete old base {}: {}", - shard_id, - base.display(), - e - ); - } - } - if incr.exists() { - if let Err(e) = std::fs::remove_file(&incr) { - warn!( - "prune_shard_files {}: failed to delete old incr {}: {}", - shard_id, - incr.display(), - e - ); - } - } - } -} - -/// Replay multi-part AOF: load base RDB then replay incremental RESP. -/// -/// Returns total keys/commands loaded. -pub fn replay_multi_part( - databases: &mut [crate::storage::Database], - manifest: &AofManifest, - engine: &dyn crate::persistence::replay::CommandReplayEngine, -) -> Result { - let mut total = 0usize; - - // Load base RDB - let base_path = manifest.base_path(); - if base_path.exists() { - match crate::persistence::rdb::load(databases, &base_path) { - Ok(n) => { - info!( - "AOF base RDB loaded: {} keys from {}", - n, - base_path.display() - ); - total += n; - } - Err(e) => { - // Base RDB is corrupt or unreadable — applying incremental - // deltas on top of missing/corrupt base gives wrong results. - error!("AOF base RDB load failed: {}", e); - return Err(e); - } - } - } else { - // Missing base is tolerable only when the incr log is also empty - // (fresh manifest from initialize(), or first boot after legacy - // upgrade). If there's incremental content but no base, replaying - // deltas (DEL, EXPIRE, HINCRBY, …) on an empty database produces - // incorrect state — fail loudly rather than silently corrupt. - let incr_path = manifest.incr_path(); - let incr_len = std::fs::metadata(&incr_path).map(|m| m.len()).unwrap_or(0); - if incr_len > 0 { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "AOF base RDB missing at {} but incr {} is {} bytes; refusing to replay incr against empty state", - base_path.display(), - incr_path.display(), - incr_len, - ), - }, - )); - } - warn!( - "AOF base RDB not found: {} (incr empty, treating as fresh init)", - base_path.display() - ); - } - - // Replay incremental RESP - let incr_path = manifest.incr_path(); - if incr_path.exists() { - let data = std::fs::read(&incr_path)?; - if !data.is_empty() { - // Pure RESP — use replay_aof_resp (no RDB preamble detection needed) - let count = replay_incr_resp(databases, &data, engine)?; - info!( - "AOF incr replayed: {} commands from {}", - count, - incr_path.display() - ); - total += count; - } - } - - Ok(total) -} - -/// Replay pure RESP commands from a byte slice. -/// -/// **Corruption handling:** On mid-stream parse errors this returns an error -/// rather than silently resyncing to the next `*` byte. Silent resync in a -/// multi-part AOF is dangerous: an undetected run of dropped commands leaves -/// the database in an inconsistent state that cannot be reconstructed. -/// Truncated tails (parser returns `Ok(None)` with bytes remaining) are -/// logged and treated as the legitimate end of the incremental log, matching -/// `replay_aof` semantics for crash-time tail truncation. -fn replay_incr_resp( - databases: &mut [crate::storage::Database], - data: &[u8], - engine: &dyn crate::persistence::replay::CommandReplayEngine, -) -> Result { - use crate::protocol::{Frame, ParseConfig, parse}; - use bytes::BytesMut; - - let total_len = data.len(); - let mut buf = BytesMut::from(data); - let config = ParseConfig::default(); - let mut selected_db: usize = 0; - let mut count: usize = 0; - - loop { - if buf.is_empty() { - break; - } - match parse::parse(&mut buf, &config) { - Ok(Some(frame)) => { - let (cmd, cmd_args) = match &frame { - Frame::Array(arr) if !arr.is_empty() => { - let name = match &arr[0] { - Frame::BulkString(s) => s.as_ref(), - Frame::SimpleString(s) => s.as_ref(), - other => { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "AOF incr command at offset {} has non-string name frame: {:?}", - total_len - buf.len(), - std::mem::discriminant(other) - ), - }, - )); - } - }; - (name as &[u8], &arr[1..]) - } - other => { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "AOF incr non-array frame at offset {}: {:?}", - total_len - buf.len(), - std::mem::discriminant(other) - ), - }, - )); - } - }; - engine.replay_command(databases, cmd, cmd_args, &mut selected_db); - count += 1; - } - Ok(None) => { - if !buf.is_empty() { - let offset = total_len - buf.len(); - warn!( - "AOF incr truncated tail: {} bytes at offset {} (treating as crash-time EOF)", - buf.len(), - offset - ); - } - break; - } - Err(e) => { - let offset = total_len - buf.len(); - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!("AOF incr parse error at offset {}: {:?}", offset, e), - }, - )); - } - } - } - - Ok(count) -} - -/// An entry that was tagged `OrderedAcrossShards` (RFC § 2 Rule 2) and -/// must be merge-replayed in global LSN order after per-shard replay -/// completes. The `shard_id` records which shard's file it came from so -/// the merge step can dispatch each entry back to its origin shard's -/// databases. -#[derive(Debug, Clone)] -pub struct OrderedEntry { - pub shard_id: u16, - pub lsn: u64, - pub bytes: bytes::Bytes, -} - -/// Replay a framed PerShard incr file: `[u64 lsn LE][u32 len LE][RESP bytes]`. -/// -/// Step 3 wrote this format; step 4 reads it. Step 5 extends the LSN field: -/// the high bit (`crate::persistence::aof::ORDERED_LSN_FLAG`) marks the -/// entry as `OrderedAcrossShards` — those entries are NOT replayed inline, -/// instead they are pushed into `ordered_buf` for the caller to merge-replay -/// in global LSN order across all shards. -/// -/// Returns `(commands_replayed, max_lsn)` — the count covers only inline -/// (non-ordered) replays. `max_lsn` is the NEXT-FREE replication offset: -/// `max(entry.lsn + entry.len)` across both inline AND ordered entries (the -/// high bit is masked out before the computation). It is the offset AFTER the -/// last byte on disk, NOT the start LSN of the last entry — because -/// `ReplicationState::issue_lsn` returns the offset BEFORE adding the entry -/// length, seeding `master_repl_offset` with a start LSN would reissue the -/// last on-disk LSN and break the lsn->entry uniqueness invariant (F5). -/// -/// **Truncated entries:** a header partly written at crash time is treated as -/// EOF (parity with `replay_incr_resp` semantics). A whole header followed by -/// a truncated payload is also EOF — the writer's invariant is that the -/// header is written first then the payload, and on partial write the most we -/// can lose is the last entry's payload tail. -/// -/// **Corruption:** a mid-stream RESP parse error inside an otherwise-complete -/// payload is fatal (same reasoning as `replay_incr_resp`). -fn replay_incr_framed( - shard_id: u16, - databases: &mut [crate::storage::Database], - data: &[u8], - engine: &dyn crate::persistence::replay::CommandReplayEngine, - ordered_buf: &mut Vec, -) -> Result<(usize, u64), crate::error::MoonError> { - use crate::protocol::{Frame, ParseConfig, parse}; - use bytes::BytesMut; - - const HEADER_LEN: usize = 12; // u64 lsn LE + u32 len LE - - let total_len = data.len(); - let mut offset: usize = 0; - let config = ParseConfig::default(); - let mut selected_db: usize = 0; - let mut count: usize = 0; - let mut max_lsn: u64 = 0; - - while offset < total_len { - if total_len - offset < HEADER_LEN { - warn!( - "AOF incr framed truncated header: {} bytes at offset {} (treating as crash-time EOF)", - total_len - offset, - offset - ); - break; - } - // SAFETY: line 1491 guarantees `total_len - offset >= HEADER_LEN` (=12), - // so the [offset..offset+8] and [offset+8..offset+12] slices are valid - // and `try_into()` to a fixed-size array cannot fail (length-matched). - #[allow(clippy::unwrap_used)] // bounds-checked above; try_into is statically length-matched - let raw_lsn = u64::from_le_bytes(data[offset..offset + 8].try_into().expect("8 bytes")); - #[allow(clippy::unwrap_used)] // same bounds-check guarantee - let len = - u32::from_le_bytes(data[offset + 8..offset + 12].try_into().expect("4 bytes")) as usize; - let payload_start = offset + HEADER_LEN; - let payload_end = payload_start.saturating_add(len); - if payload_end > total_len { - warn!( - "AOF incr framed truncated payload at offset {} (lsn {:#x}, declared len {}, have {} bytes); treating as crash-time EOF", - offset, - raw_lsn, - len, - total_len - payload_start - ); - break; - } - - // Strip the OrderedAcrossShards flag to recover the true LSN. - let is_ordered = raw_lsn & crate::persistence::aof::ORDERED_LSN_FLAG != 0; - let lsn = raw_lsn & !crate::persistence::aof::ORDERED_LSN_FLAG; - - // Ordered entries: buffer for cross-shard merge replay; do NOT - // dispatch inline. - if is_ordered { - let bytes = bytes::Bytes::copy_from_slice(&data[payload_start..payload_end]); - ordered_buf.push(OrderedEntry { - shard_id, - lsn, - bytes, - }); - // F5: track the next-free replication offset (entry end), not the - // start LSN — `issue_lsn` returns the offset BEFORE adding the - // entry length, so the seed must clear every byte already on disk. - let entry_end = lsn + len as u64; - if entry_end > max_lsn { - max_lsn = entry_end; - } - offset = payload_end; - continue; - } - - // Parse RESP from the payload slice. A standalone slice ensures one - // header maps to exactly one command — no implicit pipelining across - // headers. - let mut buf = BytesMut::from(&data[payload_start..payload_end]); - match parse::parse(&mut buf, &config) { - Ok(Some(frame)) => { - let (cmd, cmd_args) = match &frame { - Frame::Array(arr) if !arr.is_empty() => { - let name = match &arr[0] { - Frame::BulkString(s) => s.as_ref(), - Frame::SimpleString(s) => s.as_ref(), - other => { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "AOF incr framed command at offset {} (lsn {}) has non-string name frame: {:?}", - offset, - lsn, - std::mem::discriminant(other) - ), - }, - )); - } - }; - (name as &[u8], &arr[1..]) - } - other => { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "AOF incr framed non-array frame at offset {} (lsn {}): {:?}", - offset, - lsn, - std::mem::discriminant(other) - ), - }, - )); - } - }; - engine.replay_command(databases, cmd, cmd_args, &mut selected_db); - count += 1; - // F5: next-free offset = entry start LSN + RESP byte length. - let entry_end = lsn + len as u64; - if entry_end > max_lsn { - max_lsn = entry_end; - } - } - Ok(None) => { - // Header said `len` bytes of RESP, but parser can't make a - // frame from those bytes. That's corruption inside a fully - // declared payload, not a truncated tail — escalate. - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "AOF incr framed payload at offset {} (lsn {}, len {}) parsed as incomplete frame; corrupt entry", - offset, lsn, len - ), - }, - )); - } - Err(e) => { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "AOF incr framed parse error at offset {} (lsn {}, len {}): {:?}", - offset, lsn, len, e - ), - }, - )); - } - } - - offset = payload_end; - } - - Ok((count, max_lsn)) -} - -/// Replay a PerShard multi-part AOF into N parallel `Vec` buffers. -/// -/// `per_shard_databases[i]` is shard `i`'s database vector. The manifest's -/// `shards` length MUST equal `per_shard_databases.len()`; the caller is -/// expected to have run [`AofManifest::verify_shard_count`] at boot. -/// -/// Per-shard replay is fully parallel: each shard's base RDB load and incr -/// replay run in a separate OS thread via `std::thread::scope`. Shards are -/// independent (different `DashTable` instances, no shared mutable state), so -/// this is safe and correct. Parallelism delivers the RFC § 1 benefit on -/// multi-shard deployments with large AOF files. -/// -/// The `engine_factory` closure is called once per shard thread to produce an -/// independent replay engine. This is required because `CommandReplayEngine` -/// implementations (e.g., `DispatchReplayEngine` under the `graph` feature) -/// may contain non-`Sync` state (`RefCell`) that cannot be safely shared across -/// threads. Each thread owns its own engine; results (total count, max LSN, -/// ordered entries) are collected and merged in the caller thread after all -/// shard threads complete. -/// -/// Returns `(total_commands_replayed, global_max_lsn, ordered_entries)`: -/// - `total_commands_replayed` covers all inline (non-ordered) entries -/// plus the base-RDB key count. -/// - `global_max_lsn` is `max(per-shard max LSN)` across both inline and -/// ordered entries; the caller is expected to call -/// `ReplicationState::seed_master_offset(global_max_lsn)` before -/// accepting client traffic (RFC § 2 Rule 3). -/// - `ordered_entries` is the set of `OrderedAcrossShards`-tagged entries -/// across ALL shards; the caller passes them to -/// [`replay_ordered_merge`] for the cross-shard merge replay. -pub fn replay_per_shard( - per_shard_databases: &mut [&mut [crate::storage::Database]], - manifest: &AofManifest, - engine_factory: &( - dyn Fn() -> Box + Sync - ), -) -> Result<(usize, u64, Vec), crate::error::MoonError> { - debug_assert_eq!( - manifest.layout, - AofLayout::PerShard, - "replay_per_shard called on TopLevel manifest" - ); - if manifest.shards.len() != per_shard_databases.len() { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "replay_per_shard shard-count mismatch: manifest has {} shards, caller passed {} database vectors", - manifest.shards.len(), - per_shard_databases.len() - ), - }, - )); - } - - // Per-shard type alias for the thread result. - type ShardResult = Result<(usize, u64, Vec), crate::error::MoonError>; - - // Use std::thread::scope so each shard thread borrows its databases slice - // without a 'static lifetime requirement. All threads complete before scope - // exits, which satisfies the borrow checker. Errors are propagated via - // a Vec collected after join. - let shard_results: Vec = std::thread::scope(|scope| { - let mut handles = Vec::with_capacity(per_shard_databases.len()); - - for (shard_id, databases) in per_shard_databases.iter_mut().enumerate() { - let sid = shard_id as u16; - let base_path = manifest.shard_base_path(sid); - let incr_path = manifest.shard_incr_path(sid); - let engine = engine_factory(); - - handles.push(scope.spawn(move || -> ShardResult { - let mut shard_total: usize = 0; - let mut shard_max_lsn: u64 = 0; - let mut shard_ordered: Vec = Vec::new(); - - // Load this shard's base RDB. - if base_path.exists() { - match crate::persistence::rdb::load(*databases, &base_path) { - Ok(n) => { - info!( - "AOF shard-{} base RDB loaded: {} keys from {}", - sid, - n, - base_path.display() - ); - shard_total += n; - } - Err(e) => { - error!("AOF shard-{} base RDB load failed: {}", sid, e); - return Err(e); - } - } - } else { - // Missing base is tolerable only when this shard's incr file is - // empty (or absent). Same invariant as `replay_multi_part`. - let incr_len = - std::fs::metadata(&incr_path).map(|m| m.len()).unwrap_or(0); - if incr_len > 0 { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "AOF shard-{} base RDB missing at {} but incr {} is {} bytes; refusing to replay incr against empty state", - sid, - base_path.display(), - incr_path.display(), - incr_len, - ), - }, - )); - } - warn!( - "AOF shard-{} base RDB not found: {} (incr empty, treating as fresh init)", - sid, - base_path.display() - ); - } - - // Replay this shard's framed incr file. - if incr_path.exists() { - let data = std::fs::read(&incr_path).map_err(|e| { - crate::error::MoonError::from(crate::error::AofError::Io { - path: incr_path.clone(), - source: e, - }) - })?; - if !data.is_empty() { - let (count, max_lsn) = replay_incr_framed( - sid, - *databases, - &data, - engine.as_ref(), - &mut shard_ordered, - )?; - info!( - "AOF shard-{} incr replayed: {} commands from {} (max lsn {})", - sid, - count, - incr_path.display(), - max_lsn - ); - shard_total += count; - if max_lsn > shard_max_lsn { - shard_max_lsn = max_lsn; - } - } - } - - Ok((shard_total, shard_max_lsn, shard_ordered)) - })); - } - - // Collect results in shard order. - handles - .into_iter() - .map(|h| { - h.join().unwrap_or_else(|_| { - Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: "replay_per_shard worker thread panicked".to_owned(), - }, - )) - }) - }) - .collect() - }); - - // Merge per-shard results. - let mut total: usize = 0; - let mut global_max_lsn: u64 = 0; - let mut ordered_entries: Vec = Vec::new(); - - for result in shard_results { - let (shard_total, shard_max_lsn, shard_ordered) = result?; - total += shard_total; - if shard_max_lsn > global_max_lsn { - global_max_lsn = shard_max_lsn; - } - ordered_entries.extend(shard_ordered); - } - - Ok((total, global_max_lsn, ordered_entries)) -} - -/// Merge-replay `OrderedAcrossShards` entries collected across all shards -/// in global LSN order (RFC § 2 Rule 2). -/// -/// `entries` is sorted by `lsn` ascending, then each entry is dispatched -/// against its origin shard's databases — the per-shard partition is -/// preserved because each `OrderedEntry` carries the `shard_id` it was -/// read from. This guarantees that a cross-shard atomic operation -/// committed at LSN N is replayed as a coherent group (every -/// shard's portion at LSN N is applied before any shard's LSN N+1 work). -/// -/// **Crash-time atomicity:** if a cross-shard commit was mid-write at -/// crash time, some shards may have the LSN-N entry while others don't. -/// Step 5 ships the merge mechanism only; detecting partial commits and -/// performing the corresponding rollback is left to the future cross-shard -/// TXN consumer — `replay_ordered_merge` currently best-effort-applies -/// whichever entries survived. A `warn!` is emitted when the entry count -/// per LSN is uneven across shards so operators have a forensic trail. -/// -/// **Today's emitters:** none in production code. The path is exercised -/// by tests so the round-trip wiring is verified end-to-end and ready for -/// future use. -pub fn replay_ordered_merge( - per_shard_databases: &mut [&mut [crate::storage::Database]], - mut entries: Vec, - engine: &dyn crate::persistence::replay::CommandReplayEngine, -) -> Result { - use crate::protocol::{Frame, ParseConfig, parse}; - use bytes::BytesMut; - - if entries.is_empty() { - return Ok(0); - } - - entries.sort_by_key(|e| e.lsn); - - // Per-LSN cardinality audit: detect torn cross-shard commits. - // - // A "torn" commit is one where LSN N appears in fewer shard files than - // the maximum cardinality seen for any other LSN in this batch. Applying - // partial entries violates atomicity — if the write was interrupted mid- - // commit (e.g., crash between shard-0 and shard-1 writes), replaying only - // the shard-0 portion produces an inconsistent state that cannot be - // compensated. DROP the entire torn LSN instead of applying partial data. - // - // NOTE: "torn" detection is heuristic — it compares each LSN's count - // against the maximum cardinality observed. An LSN that legitimately spans - // fewer shards (e.g. single-shard ordered op) can only occur if the batch - // is heterogeneous. Production emitters (future cross-shard TXN) must - // guarantee uniform cardinality per LSN, so this heuristic is correct for - // all currently-reachable code paths. - let mut counts: std::collections::BTreeMap = std::collections::BTreeMap::new(); - for e in &entries { - *counts.entry(e.lsn).or_insert(0) += 1; - } - let max_count = counts.values().copied().max().unwrap_or(0); - let mut torn_lsns: std::collections::BTreeSet = std::collections::BTreeSet::new(); - for (&lsn, &n) in &counts { - if n < max_count { - warn!( - "OrderedAcrossShards LSN {} appears in only {} of {} shard files; \ - torn cross-shard commit detected — dropping entry for atomicity", - lsn, n, max_count - ); - torn_lsns.insert(lsn); - } - } - - let config = ParseConfig::default(); - let mut replayed: usize = 0; - - for entry in entries { - // Skip entries belonging to a torn (partially-written) commit. - if torn_lsns.contains(&entry.lsn) { - continue; - } - let shard_idx = entry.shard_id as usize; - if shard_idx >= per_shard_databases.len() { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "OrderedAcrossShards entry references shard {} but only {} shards present", - entry.shard_id, - per_shard_databases.len() - ), - }, - )); - } - let mut buf = BytesMut::from(entry.bytes.as_ref()); - match parse::parse(&mut buf, &config) { - Ok(Some(Frame::Array(arr))) if !arr.is_empty() => { - let cmd = match &arr[0] { - Frame::BulkString(s) => s.as_ref(), - Frame::SimpleString(s) => s.as_ref(), - _ => { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "OrderedAcrossShards entry at lsn {} has non-string command frame", - entry.lsn - ), - }, - )); - } - }; - let mut selected_db: usize = 0; - let databases = &mut *per_shard_databases[shard_idx]; - engine.replay_command(databases, cmd, &arr[1..], &mut selected_db); - replayed += 1; - } - other => { - return Err(crate::error::MoonError::from( - crate::error::AofError::RewriteFailed { - detail: format!( - "OrderedAcrossShards entry at lsn {} on shard {} did not parse as RESP array: {:?}", - entry.lsn, - entry.shard_id, - other.map(|_| ()).err() - ), - }, - )); - } - } - } - - Ok(replayed) -} - -#[cfg(test)] -mod tests_v2 { - //! Unit tests for the v2 (PerShard) manifest format. - //! - //! Covers the Step 1 deliverable of the per-shard AOF RFC: - //! - v1 manifests continue to load as TopLevel (single-shard, shard_id=0) - //! - v2 round-trip: write → load → equivalent struct shape - //! - shard count mismatch produces the verbatim RFC § 3 error - //! - migrate_top_level_to_per_shard performs in-place rename and rewrites - //! the manifest as v2 - //! - global_max_lsn computes max across shards - //! - is_legacy_top_level_layout detects top-level files - - use super::*; - use std::fs; - - fn temp_dir() -> PathBuf { - // Use a global atomic counter so parallel test threads (cargo test runs - // unit tests in parallel) never produce the same directory name even - // when PID and nanosecond clock resolution are the same for two threads. - static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); - let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - let d = std::env::temp_dir().join(format!( - "moon-aof-manifest-test-{}-{}", - std::process::id(), - n, - )); - fs::create_dir_all(&d).expect("temp dir create"); - d - } - - #[test] - fn v1_manifest_loads_as_top_level_single_shard() { - let dir = temp_dir(); - let m = AofManifest::initialize(&dir).expect("initialize v1"); - - assert_eq!(m.layout, AofLayout::TopLevel); - assert_eq!(m.shards.len(), 1); - assert_eq!(m.shards[0].shard_id, 0); - assert_eq!(m.shards[0].max_lsn, 0); - - // Reload from disk - let reloaded = AofManifest::load(&dir).expect("load").expect("present"); - assert_eq!(reloaded.layout, AofLayout::TopLevel); - assert_eq!(reloaded.shards.len(), 1); - assert_eq!(reloaded.seq, m.seq); - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn v2_manifest_round_trips() { - let dir = temp_dir(); - let m = AofManifest::initialize_multi(&dir, 4).expect("initialize_multi"); - - assert_eq!(m.layout, AofLayout::PerShard); - assert_eq!(m.shards.len(), 4); - for (i, s) in m.shards.iter().enumerate() { - assert_eq!(s.shard_id, i as u16); - assert_eq!(s.max_lsn, 0); - } - - // Per-shard subdirs were created with empty base + incr. - for i in 0..4u16 { - assert!(m.shard_dir(i).exists(), "shard-{} dir exists", i); - assert!(m.shard_base_path(i).exists(), "shard-{} base exists", i); - assert!(m.shard_incr_path(i).exists(), "shard-{} incr exists", i); - } - - let reloaded = AofManifest::load(&dir).expect("load").expect("present"); - assert_eq!(reloaded.layout, AofLayout::PerShard); - assert_eq!(reloaded.shards.len(), 4); - assert_eq!(reloaded.seq, m.seq); - for (i, s) in reloaded.shards.iter().enumerate() { - assert_eq!(s.shard_id, i as u16); - } - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn verify_shard_count_emits_rfc_error_verbatim() { - let m = AofManifest { - dir: PathBuf::from("/tmp/nowhere"), - seq: 1, - layout: AofLayout::PerShard, - shards: vec![ - ShardManifest { - shard_id: 0, - max_lsn: 0, - }, - ShardManifest { - shard_id: 1, - max_lsn: 0, - }, - ], - }; - let err = m.verify_shard_count(4).expect_err("should mismatch"); - assert_eq!( - err, - "ERR shard count changed (manifest=2, config=4); refusing to start to avoid data loss. See docs/runbooks/shard-count-change.md" - ); - - // Matching count succeeds. - m.verify_shard_count(2).expect("match"); - } - - #[test] - fn migrate_top_level_to_per_shard_moves_files_and_rewrites_manifest() { - let dir = temp_dir(); - let mut m = AofManifest::initialize(&dir).expect("initialize v1"); - - // Write a marker into the incr file so we can prove the contents - // survive the rename. - let original_incr = m.aof_dir().join(format!("moon.aof.{}.incr.aof", m.seq)); - fs::write(&original_incr, b"MARKER").expect("write incr marker"); - - m.migrate_top_level_to_per_shard().expect("migrate"); - - assert_eq!(m.layout, AofLayout::PerShard); - assert!(!original_incr.exists(), "old incr removed by rename"); - let new_incr = m.shard_incr_path(0); - assert!(new_incr.exists(), "new shard-0 incr exists"); - let contents = fs::read(&new_incr).expect("read new incr"); - assert_eq!(contents, b"MARKER", "incr contents preserved"); - - // Reloaded manifest is v2. - let reloaded = AofManifest::load(&dir).expect("load").expect("present"); - assert_eq!(reloaded.layout, AofLayout::PerShard); - assert_eq!(reloaded.shards.len(), 1); - - // Idempotency: second call is a no-op. - m.migrate_top_level_to_per_shard().expect("idempotent"); - assert_eq!(m.layout, AofLayout::PerShard); - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn global_max_lsn_returns_max_across_shards() { - let m = AofManifest { - dir: PathBuf::from("/tmp/nowhere"), - seq: 1, - layout: AofLayout::PerShard, - shards: vec![ - ShardManifest { - shard_id: 0, - max_lsn: 100, - }, - ShardManifest { - shard_id: 1, - max_lsn: 500, - }, - ShardManifest { - shard_id: 2, - max_lsn: 250, - }, - ], - }; - assert_eq!(m.global_max_lsn(), 500); - } - - #[test] - fn is_legacy_top_level_layout_detects_v1_files() { - let dir = temp_dir(); - // No appendonlydir yet → false. - assert!(!AofManifest::is_legacy_top_level_layout(&dir)); - - // After v1 initialize, top-level files present → true. - let _m = AofManifest::initialize(&dir).expect("init v1"); - assert!(AofManifest::is_legacy_top_level_layout(&dir)); - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn is_legacy_top_level_layout_returns_false_for_v2() { - let dir = temp_dir(); - let _m = AofManifest::initialize_multi(&dir, 2).expect("init v2"); - assert!( - !AofManifest::is_legacy_top_level_layout(&dir), - "v2 layout has no top-level moon.aof.* files" - ); - - fs::remove_dir_all(&dir).ok(); - } - - /// FIX-W3-4: v2 manifest with stray top-level .base.rdb must return false, - /// not true. The filename scan is misleading when a valid v2 manifest exists. - /// - /// Scenario: operator upgraded to v2 but left a stale `moon.aof.1.base.rdb` - /// at the top level (e.g., copied during debugging). `is_legacy_top_level_layout` - /// must check the manifest first and return false when v2 is confirmed. - #[test] - fn is_legacy_top_level_layout_ignores_stray_files_when_v2_manifest_present() { - let dir = temp_dir(); - // Initialize a genuine v2 (PerShard) layout. - let _m = AofManifest::initialize_multi(&dir, 2).expect("init v2"); - - // Plant a stale top-level base.rdb to simulate the stray-file scenario. - let stray = dir.join(AOF_DIR_NAME).join("moon.aof.1.base.rdb"); - fs::write(&stray, b"REDIS0011\xff").expect("write stray base.rdb"); - - // Even though the stray file matches the filename pattern, a valid v2 - // manifest is present, so is_legacy_top_level_layout must return false. - assert!( - !AofManifest::is_legacy_top_level_layout(&dir), - "v2 manifest + stray top-level file must still return false" - ); - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn parse_v2_rejects_shard_count_mismatch_in_file() { - let dir = temp_dir(); - let aof = dir.join(AOF_DIR_NAME); - fs::create_dir_all(&aof).unwrap(); - // Manifest claims shards 3 but only declares two shard records. - fs::write( - aof.join(MANIFEST_NAME), - "version 2\nseq 1\nshards 3\nshard 0 max_lsn 0\nshard 1 max_lsn 0\n", - ) - .unwrap(); - - let err = AofManifest::load(&dir).expect_err("should reject"); - let msg = err.to_string(); - assert!( - msg.contains("declares shards=3 but has 2 shard records"), - "got: {}", - msg - ); - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn parse_v2_rejects_non_contiguous_shard_ids() { - let dir = temp_dir(); - let aof = dir.join(AOF_DIR_NAME); - fs::create_dir_all(&aof).unwrap(); - // shards=2 but ids are {0, 2} not {0, 1}. - fs::write( - aof.join(MANIFEST_NAME), - "version 2\nseq 1\nshards 2\nshard 0 max_lsn 0\nshard 2 max_lsn 0\n", - ) - .unwrap(); - - let err = AofManifest::load(&dir).expect_err("should reject"); - let msg = err.to_string(); - assert!(msg.contains("non-contiguous shard ids"), "got: {}", msg); - - fs::remove_dir_all(&dir).ok(); - } - - // ------------------------------------------------------------------ - // Reviewer-flagged fixes: layout-aware path helpers + migration - // rollback. See the "Verify findings against current code" review - // comment on aof_manifest.rs:669-775 and :688-717. - // ------------------------------------------------------------------ - - #[test] - fn base_incr_paths_route_to_shard_zero_after_migration() { - let dir = temp_dir(); - let mut m = AofManifest::initialize(&dir).expect("init v1"); - // Pre-migration: TopLevel paths under appendonlydir/ directly. - assert_eq!(m.base_path(), m.aof_dir().join("moon.aof.1.base.rdb")); - assert_eq!(m.incr_path(), m.aof_dir().join("moon.aof.1.incr.aof")); - - m.migrate_top_level_to_per_shard().expect("migrate"); - - // Post-migration: single-file helpers must route to shard-0/ so - // replay_multi_part and advance() find the actual files. This is - // the bug the reviewer flagged for aof_manifest.rs:669-775. - let shard0 = m.aof_dir().join("shard-0"); - assert_eq!(m.base_path(), shard0.join("moon.aof.1.base.rdb")); - assert_eq!(m.incr_path(), shard0.join("moon.aof.1.incr.aof")); - assert_eq!(m.base_path_seq(7), shard0.join("moon.aof.7.base.rdb")); - assert_eq!(m.incr_path_seq(7), shard0.join("moon.aof.7.incr.aof")); - // The path the helper returns must be where the file actually lives. - assert!(m.base_path().exists(), "base file at returned path"); - assert!(m.incr_path().exists(), "incr file at returned path"); - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn migrate_rolls_back_filesystem_when_incr_rename_fails() { - // Simulate the rename(old_incr → new_incr) failure path by making - // the destination already exist as a directory (rename onto a - // non-empty directory is an error on every supported OS). - let dir = temp_dir(); - let mut m = AofManifest::initialize(&dir).expect("init v1"); - let original_base = m.aof_dir().join("moon.aof.1.base.rdb"); - let original_incr = m.aof_dir().join("moon.aof.1.incr.aof"); - fs::write(&original_incr, b"INCR_MARKER").expect("seed incr"); - let base_bytes_before = fs::read(&original_base).expect("read base"); - - // Pre-create shard-0/moon.aof.1.incr.aof as a DIRECTORY so the - // rename fails after the base rename has already succeeded. - let shard0 = m.aof_dir().join("shard-0"); - fs::create_dir_all(shard0.join("moon.aof.1.incr.aof")).expect("seed blocker"); - - let err = m - .migrate_top_level_to_per_shard() - .expect_err("incr rename should fail"); - let _ = err; // exact error kind depends on OS - - // Rollback invariants: - // 1. Layout stays TopLevel in memory. - // 2. base file restored to its original TopLevel path. - // 3. base file contents unchanged. - // 4. on-disk manifest is still v1 (load returns layout TopLevel). - assert_eq!(m.layout, AofLayout::TopLevel, "in-memory layout reverted"); - assert!(original_base.exists(), "base restored to top-level"); - let base_bytes_after = fs::read(&original_base).expect("read base"); - assert_eq!(base_bytes_after, base_bytes_before, "base contents intact"); - let reloaded = AofManifest::load(&dir).expect("load").expect("present"); - assert_eq!(reloaded.layout, AofLayout::TopLevel, "on-disk manifest v1"); - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn migrate_does_not_mutate_on_missing_base() { - let dir = temp_dir(); - let mut m = AofManifest::initialize(&dir).expect("init v1"); - let base = m.aof_dir().join("moon.aof.1.base.rdb"); - fs::remove_file(&base).expect("remove base"); - - let err = m - .migrate_top_level_to_per_shard() - .expect_err("missing base should fail"); - assert_eq!(err.kind(), std::io::ErrorKind::NotFound); - // Layout never flipped, no rollback needed. - assert_eq!(m.layout, AofLayout::TopLevel); - - fs::remove_dir_all(&dir).ok(); - } - - // -- Step 4 (per-shard replay) tests --------------------------------- - - fn frame_entry(lsn: u64, resp: &[u8]) -> Vec { - let mut buf = Vec::with_capacity(12 + resp.len()); - buf.extend_from_slice(&lsn.to_le_bytes()); - buf.extend_from_slice(&(resp.len() as u32).to_le_bytes()); - buf.extend_from_slice(resp); - buf - } - - /// Minimal `CommandReplayEngine` that records (lsn-implicit-via-order, cmd - /// name) calls without touching real storage. Tests use this to assert - /// the framed parser hands the right command sequence to the engine. - struct RecordingEngine { - calls: std::cell::RefCell>, - } - - impl RecordingEngine { - fn new() -> Self { - Self { - calls: std::cell::RefCell::new(Vec::new()), - } - } - } - - impl crate::persistence::replay::CommandReplayEngine for RecordingEngine { - fn replay_command( - &self, - _databases: &mut [crate::storage::Database], - cmd: &[u8], - _args: &[crate::protocol::Frame], - _selected_db: &mut usize, - ) { - self.calls - .borrow_mut() - .push(String::from_utf8_lossy(cmd).into_owned()); - } - } - - #[test] - fn replay_incr_framed_decodes_lsn_and_resp() { - // Two framed entries: PING and DBSIZE (no args, both small RESP arrays). - let mut bytes = frame_entry(7, b"*1\r\n$4\r\nPING\r\n"); - bytes.extend_from_slice(&frame_entry(11, b"*1\r\n$6\r\nDBSIZE\r\n")); - - let mut dbs: Vec = vec![crate::storage::Database::new()]; - let engine = RecordingEngine::new(); - let mut ordered: Vec = Vec::new(); - let (count, max_lsn) = - replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered).expect("framed replay"); - assert!(ordered.is_empty(), "no ordered entries in this stream"); - - assert_eq!(count, 2); - // F5: max_lsn is the NEXT-FREE offset = max(lsn + len) = max(7+14, 11+16) = 27. - assert_eq!(max_lsn, 27); - let calls = engine.calls.borrow(); - assert_eq!(calls.len(), 2); - assert_eq!(calls[0], "PING"); - assert_eq!(calls[1], "DBSIZE"); - } - - #[test] - fn replay_incr_framed_max_lsn_is_next_free_offset() { - // F5: replay must return the next-free replication offset (entry end = - // start LSN + RESP byte length), not the START LSN of the last entry. - // `issue_lsn` hands out the offset BEFORE adding the entry's length, so - // seeding `master_repl_offset` with a start LSN reissues the last - // pre-crash entry's LSN — breaking lsn->entry uniqueness (RFC § 2 Rule 3). - let ping = b"*1\r\n$4\r\nPING\r\n"; // 14 bytes - let dbsize = b"*1\r\n$6\r\nDBSIZE\r\n"; // 16 bytes - // Cumulative LSNs as the writer issues them: each entry starts at the - // previous entry's end. - let mut bytes = frame_entry(100, ping); - bytes.extend_from_slice(&frame_entry(100 + ping.len() as u64, dbsize)); - - let mut dbs: Vec = vec![crate::storage::Database::new()]; - let engine = RecordingEngine::new(); - let mut ordered: Vec = Vec::new(); - let (_count, max_lsn) = - replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered).expect("framed replay"); - - // Last entry: start 114 + len 16 = 130. The next write MUST get >= 130. - let expected_next_free = 100 + ping.len() as u64 + dbsize.len() as u64; - assert_eq!(expected_next_free, 130); - assert_eq!( - max_lsn, expected_next_free, - "max_lsn must be the next-free offset (entry end), not the last start LSN" - ); - } - - #[test] - fn replay_incr_framed_truncated_header_is_crash_eof() { - // One valid entry, then a partial 5-byte header (crash mid-write). - let mut bytes = frame_entry(3, b"*1\r\n$4\r\nPING\r\n"); - bytes.extend_from_slice(&[0u8; 5]); - - let mut dbs: Vec = vec![crate::storage::Database::new()]; - let engine = RecordingEngine::new(); - let mut ordered: Vec = Vec::new(); - let (count, max_lsn) = replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered) - .expect("truncated-header is EOF"); - - assert_eq!(count, 1); - // F5: next-free offset = PING entry start 3 + RESP len 14 = 17. - assert_eq!(max_lsn, 17); - } - - #[test] - fn replay_incr_framed_truncated_payload_is_crash_eof() { - // Header declares 14 bytes of RESP but only 5 actually present. - let mut bytes = Vec::new(); - bytes.extend_from_slice(&5u64.to_le_bytes()); - bytes.extend_from_slice(&14u32.to_le_bytes()); - bytes.extend_from_slice(b"*1\r\n$"); // 5 bytes, payload truncated - - let mut dbs: Vec = vec![crate::storage::Database::new()]; - let engine = RecordingEngine::new(); - let mut ordered: Vec = Vec::new(); - let (count, max_lsn) = replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered) - .expect("truncated-payload is EOF"); - - assert_eq!(count, 0); - assert_eq!(max_lsn, 0); - } - - #[test] - fn replay_incr_framed_complete_but_corrupt_payload_errors() { - // Header declares 4 bytes, payload is 4 bytes of garbage that won't - // parse as a RESP frame. - let mut bytes = Vec::new(); - bytes.extend_from_slice(&1u64.to_le_bytes()); - bytes.extend_from_slice(&4u32.to_le_bytes()); - bytes.extend_from_slice(b"XXXX"); - - let mut dbs: Vec = vec![crate::storage::Database::new()]; - let engine = RecordingEngine::new(); - let mut ordered: Vec = Vec::new(); - let err = replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered) - .expect_err("complete-but-corrupt should error"); - let msg = format!("{err}"); - assert!( - msg.contains("framed"), - "error should mention framed context, got: {msg}" - ); - } - - #[test] - fn replay_per_shard_round_trips_two_shards() { - let dir = temp_dir(); - let manifest = AofManifest::initialize_multi(&dir, 2).expect("initialize_multi 2 shards"); - - // Hand-author framed incr files: shard-0 SETs k0/v0 at lsn=10, - // shard-1 SETs k1/v1 at lsn=20. - let set_k0 = frame_entry(10, b"*3\r\n$3\r\nSET\r\n$2\r\nk0\r\n$2\r\nv0\r\n"); - let set_k1 = frame_entry(20, b"*3\r\n$3\r\nSET\r\n$2\r\nk1\r\n$2\r\nv1\r\n"); - fs::write(manifest.shard_incr_path(0), &set_k0).expect("write shard-0 incr"); - fs::write(manifest.shard_incr_path(1), &set_k1).expect("write shard-1 incr"); - - // Two independent shard database vectors. - let mut shard0: Vec = vec![crate::storage::Database::new()]; - let mut shard1: Vec = vec![crate::storage::Database::new()]; - - let (total, global_max_lsn, ordered) = { - let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0, &mut shard1]; - replay_per_shard( - &mut slices, - &manifest, - &(|| { - Box::new(crate::persistence::replay::DispatchReplayEngine::new()) - as Box - }), - ) - .expect("per-shard replay") - }; - - assert_eq!(total, 2, "two SETs replayed"); - // F5: global_max_lsn = max next-free offset across shards. shard-1 SET - // is 29 RESP bytes at lsn 20 → next-free 49 (> shard-0's 10+29=39). - assert_eq!( - global_max_lsn, 49, - "global max lsn = max(shard next-free offsets)" - ); - assert!(ordered.is_empty(), "no ordered entries in this stream"); - - // Each shard's DB now holds its key (and only its key). - assert!(shard0[0].len() >= 1, "shard 0 has k0"); - assert!(shard1[0].len() >= 1, "shard 1 has k1"); - - fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn replay_per_shard_rejects_shard_count_mismatch() { - let dir = temp_dir(); - let manifest = AofManifest::initialize_multi(&dir, 2).expect("initialize_multi 2 shards"); - - // Only one slice — manifest says 2. - let mut shard0: Vec = vec![crate::storage::Database::new()]; - let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0]; - - let err = replay_per_shard( - &mut slices, - &manifest, - &(|| { - Box::new(crate::persistence::replay::DispatchReplayEngine::new()) - as Box - }), - ) - .expect_err("shard count mismatch must error"); - let msg = format!("{err}"); - assert!( - msg.contains("shard-count mismatch"), - "error message should call out the mismatch, got: {msg}" - ); - - fs::remove_dir_all(&dir).ok(); - } - - /// FIX-W3-1: parallel per-shard replay must produce identical results to - /// sequential replay. N=4 shards, one key per shard. - /// - /// Test gate: correctness (same total/max_lsn/key distribution as sequential). - /// Wall-time comparison is flaky in CI and omitted. - #[test] - fn replay_per_shard_parallel_matches_sequential() { - let dir = temp_dir(); - let n_shards: u16 = 4; - let manifest = - AofManifest::initialize_multi(&dir, n_shards).expect("initialize_multi 4 shards"); - - // Each shard gets one SET at lsn = shard_id * 10 + 10. - for sid in 0..n_shards { - let lsn = (sid as u64 + 1) * 10; - let key = format!("k{sid}"); - let val = format!("v{sid}"); - let resp = format!( - "*3\r\n$3\r\nSET\r\n${klen}\r\n{key}\r\n${vlen}\r\n{val}\r\n", - klen = key.len(), - vlen = val.len(), - ); - let entry = frame_entry(lsn, resp.as_bytes()); - fs::write(manifest.shard_incr_path(sid), &entry).expect("write shard incr"); - } - - let mut shards: Vec> = (0..n_shards as usize) - .map(|_| vec![crate::storage::Database::new()]) - .collect(); - - let engine_factory = || { - Box::new(crate::persistence::replay::DispatchReplayEngine::new()) - as Box - }; - let (total, global_max_lsn, ordered) = { - let mut slices: Vec<&mut [crate::storage::Database]> = - shards.iter_mut().map(|s| s.as_mut_slice()).collect(); - replay_per_shard(&mut slices, &manifest, &engine_factory) - .expect("parallel per-shard replay") - }; - - assert_eq!(total, n_shards as usize, "one SET per shard = N total"); - // F5: global_max_lsn = highest shard's NEXT-FREE offset. The highest - // shard (sid=N-1) SETs at lsn N*10 with a 29-byte RESP → next-free - // N*10 + 29 (here 40 + 29 = 69), not the bare start LSN. - assert_eq!( - global_max_lsn, - n_shards as u64 * 10 + 29, - "global max lsn = highest shard next-free offset" - ); - assert!(ordered.is_empty(), "no ordered entries"); - - // Each shard must have exactly one key. - for (sid, shard) in shards.iter().enumerate() { - assert_eq!( - shard[0].len(), - 1, - "shard {} must have exactly 1 key after parallel replay", - sid - ); - } - - fs::remove_dir_all(&dir).ok(); - } - - // -- Step 5 (OrderedAcrossShards merge) tests ------------------------ - - /// Frame an ordered entry: same on-disk layout as `frame_entry`, with - /// the high bit of LSN set. - fn frame_ordered(lsn: u64, resp: &[u8]) -> Vec { - assert_eq!( - lsn & crate::persistence::aof::ORDERED_LSN_FLAG, - 0, - "test helper expects raw lsn without the ordered flag" - ); - let tagged = lsn | crate::persistence::aof::ORDERED_LSN_FLAG; - let mut buf = Vec::with_capacity(12 + resp.len()); - buf.extend_from_slice(&tagged.to_le_bytes()); - buf.extend_from_slice(&(resp.len() as u32).to_le_bytes()); - buf.extend_from_slice(resp); - buf - } - - #[test] - fn replay_incr_framed_buffers_ordered_entries() { - // Mix: normal PING, then an ordered SET, then normal DBSIZE. - let mut bytes = frame_entry(5, b"*1\r\n$4\r\nPING\r\n"); - bytes.extend_from_slice(&frame_ordered( - 8, - b"*3\r\n$3\r\nSET\r\n$1\r\nk\r\n$1\r\nv\r\n", - )); - bytes.extend_from_slice(&frame_entry(12, b"*1\r\n$6\r\nDBSIZE\r\n")); - - let mut dbs: Vec = vec![crate::storage::Database::new()]; - let engine = RecordingEngine::new(); - let mut ordered: Vec = Vec::new(); - let (count, max_lsn) = replay_incr_framed(3, &mut dbs, &bytes, &engine, &mut ordered) - .expect("framed replay with ordered"); - - assert_eq!(count, 2, "two inline entries dispatched (PING, DBSIZE)"); - // F5: max_lsn is the next-free offset across inline AND ordered. The - // ordered SET (27 RESP bytes) at lsn 8 → next-free 35, exceeding the - // PING (5+14=19) and DBSIZE (12+16=28) ends. - assert_eq!( - max_lsn, 35, - "max LSN = next-free offset across inline and ordered" - ); - assert_eq!(ordered.len(), 1, "one entry buffered as ordered"); - let buffered = &ordered[0]; - assert_eq!(buffered.shard_id, 3, "shard_id forwarded"); - assert_eq!(buffered.lsn, 8, "buffered LSN has the high bit masked off"); - let calls = engine.calls.borrow(); - assert_eq!(calls.len(), 2); - assert_eq!(calls[0], "PING"); - assert_eq!(calls[1], "DBSIZE", "ordered SET was NOT dispatched inline"); - } - - #[test] - fn replay_ordered_merge_sorts_by_lsn_across_shards() { - use crate::persistence::replay::DispatchReplayEngine; - - // Three ordered entries across two shards, deliberately out of LSN - // order on the wire so the merge step has work to do. - let entries = vec![ - OrderedEntry { - shard_id: 1, - lsn: 30, - bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\nb1\r\n$1\r\n3\r\n"), - }, - OrderedEntry { - shard_id: 0, - lsn: 10, - bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\na1\r\n$1\r\n1\r\n"), - }, - OrderedEntry { - shard_id: 0, - lsn: 20, - bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\na2\r\n$1\r\n2\r\n"), - }, - ]; - - let mut shard0: Vec = vec![crate::storage::Database::new()]; - let mut shard1: Vec = vec![crate::storage::Database::new()]; - let replayed = { - let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0, &mut shard1]; - replay_ordered_merge(&mut slices, entries, &DispatchReplayEngine::new()) - .expect("ordered merge replay") - }; - - assert_eq!(replayed, 3); - assert!(shard0[0].len() >= 2, "shard 0 received a1 + a2"); - assert!(shard1[0].len() >= 1, "shard 1 received b1"); - } - - #[test] - fn replay_ordered_merge_empty_returns_zero() { - use crate::persistence::replay::DispatchReplayEngine; - - let mut shard0: Vec = vec![crate::storage::Database::new()]; - let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0]; - let replayed = replay_ordered_merge(&mut slices, Vec::new(), &DispatchReplayEngine::new()) - .expect("empty merge ok"); - assert_eq!(replayed, 0); - } - - /// FIX-W3-3: torn cross-shard commit must be DROPPED entirely, not partially applied. - /// - /// Synthesize a 2-shard AOF where LSN 100 appears on shard 0 only (N=1 - /// of K=2 expected). After replay, shard 0 must NOT have the key written - /// by the LSN-100 entry (it was dropped for atomicity). - #[test] - fn replay_ordered_merge_drops_torn_commit() { - use crate::persistence::replay::DispatchReplayEngine; - - // Two shards, two complete entries at LSN 10 (one per shard) — these - // should succeed. LSN 100 appears only on shard 0 (torn) — must be dropped. - let entries = vec![ - // Complete pair: LSN 10 on both shards - OrderedEntry { - shard_id: 0, - lsn: 10, - bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\nc0\r\n$1\r\n1\r\n"), - }, - OrderedEntry { - shard_id: 1, - lsn: 10, - bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\nc1\r\n$1\r\n1\r\n"), - }, - // Torn entry: LSN 100 only on shard 0, not shard 1 - OrderedEntry { - shard_id: 0, - lsn: 100, - bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$5\r\ntorn0\r\n$1\r\nv\r\n"), - }, - ]; - - let mut shard0: Vec = vec![crate::storage::Database::new()]; - let mut shard1: Vec = vec![crate::storage::Database::new()]; - let replayed = { - let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0, &mut shard1]; - replay_ordered_merge(&mut slices, entries, &DispatchReplayEngine::new()) - .expect("ordered merge replay") - }; - - // The torn LSN-100 entry must NOT be applied (dropped for atomicity). - assert_eq!(replayed, 2, "only the complete LSN-10 pair is replayed"); - assert_eq!( - shard0[0].len(), - 1, - "shard-0 only has the complete LSN-10 key; torn LSN-100 entry must not be applied" - ); - // Verify the torn key is absent - assert!( - shard0[0].get(b"torn0").is_none(), - "torn shard-0 entry (LSN 100) must NOT be applied" - ); - } - - #[test] - fn ordered_entry_lsn_flag_set_via_try_send_append_ordered() { - use crate::persistence::aof::{AofMessage, AofWriterPool, ORDERED_LSN_FLAG}; - use crate::runtime::channel; - - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - // Raw lsn = 42; high bit must end up set on the receive side. - pool.try_send_append_ordered(0, 42, bytes::Bytes::from_static(b"x")); - let msg = rx0.try_recv().expect("ordered append delivered"); - match msg { - AofMessage::Append { lsn, .. } => { - assert_eq!( - lsn & ORDERED_LSN_FLAG, - ORDERED_LSN_FLAG, - "ordered flag set on lsn" - ); - assert_eq!( - lsn & !ORDERED_LSN_FLAG, - 42, - "low bits preserve the original lsn" - ); - } - _ => panic!("expected Append"), - } - } - - // ----------------------------------------------------------------------- - // FIX-W2-1: cleanup_orphans must recurse into shard-N/ subdirectories - // ----------------------------------------------------------------------- - - /// Build a minimal PerShard manifest fixture on disk at `seq` without - /// needing `advance_shard`. Directly creates the expected directory layout - /// so the test is self-contained and doesn't depend on FIX-W2-3 methods. - fn write_per_shard_manifest_at_seq(dir: &Path, num_shards: u16, seq: u64) -> AofManifest { - let aof_dir = dir.join(AOF_DIR_NAME); - fs::create_dir_all(&aof_dir).unwrap(); - let empty_rdb = crate::persistence::rdb::save_to_bytes(&[] as &[crate::storage::Database]) - .expect("empty rdb"); - let shards: Vec = (0..num_shards) - .map(|id| ShardManifest { - shard_id: id, - max_lsn: 0, - }) - .collect(); - let manifest = AofManifest { - dir: dir.to_path_buf(), - seq, - layout: AofLayout::PerShard, - shards, - }; - for shard_id in 0..num_shards { - let shard_dir = manifest.shard_dir(shard_id); - fs::create_dir_all(&shard_dir).unwrap(); - let base = manifest.shard_base_path(shard_id); - let tmp = base.with_extension("rdb.tmp"); - fs::write(&tmp, &empty_rdb).unwrap(); - fs::rename(&tmp, &base).unwrap(); - fs::write(manifest.shard_incr_path(shard_id), b"").unwrap(); - } - manifest.write_manifest().unwrap(); - manifest - } - - #[test] - fn cleanup_orphans_removes_stale_files_in_shard_subdirs() { - let dir = temp_dir(); - - // Build a 2-shard PerShard manifest at seq=2. - let manifest = write_per_shard_manifest_at_seq(&dir, 2, 2); - - // Inject orphan files in shard-0/ that a crashed BGREWRITEAOF would leave. - // seq=1 tmp (aborted write) and a seq=5 incr (future zombie). - let shard0_dir = manifest.shard_dir(0); - let orphan_tmp = shard0_dir.join("moon.aof.1.base.rdb.tmp"); - let orphan_old_incr = shard0_dir.join("moon.aof.5.incr.aof"); - fs::write(&orphan_tmp, b"").expect("write orphan tmp"); - fs::write(&orphan_old_incr, b"").expect("write orphan incr"); - - // Active files for seq=2 must survive. - let active_base = manifest.shard_base_path(0); - let active_incr = manifest.shard_incr_path(0); - assert!( - active_base.exists(), - "active base must exist before cleanup" - ); - assert!( - active_incr.exists(), - "active incr must exist before cleanup" - ); - - // Reload the manifest — this triggers cleanup_orphans. - let _reloaded = AofManifest::load(&dir).expect("load").expect("present"); - - assert!( - !orphan_tmp.exists(), - "orphan .rdb.tmp in shard-0/ must be deleted by cleanup_orphans" - ); - assert!( - !orphan_old_incr.exists(), - "orphan old incr in shard-0/ must be deleted by cleanup_orphans" - ); - assert!( - active_base.exists(), - "active seq=2 base must survive cleanup" - ); - assert!( - active_incr.exists(), - "active seq=2 incr must survive cleanup" - ); - - fs::remove_dir_all(&dir).ok(); - } - - // ----------------------------------------------------------------------- - // FIX-W2-2: initialize_multi idempotency — second call returns error - // ----------------------------------------------------------------------- - #[test] - fn initialize_multi_second_call_returns_already_initialized_error() { - let dir = temp_dir(); - - // First call must succeed. - let _m = AofManifest::initialize_multi(&dir, 4).expect("first call ok"); - - // Count files before second call. - let aof_dir = dir.join(AOF_DIR_NAME); - let count_before: usize = (0..4u16) - .map(|sid| { - let shard_dir = aof_dir.join(format!("shard-{}", sid)); - fs::read_dir(&shard_dir).map(|e| e.count()).unwrap_or(0) - }) - .sum(); - - // Second call must return an error with the manifest already present. - let result = AofManifest::initialize_multi(&dir, 4); - assert!( - result.is_err(), - "second initialize_multi must fail when manifest already exists" - ); - let err = result.unwrap_err(); - assert_eq!( - err.kind(), - std::io::ErrorKind::AlreadyExists, - "error kind must be AlreadyExists; got {:?}: {}", - err.kind(), - err - ); - - // File count must be unchanged — no files were overwritten. - let count_after: usize = (0..4u16) - .map(|sid| { - let shard_dir = aof_dir.join(format!("shard-{}", sid)); - fs::read_dir(&shard_dir).map(|e| e.count()).unwrap_or(0) - }) - .sum(); - assert_eq!( - count_before, count_after, - "second call must not create or overwrite any shard files" - ); - - fs::remove_dir_all(&dir).ok(); - } - - // ----------------------------------------------------------------------- - // F6 crash-safety ordering: advance_shard writes new base+incr but MUST - // NOT delete old files. Deleting before the manifest durably commits the - // new seq leaves a window where a crash orphans the persisted (old) seq - // whose files are already gone → recovery reads a missing base → data - // loss for completed shards. Deletion is the coordinator's job, AFTER - // write_manifest(), via prune_shard_files(). This mirrors the proven - // ordering in advance() (TopLevel), which deletes only post-commit. - // ----------------------------------------------------------------------- - #[test] - fn advance_shard_defers_delete_until_after_commit() { - let dir = temp_dir(); - - // Initialize 2-shard manifest at seq=1. - let mut manifest = AofManifest::initialize_multi(&dir, 2).expect("initialize_multi"); - assert_eq!(manifest.seq, 1); - - let empty_rdb = crate::persistence::rdb::save_to_bytes(&[] as &[crate::storage::Database]) - .expect("empty rdb"); - - // Old shard files at seq=1 must exist before advance. - let old_base_s0 = manifest.shard_base_path_seq(0, 1); - let old_incr_s0 = manifest.shard_incr_path_seq(0, 1); - let old_base_s1 = manifest.shard_base_path_seq(1, 1); - let old_incr_s1 = manifest.shard_incr_path_seq(1, 1); - assert!(old_base_s0.exists(), "seq=1 base must exist for shard 0"); - assert!(old_incr_s0.exists(), "seq=1 incr must exist for shard 0"); - - // Fan out: coordinator picks new_seq=2 once and advances every shard - // to it. No manifest write, no deletion, happens inside the fan-out. - let new_incr_s0 = manifest - .advance_shard(0, 2, &empty_rdb) - .expect("advance_shard 0 → seq=2"); - let new_incr_s1 = manifest - .advance_shard(1, 2, &empty_rdb) - .expect("advance_shard 1 → seq=2"); - assert!(new_incr_s0.exists(), "new incr file must be created (s0)"); - assert!(new_incr_s1.exists(), "new incr file must be created (s1)"); - - // PRE-COMMIT INVARIANT: new files written, OLD files NOT yet deleted. - // This is the regression guard against delete-before-commit. - assert!( - manifest.shard_base_path_seq(0, 2).exists(), - "new seq=2 base must exist for shard 0" - ); - assert!( - manifest.shard_base_path_seq(1, 2).exists(), - "new seq=2 base must exist for shard 1" - ); - assert!( - old_base_s0.exists(), - "old seq=1 base (s0) MUST survive until the manifest commits" - ); - assert!( - old_incr_s0.exists(), - "old seq=1 incr (s0) MUST survive until the manifest commits" - ); - assert!( - old_base_s1.exists(), - "old seq=1 base (s1) MUST survive until the manifest commits" - ); - - // COMMIT: coordinator bumps seq and persists the manifest atomically. - // This is the single durable commit point for the whole rewrite. - manifest.seq = 2; - manifest - .write_manifest() - .expect("write manifest after advance"); - - // POST-COMMIT: coordinator prunes old files — safe now that recovery - // resolves base/incr by the durably-committed new seq. - manifest.prune_shard_files(0, 1); - manifest.prune_shard_files(1, 1); - assert!( - !old_base_s0.exists(), - "old seq=1 base (s0) pruned post-commit" - ); - assert!( - !old_incr_s0.exists(), - "old seq=1 incr (s0) pruned post-commit" - ); - assert!( - !old_base_s1.exists(), - "old seq=1 base (s1) pruned post-commit" - ); - assert!( - !old_incr_s1.exists(), - "old seq=1 incr (s1) pruned post-commit" - ); - assert!( - manifest.shard_base_path_seq(0, 2).exists(), - "new seq=2 base (s0) must remain after prune" - ); - - // Recovery reads base by manifest.seq — must resolve to the new seq. - let reloaded = AofManifest::load(&dir).expect("load").expect("present"); - assert_eq!(reloaded.seq, 2); - - fs::remove_dir_all(&dir).ok(); - } - - // ----------------------------------------------------------------------- - // FIX-W2-7: smoke test — fsync helper consolidation did not break - // initialize_multi. Checks that the post-consolidation manifest has the - // correct PerShard layout, the expected shard count, and per-shard - // base/incr files. This is discriminating: a regression that produces a - // TopLevel manifest or wrong shard count will be caught here. - // ----------------------------------------------------------------------- - #[test] - fn initialize_multi_smoke_after_fsync_consolidation() { - let tmp = tempfile::tempdir().expect("tempdir"); - let dir = tmp.path(); - let n: u16 = 2; - let result = AofManifest::initialize_multi(dir, n); - assert!( - result.is_ok(), - "initialize_multi({n} shards) must succeed: {:?}", - result.err() - ); - let manifest = result.unwrap(); - - // Discriminating: layout must be PerShard, not TopLevel. - assert_eq!( - manifest.layout, - AofLayout::PerShard, - "initialize_multi must produce a PerShard manifest" - ); - // Discriminating: shard count must match the requested count. - assert_eq!( - manifest.shards.len() as u16, - n, - "manifest must record exactly {n} shards, got {}", - manifest.shards.len() - ); - // Discriminating: per-shard base RDB and incr files must exist on disk. - for shard_id in 0..n { - assert!( - manifest.shard_base_path(shard_id).exists(), - "shard-{shard_id} base RDB must exist at {}", - manifest.shard_base_path(shard_id).display() - ); - assert!( - manifest.shard_incr_path(shard_id).exists(), - "shard-{shard_id} incr file must exist at {}", - manifest.shard_incr_path(shard_id).display() - ); - } - // Discriminating: the on-disk manifest file must contain `version 2` - // (PerShard v2 header), not be a bare v1 file. - let manifest_path = dir.join(AOF_DIR_NAME).join("moon.aof.manifest"); - let content = - std::fs::read_to_string(&manifest_path).expect("manifest file must be readable"); - assert!( - content.contains("version 2"), - "manifest file must contain 'version 2' (PerShard v2 header); got:\n{}", - content - ); - } -} diff --git a/src/persistence/aof_manifest/mod.rs b/src/persistence/aof_manifest/mod.rs new file mode 100644 index 00000000..9f74ed6a --- /dev/null +++ b/src/persistence/aof_manifest/mod.rs @@ -0,0 +1,1388 @@ +//! Multi-part AOF manifest: tracks base (RDB) and incremental (RESP) files. +//! +//! Part of the **storage format v1** umbrella commitment — see +//! [`docs/STORAGE-FORMAT-V1.md`](../../../docs/STORAGE-FORMAT-V1.md). The +//! manifest framing is the canonical on-disk marker; the human-readable +//! "v1" umbrella also covers WAL v3 and RDB v2 sub-formats. +//! +//! Two on-disk layouts coexist (selected at manifest creation time, never mixed +//! within one directory): +//! +//! **TopLevel (manifest v1, single-shard / legacy):** +//! ```text +//! appendonlydir/ +//! moon.aof.1.base.rdb # RDB snapshot base +//! moon.aof.1.incr.aof # Incremental RESP since base +//! moon.aof.manifest # v1 text format +//! ``` +//! +//! **PerShard (manifest v2, multi-shard durability):** +//! ```text +//! appendonlydir/ +//! moon.aof.manifest # v2 text format (carries shard count + max_lsn) +//! shard-0/ +//! moon.aof.1.base.rdb +//! moon.aof.1.incr.aof +//! shard-1/ +//! moon.aof.1.base.rdb +//! moon.aof.1.incr.aof +//! … +//! ``` +//! +//! The manifest text format is line-prefix based. v1 manifests have no +//! `version` line; v2 manifests begin with `version 2`. On BGREWRITEAOF the +//! sequence increments, a new base + incr pair is created per shard (PerShard) +//! or at top level (TopLevel), and old files are deleted. + +use std::io::Write; +use std::path::{Path, PathBuf}; + +use tracing::{error, info, warn}; + +use crate::persistence::fsync::fsync_directory; + +mod shard_replay; +mod shard_rewrite; + +// Re-export the relocated public replay API so external paths +// (`crate::persistence::aof_manifest::replay_*`) keep resolving (issue #143). +pub use shard_replay::{OrderedEntry, replay_multi_part, replay_ordered_merge, replay_per_shard}; + +const MANIFEST_NAME: &str = "moon.aof.manifest"; +const AOF_DIR_NAME: &str = "appendonlydir"; + +/// Fsync the parent directory of `path` (best-effort). +/// +/// POSIX guarantees atomicity of `rename()` but does NOT guarantee that the +/// directory entry update is durable after a crash. On ext4 and XFS without +/// `data=ordered`, a crash between the rename and a directory fsync can leave +/// the old file name visible on the next boot even though the rename completed +/// in memory. Calling this after every manifest-visible rename closes that gap. +/// +/// Best-effort: logs on failure but does not propagate the error. A failed +/// dir fsync means the rename may not survive a crash — the worst case is +/// that recovery falls back to the previous manifest state, which is still +/// consistent (the atomic rename guarantees the file is either fully old or +/// fully new). Call sites that CAN propagate (i.e., are in a fallible fn that +/// returns `std::io::Result`) should call `fsync_directory(parent)?` directly. +fn fsync_parent_best_effort(path: &Path) { + let parent = match path.parent() { + Some(p) if !p.as_os_str().is_empty() => p, + _ => return, // root or no parent — nothing to fsync + }; + if let Err(e) = fsync_directory(parent) { + warn!( + "fsync_parent_best_effort: failed to fsync dir {} after rename of {}: {}", + parent.display(), + path.display(), + e + ); + } +} + +/// On-disk layout discriminator. +/// +/// `TopLevel` is the legacy single-shard layout from manifest v1. `PerShard` +/// is the multi-shard layout introduced with manifest v2 — used whenever +/// `num_shards >= 2`. A `--shards 1` deployment with an existing v1 manifest +/// stays TopLevel until explicitly migrated. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AofLayout { + /// Legacy single-shard layout: `appendonlydir/moon.aof.{seq}.{base|incr}.*`. + TopLevel, + /// Per-shard layout: `appendonlydir/shard-{N}/moon.aof.{seq}.{base|incr}.*`. + PerShard, +} + +/// Per-shard manifest entry. One per shard in `PerShard` layout. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ShardManifest { + /// Shard ID (0..num_shards). + pub shard_id: u16, + /// Max LSN persisted to this shard's incr file so far. Semantics defined + /// by step 3 (LSN tagging) of the per-shard AOF RFC — until then this is + /// 0 and recovery does not use it. Once step 3 ships, recovery seeds + /// `master_repl_offset = max(shards[*].max_lsn)` before accepting writes. + pub max_lsn: u64, +} + +/// Active AOF file set tracked by the manifest. +#[derive(Debug, Clone)] +pub struct AofManifest { + /// Base directory (parent of `appendonlydir/`) + pub dir: PathBuf, + /// Current sequence number (incremented on each rewrite). + pub seq: u64, + /// On-disk layout. Determines path computation for base/incr files. + pub layout: AofLayout, + /// Per-shard metadata. Length is 1 for `TopLevel`, `num_shards` for + /// `PerShard`. Indexed by `shard_id`. + pub shards: Vec, +} + +impl AofManifest { + /// Path to the `appendonlydir/` directory. + pub fn aof_dir(&self) -> PathBuf { + self.dir.join(AOF_DIR_NAME) + } + + /// Path to the manifest file. + pub fn manifest_path(&self) -> PathBuf { + self.aof_dir().join(MANIFEST_NAME) + } + + /// Path to the base RDB file for the current sequence. + /// + /// Layout-aware: TopLevel returns `appendonlydir/moon.aof.{seq}.base.rdb`; + /// PerShard routes to `appendonlydir/shard-0/moon.aof.{seq}.base.rdb`. + /// This single-file helper is meaningful only when there is one shard + /// (post-migration `--shards 1`); a multi-shard PerShard manifest has N + /// base files and the caller must use [`Self::shard_base_path`] instead. + /// In debug builds, calling this on a multi-shard PerShard manifest + /// asserts; in release it returns the shard-0 path so production stays + /// recoverable rather than panicking on a stale call site. + pub fn base_path(&self) -> PathBuf { + match self.layout { + AofLayout::TopLevel => self + .aof_dir() + .join(format!("moon.aof.{}.base.rdb", self.seq)), + AofLayout::PerShard => { + debug_assert!( + self.shards.len() == 1, + "base_path() called on multi-shard PerShard manifest; use shard_base_path(shard_id)", + ); + self.shard_base_path_seq(0, self.seq) + } + } + } + + /// Path to the incremental RESP file for the current sequence. + /// + /// Layout-aware — see [`Self::base_path`] for the same routing rules. + pub fn incr_path(&self) -> PathBuf { + match self.layout { + AofLayout::TopLevel => self + .aof_dir() + .join(format!("moon.aof.{}.incr.aof", self.seq)), + AofLayout::PerShard => { + debug_assert!( + self.shards.len() == 1, + "incr_path() called on multi-shard PerShard manifest; use shard_incr_path(shard_id)", + ); + self.shard_incr_path_seq(0, self.seq) + } + } + } + + /// Path to the base RDB file for a given sequence. Layout-aware — see + /// [`Self::base_path`]. + pub fn base_path_seq(&self, seq: u64) -> PathBuf { + match self.layout { + AofLayout::TopLevel => self.aof_dir().join(format!("moon.aof.{}.base.rdb", seq)), + AofLayout::PerShard => { + debug_assert!( + self.shards.len() == 1, + "base_path_seq() called on multi-shard PerShard manifest; use shard_base_path_seq(shard_id, seq)", + ); + self.shard_base_path_seq(0, seq) + } + } + } + + /// Path to the incremental RESP file for a given sequence. Layout-aware — + /// see [`Self::base_path`]. + pub fn incr_path_seq(&self, seq: u64) -> PathBuf { + match self.layout { + AofLayout::TopLevel => self.aof_dir().join(format!("moon.aof.{}.incr.aof", seq)), + AofLayout::PerShard => { + debug_assert!( + self.shards.len() == 1, + "incr_path_seq() called on multi-shard PerShard manifest; use shard_incr_path_seq(shard_id, seq)", + ); + self.shard_incr_path_seq(0, seq) + } + } + } + + /// Create the `appendonlydir/` and write the initial manifest. + /// + /// Prefer [`Self::initialize_with_base`] when the in-memory databases + /// already contain state (e.g. first upgrade from legacy single-file AOF + /// or per-shard WAL) — otherwise subsequent boots cannot reconstruct that + /// state because there is no base RDB for `replay_multi_part` to load. + /// + /// B4 fix: even on fresh install (no prior state), materialize an EMPTY + /// base RDB so the `(base + incr)` invariant always holds. Without this, + /// the recovery path refuses to replay incr-only state and the server + /// fails to restart after a graceful shutdown that only wrote incr. + pub fn initialize(dir: &Path) -> std::io::Result { + let manifest = Self { + dir: dir.to_path_buf(), + seq: 1, + layout: AofLayout::TopLevel, + shards: vec![ShardManifest { + shard_id: 0, + max_lsn: 0, + }], + }; + std::fs::create_dir_all(manifest.aof_dir())?; + + // Serialize an empty database vector to an empty base RDB so the + // (base + incr) invariant holds from the first boot. + let empty_dbs: [crate::storage::Database; 0] = []; + let empty_rdb = crate::persistence::rdb::save_to_bytes(&empty_dbs) + .map_err(|e| std::io::Error::other(format!("empty RDB serialize: {e}")))?; + let base_path = manifest.base_path(); + let tmp_path = base_path.with_extension("rdb.tmp"); + { + let mut f = std::fs::File::create(&tmp_path)?; + f.write_all(&empty_rdb)?; + f.sync_data()?; + } + std::fs::rename(&tmp_path, &base_path)?; + fsync_parent_best_effort(&base_path); + + // Create the empty incr file so the writer has a target. + std::fs::File::create(manifest.incr_path())?; + + manifest.write_manifest()?; + Ok(manifest) + } + + /// Create the `appendonlydir/` and write an initial manifest with a base RDB + /// capturing the current in-memory state. + /// + /// Used on first upgrade from legacy persistence formats: after + /// `restore_from_persistence` has loaded state from the per-shard WAL or + /// `appendonly.aof`, this call materializes that state as the seq 1 base + /// RDB. Without a base, on the next boot the multi-part replay path would + /// clear the databases and then fail (missing base with non-empty incr) + /// or silently restart from empty state. + pub fn initialize_with_base(dir: &Path, rdb_bytes: &[u8]) -> std::io::Result { + let manifest = Self { + dir: dir.to_path_buf(), + seq: 1, + layout: AofLayout::TopLevel, + shards: vec![ShardManifest { + shard_id: 0, + max_lsn: 0, + }], + }; + std::fs::create_dir_all(manifest.aof_dir())?; + + // Write base RDB atomically: tmp file + fsync + rename. + let base_path = manifest.base_path(); + let tmp_path = base_path.with_extension("rdb.tmp"); + { + let mut f = std::fs::File::create(&tmp_path)?; + f.write_all(rdb_bytes)?; + f.sync_data()?; + } + std::fs::rename(&tmp_path, &base_path)?; + fsync_parent_best_effort(&base_path); + + // Create empty incr file so the writer has something to append to. + std::fs::File::create(manifest.incr_path())?; + + manifest.write_manifest()?; + Ok(manifest) + } + + /// Load manifest from disk. + /// + /// Returns: + /// - `Ok(None)` — manifest file does not exist (fresh install or legacy single-file AOF) + /// - `Ok(Some(manifest))` — manifest loaded successfully + /// - `Err(_)` — manifest file exists but is unreadable or corrupt. + /// Callers MUST treat this as fatal: overwriting a corrupt manifest with a + /// fresh one silently destroys the reference to the real base RDB and loses data. + pub fn load(dir: &Path) -> std::io::Result> { + let aof_dir = dir.join(AOF_DIR_NAME); + let manifest_path = aof_dir.join(MANIFEST_NAME); + + if !manifest_path.exists() { + return Ok(None); + } + + let content = std::fs::read_to_string(&manifest_path)?; + + // Detect format version. v1 manifests have no `version` line and use + // line prefixes `seq`/`base`/`incr`. v2 manifests start with `version 2` + // and carry per-shard records. + let mut format_version: u8 = 1; + for line in content.lines() { + let line = line.trim(); + if let Some(val) = line.strip_prefix("version ") { + if let Ok(v) = val.parse::() { + format_version = v; + } + break; + } + if !line.is_empty() { + // First non-blank line is not a version header → v1. + break; + } + } + + let manifest = match format_version { + 1 => Self::parse_v1(&content, dir, &manifest_path)?, + 2 => Self::parse_v2(&content, dir, &manifest_path)?, + other => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} has unsupported format version {} (max supported: 2)", + manifest_path.display(), + other, + ), + )); + } + }; + + // Best-effort orphan cleanup: delete stray base/incr files from aborted + // rewrites. A crash between advance() steps 1-3 leaves a new base RDB on + // disk that the active manifest never references. Without this sweep, + // repeated crashes during rewrite can fill the disk with zombie files. + // + // Safe to call here: parse_* verified the manifest has all required + // records, so cleanup_orphans won't delete the active files. + manifest.cleanup_orphans(); + + Ok(Some(manifest)) + } + + /// Parse a v1 (TopLevel, single-shard) manifest. + fn parse_v1(content: &str, dir: &Path, manifest_path: &Path) -> std::io::Result { + let mut seq = 0u64; + let mut has_base_record = false; + let mut has_incr_record = false; + for line in content.lines() { + let line = line.trim(); + if let Some(val) = line.strip_prefix("seq ") { + if let Ok(n) = val.parse::() { + seq = n; + } + } else if line.starts_with("base ") { + has_base_record = true; + } else if line.starts_with("incr ") { + has_incr_record = true; + } + } + + if seq == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} has no valid sequence number", + manifest_path.display() + ), + )); + } + + if !has_base_record || !has_incr_record { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} is truncated: seq={} base={} incr={}", + manifest_path.display(), + seq, + has_base_record, + has_incr_record, + ), + )); + } + + Ok(Self { + dir: dir.to_path_buf(), + seq, + layout: AofLayout::TopLevel, + shards: vec![ShardManifest { + shard_id: 0, + max_lsn: 0, + }], + }) + } + + /// Parse a v2 (PerShard, multi-shard) manifest. + /// + /// Expected line format: + /// ```text + /// version 2 + /// seq N + /// shards K + /// shard 0 max_lsn LSN0 + /// shard 1 max_lsn LSN1 + /// ... + /// ``` + /// + /// Per-shard `base`/`incr` paths are derived from `shard-{N}/moon.aof.{seq}.*` + /// rather than stored explicitly — the layout is canonical, so storing + /// paths invites drift between the stored value and the computed one. + fn parse_v2(content: &str, dir: &Path, manifest_path: &Path) -> std::io::Result { + let mut seq = 0u64; + let mut num_shards: Option = None; + let mut shards: Vec = Vec::new(); + + for line in content.lines() { + let line = line.trim(); + if line.is_empty() || line.starts_with('#') { + continue; + } + if line == "version 2" { + continue; + } else if let Some(val) = line.strip_prefix("seq ") { + seq = val.parse::().map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} has invalid seq line `{}`: {}", + manifest_path.display(), + line, + e, + ), + ) + })?; + } else if let Some(val) = line.strip_prefix("shards ") { + num_shards = Some(val.parse::().map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} has invalid shards line `{}`: {}", + manifest_path.display(), + line, + e, + ), + ) + })?); + } else if let Some(rest) = line.strip_prefix("shard ") { + // Format: `shard max_lsn ` + let mut it = rest.split_whitespace(); + let id_str = it.next().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} has shard line missing id: `{}`", + manifest_path.display(), + line, + ), + ) + })?; + let id: u16 = id_str.parse().map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} has shard line invalid id `{}`: {}", + manifest_path.display(), + id_str, + e, + ), + ) + })?; + // Expect `max_lsn `. + let label = it.next().unwrap_or(""); + let val_str = it.next().unwrap_or("0"); + if label != "max_lsn" { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} shard {} expected `max_lsn`, got `{}`", + manifest_path.display(), + id, + label, + ), + )); + } + let max_lsn: u64 = val_str.parse().map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} shard {} invalid max_lsn `{}`: {}", + manifest_path.display(), + id, + val_str, + e, + ), + ) + })?; + shards.push(ShardManifest { + shard_id: id, + max_lsn, + }); + } + // Unknown lines are tolerated (forward-compat). Strict parsers can + // be added at v3 if needed. + } + + if seq == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} has no valid sequence number", + manifest_path.display() + ), + )); + } + + let expected = num_shards.ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} is missing required `shards N` line", + manifest_path.display() + ), + ) + })?; + + if shards.len() != expected as usize { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} declares shards={} but has {} shard records", + manifest_path.display(), + expected, + shards.len(), + ), + )); + } + + // Sort by shard_id and verify contiguous range [0, expected). + shards.sort_by_key(|s| s.shard_id); + for (i, s) in shards.iter().enumerate() { + if s.shard_id as usize != i { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "AOF manifest at {} has non-contiguous shard ids (expected {} at position {}, got {})", + manifest_path.display(), + i, + i, + s.shard_id, + ), + )); + } + } + + Ok(Self { + dir: dir.to_path_buf(), + seq, + layout: AofLayout::PerShard, + shards, + }) + } + + /// Delete any base/incr files in `appendonlydir/` that do not match the + /// current sequence. Best-effort — logs but does not propagate errors. + /// + /// For `PerShard` layout, also recurses into every `shard-N/` subdirectory + /// and removes stale/tmp files there. Aborted BGREWRITEAOF runs leave + /// `.rdb.tmp` files in the shard subdirs that otherwise accumulate forever. + fn cleanup_orphans(&self) { + match self.layout { + AofLayout::TopLevel => { + self.cleanup_orphans_dir(&self.aof_dir(), self.seq); + } + AofLayout::PerShard => { + // Top-level appendonlydir/ holds only the manifest — no data files + // to clean up there. All data lives in shard-N/ subdirs. + for shard in &self.shards { + self.cleanup_orphans_shard(shard.shard_id); + } + } + } + } + + /// Scan a single shard's directory for orphan base/incr/tmp files that do + /// not correspond to the current manifest sequence. Best-effort. + fn cleanup_orphans_shard(&self, shard_id: u16) { + self.cleanup_orphans_dir(&self.shard_dir(shard_id), self.seq); + } + + /// Core orphan sweep: scan `dir` and remove any `moon.aof.*` files whose + /// sequence is not `keep_seq`. Skips the manifest file itself. + fn cleanup_orphans_dir(&self, dir: &Path, keep_seq: u64) { + let entries = match std::fs::read_dir(dir) { + Ok(e) => e, + Err(_) => return, + }; + let current_base = format!("moon.aof.{}.base.rdb", keep_seq); + let current_incr = format!("moon.aof.{}.incr.aof", keep_seq); + for entry in entries.flatten() { + let name = entry.file_name(); + let name_str = match name.to_str() { + Some(s) => s, + None => continue, + }; + // Keep manifest, current base, current incr. Delete any other moon.aof.*. + if name_str == MANIFEST_NAME || name_str == current_base || name_str == current_incr { + continue; + } + let is_moon_aof = name_str.starts_with("moon.aof.") + && (name_str.ends_with(".base.rdb") + || name_str.ends_with(".incr.aof") + || name_str.ends_with(".rdb.tmp") + || name_str.ends_with(".tmp")); + if !is_moon_aof { + continue; + } + let path = entry.path(); + match std::fs::remove_file(&path) { + Ok(()) => info!("AOF orphan cleanup: removed {}", path.display()), + Err(e) => warn!( + "AOF orphan cleanup: failed to remove {}: {}", + path.display(), + e + ), + } + } + } + + /// Write the manifest file atomically (write tmp + rename). + /// + /// Emits v1 format for `TopLevel` and v2 for `PerShard`. The format is + /// selected by `self.layout`, never by callers — preserving the invariant + /// that one directory holds one layout. + pub fn write_manifest(&self) -> std::io::Result<()> { + let manifest_path = self.manifest_path(); + let tmp_path = manifest_path.with_extension("tmp"); + + let content = match self.layout { + AofLayout::TopLevel => format!( + "seq {}\nbase moon.aof.{}.base.rdb\nincr moon.aof.{}.incr.aof\n", + self.seq, self.seq, self.seq + ), + AofLayout::PerShard => { + let mut s = String::with_capacity(64 + self.shards.len() * 40); + s.push_str("version 2\n"); + s.push_str(&format!("seq {}\n", self.seq)); + s.push_str(&format!("shards {}\n", self.shards.len())); + for shard in &self.shards { + s.push_str(&format!( + "shard {} max_lsn {}\n", + shard.shard_id, shard.max_lsn + )); + } + s + } + }; + + let mut f = std::fs::File::create(&tmp_path)?; + f.write_all(content.as_bytes())?; + f.sync_data()?; + std::fs::rename(&tmp_path, &manifest_path)?; + fsync_parent_best_effort(&manifest_path); + Ok(()) + } + + // ------------------------------------------------------------------ + // Per-shard layout helpers + // ------------------------------------------------------------------ + + /// Directory holding a shard's AOF files. + /// + /// - `TopLevel`: `appendonlydir/` (the shard_id argument is asserted to be 0). + /// - `PerShard`: `appendonlydir/shard-{shard_id}/`. + pub fn shard_dir(&self, shard_id: u16) -> PathBuf { + match self.layout { + AofLayout::TopLevel => { + debug_assert_eq!(shard_id, 0, "TopLevel layout only has shard 0"); + self.aof_dir() + } + AofLayout::PerShard => self.aof_dir().join(format!("shard-{}", shard_id)), + } + } + + /// Path to a shard's base RDB file for the current sequence. + pub fn shard_base_path(&self, shard_id: u16) -> PathBuf { + self.shard_dir(shard_id) + .join(format!("moon.aof.{}.base.rdb", self.seq)) + } + + /// Path to a shard's incremental RESP file for the current sequence. + pub fn shard_incr_path(&self, shard_id: u16) -> PathBuf { + self.shard_dir(shard_id) + .join(format!("moon.aof.{}.incr.aof", self.seq)) + } + + /// Path to a shard's base RDB file for a given sequence. + pub fn shard_base_path_seq(&self, shard_id: u16, seq: u64) -> PathBuf { + self.shard_dir(shard_id) + .join(format!("moon.aof.{}.base.rdb", seq)) + } + + /// Path to a shard's incremental RESP file for a given sequence. + pub fn shard_incr_path_seq(&self, shard_id: u16, seq: u64) -> PathBuf { + self.shard_dir(shard_id) + .join(format!("moon.aof.{}.incr.aof", seq)) + } + + /// Maximum LSN persisted across all shards. + /// + /// Computed (not stored) so the stored value can never drift from + /// the per-shard records. Returns 0 if `shards` is empty (defensive; + /// constructors guarantee at least one shard). + pub fn global_max_lsn(&self) -> u64 { + self.shards.iter().map(|s| s.max_lsn).max().unwrap_or(0) + } + + /// Verify that the manifest matches the runtime shard count. + /// + /// Returns the verbatim error from RFC § 3 if the shard count differs, + /// for operator-facing consistency. Callers (typically `main.rs` boot) + /// should treat this as fatal: continuing with a mismatched shard count + /// silently drops data from shards that no longer exist or replays a + /// shard's data into the wrong DashTable. + pub fn verify_shard_count(&self, expected: u16) -> Result<(), String> { + let actual = self.shards.len() as u16; + if actual != expected { + return Err(format!( + "ERR shard count changed (manifest={}, config={}); refusing to start to avoid data loss. See docs/runbooks/shard-count-change.md", + actual, expected + )); + } + Ok(()) + } + + /// Returns true if the on-disk layout under `appendonlydir/` matches the + /// legacy TopLevel format (files at top level, no `shard-N/` subdirs). + /// + /// Used by callers to detect when a v1 single-shard deployment is being + /// upgraded to v2 multi-shard and needs explicit migration. Does NOT + /// migrate — separate from `migrate_top_level_to_per_shard` so the side + /// effect is opt-in, not hidden in a load path. + pub fn is_legacy_top_level_layout(dir: &Path) -> bool { + let aof_dir = dir.join(AOF_DIR_NAME); + if !aof_dir.exists() { + return false; + } + + // Check manifest version first. If a valid v2 (PerShard) manifest exists, + // return false regardless of stray top-level files. Operators occasionally + // leave old base.rdb / incr.aof files at the top level during debugging + // or failed upgrades; scanning filenames without reading the manifest would + // produce a misleading "legacy detected" result and trigger unwanted + // migration on an already-upgraded deployment. + if let Ok(Some(m)) = Self::load(dir) { + if m.layout == AofLayout::PerShard { + return false; + } + } + + let entries = match std::fs::read_dir(&aof_dir) { + Ok(e) => e, + Err(_) => return false, + }; + for entry in entries.flatten() { + let name = entry.file_name(); + let Some(name_str) = name.to_str() else { + continue; + }; + if name_str.starts_with("moon.aof.") + && (name_str.ends_with(".base.rdb") || name_str.ends_with(".incr.aof")) + { + return true; + } + } + false + } + + /// Migrate a single-shard TopLevel layout in place to a single-shard + /// PerShard layout. + /// + /// Moves `appendonlydir/moon.aof.{seq}.{base.rdb,incr.aof}` into + /// `appendonlydir/shard-0/`, then rewrites the manifest as v2 with + /// `shards 1`. Idempotent: a second call on an already-PerShard manifest + /// returns Ok with no filesystem changes. + /// + /// This is the RFC § 5 case 1 migration — zero data movement (rename only), + /// safe to run on first boot after upgrading from v0.1.x. Multi-shard + /// migrations from legacy AOF (case 2) use the `moon migrate-aof` + /// subcommand and are NOT handled here. + pub fn migrate_top_level_to_per_shard(&mut self) -> std::io::Result<()> { + if self.layout == AofLayout::PerShard { + return Ok(()); + } + if self.shards.len() != 1 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "migrate_top_level_to_per_shard called with {} shards; \ + only single-shard TopLevel can be migrated in place", + self.shards.len() + ), + )); + } + + // Compute paths up front. shard_dir/shard_*_path_seq for a single- + // shard target are pure path computations and do NOT depend on + // self.layout, so it is safe to derive them while layout is still + // TopLevel. + let old_base = self + .aof_dir() + .join(format!("moon.aof.{}.base.rdb", self.seq)); + let old_incr = self + .aof_dir() + .join(format!("moon.aof.{}.incr.aof", self.seq)); + let new_dir = self.aof_dir().join("shard-0"); + let new_base = new_dir.join(format!("moon.aof.{}.base.rdb", self.seq)); + let new_incr = new_dir.join(format!("moon.aof.{}.incr.aof", self.seq)); + + if !old_base.exists() { + // Pre-flight check: nothing moved yet, no rollback needed. + return Err(std::io::Error::new( + std::io::ErrorKind::NotFound, + format!( + "TopLevel→PerShard migration: source base {} not found", + old_base.display() + ), + )); + } + std::fs::create_dir_all(&new_dir)?; + + // Move base. If the rename itself fails, no on-disk mutation has + // happened yet — bail without rollback. Layout stays TopLevel until + // commit at the bottom. + std::fs::rename(&old_base, &new_base)?; + + // Fsync the target directory so the base rename is durable before we + // proceed. A crash after rename but before dir-fsync could leave the + // old filename visible on the next boot. + // + // NOTE: if this fsync fails, old_base has already moved to new_base — + // rollback the rename before returning so the manifest stays consistent. + if let Err(e) = fsync_directory(&new_dir) { + if let Err(re) = std::fs::rename(&new_base, &old_base) { + error!( + "Migration rollback: failed to restore base {} → {} after fsync_directory failure: {}", + new_base.display(), + old_base.display(), + re + ); + } + return Err(e); + } + + // Base is now durably in shard-0/. Any subsequent error must restore it. + let moved_incr: bool; + let created_incr: bool; + if old_incr.exists() { + if let Err(e) = std::fs::rename(&old_incr, &new_incr) { + if let Err(re) = std::fs::rename(&new_base, &old_base) { + error!( + "Migration rollback: failed to restore base {} → {}: {}", + new_base.display(), + old_base.display(), + re + ); + } + return Err(e); + } + // Fsync the shard directory to make the incr rename durable. + // If this fails, roll back both incr and base renames. + if let Err(e) = fsync_directory(&new_dir) { + if let Err(re) = std::fs::rename(&new_incr, &old_incr) { + error!( + "Migration rollback: failed to restore incr {} → {} after fsync_directory failure: {}", + new_incr.display(), + old_incr.display(), + re + ); + } + if let Err(re) = std::fs::rename(&new_base, &old_base) { + error!( + "Migration rollback: failed to restore base {} → {} after fsync_directory failure: {}", + new_base.display(), + old_base.display(), + re + ); + } + return Err(e); + } + moved_incr = true; + created_incr = false; + } else { + match std::fs::File::create(&new_incr) { + Ok(_) => { + moved_incr = false; + created_incr = true; + } + Err(e) => { + if let Err(re) = std::fs::rename(&new_base, &old_base) { + error!( + "Migration rollback: failed to restore base {} → {}: {}", + new_base.display(), + old_base.display(), + re + ); + } + return Err(e); + } + } + } + + // Commit: flip layout, persist as v2. If write_manifest fails, undo + // every filesystem mutation and restore layout so the next boot still + // sees a valid v1 TopLevel deployment. + self.layout = AofLayout::PerShard; + if let Err(e) = self.write_manifest() { + self.layout = AofLayout::TopLevel; + if moved_incr { + if let Err(re) = std::fs::rename(&new_incr, &old_incr) { + error!( + "Migration rollback: failed to restore incr {} → {}: {}", + new_incr.display(), + old_incr.display(), + re + ); + } + } else if created_incr { + if let Err(re) = std::fs::remove_file(&new_incr) { + warn!( + "Migration rollback: failed to remove freshly created incr {}: {}", + new_incr.display(), + re + ); + } + } + if let Err(re) = std::fs::rename(&new_base, &old_base) { + error!( + "Migration rollback: failed to restore base {} → {}: {}. \ + Manifest dir {} may be in an inconsistent state.", + new_base.display(), + old_base.display(), + re, + self.dir.display() + ); + } + return Err(e); + } + + info!( + "AOF migrated: TopLevel → PerShard (single shard) at {}", + self.aof_dir().display() + ); + Ok(()) + } + + /// Advance to the next sequence: write new base RDB, create new incr file, + /// update manifest, delete old files. + /// + /// Returns the path to the new incremental file (caller should switch writing to it). + pub fn advance(&mut self, rdb_bytes: &[u8]) -> Result { + let old_seq = self.seq; + let new_seq = old_seq + 1; + + let aof_dir = self.aof_dir(); + std::fs::create_dir_all(&aof_dir).map_err(|e| crate::error::AofError::Io { + path: aof_dir.clone(), + source: e, + })?; + + // 1. Write new base RDB (atomic: tmp + fsync + rename). + // Must fsync the data BEFORE renaming — a rename without prior fsync + // can publish a file whose contents aren't durable, so a crash leaves + // the manifest pointing at an empty/partial base RDB. + let new_base = self.base_path_seq(new_seq); + let tmp_base = new_base.with_extension("rdb.tmp"); + { + let mut f = + std::fs::File::create(&tmp_base).map_err(|e| crate::error::AofError::Io { + path: tmp_base.clone(), + source: e, + })?; + f.write_all(rdb_bytes) + .map_err(|e| crate::error::AofError::Io { + path: tmp_base.clone(), + source: e, + })?; + f.sync_data().map_err(|e| crate::error::AofError::Io { + path: tmp_base.clone(), + source: e, + })?; + } + std::fs::rename(&tmp_base, &new_base).map_err(|e| { + crate::error::AofError::RewriteFailed { + detail: format!("rename base: {}", e), + } + })?; + fsync_parent_best_effort(&new_base); + + // 2. Create empty new incremental file + let new_incr = self.incr_path_seq(new_seq); + std::fs::File::create(&new_incr).map_err(|e| crate::error::AofError::Io { + path: new_incr.clone(), + source: e, + })?; + + // 3. Update manifest (atomic) + self.seq = new_seq; + self.write_manifest() + .map_err(|e| crate::error::AofError::Io { + path: self.manifest_path(), + source: e, + })?; + + // 4. Delete old files (best-effort) + let old_base = self.base_path_seq(old_seq); + let old_incr = self.incr_path_seq(old_seq); + if old_base.exists() { + if let Err(e) = std::fs::remove_file(&old_base) { + warn!("Failed to delete old base {}: {}", old_base.display(), e); + } + } + if old_incr.exists() { + if let Err(e) = std::fs::remove_file(&old_incr) { + warn!("Failed to delete old incr {}: {}", old_incr.display(), e); + } + } + + info!( + "AOF advanced to seq {}: base={} bytes, incr={}", + new_seq, + rdb_bytes.len(), + new_incr.display() + ); + + Ok(new_incr) + } +} + +#[cfg(test)] +mod tests_v2 { + //! Unit tests for the v2 (PerShard) manifest format. + //! + //! Covers the Step 1 deliverable of the per-shard AOF RFC: + //! - v1 manifests continue to load as TopLevel (single-shard, shard_id=0) + //! - v2 round-trip: write → load → equivalent struct shape + //! - shard count mismatch produces the verbatim RFC § 3 error + //! - migrate_top_level_to_per_shard performs in-place rename and rewrites + //! the manifest as v2 + //! - global_max_lsn computes max across shards + //! - is_legacy_top_level_layout detects top-level files + + use super::*; + use std::fs; + + fn temp_dir() -> PathBuf { + // Use a global atomic counter so parallel test threads (cargo test runs + // unit tests in parallel) never produce the same directory name even + // when PID and nanosecond clock resolution are the same for two threads. + static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let d = std::env::temp_dir().join(format!( + "moon-aof-manifest-test-{}-{}", + std::process::id(), + n, + )); + fs::create_dir_all(&d).expect("temp dir create"); + d + } + + #[test] + fn v1_manifest_loads_as_top_level_single_shard() { + let dir = temp_dir(); + let m = AofManifest::initialize(&dir).expect("initialize v1"); + + assert_eq!(m.layout, AofLayout::TopLevel); + assert_eq!(m.shards.len(), 1); + assert_eq!(m.shards[0].shard_id, 0); + assert_eq!(m.shards[0].max_lsn, 0); + + // Reload from disk + let reloaded = AofManifest::load(&dir).expect("load").expect("present"); + assert_eq!(reloaded.layout, AofLayout::TopLevel); + assert_eq!(reloaded.shards.len(), 1); + assert_eq!(reloaded.seq, m.seq); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn v2_manifest_round_trips() { + let dir = temp_dir(); + let m = AofManifest::initialize_multi(&dir, 4).expect("initialize_multi"); + + assert_eq!(m.layout, AofLayout::PerShard); + assert_eq!(m.shards.len(), 4); + for (i, s) in m.shards.iter().enumerate() { + assert_eq!(s.shard_id, i as u16); + assert_eq!(s.max_lsn, 0); + } + + // Per-shard subdirs were created with empty base + incr. + for i in 0..4u16 { + assert!(m.shard_dir(i).exists(), "shard-{} dir exists", i); + assert!(m.shard_base_path(i).exists(), "shard-{} base exists", i); + assert!(m.shard_incr_path(i).exists(), "shard-{} incr exists", i); + } + + let reloaded = AofManifest::load(&dir).expect("load").expect("present"); + assert_eq!(reloaded.layout, AofLayout::PerShard); + assert_eq!(reloaded.shards.len(), 4); + assert_eq!(reloaded.seq, m.seq); + for (i, s) in reloaded.shards.iter().enumerate() { + assert_eq!(s.shard_id, i as u16); + } + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn verify_shard_count_emits_rfc_error_verbatim() { + let m = AofManifest { + dir: PathBuf::from("/tmp/nowhere"), + seq: 1, + layout: AofLayout::PerShard, + shards: vec![ + ShardManifest { + shard_id: 0, + max_lsn: 0, + }, + ShardManifest { + shard_id: 1, + max_lsn: 0, + }, + ], + }; + let err = m.verify_shard_count(4).expect_err("should mismatch"); + assert_eq!( + err, + "ERR shard count changed (manifest=2, config=4); refusing to start to avoid data loss. See docs/runbooks/shard-count-change.md" + ); + + // Matching count succeeds. + m.verify_shard_count(2).expect("match"); + } + + #[test] + fn migrate_top_level_to_per_shard_moves_files_and_rewrites_manifest() { + let dir = temp_dir(); + let mut m = AofManifest::initialize(&dir).expect("initialize v1"); + + // Write a marker into the incr file so we can prove the contents + // survive the rename. + let original_incr = m.aof_dir().join(format!("moon.aof.{}.incr.aof", m.seq)); + fs::write(&original_incr, b"MARKER").expect("write incr marker"); + + m.migrate_top_level_to_per_shard().expect("migrate"); + + assert_eq!(m.layout, AofLayout::PerShard); + assert!(!original_incr.exists(), "old incr removed by rename"); + let new_incr = m.shard_incr_path(0); + assert!(new_incr.exists(), "new shard-0 incr exists"); + let contents = fs::read(&new_incr).expect("read new incr"); + assert_eq!(contents, b"MARKER", "incr contents preserved"); + + // Reloaded manifest is v2. + let reloaded = AofManifest::load(&dir).expect("load").expect("present"); + assert_eq!(reloaded.layout, AofLayout::PerShard); + assert_eq!(reloaded.shards.len(), 1); + + // Idempotency: second call is a no-op. + m.migrate_top_level_to_per_shard().expect("idempotent"); + assert_eq!(m.layout, AofLayout::PerShard); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn global_max_lsn_returns_max_across_shards() { + let m = AofManifest { + dir: PathBuf::from("/tmp/nowhere"), + seq: 1, + layout: AofLayout::PerShard, + shards: vec![ + ShardManifest { + shard_id: 0, + max_lsn: 100, + }, + ShardManifest { + shard_id: 1, + max_lsn: 500, + }, + ShardManifest { + shard_id: 2, + max_lsn: 250, + }, + ], + }; + assert_eq!(m.global_max_lsn(), 500); + } + + #[test] + fn is_legacy_top_level_layout_detects_v1_files() { + let dir = temp_dir(); + // No appendonlydir yet → false. + assert!(!AofManifest::is_legacy_top_level_layout(&dir)); + + // After v1 initialize, top-level files present → true. + let _m = AofManifest::initialize(&dir).expect("init v1"); + assert!(AofManifest::is_legacy_top_level_layout(&dir)); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn is_legacy_top_level_layout_returns_false_for_v2() { + let dir = temp_dir(); + let _m = AofManifest::initialize_multi(&dir, 2).expect("init v2"); + assert!( + !AofManifest::is_legacy_top_level_layout(&dir), + "v2 layout has no top-level moon.aof.* files" + ); + + fs::remove_dir_all(&dir).ok(); + } + + /// FIX-W3-4: v2 manifest with stray top-level .base.rdb must return false, + /// not true. The filename scan is misleading when a valid v2 manifest exists. + /// + /// Scenario: operator upgraded to v2 but left a stale `moon.aof.1.base.rdb` + /// at the top level (e.g., copied during debugging). `is_legacy_top_level_layout` + /// must check the manifest first and return false when v2 is confirmed. + #[test] + fn is_legacy_top_level_layout_ignores_stray_files_when_v2_manifest_present() { + let dir = temp_dir(); + // Initialize a genuine v2 (PerShard) layout. + let _m = AofManifest::initialize_multi(&dir, 2).expect("init v2"); + + // Plant a stale top-level base.rdb to simulate the stray-file scenario. + let stray = dir.join(AOF_DIR_NAME).join("moon.aof.1.base.rdb"); + fs::write(&stray, b"REDIS0011\xff").expect("write stray base.rdb"); + + // Even though the stray file matches the filename pattern, a valid v2 + // manifest is present, so is_legacy_top_level_layout must return false. + assert!( + !AofManifest::is_legacy_top_level_layout(&dir), + "v2 manifest + stray top-level file must still return false" + ); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn parse_v2_rejects_shard_count_mismatch_in_file() { + let dir = temp_dir(); + let aof = dir.join(AOF_DIR_NAME); + fs::create_dir_all(&aof).unwrap(); + // Manifest claims shards 3 but only declares two shard records. + fs::write( + aof.join(MANIFEST_NAME), + "version 2\nseq 1\nshards 3\nshard 0 max_lsn 0\nshard 1 max_lsn 0\n", + ) + .unwrap(); + + let err = AofManifest::load(&dir).expect_err("should reject"); + let msg = err.to_string(); + assert!( + msg.contains("declares shards=3 but has 2 shard records"), + "got: {}", + msg + ); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn parse_v2_rejects_non_contiguous_shard_ids() { + let dir = temp_dir(); + let aof = dir.join(AOF_DIR_NAME); + fs::create_dir_all(&aof).unwrap(); + // shards=2 but ids are {0, 2} not {0, 1}. + fs::write( + aof.join(MANIFEST_NAME), + "version 2\nseq 1\nshards 2\nshard 0 max_lsn 0\nshard 2 max_lsn 0\n", + ) + .unwrap(); + + let err = AofManifest::load(&dir).expect_err("should reject"); + let msg = err.to_string(); + assert!(msg.contains("non-contiguous shard ids"), "got: {}", msg); + + fs::remove_dir_all(&dir).ok(); + } + + // ------------------------------------------------------------------ + // Reviewer-flagged fixes: layout-aware path helpers + migration + // rollback. See the "Verify findings against current code" review + // comment on aof_manifest.rs:669-775 and :688-717. + // ------------------------------------------------------------------ + + #[test] + fn base_incr_paths_route_to_shard_zero_after_migration() { + let dir = temp_dir(); + let mut m = AofManifest::initialize(&dir).expect("init v1"); + // Pre-migration: TopLevel paths under appendonlydir/ directly. + assert_eq!(m.base_path(), m.aof_dir().join("moon.aof.1.base.rdb")); + assert_eq!(m.incr_path(), m.aof_dir().join("moon.aof.1.incr.aof")); + + m.migrate_top_level_to_per_shard().expect("migrate"); + + // Post-migration: single-file helpers must route to shard-0/ so + // replay_multi_part and advance() find the actual files. This is + // the bug the reviewer flagged for aof_manifest.rs:669-775. + let shard0 = m.aof_dir().join("shard-0"); + assert_eq!(m.base_path(), shard0.join("moon.aof.1.base.rdb")); + assert_eq!(m.incr_path(), shard0.join("moon.aof.1.incr.aof")); + assert_eq!(m.base_path_seq(7), shard0.join("moon.aof.7.base.rdb")); + assert_eq!(m.incr_path_seq(7), shard0.join("moon.aof.7.incr.aof")); + // The path the helper returns must be where the file actually lives. + assert!(m.base_path().exists(), "base file at returned path"); + assert!(m.incr_path().exists(), "incr file at returned path"); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn migrate_rolls_back_filesystem_when_incr_rename_fails() { + // Simulate the rename(old_incr → new_incr) failure path by making + // the destination already exist as a directory (rename onto a + // non-empty directory is an error on every supported OS). + let dir = temp_dir(); + let mut m = AofManifest::initialize(&dir).expect("init v1"); + let original_base = m.aof_dir().join("moon.aof.1.base.rdb"); + let original_incr = m.aof_dir().join("moon.aof.1.incr.aof"); + fs::write(&original_incr, b"INCR_MARKER").expect("seed incr"); + let base_bytes_before = fs::read(&original_base).expect("read base"); + + // Pre-create shard-0/moon.aof.1.incr.aof as a DIRECTORY so the + // rename fails after the base rename has already succeeded. + let shard0 = m.aof_dir().join("shard-0"); + fs::create_dir_all(shard0.join("moon.aof.1.incr.aof")).expect("seed blocker"); + + let err = m + .migrate_top_level_to_per_shard() + .expect_err("incr rename should fail"); + let _ = err; // exact error kind depends on OS + + // Rollback invariants: + // 1. Layout stays TopLevel in memory. + // 2. base file restored to its original TopLevel path. + // 3. base file contents unchanged. + // 4. on-disk manifest is still v1 (load returns layout TopLevel). + assert_eq!(m.layout, AofLayout::TopLevel, "in-memory layout reverted"); + assert!(original_base.exists(), "base restored to top-level"); + let base_bytes_after = fs::read(&original_base).expect("read base"); + assert_eq!(base_bytes_after, base_bytes_before, "base contents intact"); + let reloaded = AofManifest::load(&dir).expect("load").expect("present"); + assert_eq!(reloaded.layout, AofLayout::TopLevel, "on-disk manifest v1"); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn migrate_does_not_mutate_on_missing_base() { + let dir = temp_dir(); + let mut m = AofManifest::initialize(&dir).expect("init v1"); + let base = m.aof_dir().join("moon.aof.1.base.rdb"); + fs::remove_file(&base).expect("remove base"); + + let err = m + .migrate_top_level_to_per_shard() + .expect_err("missing base should fail"); + assert_eq!(err.kind(), std::io::ErrorKind::NotFound); + // Layout never flipped, no rollback needed. + assert_eq!(m.layout, AofLayout::TopLevel); + + fs::remove_dir_all(&dir).ok(); + } +} diff --git a/src/persistence/aof_manifest/shard_replay.rs b/src/persistence/aof_manifest/shard_replay.rs new file mode 100644 index 00000000..d3b3d5b1 --- /dev/null +++ b/src/persistence/aof_manifest/shard_replay.rs @@ -0,0 +1,1157 @@ +//! Multi-part / per-shard AOF replay: base RDB load + framed/RESP incr replay. +//! +//! Split out of the `aof_manifest` parent module (issue #143) to keep each file +//! under the 1500-line cap. `use super::*` pulls the manifest types, std/tracing +//! imports, and private helpers down from the parent module. + +use super::*; + +/// Replay multi-part AOF: load base RDB then replay incremental RESP. +/// +/// Returns total keys/commands loaded. +pub fn replay_multi_part( + databases: &mut [crate::storage::Database], + manifest: &AofManifest, + engine: &dyn crate::persistence::replay::CommandReplayEngine, +) -> Result { + let mut total = 0usize; + + // Load base RDB + let base_path = manifest.base_path(); + if base_path.exists() { + match crate::persistence::rdb::load(databases, &base_path) { + Ok(n) => { + info!( + "AOF base RDB loaded: {} keys from {}", + n, + base_path.display() + ); + total += n; + } + Err(e) => { + // Base RDB is corrupt or unreadable — applying incremental + // deltas on top of missing/corrupt base gives wrong results. + error!("AOF base RDB load failed: {}", e); + return Err(e); + } + } + } else { + // Missing base is tolerable only when the incr log is also empty + // (fresh manifest from initialize(), or first boot after legacy + // upgrade). If there's incremental content but no base, replaying + // deltas (DEL, EXPIRE, HINCRBY, …) on an empty database produces + // incorrect state — fail loudly rather than silently corrupt. + let incr_path = manifest.incr_path(); + let incr_len = std::fs::metadata(&incr_path).map(|m| m.len()).unwrap_or(0); + if incr_len > 0 { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "AOF base RDB missing at {} but incr {} is {} bytes; refusing to replay incr against empty state", + base_path.display(), + incr_path.display(), + incr_len, + ), + }, + )); + } + warn!( + "AOF base RDB not found: {} (incr empty, treating as fresh init)", + base_path.display() + ); + } + + // Replay incremental RESP + let incr_path = manifest.incr_path(); + if incr_path.exists() { + let data = std::fs::read(&incr_path)?; + if !data.is_empty() { + // Pure RESP — use replay_aof_resp (no RDB preamble detection needed) + let count = replay_incr_resp(databases, &data, engine)?; + info!( + "AOF incr replayed: {} commands from {}", + count, + incr_path.display() + ); + total += count; + } + } + + Ok(total) +} + +/// Replay pure RESP commands from a byte slice. +/// +/// **Corruption handling:** On mid-stream parse errors this returns an error +/// rather than silently resyncing to the next `*` byte. Silent resync in a +/// multi-part AOF is dangerous: an undetected run of dropped commands leaves +/// the database in an inconsistent state that cannot be reconstructed. +/// Truncated tails (parser returns `Ok(None)` with bytes remaining) are +/// logged and treated as the legitimate end of the incremental log, matching +/// `replay_aof` semantics for crash-time tail truncation. +fn replay_incr_resp( + databases: &mut [crate::storage::Database], + data: &[u8], + engine: &dyn crate::persistence::replay::CommandReplayEngine, +) -> Result { + use crate::protocol::{Frame, ParseConfig, parse}; + use bytes::BytesMut; + + let total_len = data.len(); + let mut buf = BytesMut::from(data); + let config = ParseConfig::default(); + let mut selected_db: usize = 0; + let mut count: usize = 0; + + loop { + if buf.is_empty() { + break; + } + match parse::parse(&mut buf, &config) { + Ok(Some(frame)) => { + let (cmd, cmd_args) = match &frame { + Frame::Array(arr) if !arr.is_empty() => { + let name = match &arr[0] { + Frame::BulkString(s) => s.as_ref(), + Frame::SimpleString(s) => s.as_ref(), + other => { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "AOF incr command at offset {} has non-string name frame: {:?}", + total_len - buf.len(), + std::mem::discriminant(other) + ), + }, + )); + } + }; + (name as &[u8], &arr[1..]) + } + other => { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "AOF incr non-array frame at offset {}: {:?}", + total_len - buf.len(), + std::mem::discriminant(other) + ), + }, + )); + } + }; + engine.replay_command(databases, cmd, cmd_args, &mut selected_db); + count += 1; + } + Ok(None) => { + if !buf.is_empty() { + let offset = total_len - buf.len(); + warn!( + "AOF incr truncated tail: {} bytes at offset {} (treating as crash-time EOF)", + buf.len(), + offset + ); + } + break; + } + Err(e) => { + let offset = total_len - buf.len(); + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!("AOF incr parse error at offset {}: {:?}", offset, e), + }, + )); + } + } + } + + Ok(count) +} + +/// An entry that was tagged `OrderedAcrossShards` (RFC § 2 Rule 2) and +/// must be merge-replayed in global LSN order after per-shard replay +/// completes. The `shard_id` records which shard's file it came from so +/// the merge step can dispatch each entry back to its origin shard's +/// databases. +#[derive(Debug, Clone)] +pub struct OrderedEntry { + pub shard_id: u16, + pub lsn: u64, + pub bytes: bytes::Bytes, +} + +/// Replay a framed PerShard incr file: `[u64 lsn LE][u32 len LE][RESP bytes]`. +/// +/// Step 3 wrote this format; step 4 reads it. Step 5 extends the LSN field: +/// the high bit (`crate::persistence::aof::ORDERED_LSN_FLAG`) marks the +/// entry as `OrderedAcrossShards` — those entries are NOT replayed inline, +/// instead they are pushed into `ordered_buf` for the caller to merge-replay +/// in global LSN order across all shards. +/// +/// Returns `(commands_replayed, max_lsn)` — the count covers only inline +/// (non-ordered) replays. `max_lsn` is the NEXT-FREE replication offset: +/// `max(entry.lsn + entry.len)` across both inline AND ordered entries (the +/// high bit is masked out before the computation). It is the offset AFTER the +/// last byte on disk, NOT the start LSN of the last entry — because +/// `ReplicationState::issue_lsn` returns the offset BEFORE adding the entry +/// length, seeding `master_repl_offset` with a start LSN would reissue the +/// last on-disk LSN and break the lsn->entry uniqueness invariant (F5). +/// +/// **Truncated entries:** a header partly written at crash time is treated as +/// EOF (parity with `replay_incr_resp` semantics). A whole header followed by +/// a truncated payload is also EOF — the writer's invariant is that the +/// header is written first then the payload, and on partial write the most we +/// can lose is the last entry's payload tail. +/// +/// **Corruption:** a mid-stream RESP parse error inside an otherwise-complete +/// payload is fatal (same reasoning as `replay_incr_resp`). +fn replay_incr_framed( + shard_id: u16, + databases: &mut [crate::storage::Database], + data: &[u8], + engine: &dyn crate::persistence::replay::CommandReplayEngine, + ordered_buf: &mut Vec, +) -> Result<(usize, u64), crate::error::MoonError> { + use crate::protocol::{Frame, ParseConfig, parse}; + use bytes::BytesMut; + + const HEADER_LEN: usize = 12; // u64 lsn LE + u32 len LE + + let total_len = data.len(); + let mut offset: usize = 0; + let config = ParseConfig::default(); + let mut selected_db: usize = 0; + let mut count: usize = 0; + let mut max_lsn: u64 = 0; + + while offset < total_len { + if total_len - offset < HEADER_LEN { + warn!( + "AOF incr framed truncated header: {} bytes at offset {} (treating as crash-time EOF)", + total_len - offset, + offset + ); + break; + } + // SAFETY: line 1491 guarantees `total_len - offset >= HEADER_LEN` (=12), + // so the [offset..offset+8] and [offset+8..offset+12] slices are valid + // and `try_into()` to a fixed-size array cannot fail (length-matched). + #[allow(clippy::unwrap_used)] // bounds-checked above; try_into is statically length-matched + let raw_lsn = u64::from_le_bytes(data[offset..offset + 8].try_into().expect("8 bytes")); + #[allow(clippy::unwrap_used)] // same bounds-check guarantee + let len = + u32::from_le_bytes(data[offset + 8..offset + 12].try_into().expect("4 bytes")) as usize; + let payload_start = offset + HEADER_LEN; + let payload_end = payload_start.saturating_add(len); + if payload_end > total_len { + warn!( + "AOF incr framed truncated payload at offset {} (lsn {:#x}, declared len {}, have {} bytes); treating as crash-time EOF", + offset, + raw_lsn, + len, + total_len - payload_start + ); + break; + } + + // Strip the OrderedAcrossShards flag to recover the true LSN. + let is_ordered = raw_lsn & crate::persistence::aof::ORDERED_LSN_FLAG != 0; + let lsn = raw_lsn & !crate::persistence::aof::ORDERED_LSN_FLAG; + + // Ordered entries: buffer for cross-shard merge replay; do NOT + // dispatch inline. + if is_ordered { + let bytes = bytes::Bytes::copy_from_slice(&data[payload_start..payload_end]); + ordered_buf.push(OrderedEntry { + shard_id, + lsn, + bytes, + }); + // F5: track the next-free replication offset (entry end), not the + // start LSN — `issue_lsn` returns the offset BEFORE adding the + // entry length, so the seed must clear every byte already on disk. + let entry_end = lsn + len as u64; + if entry_end > max_lsn { + max_lsn = entry_end; + } + offset = payload_end; + continue; + } + + // Parse RESP from the payload slice. A standalone slice ensures one + // header maps to exactly one command — no implicit pipelining across + // headers. + let mut buf = BytesMut::from(&data[payload_start..payload_end]); + match parse::parse(&mut buf, &config) { + Ok(Some(frame)) => { + let (cmd, cmd_args) = match &frame { + Frame::Array(arr) if !arr.is_empty() => { + let name = match &arr[0] { + Frame::BulkString(s) => s.as_ref(), + Frame::SimpleString(s) => s.as_ref(), + other => { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "AOF incr framed command at offset {} (lsn {}) has non-string name frame: {:?}", + offset, + lsn, + std::mem::discriminant(other) + ), + }, + )); + } + }; + (name as &[u8], &arr[1..]) + } + other => { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "AOF incr framed non-array frame at offset {} (lsn {}): {:?}", + offset, + lsn, + std::mem::discriminant(other) + ), + }, + )); + } + }; + engine.replay_command(databases, cmd, cmd_args, &mut selected_db); + count += 1; + // F5: next-free offset = entry start LSN + RESP byte length. + let entry_end = lsn + len as u64; + if entry_end > max_lsn { + max_lsn = entry_end; + } + } + Ok(None) => { + // Header said `len` bytes of RESP, but parser can't make a + // frame from those bytes. That's corruption inside a fully + // declared payload, not a truncated tail — escalate. + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "AOF incr framed payload at offset {} (lsn {}, len {}) parsed as incomplete frame; corrupt entry", + offset, lsn, len + ), + }, + )); + } + Err(e) => { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "AOF incr framed parse error at offset {} (lsn {}, len {}): {:?}", + offset, lsn, len, e + ), + }, + )); + } + } + + offset = payload_end; + } + + Ok((count, max_lsn)) +} + +/// Replay a PerShard multi-part AOF into N parallel `Vec` buffers. +/// +/// `per_shard_databases[i]` is shard `i`'s database vector. The manifest's +/// `shards` length MUST equal `per_shard_databases.len()`; the caller is +/// expected to have run [`AofManifest::verify_shard_count`] at boot. +/// +/// Per-shard replay is fully parallel: each shard's base RDB load and incr +/// replay run in a separate OS thread via `std::thread::scope`. Shards are +/// independent (different `DashTable` instances, no shared mutable state), so +/// this is safe and correct. Parallelism delivers the RFC § 1 benefit on +/// multi-shard deployments with large AOF files. +/// +/// The `engine_factory` closure is called once per shard thread to produce an +/// independent replay engine. This is required because `CommandReplayEngine` +/// implementations (e.g., `DispatchReplayEngine` under the `graph` feature) +/// may contain non-`Sync` state (`RefCell`) that cannot be safely shared across +/// threads. Each thread owns its own engine; results (total count, max LSN, +/// ordered entries) are collected and merged in the caller thread after all +/// shard threads complete. +/// +/// Returns `(total_commands_replayed, global_max_lsn, ordered_entries)`: +/// - `total_commands_replayed` covers all inline (non-ordered) entries +/// plus the base-RDB key count. +/// - `global_max_lsn` is `max(per-shard max LSN)` across both inline and +/// ordered entries; the caller is expected to call +/// `ReplicationState::seed_master_offset(global_max_lsn)` before +/// accepting client traffic (RFC § 2 Rule 3). +/// - `ordered_entries` is the set of `OrderedAcrossShards`-tagged entries +/// across ALL shards; the caller passes them to +/// [`replay_ordered_merge`] for the cross-shard merge replay. +pub fn replay_per_shard( + per_shard_databases: &mut [&mut [crate::storage::Database]], + manifest: &AofManifest, + engine_factory: &( + dyn Fn() -> Box + Sync + ), +) -> Result<(usize, u64, Vec), crate::error::MoonError> { + debug_assert_eq!( + manifest.layout, + AofLayout::PerShard, + "replay_per_shard called on TopLevel manifest" + ); + if manifest.shards.len() != per_shard_databases.len() { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "replay_per_shard shard-count mismatch: manifest has {} shards, caller passed {} database vectors", + manifest.shards.len(), + per_shard_databases.len() + ), + }, + )); + } + + // Per-shard type alias for the thread result. + type ShardResult = Result<(usize, u64, Vec), crate::error::MoonError>; + + // Use std::thread::scope so each shard thread borrows its databases slice + // without a 'static lifetime requirement. All threads complete before scope + // exits, which satisfies the borrow checker. Errors are propagated via + // a Vec collected after join. + let shard_results: Vec = std::thread::scope(|scope| { + let mut handles = Vec::with_capacity(per_shard_databases.len()); + + for (shard_id, databases) in per_shard_databases.iter_mut().enumerate() { + let sid = shard_id as u16; + let base_path = manifest.shard_base_path(sid); + let incr_path = manifest.shard_incr_path(sid); + let engine = engine_factory(); + + handles.push(scope.spawn(move || -> ShardResult { + let mut shard_total: usize = 0; + let mut shard_max_lsn: u64 = 0; + let mut shard_ordered: Vec = Vec::new(); + + // Load this shard's base RDB. + if base_path.exists() { + match crate::persistence::rdb::load(*databases, &base_path) { + Ok(n) => { + info!( + "AOF shard-{} base RDB loaded: {} keys from {}", + sid, + n, + base_path.display() + ); + shard_total += n; + } + Err(e) => { + error!("AOF shard-{} base RDB load failed: {}", sid, e); + return Err(e); + } + } + } else { + // Missing base is tolerable only when this shard's incr file is + // empty (or absent). Same invariant as `replay_multi_part`. + let incr_len = + std::fs::metadata(&incr_path).map(|m| m.len()).unwrap_or(0); + if incr_len > 0 { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "AOF shard-{} base RDB missing at {} but incr {} is {} bytes; refusing to replay incr against empty state", + sid, + base_path.display(), + incr_path.display(), + incr_len, + ), + }, + )); + } + warn!( + "AOF shard-{} base RDB not found: {} (incr empty, treating as fresh init)", + sid, + base_path.display() + ); + } + + // Replay this shard's framed incr file. + if incr_path.exists() { + let data = std::fs::read(&incr_path).map_err(|e| { + crate::error::MoonError::from(crate::error::AofError::Io { + path: incr_path.clone(), + source: e, + }) + })?; + if !data.is_empty() { + let (count, max_lsn) = replay_incr_framed( + sid, + *databases, + &data, + engine.as_ref(), + &mut shard_ordered, + )?; + info!( + "AOF shard-{} incr replayed: {} commands from {} (max lsn {})", + sid, + count, + incr_path.display(), + max_lsn + ); + shard_total += count; + if max_lsn > shard_max_lsn { + shard_max_lsn = max_lsn; + } + } + } + + Ok((shard_total, shard_max_lsn, shard_ordered)) + })); + } + + // Collect results in shard order. + handles + .into_iter() + .map(|h| { + h.join().unwrap_or_else(|_| { + Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: "replay_per_shard worker thread panicked".to_owned(), + }, + )) + }) + }) + .collect() + }); + + // Merge per-shard results. + let mut total: usize = 0; + let mut global_max_lsn: u64 = 0; + let mut ordered_entries: Vec = Vec::new(); + + for result in shard_results { + let (shard_total, shard_max_lsn, shard_ordered) = result?; + total += shard_total; + if shard_max_lsn > global_max_lsn { + global_max_lsn = shard_max_lsn; + } + ordered_entries.extend(shard_ordered); + } + + Ok((total, global_max_lsn, ordered_entries)) +} + +/// Merge-replay `OrderedAcrossShards` entries collected across all shards +/// in global LSN order (RFC § 2 Rule 2). +/// +/// `entries` is sorted by `lsn` ascending, then each entry is dispatched +/// against its origin shard's databases — the per-shard partition is +/// preserved because each `OrderedEntry` carries the `shard_id` it was +/// read from. This guarantees that a cross-shard atomic operation +/// committed at LSN N is replayed as a coherent group (every +/// shard's portion at LSN N is applied before any shard's LSN N+1 work). +/// +/// **Crash-time atomicity:** if a cross-shard commit was mid-write at +/// crash time, some shards may have the LSN-N entry while others don't. +/// Step 5 ships the merge mechanism only; detecting partial commits and +/// performing the corresponding rollback is left to the future cross-shard +/// TXN consumer — `replay_ordered_merge` currently best-effort-applies +/// whichever entries survived. A `warn!` is emitted when the entry count +/// per LSN is uneven across shards so operators have a forensic trail. +/// +/// **Today's emitters:** none in production code. The path is exercised +/// by tests so the round-trip wiring is verified end-to-end and ready for +/// future use. +pub fn replay_ordered_merge( + per_shard_databases: &mut [&mut [crate::storage::Database]], + mut entries: Vec, + engine: &dyn crate::persistence::replay::CommandReplayEngine, +) -> Result { + use crate::protocol::{Frame, ParseConfig, parse}; + use bytes::BytesMut; + + if entries.is_empty() { + return Ok(0); + } + + entries.sort_by_key(|e| e.lsn); + + // Per-LSN cardinality audit: detect torn cross-shard commits. + // + // A "torn" commit is one where LSN N appears in fewer shard files than + // the maximum cardinality seen for any other LSN in this batch. Applying + // partial entries violates atomicity — if the write was interrupted mid- + // commit (e.g., crash between shard-0 and shard-1 writes), replaying only + // the shard-0 portion produces an inconsistent state that cannot be + // compensated. DROP the entire torn LSN instead of applying partial data. + // + // NOTE: "torn" detection is heuristic — it compares each LSN's count + // against the maximum cardinality observed. An LSN that legitimately spans + // fewer shards (e.g. single-shard ordered op) can only occur if the batch + // is heterogeneous. Production emitters (future cross-shard TXN) must + // guarantee uniform cardinality per LSN, so this heuristic is correct for + // all currently-reachable code paths. + let mut counts: std::collections::BTreeMap = std::collections::BTreeMap::new(); + for e in &entries { + *counts.entry(e.lsn).or_insert(0) += 1; + } + let max_count = counts.values().copied().max().unwrap_or(0); + let mut torn_lsns: std::collections::BTreeSet = std::collections::BTreeSet::new(); + for (&lsn, &n) in &counts { + if n < max_count { + warn!( + "OrderedAcrossShards LSN {} appears in only {} of {} shard files; \ + torn cross-shard commit detected — dropping entry for atomicity", + lsn, n, max_count + ); + torn_lsns.insert(lsn); + } + } + + let config = ParseConfig::default(); + let mut replayed: usize = 0; + + for entry in entries { + // Skip entries belonging to a torn (partially-written) commit. + if torn_lsns.contains(&entry.lsn) { + continue; + } + let shard_idx = entry.shard_id as usize; + if shard_idx >= per_shard_databases.len() { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "OrderedAcrossShards entry references shard {} but only {} shards present", + entry.shard_id, + per_shard_databases.len() + ), + }, + )); + } + let mut buf = BytesMut::from(entry.bytes.as_ref()); + match parse::parse(&mut buf, &config) { + Ok(Some(Frame::Array(arr))) if !arr.is_empty() => { + let cmd = match &arr[0] { + Frame::BulkString(s) => s.as_ref(), + Frame::SimpleString(s) => s.as_ref(), + _ => { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "OrderedAcrossShards entry at lsn {} has non-string command frame", + entry.lsn + ), + }, + )); + } + }; + let mut selected_db: usize = 0; + let databases = &mut *per_shard_databases[shard_idx]; + engine.replay_command(databases, cmd, &arr[1..], &mut selected_db); + replayed += 1; + } + other => { + return Err(crate::error::MoonError::from( + crate::error::AofError::RewriteFailed { + detail: format!( + "OrderedAcrossShards entry at lsn {} on shard {} did not parse as RESP array: {:?}", + entry.lsn, + entry.shard_id, + other.map(|_| ()).err() + ), + }, + )); + } + } + } + + Ok(replayed) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + fn temp_dir() -> PathBuf { + // Use a global atomic counter so parallel test threads (cargo test runs + // unit tests in parallel) never produce the same directory name even + // when PID and nanosecond clock resolution are the same for two threads. + static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let d = std::env::temp_dir().join(format!( + "moon-aof-manifest-test-replay-{}-{}", + std::process::id(), + n, + )); + fs::create_dir_all(&d).expect("temp dir create"); + d + } + + // -- Step 4 (per-shard replay) tests --------------------------------- + + fn frame_entry(lsn: u64, resp: &[u8]) -> Vec { + let mut buf = Vec::with_capacity(12 + resp.len()); + buf.extend_from_slice(&lsn.to_le_bytes()); + buf.extend_from_slice(&(resp.len() as u32).to_le_bytes()); + buf.extend_from_slice(resp); + buf + } + + /// Minimal `CommandReplayEngine` that records (lsn-implicit-via-order, cmd + /// name) calls without touching real storage. Tests use this to assert + /// the framed parser hands the right command sequence to the engine. + struct RecordingEngine { + calls: std::cell::RefCell>, + } + + impl RecordingEngine { + fn new() -> Self { + Self { + calls: std::cell::RefCell::new(Vec::new()), + } + } + } + + impl crate::persistence::replay::CommandReplayEngine for RecordingEngine { + fn replay_command( + &self, + _databases: &mut [crate::storage::Database], + cmd: &[u8], + _args: &[crate::protocol::Frame], + _selected_db: &mut usize, + ) { + self.calls + .borrow_mut() + .push(String::from_utf8_lossy(cmd).into_owned()); + } + } + + #[test] + fn replay_incr_framed_decodes_lsn_and_resp() { + // Two framed entries: PING and DBSIZE (no args, both small RESP arrays). + let mut bytes = frame_entry(7, b"*1\r\n$4\r\nPING\r\n"); + bytes.extend_from_slice(&frame_entry(11, b"*1\r\n$6\r\nDBSIZE\r\n")); + + let mut dbs: Vec = vec![crate::storage::Database::new()]; + let engine = RecordingEngine::new(); + let mut ordered: Vec = Vec::new(); + let (count, max_lsn) = + replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered).expect("framed replay"); + assert!(ordered.is_empty(), "no ordered entries in this stream"); + + assert_eq!(count, 2); + // F5: max_lsn is the NEXT-FREE offset = max(lsn + len) = max(7+14, 11+16) = 27. + assert_eq!(max_lsn, 27); + let calls = engine.calls.borrow(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0], "PING"); + assert_eq!(calls[1], "DBSIZE"); + } + + #[test] + fn replay_incr_framed_max_lsn_is_next_free_offset() { + // F5: replay must return the next-free replication offset (entry end = + // start LSN + RESP byte length), not the START LSN of the last entry. + // `issue_lsn` hands out the offset BEFORE adding the entry's length, so + // seeding `master_repl_offset` with a start LSN reissues the last + // pre-crash entry's LSN — breaking lsn->entry uniqueness (RFC § 2 Rule 3). + let ping = b"*1\r\n$4\r\nPING\r\n"; // 14 bytes + let dbsize = b"*1\r\n$6\r\nDBSIZE\r\n"; // 16 bytes + // Cumulative LSNs as the writer issues them: each entry starts at the + // previous entry's end. + let mut bytes = frame_entry(100, ping); + bytes.extend_from_slice(&frame_entry(100 + ping.len() as u64, dbsize)); + + let mut dbs: Vec = vec![crate::storage::Database::new()]; + let engine = RecordingEngine::new(); + let mut ordered: Vec = Vec::new(); + let (_count, max_lsn) = + replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered).expect("framed replay"); + + // Last entry: start 114 + len 16 = 130. The next write MUST get >= 130. + let expected_next_free = 100 + ping.len() as u64 + dbsize.len() as u64; + assert_eq!(expected_next_free, 130); + assert_eq!( + max_lsn, expected_next_free, + "max_lsn must be the next-free offset (entry end), not the last start LSN" + ); + } + + #[test] + fn replay_incr_framed_truncated_header_is_crash_eof() { + // One valid entry, then a partial 5-byte header (crash mid-write). + let mut bytes = frame_entry(3, b"*1\r\n$4\r\nPING\r\n"); + bytes.extend_from_slice(&[0u8; 5]); + + let mut dbs: Vec = vec![crate::storage::Database::new()]; + let engine = RecordingEngine::new(); + let mut ordered: Vec = Vec::new(); + let (count, max_lsn) = replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered) + .expect("truncated-header is EOF"); + + assert_eq!(count, 1); + // F5: next-free offset = PING entry start 3 + RESP len 14 = 17. + assert_eq!(max_lsn, 17); + } + + #[test] + fn replay_incr_framed_truncated_payload_is_crash_eof() { + // Header declares 14 bytes of RESP but only 5 actually present. + let mut bytes = Vec::new(); + bytes.extend_from_slice(&5u64.to_le_bytes()); + bytes.extend_from_slice(&14u32.to_le_bytes()); + bytes.extend_from_slice(b"*1\r\n$"); // 5 bytes, payload truncated + + let mut dbs: Vec = vec![crate::storage::Database::new()]; + let engine = RecordingEngine::new(); + let mut ordered: Vec = Vec::new(); + let (count, max_lsn) = replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered) + .expect("truncated-payload is EOF"); + + assert_eq!(count, 0); + assert_eq!(max_lsn, 0); + } + + #[test] + fn replay_incr_framed_complete_but_corrupt_payload_errors() { + // Header declares 4 bytes, payload is 4 bytes of garbage that won't + // parse as a RESP frame. + let mut bytes = Vec::new(); + bytes.extend_from_slice(&1u64.to_le_bytes()); + bytes.extend_from_slice(&4u32.to_le_bytes()); + bytes.extend_from_slice(b"XXXX"); + + let mut dbs: Vec = vec![crate::storage::Database::new()]; + let engine = RecordingEngine::new(); + let mut ordered: Vec = Vec::new(); + let err = replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered) + .expect_err("complete-but-corrupt should error"); + let msg = format!("{err}"); + assert!( + msg.contains("framed"), + "error should mention framed context, got: {msg}" + ); + } + + #[test] + fn replay_per_shard_round_trips_two_shards() { + let dir = temp_dir(); + let manifest = AofManifest::initialize_multi(&dir, 2).expect("initialize_multi 2 shards"); + + // Hand-author framed incr files: shard-0 SETs k0/v0 at lsn=10, + // shard-1 SETs k1/v1 at lsn=20. + let set_k0 = frame_entry(10, b"*3\r\n$3\r\nSET\r\n$2\r\nk0\r\n$2\r\nv0\r\n"); + let set_k1 = frame_entry(20, b"*3\r\n$3\r\nSET\r\n$2\r\nk1\r\n$2\r\nv1\r\n"); + fs::write(manifest.shard_incr_path(0), &set_k0).expect("write shard-0 incr"); + fs::write(manifest.shard_incr_path(1), &set_k1).expect("write shard-1 incr"); + + // Two independent shard database vectors. + let mut shard0: Vec = vec![crate::storage::Database::new()]; + let mut shard1: Vec = vec![crate::storage::Database::new()]; + + let (total, global_max_lsn, ordered) = { + let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0, &mut shard1]; + replay_per_shard( + &mut slices, + &manifest, + &(|| { + Box::new(crate::persistence::replay::DispatchReplayEngine::new()) + as Box + }), + ) + .expect("per-shard replay") + }; + + assert_eq!(total, 2, "two SETs replayed"); + // F5: global_max_lsn = max next-free offset across shards. shard-1 SET + // is 29 RESP bytes at lsn 20 → next-free 49 (> shard-0's 10+29=39). + assert_eq!( + global_max_lsn, 49, + "global max lsn = max(shard next-free offsets)" + ); + assert!(ordered.is_empty(), "no ordered entries in this stream"); + + // Each shard's DB now holds its key (and only its key). + assert!(shard0[0].len() >= 1, "shard 0 has k0"); + assert!(shard1[0].len() >= 1, "shard 1 has k1"); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn replay_per_shard_rejects_shard_count_mismatch() { + let dir = temp_dir(); + let manifest = AofManifest::initialize_multi(&dir, 2).expect("initialize_multi 2 shards"); + + // Only one slice — manifest says 2. + let mut shard0: Vec = vec![crate::storage::Database::new()]; + let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0]; + + let err = replay_per_shard( + &mut slices, + &manifest, + &(|| { + Box::new(crate::persistence::replay::DispatchReplayEngine::new()) + as Box + }), + ) + .expect_err("shard count mismatch must error"); + let msg = format!("{err}"); + assert!( + msg.contains("shard-count mismatch"), + "error message should call out the mismatch, got: {msg}" + ); + + fs::remove_dir_all(&dir).ok(); + } + + /// FIX-W3-1: parallel per-shard replay must produce identical results to + /// sequential replay. N=4 shards, one key per shard. + /// + /// Test gate: correctness (same total/max_lsn/key distribution as sequential). + /// Wall-time comparison is flaky in CI and omitted. + #[test] + fn replay_per_shard_parallel_matches_sequential() { + let dir = temp_dir(); + let n_shards: u16 = 4; + let manifest = + AofManifest::initialize_multi(&dir, n_shards).expect("initialize_multi 4 shards"); + + // Each shard gets one SET at lsn = shard_id * 10 + 10. + for sid in 0..n_shards { + let lsn = (sid as u64 + 1) * 10; + let key = format!("k{sid}"); + let val = format!("v{sid}"); + let resp = format!( + "*3\r\n$3\r\nSET\r\n${klen}\r\n{key}\r\n${vlen}\r\n{val}\r\n", + klen = key.len(), + vlen = val.len(), + ); + let entry = frame_entry(lsn, resp.as_bytes()); + fs::write(manifest.shard_incr_path(sid), &entry).expect("write shard incr"); + } + + let mut shards: Vec> = (0..n_shards as usize) + .map(|_| vec![crate::storage::Database::new()]) + .collect(); + + let engine_factory = || { + Box::new(crate::persistence::replay::DispatchReplayEngine::new()) + as Box + }; + let (total, global_max_lsn, ordered) = { + let mut slices: Vec<&mut [crate::storage::Database]> = + shards.iter_mut().map(|s| s.as_mut_slice()).collect(); + replay_per_shard(&mut slices, &manifest, &engine_factory) + .expect("parallel per-shard replay") + }; + + assert_eq!(total, n_shards as usize, "one SET per shard = N total"); + // F5: global_max_lsn = highest shard's NEXT-FREE offset. The highest + // shard (sid=N-1) SETs at lsn N*10 with a 29-byte RESP → next-free + // N*10 + 29 (here 40 + 29 = 69), not the bare start LSN. + assert_eq!( + global_max_lsn, + n_shards as u64 * 10 + 29, + "global max lsn = highest shard next-free offset" + ); + assert!(ordered.is_empty(), "no ordered entries"); + + // Each shard must have exactly one key. + for (sid, shard) in shards.iter().enumerate() { + assert_eq!( + shard[0].len(), + 1, + "shard {} must have exactly 1 key after parallel replay", + sid + ); + } + + fs::remove_dir_all(&dir).ok(); + } + + // -- Step 5 (OrderedAcrossShards merge) tests ------------------------ + + /// Frame an ordered entry: same on-disk layout as `frame_entry`, with + /// the high bit of LSN set. + fn frame_ordered(lsn: u64, resp: &[u8]) -> Vec { + assert_eq!( + lsn & crate::persistence::aof::ORDERED_LSN_FLAG, + 0, + "test helper expects raw lsn without the ordered flag" + ); + let tagged = lsn | crate::persistence::aof::ORDERED_LSN_FLAG; + let mut buf = Vec::with_capacity(12 + resp.len()); + buf.extend_from_slice(&tagged.to_le_bytes()); + buf.extend_from_slice(&(resp.len() as u32).to_le_bytes()); + buf.extend_from_slice(resp); + buf + } + + #[test] + fn replay_incr_framed_buffers_ordered_entries() { + // Mix: normal PING, then an ordered SET, then normal DBSIZE. + let mut bytes = frame_entry(5, b"*1\r\n$4\r\nPING\r\n"); + bytes.extend_from_slice(&frame_ordered( + 8, + b"*3\r\n$3\r\nSET\r\n$1\r\nk\r\n$1\r\nv\r\n", + )); + bytes.extend_from_slice(&frame_entry(12, b"*1\r\n$6\r\nDBSIZE\r\n")); + + let mut dbs: Vec = vec![crate::storage::Database::new()]; + let engine = RecordingEngine::new(); + let mut ordered: Vec = Vec::new(); + let (count, max_lsn) = replay_incr_framed(3, &mut dbs, &bytes, &engine, &mut ordered) + .expect("framed replay with ordered"); + + assert_eq!(count, 2, "two inline entries dispatched (PING, DBSIZE)"); + // F5: max_lsn is the next-free offset across inline AND ordered. The + // ordered SET (27 RESP bytes) at lsn 8 → next-free 35, exceeding the + // PING (5+14=19) and DBSIZE (12+16=28) ends. + assert_eq!( + max_lsn, 35, + "max LSN = next-free offset across inline and ordered" + ); + assert_eq!(ordered.len(), 1, "one entry buffered as ordered"); + let buffered = &ordered[0]; + assert_eq!(buffered.shard_id, 3, "shard_id forwarded"); + assert_eq!(buffered.lsn, 8, "buffered LSN has the high bit masked off"); + let calls = engine.calls.borrow(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0], "PING"); + assert_eq!(calls[1], "DBSIZE", "ordered SET was NOT dispatched inline"); + } + + #[test] + fn replay_ordered_merge_sorts_by_lsn_across_shards() { + use crate::persistence::replay::DispatchReplayEngine; + + // Three ordered entries across two shards, deliberately out of LSN + // order on the wire so the merge step has work to do. + let entries = vec![ + OrderedEntry { + shard_id: 1, + lsn: 30, + bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\nb1\r\n$1\r\n3\r\n"), + }, + OrderedEntry { + shard_id: 0, + lsn: 10, + bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\na1\r\n$1\r\n1\r\n"), + }, + OrderedEntry { + shard_id: 0, + lsn: 20, + bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\na2\r\n$1\r\n2\r\n"), + }, + ]; + + let mut shard0: Vec = vec![crate::storage::Database::new()]; + let mut shard1: Vec = vec![crate::storage::Database::new()]; + let replayed = { + let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0, &mut shard1]; + replay_ordered_merge(&mut slices, entries, &DispatchReplayEngine::new()) + .expect("ordered merge replay") + }; + + assert_eq!(replayed, 3); + assert!(shard0[0].len() >= 2, "shard 0 received a1 + a2"); + assert!(shard1[0].len() >= 1, "shard 1 received b1"); + } + + #[test] + fn replay_ordered_merge_empty_returns_zero() { + use crate::persistence::replay::DispatchReplayEngine; + + let mut shard0: Vec = vec![crate::storage::Database::new()]; + let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0]; + let replayed = replay_ordered_merge(&mut slices, Vec::new(), &DispatchReplayEngine::new()) + .expect("empty merge ok"); + assert_eq!(replayed, 0); + } + + /// FIX-W3-3: torn cross-shard commit must be DROPPED entirely, not partially applied. + /// + /// Synthesize a 2-shard AOF where LSN 100 appears on shard 0 only (N=1 + /// of K=2 expected). After replay, shard 0 must NOT have the key written + /// by the LSN-100 entry (it was dropped for atomicity). + #[test] + fn replay_ordered_merge_drops_torn_commit() { + use crate::persistence::replay::DispatchReplayEngine; + + // Two shards, two complete entries at LSN 10 (one per shard) — these + // should succeed. LSN 100 appears only on shard 0 (torn) — must be dropped. + let entries = vec![ + // Complete pair: LSN 10 on both shards + OrderedEntry { + shard_id: 0, + lsn: 10, + bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\nc0\r\n$1\r\n1\r\n"), + }, + OrderedEntry { + shard_id: 1, + lsn: 10, + bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$2\r\nc1\r\n$1\r\n1\r\n"), + }, + // Torn entry: LSN 100 only on shard 0, not shard 1 + OrderedEntry { + shard_id: 0, + lsn: 100, + bytes: bytes::Bytes::from_static(b"*3\r\n$3\r\nSET\r\n$5\r\ntorn0\r\n$1\r\nv\r\n"), + }, + ]; + + let mut shard0: Vec = vec![crate::storage::Database::new()]; + let mut shard1: Vec = vec![crate::storage::Database::new()]; + let replayed = { + let mut slices: Vec<&mut [crate::storage::Database]> = vec![&mut shard0, &mut shard1]; + replay_ordered_merge(&mut slices, entries, &DispatchReplayEngine::new()) + .expect("ordered merge replay") + }; + + // The torn LSN-100 entry must NOT be applied (dropped for atomicity). + assert_eq!(replayed, 2, "only the complete LSN-10 pair is replayed"); + assert_eq!( + shard0[0].len(), + 1, + "shard-0 only has the complete LSN-10 key; torn LSN-100 entry must not be applied" + ); + // Verify the torn key is absent + assert!( + shard0[0].get(b"torn0").is_none(), + "torn shard-0 entry (LSN 100) must NOT be applied" + ); + } + + #[test] + fn ordered_entry_lsn_flag_set_via_try_send_append_ordered() { + use crate::persistence::aof::{AofMessage, AofWriterPool, ORDERED_LSN_FLAG}; + use crate::runtime::channel; + + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + // Raw lsn = 42; high bit must end up set on the receive side. + pool.try_send_append_ordered(0, 42, bytes::Bytes::from_static(b"x")); + let msg = rx0.try_recv().expect("ordered append delivered"); + match msg { + AofMessage::Append { lsn, .. } => { + assert_eq!( + lsn & ORDERED_LSN_FLAG, + ORDERED_LSN_FLAG, + "ordered flag set on lsn" + ); + assert_eq!( + lsn & !ORDERED_LSN_FLAG, + 42, + "low bits preserve the original lsn" + ); + } + _ => panic!("expected Append"), + } + } + + // ----------------------------------------------------------------------- + // FIX-W2-1: cleanup_orphans must recurse into shard-N/ subdirectories + // ----------------------------------------------------------------------- +} diff --git a/src/persistence/aof_manifest/shard_rewrite.rs b/src/persistence/aof_manifest/shard_rewrite.rs new file mode 100644 index 00000000..cb30cbf1 --- /dev/null +++ b/src/persistence/aof_manifest/shard_rewrite.rs @@ -0,0 +1,577 @@ +//! Per-shard manifest rewrite helpers: `initialize_multi` / +//! `try_initialize_multi` and `advance_shard` / `prune_shard_files`. +//! +//! Split out of the `aof_manifest` parent module (issue #143) to keep each file +//! under the 1500-line cap. These are inherent `AofManifest` methods; `use +//! super::*` brings the type, its private fields (a child module sees parent +//! privates), and the std/tracing imports into scope. + +use super::*; + +impl AofManifest { + /// Create the `appendonlydir/` and write an initial v2 manifest for the + /// given shard count. + /// + /// Each shard gets its own `shard-{N}/` subdirectory with an empty base + /// RDB and an empty incr file. Mirrors `initialize()` semantics: the + /// `(base + incr)` invariant holds from the first boot, so recovery can + /// replay incr-only state without complaint. + /// + /// **Idempotency pre-flight:** if `appendonlydir/moon.aof.manifest` already + /// exists, returns `Err(AlreadyExists)` without modifying any files. A + /// mid-loop crash followed by a retry would otherwise overwrite the already- + /// written shard-0 base RDB with an empty RDB, losing state. Callers that + /// want resume-or-skip semantics should use [`Self::try_initialize_multi`]. + /// + /// **Rollback on partial failure:** if the per-shard loop fails mid-way (e.g. + /// shard-1 write fails after shard-0 succeeded), all already-created shard + /// base RDB files are deleted before returning the error. + pub fn initialize_multi(dir: &Path, num_shards: u16) -> std::io::Result { + if num_shards == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "initialize_multi requires num_shards >= 1", + )); + } + let manifest = Self { + dir: dir.to_path_buf(), + seq: 1, + layout: AofLayout::PerShard, + shards: (0..num_shards) + .map(|id| ShardManifest { + shard_id: id, + max_lsn: 0, + }) + .collect(), + }; + std::fs::create_dir_all(manifest.aof_dir())?; + + // Pre-flight: refuse if manifest already exists to avoid overwriting + // already-written shard base RDB files (idempotency guard). + let manifest_path = manifest.manifest_path(); + if manifest_path.exists() { + return Err(std::io::Error::new( + std::io::ErrorKind::AlreadyExists, + format!( + "initialize_multi: manifest already exists at {}; \ + use try_initialize_multi() for idempotent initialization", + manifest_path.display() + ), + )); + } + + // Per-shard empty RDB. Single Database::default() inside a 1-element + // slice matches `initialize()`'s empty-RDB shape for each shard. + let empty_dbs: [crate::storage::Database; 0] = []; + let empty_rdb = crate::persistence::rdb::save_to_bytes(&empty_dbs) + .map_err(|e| std::io::Error::other(format!("empty RDB serialize: {e}")))?; + + // Track which shard directories were successfully created so we can + // roll them back on partial failure. + let mut created_shards: Vec = Vec::with_capacity(num_shards as usize); + + let loop_result = (|| -> std::io::Result<()> { + for shard_id in 0..num_shards { + let shard_dir = manifest.shard_dir(shard_id); + std::fs::create_dir_all(&shard_dir)?; + + let base_path = manifest.shard_base_path(shard_id); + let tmp_path = base_path.with_extension("rdb.tmp"); + { + let mut f = std::fs::File::create(&tmp_path)?; + f.write_all(&empty_rdb)?; + f.sync_data()?; + } + std::fs::rename(&tmp_path, &base_path)?; + fsync_parent_best_effort(&base_path); + std::fs::File::create(manifest.shard_incr_path(shard_id))?; + created_shards.push(shard_id); + } + Ok(()) + })(); + + if let Err(e) = loop_result { + // Rollback: remove base RDB files for all successfully-created shards. + for sid in created_shards { + let base = manifest.shard_base_path(sid); + if let Err(re) = std::fs::remove_file(&base) { + warn!( + "initialize_multi rollback: failed to remove {}: {}", + base.display(), + re + ); + } + } + return Err(e); + } + + manifest.write_manifest()?; + Ok(manifest) + } + /// Initialize a v2 multi-shard manifest only if one does not already exist. + /// + /// Returns `Ok(Some(manifest))` on successful creation, or `Ok(None)` if the + /// manifest file already existed (already initialized — no files modified). + /// Returns `Err(_)` only on actual I/O failures. + pub fn try_initialize_multi(dir: &Path, num_shards: u16) -> std::io::Result> { + match Self::initialize_multi(dir, num_shards) { + Ok(m) => Ok(Some(m)), + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(None), + Err(e) => Err(e), + } + } + /// Advance a single shard to a new sequence: write the shard's new base RDB, + /// create a new empty incr file, then update the shard's `max_lsn` in the + /// in-memory manifest. + /// + /// **Does NOT delete the old generation's files.** Deletion is deferred to + /// the coordinator via [`prune_shard_files`](Self::prune_shard_files), + /// called only AFTER `write_manifest()` durably commits the new seq. + /// Deleting before the commit would leave a crash window where the + /// persisted (old) seq points at files that are already gone — recovery + /// resolves base/incr by `self.seq`, so a crash mid-fan-out would lose data + /// for any shard that had already advanced. This matches the post-commit + /// deletion ordering in [`advance`](Self::advance) (TopLevel layout). + /// + /// **Caller MUST call `write_manifest()` after all shards have been advanced** + /// (and set `self.seq` to the new seq) to persist the updated manifest + /// atomically — this is the single durable commit point for the rewrite. + /// Advancing shards one at a time and writing the manifest per-shard would + /// leave the manifest in an inconsistent state between calls. + /// + /// For `TopLevel` layout, `shard_id` must be 0 and this delegates to + /// `advance()` (which deletes post-commit internally). For `PerShard` + /// layout, files are written to `shard_dir(shard_id)/`. + /// + /// Returns the path to the new incremental file for this shard. + pub fn advance_shard( + &mut self, + shard_id: u16, + new_seq: u64, + rdb_bytes: &[u8], + ) -> Result { + if self.layout == AofLayout::TopLevel { + debug_assert_eq!(shard_id, 0, "TopLevel layout only has shard 0"); + return self.advance(rdb_bytes); + } + + // Validate shard_id is known in this manifest. + let shard_idx = self + .shards + .iter() + .position(|s| s.shard_id == shard_id) + .ok_or_else(|| crate::error::AofError::RewriteFailed { + detail: format!( + "advance_shard: shard_id {} not in manifest (shards: {})", + shard_id, + self.shards.len() + ), + })?; + + let shard_dir = self.shard_dir(shard_id); + std::fs::create_dir_all(&shard_dir).map_err(|e| crate::error::AofError::Io { + path: shard_dir.clone(), + source: e, + })?; + + // 1. Write new base RDB atomically: tmp + fsync + rename. + let new_base = self.shard_base_path_seq(shard_id, new_seq); + let tmp_base = new_base.with_extension("rdb.tmp"); + { + let mut f = + std::fs::File::create(&tmp_base).map_err(|e| crate::error::AofError::Io { + path: tmp_base.clone(), + source: e, + })?; + f.write_all(rdb_bytes) + .map_err(|e| crate::error::AofError::Io { + path: tmp_base.clone(), + source: e, + })?; + f.sync_data().map_err(|e| crate::error::AofError::Io { + path: tmp_base.clone(), + source: e, + })?; + } + std::fs::rename(&tmp_base, &new_base).map_err(|e| { + crate::error::AofError::RewriteFailed { + detail: format!( + "advance_shard {}: rename base {}: {}", + shard_id, + tmp_base.display(), + e + ), + } + })?; + fsync_parent_best_effort(&new_base); + + // 2. Create empty new incremental file. + let new_incr = self.shard_incr_path_seq(shard_id, new_seq); + std::fs::File::create(&new_incr).map_err(|e| crate::error::AofError::Io { + path: new_incr.clone(), + source: e, + })?; + + // 3. Update per-shard LSN in-memory (manifest write is the caller's job). + // Old-generation files are intentionally NOT deleted here — the + // coordinator prunes them via `prune_shard_files` only after + // `write_manifest()` durably commits the new seq (see the fn doc; + // delete-before-commit would lose data on a mid-fan-out crash). + self.shards[shard_idx].max_lsn = self.shards[shard_idx].max_lsn.max(new_seq); + + info!( + "AOF shard {} advanced to seq {}: base={} bytes, incr={}", + shard_id, + new_seq, + rdb_bytes.len(), + new_incr.display() + ); + + Ok(new_incr) + } + /// Delete a shard's base + incr files for a specific `seq`. Best-effort. + /// + /// **Crash-safety contract:** the rewrite coordinator MUST call this only + /// AFTER `write_manifest()` has durably committed the new seq. Deleting an + /// old generation's files before the manifest flips would orphan the + /// persisted (old) seq whose files are already gone — recovery resolves + /// base/incr by `self.seq`, so it would read a missing base and lose data + /// for any shard that completed before the crash. This mirrors the + /// post-commit deletion ordering in `advance()` (TopLevel layout). + pub fn prune_shard_files(&self, shard_id: u16, seq: u64) { + let base = self.shard_base_path_seq(shard_id, seq); + let incr = self.shard_incr_path_seq(shard_id, seq); + if base.exists() { + if let Err(e) = std::fs::remove_file(&base) { + warn!( + "prune_shard_files {}: failed to delete old base {}: {}", + shard_id, + base.display(), + e + ); + } + } + if incr.exists() { + if let Err(e) = std::fs::remove_file(&incr) { + warn!( + "prune_shard_files {}: failed to delete old incr {}: {}", + shard_id, + incr.display(), + e + ); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + fn temp_dir() -> PathBuf { + // Use a global atomic counter so parallel test threads (cargo test runs + // unit tests in parallel) never produce the same directory name even + // when PID and nanosecond clock resolution are the same for two threads. + static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let d = std::env::temp_dir().join(format!( + "moon-aof-manifest-test-rewrite-{}-{}", + std::process::id(), + n, + )); + fs::create_dir_all(&d).expect("temp dir create"); + d + } + + /// Build a minimal PerShard manifest fixture on disk at `seq` without + /// needing `advance_shard`. Directly creates the expected directory layout + /// so the test is self-contained and doesn't depend on FIX-W2-3 methods. + fn write_per_shard_manifest_at_seq(dir: &Path, num_shards: u16, seq: u64) -> AofManifest { + let aof_dir = dir.join(AOF_DIR_NAME); + fs::create_dir_all(&aof_dir).unwrap(); + let empty_rdb = crate::persistence::rdb::save_to_bytes(&[] as &[crate::storage::Database]) + .expect("empty rdb"); + let shards: Vec = (0..num_shards) + .map(|id| ShardManifest { + shard_id: id, + max_lsn: 0, + }) + .collect(); + let manifest = AofManifest { + dir: dir.to_path_buf(), + seq, + layout: AofLayout::PerShard, + shards, + }; + for shard_id in 0..num_shards { + let shard_dir = manifest.shard_dir(shard_id); + fs::create_dir_all(&shard_dir).unwrap(); + let base = manifest.shard_base_path(shard_id); + let tmp = base.with_extension("rdb.tmp"); + fs::write(&tmp, &empty_rdb).unwrap(); + fs::rename(&tmp, &base).unwrap(); + fs::write(manifest.shard_incr_path(shard_id), b"").unwrap(); + } + manifest.write_manifest().unwrap(); + manifest + } + + #[test] + fn cleanup_orphans_removes_stale_files_in_shard_subdirs() { + let dir = temp_dir(); + + // Build a 2-shard PerShard manifest at seq=2. + let manifest = write_per_shard_manifest_at_seq(&dir, 2, 2); + + // Inject orphan files in shard-0/ that a crashed BGREWRITEAOF would leave. + // seq=1 tmp (aborted write) and a seq=5 incr (future zombie). + let shard0_dir = manifest.shard_dir(0); + let orphan_tmp = shard0_dir.join("moon.aof.1.base.rdb.tmp"); + let orphan_old_incr = shard0_dir.join("moon.aof.5.incr.aof"); + fs::write(&orphan_tmp, b"").expect("write orphan tmp"); + fs::write(&orphan_old_incr, b"").expect("write orphan incr"); + + // Active files for seq=2 must survive. + let active_base = manifest.shard_base_path(0); + let active_incr = manifest.shard_incr_path(0); + assert!( + active_base.exists(), + "active base must exist before cleanup" + ); + assert!( + active_incr.exists(), + "active incr must exist before cleanup" + ); + + // Reload the manifest — this triggers cleanup_orphans. + let _reloaded = AofManifest::load(&dir).expect("load").expect("present"); + + assert!( + !orphan_tmp.exists(), + "orphan .rdb.tmp in shard-0/ must be deleted by cleanup_orphans" + ); + assert!( + !orphan_old_incr.exists(), + "orphan old incr in shard-0/ must be deleted by cleanup_orphans" + ); + assert!( + active_base.exists(), + "active seq=2 base must survive cleanup" + ); + assert!( + active_incr.exists(), + "active seq=2 incr must survive cleanup" + ); + + fs::remove_dir_all(&dir).ok(); + } + + // ----------------------------------------------------------------------- + // FIX-W2-2: initialize_multi idempotency — second call returns error + // ----------------------------------------------------------------------- + #[test] + fn initialize_multi_second_call_returns_already_initialized_error() { + let dir = temp_dir(); + + // First call must succeed. + let _m = AofManifest::initialize_multi(&dir, 4).expect("first call ok"); + + // Count files before second call. + let aof_dir = dir.join(AOF_DIR_NAME); + let count_before: usize = (0..4u16) + .map(|sid| { + let shard_dir = aof_dir.join(format!("shard-{}", sid)); + fs::read_dir(&shard_dir).map(|e| e.count()).unwrap_or(0) + }) + .sum(); + + // Second call must return an error with the manifest already present. + let result = AofManifest::initialize_multi(&dir, 4); + assert!( + result.is_err(), + "second initialize_multi must fail when manifest already exists" + ); + let err = result.unwrap_err(); + assert_eq!( + err.kind(), + std::io::ErrorKind::AlreadyExists, + "error kind must be AlreadyExists; got {:?}: {}", + err.kind(), + err + ); + + // File count must be unchanged — no files were overwritten. + let count_after: usize = (0..4u16) + .map(|sid| { + let shard_dir = aof_dir.join(format!("shard-{}", sid)); + fs::read_dir(&shard_dir).map(|e| e.count()).unwrap_or(0) + }) + .sum(); + assert_eq!( + count_before, count_after, + "second call must not create or overwrite any shard files" + ); + + fs::remove_dir_all(&dir).ok(); + } + + // ----------------------------------------------------------------------- + // F6 crash-safety ordering: advance_shard writes new base+incr but MUST + // NOT delete old files. Deleting before the manifest durably commits the + // new seq leaves a window where a crash orphans the persisted (old) seq + // whose files are already gone → recovery reads a missing base → data + // loss for completed shards. Deletion is the coordinator's job, AFTER + // write_manifest(), via prune_shard_files(). This mirrors the proven + // ordering in advance() (TopLevel), which deletes only post-commit. + // ----------------------------------------------------------------------- + #[test] + fn advance_shard_defers_delete_until_after_commit() { + let dir = temp_dir(); + + // Initialize 2-shard manifest at seq=1. + let mut manifest = AofManifest::initialize_multi(&dir, 2).expect("initialize_multi"); + assert_eq!(manifest.seq, 1); + + let empty_rdb = crate::persistence::rdb::save_to_bytes(&[] as &[crate::storage::Database]) + .expect("empty rdb"); + + // Old shard files at seq=1 must exist before advance. + let old_base_s0 = manifest.shard_base_path_seq(0, 1); + let old_incr_s0 = manifest.shard_incr_path_seq(0, 1); + let old_base_s1 = manifest.shard_base_path_seq(1, 1); + let old_incr_s1 = manifest.shard_incr_path_seq(1, 1); + assert!(old_base_s0.exists(), "seq=1 base must exist for shard 0"); + assert!(old_incr_s0.exists(), "seq=1 incr must exist for shard 0"); + + // Fan out: coordinator picks new_seq=2 once and advances every shard + // to it. No manifest write, no deletion, happens inside the fan-out. + let new_incr_s0 = manifest + .advance_shard(0, 2, &empty_rdb) + .expect("advance_shard 0 → seq=2"); + let new_incr_s1 = manifest + .advance_shard(1, 2, &empty_rdb) + .expect("advance_shard 1 → seq=2"); + assert!(new_incr_s0.exists(), "new incr file must be created (s0)"); + assert!(new_incr_s1.exists(), "new incr file must be created (s1)"); + + // PRE-COMMIT INVARIANT: new files written, OLD files NOT yet deleted. + // This is the regression guard against delete-before-commit. + assert!( + manifest.shard_base_path_seq(0, 2).exists(), + "new seq=2 base must exist for shard 0" + ); + assert!( + manifest.shard_base_path_seq(1, 2).exists(), + "new seq=2 base must exist for shard 1" + ); + assert!( + old_base_s0.exists(), + "old seq=1 base (s0) MUST survive until the manifest commits" + ); + assert!( + old_incr_s0.exists(), + "old seq=1 incr (s0) MUST survive until the manifest commits" + ); + assert!( + old_base_s1.exists(), + "old seq=1 base (s1) MUST survive until the manifest commits" + ); + + // COMMIT: coordinator bumps seq and persists the manifest atomically. + // This is the single durable commit point for the whole rewrite. + manifest.seq = 2; + manifest + .write_manifest() + .expect("write manifest after advance"); + + // POST-COMMIT: coordinator prunes old files — safe now that recovery + // resolves base/incr by the durably-committed new seq. + manifest.prune_shard_files(0, 1); + manifest.prune_shard_files(1, 1); + assert!( + !old_base_s0.exists(), + "old seq=1 base (s0) pruned post-commit" + ); + assert!( + !old_incr_s0.exists(), + "old seq=1 incr (s0) pruned post-commit" + ); + assert!( + !old_base_s1.exists(), + "old seq=1 base (s1) pruned post-commit" + ); + assert!( + !old_incr_s1.exists(), + "old seq=1 incr (s1) pruned post-commit" + ); + assert!( + manifest.shard_base_path_seq(0, 2).exists(), + "new seq=2 base (s0) must remain after prune" + ); + + // Recovery reads base by manifest.seq — must resolve to the new seq. + let reloaded = AofManifest::load(&dir).expect("load").expect("present"); + assert_eq!(reloaded.seq, 2); + + fs::remove_dir_all(&dir).ok(); + } + + // ----------------------------------------------------------------------- + // FIX-W2-7: smoke test — fsync helper consolidation did not break + // initialize_multi. Checks that the post-consolidation manifest has the + // correct PerShard layout, the expected shard count, and per-shard + // base/incr files. This is discriminating: a regression that produces a + // TopLevel manifest or wrong shard count will be caught here. + // ----------------------------------------------------------------------- + #[test] + fn initialize_multi_smoke_after_fsync_consolidation() { + let tmp = tempfile::tempdir().expect("tempdir"); + let dir = tmp.path(); + let n: u16 = 2; + let result = AofManifest::initialize_multi(dir, n); + assert!( + result.is_ok(), + "initialize_multi({n} shards) must succeed: {:?}", + result.err() + ); + let manifest = result.unwrap(); + + // Discriminating: layout must be PerShard, not TopLevel. + assert_eq!( + manifest.layout, + AofLayout::PerShard, + "initialize_multi must produce a PerShard manifest" + ); + // Discriminating: shard count must match the requested count. + assert_eq!( + manifest.shards.len() as u16, + n, + "manifest must record exactly {n} shards, got {}", + manifest.shards.len() + ); + // Discriminating: per-shard base RDB and incr files must exist on disk. + for shard_id in 0..n { + assert!( + manifest.shard_base_path(shard_id).exists(), + "shard-{shard_id} base RDB must exist at {}", + manifest.shard_base_path(shard_id).display() + ); + assert!( + manifest.shard_incr_path(shard_id).exists(), + "shard-{shard_id} incr file must exist at {}", + manifest.shard_incr_path(shard_id).display() + ); + } + // Discriminating: the on-disk manifest file must contain `version 2` + // (PerShard v2 header), not be a bare v1 file. + let manifest_path = dir.join(AOF_DIR_NAME).join("moon.aof.manifest"); + let content = + std::fs::read_to_string(&manifest_path).expect("manifest file must be readable"); + assert!( + content.contains("version 2"), + "manifest file must contain 'version 2' (PerShard v2 header); got:\n{}", + content + ); + } +} From 566869d02115ef39d252268e3c274511f37cb367 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 13:57:49 +0700 Subject: [PATCH 09/15] chore(storage): correct stale async-spill doc, drop dead remove-first evict path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit's PR #144 review raised a 🟠 Major data-loss finding against the async-spill eviction path (inline on spill_thread.rs:476, citing eviction.rs:230-239): it claimed the eviction code removes a key from the hot tier BEFORE sending the SpillRequest, so a full/disconnected spill channel drops the request and loses data. Verified against the live code: the finding is STALE, not a real defect. The production path `evict_one_async_spill` (eviction.rs:520-526) is fail-safe send-before-remove: if sender.try_send(req).is_err() { return false; } // key stays in RAM db.remove(key.as_bytes()); // only after success On a full/disconnected channel the send fails, the entry stays resident, and the eviction loop bails to retry next tick — no acknowledged write is lost. The worst case under sustained backpressure is keys remaining resident (eventually an OOM write-rejection once at budget), which is correct fail-safe behaviour. All live callers (handler_monoio, handler_sharded, persistence_tick) route through this path. What actually misled the reviewer were two stale artefacts, both fixed here — no behaviour change to the live path: 1. A stale doc comment that had drifted off `try_evict_if_needed_async_spill` and silently merged into `next_spill_file_id_seed`'s doc (missing blank line). It still described the OLD "remove from RAM, then best-effort send, drop request if channel full" ordering. Replaced with an accurate doc on the real function describing the fail-safe send-before-remove contract. 2. Dead code `try_evict_deferred` + its private helper `find_victim_for_policy` (zero callers in src/, tests/, benches/, examples/ — verified by repo-root grep and `cargo check --all-targets` on both feature sets). This was the only code in the tree that actually embodied the dangerous remove-first pattern the reviewer feared (db.remove before any send). Removed so the live tree has exactly one eviction-spill path, and it is the fail-safe one. Validated: clippy -D warnings clean on default (monoio) and tokio,jemalloc; 15 storage::eviction unit tests green; cargo check --all-targets green on both feature sets; fmt clean. The review's second 🟠 Major (conn_accept.rs:277-320, gate spawn_migrated_tokio_connection behind cfg(unix)) is skipped: moon targets only Linux + macOS (both Unix), the function compiles and clippy-passes under runtime-tokio on every supported target, and the gap only affects unsupported Windows builds — gating it would touch the event_loop.rs call site for zero supported-platform benefit. author: Tin Dang --- src/storage/eviction.rs | 95 +++++++---------------------------------- 1 file changed, 15 insertions(+), 80 deletions(-) diff --git a/src/storage/eviction.rs b/src/storage/eviction.rs index b5f1404e..6e05ca5f 100644 --- a/src/storage/eviction.rs +++ b/src/storage/eviction.rs @@ -227,16 +227,6 @@ pub fn try_evict_if_needed_with_spill_and_total( Ok(()) } -/// Check if eviction is needed, spilling evicted entries asynchronously via -/// a background `SpillThread` instead of doing synchronous pwrite. -/// -/// The async path: extracts key/value bytes, removes entry from DashTable -/// (freeing RAM immediately), then sends a `SpillRequest` to the background -/// thread. The pwrite is best-effort -- if the channel is full, the request -/// is dropped (entry already removed from RAM). -/// -/// Callers must poll `SpillThread::drain_completions()` to apply manifest -/// and ColdIndex updates from completed spills. /// Seed value for a shard's spill `file_id` counter after recovery. /// /// On restart, AOF/RDB replay re-populates the hot tier and the persistence @@ -296,6 +286,21 @@ pub fn next_spill_file_id_seed(shard_dir: Option<&Path>) -> u64 { } } +/// Evict over-budget entries, spilling them to disk asynchronously via a +/// background `SpillThread` instead of doing a synchronous pwrite on the event +/// loop. +/// +/// Per victim (`evict_one_async_spill`): serialize the key/value into a +/// `SpillRequest`, **`try_send` it to the spill channel BEFORE removing the +/// entry from the hot tier, and only remove from RAM once the send succeeds.** +/// This ordering is fail-safe: if the channel is full or disconnected the send +/// fails, the key stays in RAM, and the eviction loop bails out to retry on the +/// next tick — no acknowledged data is lost. The worst case under sustained +/// backpressure is keys remaining resident (eventually an OOM write-rejection +/// once at budget), which is correct behaviour, not data loss. +/// +/// Callers must poll `SpillThread::drain_completions()` to apply the manifest +/// and `ColdIndex` updates produced by completed spills. pub fn try_evict_if_needed_async_spill( db: &mut Database, config: &RuntimeConfig, @@ -357,76 +362,6 @@ pub fn try_evict_if_needed_async_spill_with_total( Ok(()) } -/// Evict entries to bring memory under maxmemory, returning removed -/// (key, Entry) pairs for deferred spill OUTSIDE the write lock. -/// -/// Inside the lock: only find_victim + db.remove (~600ns per eviction). -/// The caller extracts value bytes from the owned Entry after releasing -/// the lock, then sends SpillRequests to the background thread. -pub fn try_evict_deferred( - db: &mut Database, - config: &RuntimeConfig, -) -> Result, Frame> { - if config.maxmemory == 0 { - return Ok(smallvec::SmallVec::new()); - } - - // Per-shard budget (see `try_evict_if_needed_with_spill_and_total`). - let budget = config.maxmemory_per_shard(); - let total_memory = db.estimated_memory(); - if total_memory <= budget { - return Ok(smallvec::SmallVec::new()); - } - - let policy = EvictionPolicy::from_str(&config.maxmemory_policy); - let mut evicted = smallvec::SmallVec::new(); - let mut current_total = total_memory; - - while current_total > budget { - if policy == EvictionPolicy::NoEviction { - return Err(oom_error()); - } - - let victim = find_victim_for_policy(db, config, &policy); - let key = match victim { - Some(k) => k, - None => return Err(oom_error()), - }; - - let before = db.estimated_memory(); - let key_bytes = Bytes::copy_from_slice(key.as_bytes()); - if let Some(entry) = db.remove(key.as_bytes()) { - evicted.push((key_bytes, entry)); - } - let after = db.estimated_memory(); - current_total = current_total.saturating_sub(before.saturating_sub(after)); - } - - Ok(evicted) -} - -/// Find a victim key using the given eviction policy. -fn find_victim_for_policy( - db: &Database, - config: &RuntimeConfig, - policy: &EvictionPolicy, -) -> Option { - match policy { - EvictionPolicy::NoEviction => None, - EvictionPolicy::AllKeysLru => find_victim_lru(db, config.maxmemory_samples, false), - EvictionPolicy::AllKeysLfu => { - find_victim_lfu(db, config.maxmemory_samples, config.lfu_decay_time, false) - } - EvictionPolicy::AllKeysRandom => find_victim_random(db, false), - EvictionPolicy::VolatileLru => find_victim_lru(db, config.maxmemory_samples, true), - EvictionPolicy::VolatileLfu => { - find_victim_lfu(db, config.maxmemory_samples, config.lfu_decay_time, true) - } - EvictionPolicy::VolatileRandom => find_victim_random(db, true), - EvictionPolicy::VolatileTtl => find_victim_volatile_ttl(db, config.maxmemory_samples), - } -} - /// Evict a single key via the async spill path. /// /// Extracts the entry, removes it from DashTable (immediate RAM relief), From 798ce096eadd2d183dd92ac35353c9aa1d67ee64 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 14:41:14 +0700 Subject: [PATCH 10/15] fix(shard): gate migrated-connection spawn fns behind cfg(unix) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit's PR #144 review (outside-diff 🟠 Major, conn_accept.rs:277-320) flagged that spawn_migrated_tokio_connection is gated only by #[cfg(feature = "runtime-tokio")] yet its signature takes std::os::unix::io::RawFd and reconstructs the socket via from_raw_fd — both Unix-only — so a runtime-tokio build on a non-Unix target would not compile. Make the platform coupling explicit and consistent: - conn_accept.rs: spawn_migrated_tokio_connection -> cfg(all(runtime-tokio, unix)) and its twin spawn_migrated_monoio_connection -> cfg(all(runtime-monoio, unix)). - event_loop.rs: the three migration-drain call sites gated to match (the two runtime-agnostic blocks' tokio/monoio wrappers, plus the monoio-only third site via #[cfg(unix)]). On every supported target (Linux + macOS, both Unix) `unix` is always true, so this is a no-op for behaviour and codegen — it only documents intent and stops a hypothetical non-Unix build from failing inside these functions. Verified: clippy -D warnings clean and cargo check --all-targets green on both default (monoio) and tokio,jemalloc; fmt clean. Scope note: this resolves the flagged surface but does NOT by itself make moon compile on a non-Unix target. The migration subsystem is Unix-coupled at the data-type level — MigrateConnectionPayload.fd (src/shard/dispatch.rs:295) is a RawFd, and the ShardMessage::MigrateConnection variant, pending_migrations Vec, and spsc handler all carry it ungated. Full non-Unix support would require cfg(unix)-gating that entire subsystem across dispatch.rs / spsc_handler.rs / handler_sharded — deliberately left out: moon targets only Linux + macOS, there is no Windows CI to verify against, and it touches hot-path dispatch for zero supported-platform benefit. author: Tin Dang --- src/shard/conn_accept.rs | 13 +++++++++++-- src/shard/event_loop.rs | 12 ++++++++---- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/src/shard/conn_accept.rs b/src/shard/conn_accept.rs index 351ad40f..173b5ca2 100644 --- a/src/shard/conn_accept.rs +++ b/src/shard/conn_accept.rs @@ -272,7 +272,13 @@ pub(crate) fn spawn_tokio_connection( /// /// TLS connections cannot be migrated because TLS session state lives in userspace and /// cannot be reconstructed from a raw FD. Only plain TCP connections should be migrated. -#[cfg(feature = "runtime-tokio")] +// `unix`-gated alongside `runtime-tokio`: the signature takes a +// `std::os::unix::io::RawFd` and reconstructs the socket via `from_raw_fd`, +// both of which only exist on Unix. Moon targets Linux + macOS (both Unix), so +// on every supported build `unix` is always true; the extra predicate just makes +// the platform coupling explicit (CodeRabbit PR #144). Full non-Unix support +// would also require gating `MigrateConnectionPayload.fd` — out of scope. +#[cfg(all(feature = "runtime-tokio", unix))] #[allow(clippy::too_many_arguments)] pub(crate) fn spawn_migrated_tokio_connection( fd: std::os::unix::io::RawFd, @@ -749,7 +755,10 @@ pub(crate) fn spawn_monoio_connection( /// # Limitations /// /// TLS connections cannot be migrated (TLS session state is in userspace). -#[cfg(feature = "runtime-monoio")] +// `unix`-gated alongside `runtime-monoio` for the same reason as the tokio twin +// above: the `RawFd` signature + `from_raw_fd` are Unix-only. Always true on +// supported targets (Linux + macOS). +#[cfg(all(feature = "runtime-monoio", unix))] #[allow(clippy::too_many_arguments)] pub(crate) fn spawn_migrated_monoio_connection( fd: std::os::unix::io::RawFd, diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 34d9f239..b393b544 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -1233,7 +1233,7 @@ impl super::Shard { "Shard {}: accepting migrated connection (fd={}, client_id={}, from={})", shard_id, fd, state.client_id, state.peer_addr ); - #[cfg(feature = "runtime-tokio")] + #[cfg(all(feature = "runtime-tokio", unix))] { conn_accept::spawn_migrated_tokio_connection( fd, state, @@ -1247,7 +1247,7 @@ impl super::Shard { &spill_sender, &spill_file_id, &disk_offload_dir, ); } - #[cfg(feature = "runtime-monoio")] + #[cfg(all(feature = "runtime-monoio", unix))] { conn_accept::spawn_migrated_monoio_connection( fd, state, @@ -1333,7 +1333,7 @@ impl super::Shard { "Shard {}: accepting migrated connection (fd={}, client_id={}, from={})", shard_id, fd, state.client_id, state.peer_addr ); - #[cfg(feature = "runtime-tokio")] + #[cfg(all(feature = "runtime-tokio", unix))] { conn_accept::spawn_migrated_tokio_connection( fd, state, @@ -1347,7 +1347,7 @@ impl super::Shard { &spill_sender, &spill_file_id, &disk_offload_dir, ); } - #[cfg(feature = "runtime-monoio")] + #[cfg(all(feature = "runtime-monoio", unix))] { conn_accept::spawn_migrated_monoio_connection( fd, state, @@ -1977,6 +1977,10 @@ impl super::Shard { state.client_id, state.peer_addr ); + // `unix`-gated to match `spawn_migrated_monoio_connection`'s + // Unix-only `RawFd` signature (CodeRabbit PR #144). Always + // compiled on supported targets (Linux + macOS). + #[cfg(unix)] conn_accept::spawn_migrated_monoio_connection( fd, state, From 1a10aab9808cc0faf84a3ed87b9cda59c0b22d79 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 14:44:32 +0700 Subject: [PATCH 11/15] test(storage): lock fail-safe send-before-remove async-spill eviction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit's PR #144 review raised a 🟠 Major claiming the async-spill eviction path removes a key from the hot tier BEFORE enqueuing its SpillRequest, so a full/disconnected spill channel would drop the request and lose data. The live code (evict_one_async_spill) is already fail-safe — it try_sends BEFORE db.remove, and a failed send leaves the key resident for retry — but that invariant had no direct regression guard. This adds one. async_spill_full_channel_keeps_victim_in_ram_no_data_loss: - pre-fills a bounded(1) spill channel so the eviction's try_send observes a FULL channel (the reviewer's exact backpressure scenario); - asserts try_evict_if_needed_async_spill surfaces OOM (no false success) AND the victim key is still resident (db.len()==1, key present) — no data loss. Red/green verified: GREEN on the live send-before-remove code; flipping the live ordering to remove-before-send locally turns it RED on the data-loss assertion (db.len() left:0, right:1), proving the test catches the regression. Live code restored byte-identical (test-only addition). No behaviour change. clippy -D warnings clean and 16 storage::eviction unit tests green on both default (monoio) and tokio,jemalloc; fmt clean. author: Tin Dang --- src/storage/eviction.rs | 56 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/src/storage/eviction.rs b/src/storage/eviction.rs index 6e05ca5f..54dbeee2 100644 --- a/src/storage/eviction.rs +++ b/src/storage/eviction.rs @@ -983,6 +983,62 @@ mod tests { assert_eq!(next_file_id, 2); } + /// Fail-safe send-before-remove: when the async-spill request channel is + /// FULL, eviction must NOT remove the victim from the hot tier (no data + /// loss) and must surface OOM rather than pretend success. The live path + /// (`evict_one_async_spill`) `try_send`s the `SpillRequest` BEFORE + /// `db.remove`, so a failed send leaves the key resident for retry on the + /// next eviction tick. This regression guard locks that ordering — the exact + /// request-side data-loss scenario CodeRabbit raised on PR #144 (verified + /// stale against the live code; this test keeps it that way). + #[test] + fn async_spill_full_channel_keeps_victim_in_ram_no_data_loss() { + let tmp = tempfile::tempdir().unwrap(); + let shard_dir = tmp.path(); + let mut next_file_id = 1u64; + + let mut db = Database::new(); + db.set_string(Bytes::from_static(b"k1"), Bytes::from_static(b"v1")); + + // Per-shard budget of 1 byte forces an eviction attempt for k1. + let config = make_config(1, "allkeys-lru"); + + // bounded(1) PRE-FILLED so the eviction's `try_send` observes a FULL + // channel (the reviewer's backpressure scenario) rather than a + // disconnect — both fail `try_send`, but full is the precise concern. + let (tx, _rx) = flume::bounded::(1); + tx.try_send(SpillRequest { + key: Bytes::from_static(b"filler"), + db_index: 0, + value_bytes: Bytes::from_static(b"x"), + value_type: ValueType::String, + flags: 0, + ttl_ms: None, + file_id: 0, + shard_dir: shard_dir.to_path_buf(), + }) + .expect("filler occupies the single channel slot"); + + let result = + try_evict_if_needed_async_spill(&mut db, &config, &tx, shard_dir, &mut next_file_id, 0); + + // Could not enqueue the spill → OOM surfaced, never a false success. + assert!( + result.is_err(), + "a full spill channel must surface OOM, never silently succeed" + ); + // CRUCIAL: the victim is still resident — no acknowledged-write data loss. + assert_eq!( + db.len(), + 1, + "victim must NOT be removed from RAM when the spill send fails" + ); + assert!( + db.data().get(&b"k1"[..]).is_some(), + "k1 must remain in the hot tier for retry on the next eviction tick" + ); + } + #[test] fn test_evict_without_spill_unchanged() { let mut db = Database::new(); From 5f5ecb35c3a09ea367b044788e132440d51d4a44 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 15:09:47 +0700 Subject: [PATCH 12/15] refactor(persistence): split aof.rs into submodules under 1500-line cap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit aof.rs had grown to 4379 lines — ~3x the project's 1500-line ceiling. Decompose it into a directory module following the split convention (mod.rs + subfiles, re-export via `pub use`, children pull parent privates via `use super::*`). Pure code relocation: no logic, control flow, or signatures changed. Layout (all under the cap): - mod.rs (1075): module doc + imports + core types (AofAck, AckOutcome, FsyncPolicy, AofMessage, PerShardRewriteCoord, ShardDoneGuard, AofPoolSendError, consts) + the codec (serialize_command, replay_aof) + codec/fsync tests. The codec stays in the parent so children reach it via `use super::*` with no sibling import (CLAUDE.md: mod.rs "holds shared helpers + tests"). - pool.rs (1379): AofWriterPool + impl + pool_tests. - writer_task.rs (1014): aof_writer_task, per_shard_aof_writer_task, TEST_FAIL_WRITE_AT. - rewrite.rs (941): generate_rewrite_commands .. rewrite_aof, drain/ do_rewrite paths + rewrite tests. Cross-module wiring: - Re-export only the moved public items from mod.rs (AofWriterPool, generate_rewrite_commands, rewrite_aof, aof_writer_task, per_shard_aof_writer_task) so external `persistence::aof::X` paths are unchanged. serialize_command/replay_aof keep their paths (still in mod.rs). - Widen to pub(crate) the rewrite items called across the new boundary by pool_tests / writer_task: do_rewrite_per_shard/single/sharded, drain_pending_appends{,_framed}, rewrite_aof_sharded_sync, and DrainOutcome (struct + drained/pending_acks fields + fulfill_acks — pool_tests reads them). - cfg-gate the imports and re-exports that are runtime-specific: rewrite_aof and TEST_FAIL_WRITE_AT (tokio-only), do_rewrite_single/sharded and drain_pending_appends (monoio-only). Matches each item's existing cfg so both runtimes compile. Verification (both runtimes, code-move so behaviour is unchanged): - cargo check --all-targets: clean on default (monoio) and tokio,jemalloc. - cargo clippy -- -D warnings: clean on both (CI-exact, no new lints in aof/). - cargo fmt --check: clean. - Unit tests: persistence::aof 66 pass (monoio) / 69 pass (tokio); 40 #[test] fns preserved exactly (no test dropped or duplicated in redistribution). - Integration (OrbStack moon-dev VM): crash_matrix_per_shard_bgrewriteaof 4 pass (do_rewrite_per_shard + replay_aof end-to-end), aof_fsync_err_subscribe_ordering 3 pass. No failures. author: Tin Dang --- src/persistence/aof.rs | 4379 ---------------------------- src/persistence/aof/mod.rs | 1078 +++++++ src/persistence/aof/pool.rs | 1382 +++++++++ src/persistence/aof/rewrite.rs | 941 ++++++ src/persistence/aof/writer_task.rs | 1017 +++++++ 5 files changed, 4418 insertions(+), 4379 deletions(-) delete mode 100644 src/persistence/aof.rs create mode 100644 src/persistence/aof/mod.rs create mode 100644 src/persistence/aof/pool.rs create mode 100644 src/persistence/aof/rewrite.rs create mode 100644 src/persistence/aof/writer_task.rs diff --git a/src/persistence/aof.rs b/src/persistence/aof.rs deleted file mode 100644 index 5ed1b67c..00000000 --- a/src/persistence/aof.rs +++ /dev/null @@ -1,4379 +0,0 @@ -//! Append-Only File (AOF) persistence: logs every write command in RESP format -//! for crash recovery. Supports three fsync policies and AOF rewriting for compaction. -//! -//! ## Unwrap Classification -//! -//! | Context | Classification | Rationale | -//! |---------|---------------|-----------| -//! | `AofWriter::append` (hot path) | **fire-and-forget** | Channel send; no Result needed | -//! | `aof_writer_task` | **must-panic** | Writer task; errors logged inline | -//! | `replay_aof` | **should-recover** (`Result<_, MoonError>`) | Startup replay; log+skip on corruption | -//! | `rewrite_aof` | **should-recover** (`Result<_, MoonError>`) | Background rewrite; caller logs error | -//! | `#[cfg(test)]` code (55 unwraps) | **test-only** | Panics are appropriate in tests | -// Suppressions narrowed: only keep what's needed for conditional compilation -#![allow(unused_imports, unused_variables, unreachable_code, clippy::empty_loop)] - -use std::path::{Path, PathBuf}; -use std::sync::Arc; -use std::time::{Duration, Instant}; - -use crate::runtime::cancel::CancellationToken; -use crate::runtime::channel; -use bytes::{Bytes, BytesMut}; -use tracing::{error, info, warn}; - -use crate::error::{AofError, MoonError}; -use crate::framevec; -use crate::persistence::replay::CommandReplayEngine; -use crate::protocol::{Frame, ParseConfig, parse, serialize}; -use crate::storage::compact_key::CompactKey; -use crate::storage::compact_value::RedisValueRef; -use crate::storage::db::Database; -use crate::storage::entry::{Entry, current_time_ms}; -/// Type alias for the per-database RwLock container. -type SharedDatabases = Arc>>; - -/// Canonical AOF fsync failure error string sent to the client as a -/// `Frame::Error` when `appendfsync=always` and the writer task does not -/// confirm durability before the response. -/// -/// All handler variants (handler_single, handler_monoio, handler_sharded) -/// MUST use this constant so operators see a consistent error regardless of -/// which connection path handles the request. -/// -/// Redis convention: errors begin with a single-word code (`ERR` for generic -/// failures) followed by a space and a human-readable message. -pub const AOF_FSYNC_ERR: &[u8] = b"ERR AOF fsync failed; write not durable"; - -/// High bit of the per-entry LSN reserved for `OrderedAcrossShards` -/// (RFC § 2 Rule 2). When set on a per-shard AOF entry, recovery treats -/// the entry as participating in a cross-shard atomic operation and -/// buffers it for the cross-shard merge replay after per-shard replay -/// completes. -/// -/// Practical LSN ceilings (even at 10 M writes/s sustained for a century) -/// sit near 2^58, so reserving bit 63 has no observable effect on normal -/// writes — the bit is always 0 in entries written by `try_send_append`. -/// Only `try_send_append_ordered` sets it. -pub const ORDERED_LSN_FLAG: u64 = 1u64 << 63; - -/// Outcome reported by the writer task back to an `AppendSync` caller -/// once the rendezvous completes. -/// -/// `Synced` is sent AFTER `sync_data()` returns successfully — the -/// caller may safely `+OK` the client. `WriteFailed`/`FsyncFailed` -/// surface the failure mode so the caller can return a specific error -/// frame; either way, durability was NOT achieved. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AofAck { - /// Bytes were written and fsynced. Durability guaranteed. - Synced, - /// `write_all()` returned an error. The entry may be partially on - /// disk; recovery handles partial-payload truncation as crash EOF. - WriteFailed, - /// `write_all()` succeeded but `sync_data()` returned an error. The - /// entry is in the kernel buffer but NOT on durable storage. - FsyncFailed, - /// The writer channel was full at the time of the send — the entry - /// was **not** enqueued. This is a backpressure signal: the writer - /// is unable to keep up with the current write rate. Callers MUST - /// treat this as a hard failure (same as `WriteFailed`) under - /// `appendfsync=always`; for `everysec`/`no` it is logged and counted. - ChannelFull, -} - -/// Global counter incremented each time an AOF `AppendSync` (or fire-and- -/// forget `Append`) is dropped because the writer channel was at capacity. -/// -/// Exposed under `# Persistence` in the INFO command as -/// `aof_backpressure_dropped`. A persistently non-zero value indicates the -/// writer is a bottleneck and the operator should investigate disk I/O or -/// switch to `appendfsync=everysec`. -pub static AOF_BACKPRESSURE_DROPPED: std::sync::atomic::AtomicU64 = - std::sync::atomic::AtomicU64::new(0); - -/// Result of awaiting an `AppendSync` ack under a bounded timeout (F2). -/// -/// Distinguishes the three terminal states the `Always` durability path -/// can reach so the caller can map each to the correct client-facing -/// outcome: -/// - `Ack(_)` — the writer reported back; inspect the `AofAck`. -/// - `Disconnected` — the writer task is gone / channel dropped (no ack -/// will ever arrive). Treated as `WriteFailed`. -/// - `TimedOut` — the fsync did not confirm within the configured -/// bound. Durability is unconfirmed; treated as `FsyncFailed`. The -/// entry may still reach disk later, so the caller must NOT report -/// success but also must not assume the write was rejected. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum AckOutcome { - Ack(AofAck), - Disconnected, - TimedOut, -} - -/// AOF fsync policy controlling when data is flushed to disk. -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum FsyncPolicy { - /// Fsync after every write command (safest, slowest). - Always, - /// Fsync once per second in the background (good balance). - EverySec, - /// Let the OS decide when to flush (fastest, least safe). - No, -} - -impl FsyncPolicy { - /// Parse a policy string (as from config). Defaults to EverySec for unknown values. - pub fn from_str(s: &str) -> Self { - match s { - "always" => FsyncPolicy::Always, - "no" => FsyncPolicy::No, - _ => FsyncPolicy::EverySec, - } - } -} - -/// Messages sent to the AOF writer task via mpsc channel. -pub enum AofMessage { - /// Append serialized RESP command bytes to the AOF file, tagged with the - /// LSN that was issued for this write (`ReplicationState::issue_lsn`). - /// - /// `lsn` semantics by writer task: - /// - **TopLevel** (`aof_writer_task`): `lsn` is **ignored**; the legacy - /// v1 disk format is plain RESP bytes with no per-entry framing. - /// - **PerShard** (`per_shard_aof_writer_task`): `lsn` is **written** as - /// a u64 header per RFC § 2 Rule 1. Disk format per entry: - /// `[u64 lsn LE][u32 len LE][RESP bytes of length len]`. - /// Recovery reads `(lsn, cmd)` pairs and merges cross-shard - /// `OrderedAcrossShards` writes by LSN (RFC § 2 Rule 2). - /// - /// Construction sites that issue a real LSN call - /// `ReplicationState::issue_lsn(shard_id, bytes.len() as u64)` and pass - /// the returned value. Sites with no replication state available pass 0 - /// (TopLevel ignores it; PerShard treats 0 as "no ordering hint"). - Append { lsn: u64, bytes: Bytes }, - /// Append + fsync + ack rendezvous (RFC § 4 — Fix 2 for the H1 - /// data-loss vector exposed by `appendfsync=always`). - /// - /// Same encoding as [`AofMessage::Append`], but the writer task ALWAYS - /// fsyncs after writing the payload and signals `ack` ONCE the - /// `sync_data()` syscall returns. The caller is expected to await - /// `ack` before responding `+OK` to the client so the durability - /// contract of `appendfsync=always` is honoured end-to-end. - /// - /// Failure semantics: on write or fsync error the writer drops `ack` - /// without sending — the caller's `OneshotReceiver` resolves with - /// `RecvError`, which it must treat as a hard failure (return an - /// error frame to the client, do NOT silently +OK). - /// - /// Production callers: none in step 7 — this commit ships the - /// mechanism plus tests. Per-handler integration (which sites use - /// AppendSync vs Append) is wired in step 9 before lifting the - /// `--unsafe-multishard-aof` gate. - AppendSync { - lsn: u64, - bytes: Bytes, - ack: crate::runtime::channel::OneshotSender, - }, - /// Trigger a full AOF rewrite (compaction) using current database state. - Rewrite(SharedDatabases), - /// Trigger AOF rewrite in sharded mode (all shards' databases). - RewriteSharded(Arc), - /// [F6] Trigger a per-shard AOF rewrite (compaction) in the PerShard - /// layout. Sent to EVERY per-shard writer at once. Each writer folds its - /// own shard (drain → lock → snapshot → write new base+incr at the - /// coordinator's `new_seq` → reopen), then decrements the shared - /// `PerShardRewriteCoord`; the last writer commits the manifest once - /// (single seq flip) and prunes the old generation. The synchronized seq - /// + single commit are what make multi-shard BGREWRITEAOF crash-safe. - RewritePerShard { - shard_dbs: Arc, - coord: Arc, - }, - /// Shut down the AOF writer task gracefully. - Shutdown, -} - -/// Coordinator shared by all per-shard writers participating in one -/// BGREWRITEAOF fan-out (F6). -/// -/// Crash-safety contract (mirrors `AofManifest::advance` ordering, but -/// distributed across N writer threads): -/// -/// 1. Each writer writes its new base+incr at `new_seq` via -/// `manifest.advance_shard(shard_id, new_seq, rdb)` — which does NOT bump -/// `manifest.seq` or rewrite the manifest. So until the final commit, the -/// on-disk manifest still resolves to `old_seq`; a crash here recovers the -/// intact old generation (no loss, no double-apply). -/// 2. The LAST writer to finish (countdown reaches zero) performs the single -/// durable commit: `manifest.seq = new_seq; write_manifest()`. This is the -/// atomic point at which recovery flips to the new generation. -/// 3. Only AFTER the commit are the old-generation files pruned. -/// -/// The manifest is shared via `Arc>` and locked ONLY for the brief, -/// await-free `advance_shard` and final-commit critical sections — never held -/// across a blocking disk write of the base RDB (that happens before the lock) -/// nor across `.await`. -pub struct PerShardRewriteCoord { - /// Writers still to finish. Starts at the shard count; the writer that - /// decrements it to zero performs the commit + prune. - remaining: std::sync::atomic::AtomicUsize, - /// Shared manifest, loaded fresh from disk by the BGREWRITEAOF handler at - /// rewrite time (normal appends never touch the manifest, and BGREWRITEAOF - /// is CAS-serialized, so a fresh load is the authoritative current state). - manifest: Arc>, - /// The generation every writer advances to. Computed once = old_seq + 1. - new_seq: u64, - /// The generation being retired; pruned only after the commit. - old_seq: u64, - /// Number of shards participating (= initial `remaining`). - n_shards: usize, - /// Set by any shard whose fold fails. The final writer checks this and - /// ABORTS the commit if set — committing `new_seq` while a shard never - /// wrote its new base would make recovery look for a missing base and - /// refuse to start. On abort the old generation (`old_seq`) stays the - /// committed state for all shards (crash-safe). - failed: std::sync::atomic::AtomicBool, - /// The committed generation, published exactly once by the terminal writer - /// (`new_seq` on success, `old_seq` on abort/commit-failure). Every folded - /// writer blocks on this in `await_outcome` after `shard_done` so it can - /// reopen its append file onto the COMMITTED generation before resuming - /// appends — without this barrier a writer that reopened onto `new_seq` in - /// phase 6 keeps appending there after an abort, into a generation recovery - /// ignores (silent data loss; the old "RESTART recommended" hazard). - outcome: parking_lot::Mutex>, - outcome_cv: parking_lot::Condvar, -} - -impl PerShardRewriteCoord { - /// Construct a coordinator for an `n_shards`-way rewrite advancing the - /// shared `manifest` from its current seq to `current_seq + 1`. - pub fn new( - manifest: Arc>, - current_seq: u64, - n_shards: usize, - ) -> Arc { - Arc::new(Self { - remaining: std::sync::atomic::AtomicUsize::new(n_shards), - manifest, - new_seq: current_seq + 1, - old_seq: current_seq, - n_shards, - failed: std::sync::atomic::AtomicBool::new(false), - outcome: parking_lot::Mutex::new(None), - outcome_cv: parking_lot::Condvar::new(), - }) - } - - /// Publish the committed generation to every writer blocked in - /// `await_outcome`. Called exactly once, by the terminal `shard_done`. - #[inline] - fn publish_outcome(&self, committed_seq: u64) { - let mut slot = self.outcome.lock(); - *slot = Some(committed_seq); - self.outcome_cv.notify_all(); - } - - /// Block until the terminal writer publishes the committed generation, then - /// return it (`new_seq` on success, `old_seq` on abort/commit-failure). Each - /// folded writer calls this AFTER `shard_done` so it can reopen its append - /// file onto the committed generation. Safe against missed wake-ups: the - /// outcome is set under the same lock the condvar waits on, so a writer that - /// arrives after the publish observes `Some` and skips the wait. - /// - /// Deadlock-free: all `n_shards` writers call `shard_done` exactly once - /// (success via phase 7, fold-error/send-failure via the caller), so the - /// countdown always reaches zero and the terminal branch always publishes. - /// Each per-shard writer runs on its own dedicated OS thread, so blocking - /// one here never starves the thread that must publish. - #[must_use] - fn await_outcome(&self) -> u64 { - let mut slot = self.outcome.lock(); - while slot.is_none() { - self.outcome_cv.wait(&mut slot); - } - slot.expect("outcome is Some after the wait loop") - } - - /// The generation writers advance to. - #[inline] - pub fn new_seq(&self) -> u64 { - self.new_seq - } - - /// Mark the whole rewrite as failed (called by a shard whose fold errored). - /// The final writer will abort the commit, leaving `old_seq` authoritative. - #[inline] - pub fn mark_failed(&self) { - self.failed - .store(true, std::sync::atomic::Ordering::Release); - } - - /// Called by each writer AFTER it has durably written its new base+incr at - /// `new_seq` and reopened its append file. Decrements the countdown; the - /// final caller commits the manifest (single seq flip) and prunes the old - /// generation, then clears the global in-progress flag. - /// - /// Crash-safety: the commit (`write_manifest`) is the atomic flip point; - /// pruning runs strictly after it, so a crash mid-prune only orphans - /// already-superseded files (recovery uses `new_seq`). - pub fn shard_done(&self) { - use std::sync::atomic::Ordering; - // AcqRel: the decrement-to-zero must observe all prior writers' - // advance_shard manifest mutations before committing. - if self.remaining.fetch_sub(1, Ordering::AcqRel) == 1 { - // Abort if any shard failed to fold: committing new_seq while a - // shard lacks its new base would break recovery. Keep old_seq. - if self.failed.load(Ordering::Acquire) { - let m = self.manifest.lock(); - // Best-effort: prune the orphaned new-seq files written by the - // shards that DID fold, so they don't linger. - for sid in 0..self.n_shards { - m.prune_shard_files(sid as u16, self.new_seq); - } - drop(m); - error!( - "F6 per-shard rewrite ABORTED: a shard failed to fold; seq stays {}. \ - Old generation remains authoritative (crash-safe). Folded writers \ - roll their append files back to the old generation at the \ - barrier — no restart required.", - self.old_seq - ); - // Publish BEFORE clearing the in-progress flag so every blocked - // writer wakes and reopens onto the (still-authoritative) old gen. - self.publish_outcome(self.old_seq); - crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); - return; - } - let mut m = self.manifest.lock(); - m.seq = self.new_seq; - if let Err(e) = m.write_manifest() { - error!( - "F6 per-shard rewrite: final manifest commit (seq {}) failed: {}. \ - Old generation remains authoritative; rewrite did not take effect.", - self.new_seq, e - ); - // Do NOT prune — old generation is still the committed state. - drop(m); - // Commit failed: old_seq is still authoritative, so folded - // writers must roll back to it (their phase-6 reopen targeted - // new_seq, which never committed). - self.publish_outcome(self.old_seq); - crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); - return; - } - for sid in 0..self.n_shards { - m.prune_shard_files(sid as u16, self.old_seq); - } - drop(m); - info!( - "F6 per-shard rewrite complete: committed seq {} across {} shards, pruned seq {}", - self.new_seq, self.n_shards, self.old_seq - ); - // Success: new_seq is committed. Folded writers already reopened onto - // new_seq in phase 6, so the barrier is a no-op for them — but it must - // still unblock them. - self.publish_outcome(self.new_seq); - crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); - } - } -} - -/// Panic-safety guard for a per-shard fold's `shard_done` obligation. -/// -/// The `await_outcome` barrier in `do_rewrite_per_shard` turns "every shard -/// calls `shard_done` exactly once" from a tidiness rule into a LIVENESS -/// requirement: any folded writer that decrements then blocks in `await_outcome` -/// hangs forever if some other shard never decrements. The one path that would -/// skip `shard_done` is a panic mid-fold (e.g. `save_snapshot_to_bytes` -/// OOM-unwinding — issue #138), which under `panic = "unwind"` would otherwise -/// hang every healthy shard's AOF writer thread. -/// -/// This guard makes the decrement unconditional. The fold creates it on entry -/// and calls [`complete`](Self::complete) on the success path (clean -/// `shard_done`). On ANY other exit — `?` error or panic unwind — `Drop` fires -/// `mark_failed` + `shard_done`, so the countdown still reaches zero, the -/// terminal writer publishes `old_seq`, and every barrier waiter wakes and rolls -/// back. Because the guard owns the decrement for ALL exits, callers MUST NOT -/// call `shard_done` again after invoking `do_rewrite_per_shard`. -struct ShardDoneGuard<'a> { - coord: &'a PerShardRewriteCoord, - done: bool, -} - -impl<'a> ShardDoneGuard<'a> { - #[inline] - fn new(coord: &'a PerShardRewriteCoord) -> Self { - Self { coord, done: false } - } - - /// Success-path completion: a single clean `shard_done` (no `mark_failed`). - /// Consumes the guard so the subsequent `Drop` is a no-op. - #[inline] - fn complete(mut self) { - self.done = true; - self.coord.shard_done(); - } -} - -impl Drop for ShardDoneGuard<'_> { - fn drop(&mut self) { - if !self.done { - // Unwind or early-error before completion: abort the rewrite (keep - // old_seq) and still decrement so the barrier releases every waiter. - self.coord.mark_failed(); - self.coord.shard_done(); - } - } -} - -/// Reasons a pool send may be refused without queueing. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum AofPoolSendError { - /// `Rewrite`/`RewriteSharded` sent to a `PerShard` pool. BGREWRITEAOF must - /// be issued per shard in the per-shard layout; the legacy single-writer - /// rewrite path is not applicable. - RewriteUnsupportedInPerShard, - /// Underlying channel send failed (writer task dead or channel full). - SendFailed, -} - -/// Bundle of per-shard AOF writer senders. -/// -/// The pool keeps the call-site API uniform regardless of layout: -/// - **TopLevel** (legacy v1, single-shard, also used for `--shards 1` v2): -/// exactly one writer thread; every `sender(shard_id)` returns the same -/// sender so all shards multiplex onto one file. -/// - **PerShard** (v2 multi-shard): one writer per shard; `sender(shard_id)` -/// returns the writer that owns `appendonlydir/shard-{shard_id}/`. -/// -/// Step 2a is additive — this type is defined here but no call site is wired -/// to it yet. Step 2c performs the type plumbing in `conn_state` and -/// `conn/core`; steps 2d/2e/2f update the call sites and spawn paths. -/// Default bound for the `appendfsync=always` fsync-ack await. Mirrors the -/// `--aof-fsync-timeout-ms` config default; used by constructors that don't -/// take an explicit timeout (non-production / test helpers). -pub const DEFAULT_AOF_FSYNC_TIMEOUT: Duration = Duration::from_millis(2000); - -#[derive(Clone)] -pub struct AofWriterPool { - senders: Vec>, - layout: crate::persistence::aof_manifest::AofLayout, - /// Fsync policy configured at writer-task construction. Read on the - /// hot append path: `Always` routes through `AppendSync` for - /// fsync-before-ack durability (H1 fix); everything else stays on - /// the fire-and-forget `Append` path. - fsync_policy: FsyncPolicy, - /// F2: max time `try_send_append_durable` waits for the `Always` fsync - /// ack before failing the write. `Duration::ZERO` means unbounded - /// (legacy behavior). Prevents a stalled disk from parking write - /// connections forever (design-for-failure). - fsync_timeout: Duration, - /// F6: persistence base dir (the parent of `appendonlydir/`), set only for - /// PerShard pools that may service a per-shard BGREWRITEAOF. Needed to load - /// the authoritative manifest fresh at rewrite time. `None` for TopLevel - /// pools and test pools that never rewrite. - base_dir: Option, -} - -impl AofWriterPool { - /// Build a TopLevel pool from a single existing writer sender. Used for - /// legacy v1 deployments and `--shards 1` v2 deployments where one writer - /// thread services every shard. - pub fn top_level(sender: channel::MpscSender) -> Arc { - Self::top_level_with_policy(sender, FsyncPolicy::EverySec, DEFAULT_AOF_FSYNC_TIMEOUT) - } - - /// Same as [`Self::top_level`] but with an explicit fsync policy. The - /// policy controls whether [`Self::try_send_append_durable`] takes the - /// fast (fire-and-forget) or rendezvous (`AppendSync`) path. - /// `fsync_timeout` bounds the `Always` ack await (F2); `Duration::ZERO` - /// = unbounded. - pub fn top_level_with_policy( - sender: channel::MpscSender, - fsync_policy: FsyncPolicy, - fsync_timeout: Duration, - ) -> Arc { - Arc::new(Self { - senders: vec![sender], - layout: crate::persistence::aof_manifest::AofLayout::TopLevel, - fsync_policy, - fsync_timeout, - base_dir: None, - }) - } - - /// Build a PerShard pool from N senders. `senders[i]` MUST be the writer - /// task that owns `appendonlydir/shard-{i}/`. The vector's length is the - /// shard count; passing a length-1 vector here is a bug — use - /// [`AofWriterPool::top_level`] instead. - pub fn per_shard(senders: Vec>) -> Arc { - Self::per_shard_with_policy(senders, FsyncPolicy::EverySec, DEFAULT_AOF_FSYNC_TIMEOUT) - } - - /// Same as [`Self::per_shard`] but with an explicit fsync policy. - /// `fsync_timeout` bounds the `Always` ack await (F2); `Duration::ZERO` - /// = unbounded. - pub fn per_shard_with_policy( - senders: Vec>, - fsync_policy: FsyncPolicy, - fsync_timeout: Duration, - ) -> Arc { - debug_assert!( - senders.len() >= 2, - "per_shard pool needs >=2 writers; use top_level for single-writer" - ); - Arc::new(Self { - senders, - layout: crate::persistence::aof_manifest::AofLayout::PerShard, - fsync_policy, - fsync_timeout, - base_dir: None, - }) - } - - /// F6: same as [`Self::per_shard_with_policy`] but records the persistence - /// `base_dir` so a per-shard BGREWRITEAOF can load the authoritative - /// manifest fresh at rewrite time. This is the production constructor used - /// by `main.rs` for the PerShard layout. - pub fn per_shard_with_base_dir( - senders: Vec>, - fsync_policy: FsyncPolicy, - fsync_timeout: Duration, - base_dir: PathBuf, - ) -> Arc { - debug_assert!( - senders.len() >= 2, - "per_shard pool needs >=2 writers; use top_level for single-writer" - ); - Arc::new(Self { - senders, - layout: crate::persistence::aof_manifest::AofLayout::PerShard, - fsync_policy, - fsync_timeout, - base_dir: Some(base_dir), - }) - } - - /// Returns the configured fsync policy. Hot-path callers read this to - /// decide between the fast (`try_send_append`) and durable - /// (`try_send_append_sync`) write paths. - #[inline] - pub fn fsync_policy(&self) -> FsyncPolicy { - self.fsync_policy - } - - /// Policy-aware AOF append. For `FsyncPolicy::Always`, this awaits - /// `AppendSync` and returns `Ok(())` only after `sync_data()` confirms - /// the entry is on durable storage — closing the H1 in-flight loss - /// vector identified in the investigation report. For `EverySec` and - /// `No`, it stays on the fire-and-forget path (zero new latency). - /// - /// Returns `Err(AofAck)` only on the Always path when the write or - /// fsync failed (or the writer task is gone). Callers MUST treat - /// `Err(_)` as a hard failure — return an error frame to the client, - /// do NOT respond `+OK`. - /// - /// Async because the Always branch awaits a oneshot receiver. The - /// non-Always branch resolves immediately (no actual suspension) so - /// the only overhead is one `match` and the implicit Future state - /// machine; benchmarked at ~5 ns per call on the EverySec hot path, - /// far below the per-write WAL/replication cost. - #[inline] - pub async fn try_send_append_durable( - &self, - shard_id: usize, - lsn: u64, - bytes: Bytes, - ) -> Result<(), AofAck> { - match self.fsync_policy { - FsyncPolicy::Always => { - let rx = self.try_send_append_sync(shard_id, lsn, bytes); - // F2 (design-for-failure): bound the wait so a stalled disk - // can't park this connection forever. On elapse the write is - // failed — the entry may still land on disk later, but - // durability is NOT confirmed, so the caller must not report - // success. `Duration::ZERO` keeps the legacy unbounded await. - match Self::await_ack(rx, self.fsync_timeout).await { - AckOutcome::Ack(AofAck::Synced) => Ok(()), - AckOutcome::Ack(other) => Err(other), - // Writer task gone / channel disconnected. - AckOutcome::Disconnected => Err(AofAck::WriteFailed), - // Fsync did not confirm within the bound. - AckOutcome::TimedOut => Err(AofAck::FsyncFailed), - } - } - FsyncPolicy::EverySec | FsyncPolicy::No => { - self.try_send_append(shard_id, lsn, bytes); - Ok(()) - } - } - } - - /// Await an `AppendSync` ack receiver under a bounded timeout (F2). - /// - /// `timeout == Duration::ZERO` preserves the legacy unbounded await - /// (used when the operator explicitly opts out via - /// `--aof-fsync-timeout-ms 0`). Otherwise the await is capped by the - /// runtime-appropriate timer; on elapse the in-flight fsync is - /// abandoned (the receiver is dropped) and `TimedOut` is returned. - /// - /// Runtime-agnostic: monoio uses `select! { rx, sleep }` (matching the - /// established `cluster::failover` pattern — monoio 0.2 has no - /// `time::timeout`); tokio uses `tokio::time::timeout`. Both resolve the - /// ack first if it arrives within the bound, otherwise `TimedOut`. - async fn await_ack( - rx: crate::runtime::channel::OneshotReceiver, - timeout: Duration, - ) -> AckOutcome { - if timeout.is_zero() { - return match rx.await { - Ok(ack) => AckOutcome::Ack(ack), - Err(_) => AckOutcome::Disconnected, - }; - } - - #[cfg(feature = "runtime-monoio")] - { - monoio::select! { - res = rx => match res { - Ok(ack) => AckOutcome::Ack(ack), - Err(_) => AckOutcome::Disconnected, - }, - _ = monoio::time::sleep(timeout) => AckOutcome::TimedOut, - } - } - #[cfg(all(feature = "runtime-tokio", not(feature = "runtime-monoio")))] - { - match tokio::time::timeout(timeout, rx).await { - Ok(Ok(ack)) => AckOutcome::Ack(ack), - Ok(Err(_)) => AckOutcome::Disconnected, - Err(_) => AckOutcome::TimedOut, - } - } - } - - /// Return the writer sender that owns the given shard's AOF file. - /// - /// For TopLevel pools, `shard_id` is ignored — all shards multiplex onto - /// the single sender. For PerShard pools, `shard_id` MUST be in range - /// `[0, num_writers())`; an out-of-range id is a programmer error and - /// panics in debug builds. - #[inline] - pub fn sender(&self, shard_id: usize) -> &channel::MpscSender { - use crate::persistence::aof_manifest::AofLayout; - match self.layout { - AofLayout::TopLevel => &self.senders[0], - AofLayout::PerShard => { - debug_assert!( - shard_id < self.senders.len(), - "shard_id {} out of range for per-shard pool of size {}", - shard_id, - self.senders.len() - ); - &self.senders[shard_id] - } - } - } - - /// Fire-and-forget append for the given shard, tagged with the LSN that - /// was issued for this write (see [`AofMessage::Append`] docs for LSN - /// semantics per layout). Call sites must source `lsn` from - /// `ReplicationState::issue_lsn(shard_id, bytes.len() as u64)` for writes - /// that participate in replication ordering; sites without a - /// replication-state handle pass 0. - #[inline] - pub fn try_send_append(&self, shard_id: usize, lsn: u64, bytes: Bytes) { - let _ = self - .sender(shard_id) - .try_send(AofMessage::Append { lsn, bytes }); - } - - /// Synchronous (fsync-before-ack) append for `appendfsync=always` - /// durability (RFC § 4 — Fix 2). Returns a receiver the caller MUST - /// await before responding to the client; `AofAck::Synced` means the - /// entry is on durable storage. - /// - /// **Failure handling:** if the write or fsync fails, the receiver - /// resolves with `AofAck::WriteFailed` / `AofAck::FsyncFailed`. If - /// the writer task is gone (shutdown / channel disconnect), the - /// receiver resolves with `Err(RecvError)`. In every failure mode the - /// caller MUST return an error frame to the client, NOT `+OK`. - /// - /// **Performance:** every call adds a writer round-trip plus an - /// fsync syscall on the critical path. This is the explicit Redis - /// contract for `appendfsync=always`; callers should gate on the - /// configured policy and prefer [`Self::try_send_append`] for - /// `everysec`/`no`. - /// - /// **`shard_id` semantics:** matches [`Self::try_send_append`] — for - /// TopLevel the parameter is ignored, for PerShard it routes to - /// `senders[shard_id]`. - pub fn try_send_append_sync( - &self, - shard_id: usize, - lsn: u64, - bytes: Bytes, - ) -> crate::runtime::channel::OneshotReceiver { - let (ack_tx, ack_rx) = crate::runtime::channel::oneshot::(); - match self.sender(shard_id).try_send(AofMessage::AppendSync { - lsn, - bytes, - ack: ack_tx, - }) { - Ok(()) => {} - Err(flume::TrySendError::Full(_)) => { - // Writer channel is at capacity — count the dropped entry and - // signal ChannelFull back to the caller via a pre-filled - // oneshot so the caller's `.await` resolves immediately to - // Err(AofAck::ChannelFull) without a writer round-trip. - AOF_BACKPRESSURE_DROPPED.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - warn!( - "AOF writer channel full (shard {}): AppendSync dropped; \ - backpressure_dropped={}", - shard_id, - AOF_BACKPRESSURE_DROPPED.load(std::sync::atomic::Ordering::Relaxed), - ); - // Pre-send ChannelFull into a fresh oneshot pair; the - // caller's `ack_rx` was already returned — we create a - // new pair and use its sender to pre-fill what the caller - // will receive. The original ack_tx (inside the dropped - // AppendSync) is dropped, causing its ack_rx to yield - // RecvError. We send ChannelFull via the *returned* ack_rx - // by using a second oneshot whose sender is immediately - // fulfilled, then return that receiver instead. - let (pre_tx, pre_rx) = crate::runtime::channel::oneshot::(); - let _ = pre_tx.send(AofAck::ChannelFull); - return pre_rx; - } - Err(flume::TrySendError::Disconnected(_)) => { - // Writer task is dead — let caller handle RecvError on ack_rx. - // ack_tx was dropped inside the Err value; ack_rx will - // resolve with RecvError, which try_send_append_durable maps - // to Err(AofAck::WriteFailed). - } - } - ack_rx - } - - /// Fire-and-forget append for a cross-shard atomic operation (RFC § 2 - /// Rule 2 — `OrderedAcrossShards` tagging). - /// - /// The high bit of `lsn` (`1 << 63`) is set before the entry is queued. - /// Recovery uses this bit to recognize cross-shard atomic entries, - /// buffer them per-shard, and replay them globally in LSN order after - /// per-shard replay completes — guaranteeing TXN/SCRIPT atomicity - /// survives a crash even when multiple shards participated. - /// - /// **Caller contract:** `lsn` MUST be < `1 << 63` (i.e. the high bit - /// MUST be clear when passed in). Practical LSN ceilings — even at - /// 10 M writes/s sustained for a century — sit around 2^58, so any - /// real LSN satisfies this. Debug builds assert; release builds mask - /// the input to keep the wire format well-formed rather than - /// corrupt-by-zero-extending. - /// - /// **Production callers today:** none. Step 5 ships the infrastructure - /// (writer, framing flag, recovery merge) so a future cross-shard TXN - /// or replicated SCRIPT command has a place to land. Until that - /// consumer exists, only test code emits ordered entries. - #[inline] - pub fn try_send_append_ordered(&self, shard_id: usize, lsn: u64, bytes: Bytes) { - debug_assert_eq!( - lsn & ORDERED_LSN_FLAG, - 0, - "try_send_append_ordered: lsn must not have the high bit set; got {:#x}", - lsn, - ); - let tagged_lsn = (lsn & !ORDERED_LSN_FLAG) | ORDERED_LSN_FLAG; - let _ = self.sender(shard_id).try_send(AofMessage::Append { - lsn: tagged_lsn, - bytes, - }); - } - - /// Issue an LSN for an AOF append at every call site that has the - /// `Option>>` shape. Wraps - /// `ReplicationState::issue_lsn` so handler call sites collapse to a - /// single line. - /// - /// Returns 0 when: - /// - `repl_state` is None (test fixtures or shutdown paths) - /// - the `RwLock` is poisoned (shouldn't happen in production — - /// ReplicationState is only `write()`-locked under known-safe paths) - /// - /// 0 is a sentinel meaning "no replication ordering for this write". - /// TopLevel writers ignore the LSN entirely so 0 is harmless there; - /// PerShard writers treat 0 the same as any other LSN (per-shard order - /// is preserved by write order, not by LSN value). The LSN only matters - /// for the cross-shard `OrderedAcrossShards` merge in RFC step 5. - #[inline] - pub fn issue_append_lsn( - repl_state: &Option>>, - shard_id: usize, - delta: usize, - ) -> u64 { - repl_state - .as_ref() - .and_then(|rs| rs.read().ok().map(|g| g.issue_lsn(shard_id, delta as u64))) - .unwrap_or(0) - } - - /// Submit a Rewrite/RewriteSharded message. Only legal for TopLevel pools; - /// PerShard rewrites are per-shard operations and must be initiated by - /// the BGREWRITEAOF code path in step 6, not via this enum variant. - pub fn try_send_rewrite(&self, msg: AofMessage) -> Result<(), AofPoolSendError> { - use crate::persistence::aof_manifest::AofLayout; - debug_assert!( - matches!(msg, AofMessage::Rewrite(_) | AofMessage::RewriteSharded(_)), - "try_send_rewrite called with a non-Rewrite variant", - ); - if self.layout == AofLayout::PerShard { - return Err(AofPoolSendError::RewriteUnsupportedInPerShard); - } - self.senders[0] - .try_send(msg) - .map_err(|_| AofPoolSendError::SendFailed) - } - - /// [F6] Initiate a per-shard BGREWRITEAOF across every writer in a - /// PerShard pool. - /// - /// Loads the authoritative manifest fresh from `base_dir` (normal appends - /// never mutate the manifest, and BGREWRITEAOF is CAS-serialized by - /// `AOF_REWRITE_IN_PROGRESS`, so a fresh load is the current committed - /// state), builds a shared [`PerShardRewriteCoord`] that advances the - /// generation by one, and hands every writer the same `coord` + a cheap - /// `Arc` clone of `shard_dbs`. - /// - /// **Reliable delivery (design-for-failure):** the fan-out uses the - /// *blocking* `send` rather than `try_send`. A dropped rewrite message - /// would leave the countdown unable to reach zero — folded writers would - /// have reopened to new-seq files that the manifest never commits, silently - /// losing their post-rewrite appends. The writers run on dedicated threads - /// draining continuously, so `send` blocks only until a channel slot frees - /// (sub-millisecond), which is acceptable for a rare admin command. - /// - /// Returns `SendFailed` if `base_dir` is unset, the manifest can't be - /// loaded, or a writer thread is gone (disconnected channel). On the last - /// case the rewrite aborts WITHOUT committing — the old generation stays - /// authoritative (crash-safe), but a dead writer already means that shard's - /// persistence was compromised before this call. - pub fn try_send_rewrite_per_shard( - &self, - shard_dbs: Arc, - ) -> Result<(), AofPoolSendError> { - use crate::persistence::aof_manifest::{AofLayout, AofManifest}; - if self.layout != AofLayout::PerShard { - // A TopLevel pool rewrites via try_send_rewrite; this entry point - // is PerShard-only. - return Err(AofPoolSendError::RewriteUnsupportedInPerShard); - } - let base_dir = self.base_dir.as_ref().ok_or(AofPoolSendError::SendFailed)?; - let manifest = match AofManifest::load(base_dir) { - Ok(Some(m)) if m.layout == AofLayout::PerShard => m, - Ok(_) => { - error!( - "F6 per-shard rewrite: manifest at {} missing or not PerShard; aborting", - base_dir.display() - ); - return Err(AofPoolSendError::SendFailed); - } - Err(e) => { - error!( - "F6 per-shard rewrite: failed to load manifest at {}: {}", - base_dir.display(), - e - ); - return Err(AofPoolSendError::SendFailed); - } - }; - let current_seq = manifest.seq; - let n_shards = self.senders.len(); - let shared_manifest = Arc::new(parking_lot::Mutex::new(manifest)); - let coord = PerShardRewriteCoord::new(shared_manifest, current_seq, n_shards); - for (idx, s) in self.senders.iter().enumerate() { - // Blocking send for guaranteed delivery — see the doc comment. - if s.send(AofMessage::RewritePerShard { - shard_dbs: shard_dbs.clone(), - coord: coord.clone(), - }) - .is_err() - { - error!( - "F6 per-shard rewrite: writer {} channel disconnected; \ - rewrite aborted (no manifest commit, old generation remains \ - authoritative). Inspect AOF writer threads.", - idx - ); - // Account for the shards that did NOT receive the message (this - // one plus the unsent tail). The writers that DID receive it will - // fold and call `shard_done`; decrementing for the rest here lets - // the countdown still reach zero, so the final `shard_done` runs - // the abort path (keeps `old_seq`, clears AOF_REWRITE_IN_PROGRESS) - // instead of wedging the rewrite flag forever. - coord.mark_failed(); - for _ in idx..n_shards { - coord.shard_done(); - } - return Err(AofPoolSendError::SendFailed); - } - } - info!( - "F6 per-shard rewrite dispatched: seq {} -> {} across {} shards", - current_seq, - current_seq + 1, - n_shards - ); - Ok(()) - } - - /// Broadcast `Shutdown` to every writer. Used by orchestrated shutdown - /// paths in `main.rs`/`embedded.rs`. Each writer drains its channel and - /// fsyncs before exiting. - pub fn broadcast_shutdown(&self) { - for s in &self.senders { - let _ = s.try_send(AofMessage::Shutdown); - } - } - - /// Number of underlying writer senders. 1 for TopLevel, num_shards for - /// PerShard. - #[inline] - pub fn num_writers(&self) -> usize { - self.senders.len() - } - - /// Reports the pool's layout. Useful for places that need to refuse - /// PerShard-incompatible legacy code paths with a clear error. - #[inline] - pub fn layout(&self) -> crate::persistence::aof_manifest::AofLayout { - self.layout - } -} - -#[cfg(test)] -mod pool_tests { - use super::*; - use crate::persistence::aof_manifest::AofLayout; - use crate::runtime::channel; - - /// A per-shard BGREWRITEAOF fan-out that fails partway (a writer channel is - /// disconnected) must NOT leave `AOF_REWRITE_IN_PROGRESS` stuck true. The - /// failed/unsent shards are accounted for (`mark_failed` + `shard_done`) so - /// the countdown still reaches zero and the abort path clears the flag, - /// leaving the old generation authoritative. - #[test] - fn rewrite_fan_out_partial_failure_clears_in_progress_flag() { - use crate::command::persistence::AOF_REWRITE_IN_PROGRESS; - use std::sync::atomic::Ordering; - - let tmp = tempfile::tempdir().unwrap(); - // PerShard manifest on disk so the rewrite reaches the fan-out loop. - crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); - - // Disconnect shard-0's writer so the FIRST send fails (before any - // successful send), making the abort accounting deterministic. - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - drop(rx0); - let pool = AofWriterPool::per_shard_with_base_dir( - vec![tx0, tx1], - FsyncPolicy::EverySec, - DEFAULT_AOF_FSYNC_TIMEOUT, - tmp.path().to_path_buf(), - ); - - let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![ - vec![crate::storage::Database::new()], - vec![crate::storage::Database::new()], - ]); - - // The command handler sets this before dispatching the rewrite. - AOF_REWRITE_IN_PROGRESS.store(true, Ordering::SeqCst); - let res = pool.try_send_rewrite_per_shard(shard_dbs); - assert!(res.is_err(), "disconnected writer must fail the fan-out"); - assert!( - !AOF_REWRITE_IN_PROGRESS.load(Ordering::SeqCst), - "partial fan-out failure must clear AOF_REWRITE_IN_PROGRESS, not wedge it" - ); - } - - /// A per-shard rewrite that ABORTS (another shard failed to fold) must roll - /// THIS shard's append file back onto the committed old generation. Phase 6 - /// reopens `*file` onto the new-seq incr; if the rewrite then aborts, the - /// manifest keeps `old_seq` and prunes the new-seq files — so without a - /// barrier-before-resume the writer keeps appending into a discarded incr - /// that recovery ignores (silent data loss). This test drives one shard's - /// real fold to completion with a second shard pre-marked failed, then proves - /// post-abort appends land in the COMMITTED old-gen incr. - #[test] - fn rewrite_abort_reopens_writer_onto_committed_old_generation() { - use crate::command::persistence::AOF_REWRITE_IN_PROGRESS; - use std::io::Write; - use std::sync::atomic::Ordering; - - let tmp = tempfile::tempdir().unwrap(); - // 2-shard manifest at seq=1 on disk; share it through the coord's Arc. - let manifest = - crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); - let old_seq = manifest.seq; - let old_incr_s0 = manifest.shard_incr_path_seq(0, old_seq); - assert!(old_incr_s0.exists(), "old-gen incr must exist pre-rewrite"); - let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); - - let coord = PerShardRewriteCoord::new(manifest.clone(), old_seq, 2); - let new_seq = coord.new_seq(); - assert_ne!(new_seq, old_seq); - - // Shard 1 "fails to fold": mark_failed + shard_done (countdown 2 -> 1). - // Shard 0's shard_done (inside do_rewrite_per_shard) will then be the - // terminal decrement and take the abort path. - coord.mark_failed(); - coord.shard_done(); - - // Shard 0's append file starts on the OLD incr (where the live writer - // appends before a rewrite). Empty databases keep the fold trivial; the - // reopen behaviour is independent of key count. - let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![ - vec![crate::storage::Database::new()], - vec![crate::storage::Database::new()], - ]); - let mut file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&old_incr_s0) - .unwrap(); - let (_tx, rx) = channel::mpsc_bounded::(4); - - AOF_REWRITE_IN_PROGRESS.store(true, Ordering::SeqCst); - // The fold itself succeeds; aborting is a coordinator decision, so this - // returns Ok even though the rewrite is discarded. - do_rewrite_per_shard(0, &shard_dbs, &mut file, &rx, &coord).unwrap(); - - // Abort kept the old generation committed and pruned the new-gen incr. - assert_eq!(manifest.lock().seq, old_seq, "abort must keep old_seq"); - let new_incr_s0 = manifest.lock().shard_incr_path_seq(0, new_seq); - assert!( - !new_incr_s0.exists(), - "aborted new-gen incr must be pruned (the dangling target)" - ); - - // The writer must have been rolled back onto the old-gen incr: appends - // through `*file` after the rewrite must land in the committed file, not - // the pruned/unlinked new-gen inode. - let marker = b"*1\r\n$4\r\nPING\r\n"; - file.write_all(marker).unwrap(); - file.sync_data().unwrap(); - drop(file); - let committed = std::fs::read(&old_incr_s0).unwrap(); - assert!( - committed.windows(marker.len()).any(|w| w == marker), - "post-abort appends must land in the committed old-gen incr (no silent loss)" - ); - } - - /// Cross-thread barrier wakeup: a folded writer that decrements then blocks - /// in `await_outcome` on one thread must be woken with the committed - /// generation by the terminal `shard_done` running on ANOTHER thread. The - /// single-threaded behavioural test above never actually blocks (its - /// `await_outcome` returns immediately), so this exercises the real condvar - /// wait/notify path across threads. - #[test] - fn rewrite_abort_wakes_waiter_cross_thread() { - let tmp = tempfile::tempdir().unwrap(); - let manifest = - crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); - let old_seq = manifest.seq; - let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); - let coord = PerShardRewriteCoord::new(manifest, old_seq, 2); - - // Shard A: decrement (countdown 2 -> 1, non-terminal) then block. - let c_a = coord.clone(); - let waiter = std::thread::spawn(move || { - c_a.shard_done(); - c_a.await_outcome() - }); - - // Shard B fails: mark_failed + the terminal decrement (-> 0). Whichever - // of the two `shard_done` calls is terminal sees `failed` set (B sets it - // first) and publishes old_seq. - coord.mark_failed(); - coord.shard_done(); - - let observed = waiter.join().expect("waiter must wake, not hang or panic"); - assert_eq!( - observed, old_seq, - "aborted rewrite must publish old_seq to the cross-thread waiter" - ); - } - - /// Panic-safety / liveness: a fold that PANICS mid-flight (the OOM-unwind - /// hazard of issue #138) must not hang the other shards' writers at the - /// barrier. The `ShardDoneGuard`'s `Drop` must fire `mark_failed` + - /// `shard_done` on unwind so the countdown still closes and every waiter - /// wakes. Without the guard this test hangs forever (the panicking shard - /// never decrements). - #[test] - fn rewrite_fold_panic_releases_barrier_for_other_writers() { - use std::sync::atomic::Ordering; - - let tmp = tempfile::tempdir().unwrap(); - let manifest = - crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); - let old_seq = manifest.seq; - let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); - let coord = PerShardRewriteCoord::new(manifest, old_seq, 2); - - // Shard A folds, decrements, then blocks on the barrier in a thread. - let c_a = coord.clone(); - let waiter = std::thread::spawn(move || { - c_a.shard_done(); - c_a.await_outcome() - }); - - // Shard B "panics mid-fold": a ShardDoneGuard dropped during unwind. The - // guard's Drop must abort + decrement so A's barrier releases. - let c_b = coord.clone(); - let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - let _guard = ShardDoneGuard::new(&c_b); - panic!("simulated OOM unwind during fold"); - })); - assert!(panicked.is_err(), "the simulated fold must have panicked"); - - let observed = waiter - .join() - .expect("waiter must wake after the panicking shard's guard fires"); - assert_eq!( - observed, old_seq, - "panic-aborted rewrite must publish old_seq, not hang" - ); - assert!( - coord.failed.load(Ordering::Acquire), - "the dropped guard must have marked the rewrite failed" - ); - } - - #[test] - fn top_level_pool_routes_all_shards_to_writer_zero() { - let (tx, rx) = channel::mpsc_bounded::(8); - let pool = AofWriterPool::top_level(tx); - assert_eq!(pool.num_writers(), 1); - assert_eq!(pool.layout(), AofLayout::TopLevel); - - pool.try_send_append(0, 0, Bytes::from_static(b"a")); - pool.try_send_append(7, 0, Bytes::from_static(b"b")); - pool.try_send_append(42, 0, Bytes::from_static(b"c")); - - let mut seen = 0; - while rx.try_recv().is_ok() { - seen += 1; - } - assert_eq!(seen, 3, "all 3 appends should land on writer 0"); - } - - #[test] - fn per_shard_pool_routes_each_shard_to_its_own_writer() { - let (tx0, rx0) = channel::mpsc_bounded::(8); - let (tx1, rx1) = channel::mpsc_bounded::(8); - let (tx2, rx2) = channel::mpsc_bounded::(8); - let pool = AofWriterPool::per_shard(vec![tx0, tx1, tx2]); - assert_eq!(pool.num_writers(), 3); - assert_eq!(pool.layout(), AofLayout::PerShard); - - pool.try_send_append(0, 100, Bytes::from_static(b"shard0")); - pool.try_send_append(1, 200, Bytes::from_static(b"shard1a")); - pool.try_send_append(1, 300, Bytes::from_static(b"shard1b")); - pool.try_send_append(2, 400, Bytes::from_static(b"shard2")); - - let count = |rx: &channel::MpscReceiver| -> usize { - let mut n = 0; - while rx.try_recv().is_ok() { - n += 1; - } - n - }; - assert_eq!(count(&rx0), 1, "shard 0 writer should receive exactly 1"); - assert_eq!(count(&rx1), 2, "shard 1 writer should receive exactly 2"); - assert_eq!(count(&rx2), 1, "shard 2 writer should receive exactly 1"); - } - - #[test] - fn per_shard_pool_rejects_rewrite_with_explicit_error() { - let (tx0, _rx0) = channel::mpsc_bounded::(8); - let (tx1, _rx1) = channel::mpsc_bounded::(8); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - let dummies: SharedDatabases = Arc::new(vec![]); - let err = pool - .try_send_rewrite(AofMessage::Rewrite(dummies)) - .unwrap_err(); - assert_eq!(err, AofPoolSendError::RewriteUnsupportedInPerShard); - } - - #[test] - fn top_level_pool_accepts_rewrite() { - let (tx, rx) = channel::mpsc_bounded::(8); - let pool = AofWriterPool::top_level(tx); - - let dummies: SharedDatabases = Arc::new(vec![]); - pool.try_send_rewrite(AofMessage::Rewrite(dummies)).unwrap(); - assert!(matches!(rx.try_recv(), Ok(AofMessage::Rewrite(_)))); - } - - #[test] - fn per_shard_pool_threads_lsn_field_to_each_writer() { - // Step 3 wire-format contract: try_send_append carries the issued LSN - // through to the writer task, which writes it as the per-entry header - // under PerShard layout. This unit test pins the channel-side contract - // (the disk-side framing is covered by writer-task integration). - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - pool.try_send_append(0, 42, Bytes::from_static(b"set foo 1")); - pool.try_send_append(1, 43, Bytes::from_static(b"set bar 2")); - pool.try_send_append(0, 44, Bytes::from_static(b"del foo")); - - // Shard 0 should see (42, "set foo 1") then (44, "del foo"). - match rx0.try_recv() { - Ok(AofMessage::Append { lsn, bytes }) => { - assert_eq!(lsn, 42, "shard 0 first entry lsn"); - assert_eq!(bytes.as_ref(), b"set foo 1"); - } - other => panic!( - "shard 0 first recv expected Append, got {:?}", - other.is_ok() - ), - } - match rx0.try_recv() { - Ok(AofMessage::Append { lsn, bytes }) => { - assert_eq!(lsn, 44, "shard 0 second entry lsn"); - assert_eq!(bytes.as_ref(), b"del foo"); - } - other => panic!( - "shard 0 second recv expected Append, got {:?}", - other.is_ok() - ), - } - // Shard 1 should see (43, "set bar 2") only. - match rx1.try_recv() { - Ok(AofMessage::Append { lsn, bytes }) => { - assert_eq!(lsn, 43, "shard 1 entry lsn"); - assert_eq!(bytes.as_ref(), b"set bar 2"); - } - other => panic!("shard 1 recv expected Append, got {:?}", other.is_ok()), - } - } - - #[test] - fn try_send_append_sync_queues_appendsync_with_ack() { - // Channel-level wiring contract for the H1 fix: `try_send_append_sync` - // queues `AofMessage::AppendSync { lsn, bytes, ack }`, and the - // returned receiver resolves to whatever value the (mocked) writer - // sends on `ack`. End-to-end durability is covered by step 8 - // (CRASH-01-LITE); this pins the API contract. - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - let recv = pool.try_send_append_sync(0, 99, Bytes::from_static(b"SET k v")); - - // Drain the queue; the writer would normally do this. Capture the - // ack sender, do the (mock) durable write, then ack Synced. - let ack = match rx0.try_recv() { - Ok(AofMessage::AppendSync { lsn, bytes, ack }) => { - assert_eq!(lsn, 99, "lsn forwarded through the channel"); - assert_eq!(bytes.as_ref(), b"SET k v", "bytes forwarded"); - ack - } - other => panic!("expected AppendSync, got {:?}", other.is_ok()), - }; - - // Writer reports Synced — caller observes Synced. - let _ = ack.send(AofAck::Synced); - let result = recv.recv_blocking().expect("receiver resolves"); - assert_eq!(result, AofAck::Synced); - } - - #[test] - fn append_sync_writer_dropped_resolves_recv_error() { - // If the writer task is dead or the channel disconnects between - // queueing and the ack send, the receiver MUST resolve with an - // error rather than hang. Callers treat that as a hard failure - // (return an error frame, do not +OK). - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - let recv = pool.try_send_append_sync(0, 7, Bytes::from_static(b"x")); - - // Drain the message but DROP the ack sender without sending. - match rx0.try_recv() { - Ok(AofMessage::AppendSync { ack, .. }) => drop(ack), - other => panic!("expected AppendSync, got {:?}", other.is_ok()), - } - - let err = recv.recv_blocking().expect_err("dropped ack -> RecvError"); - // Crash-safe: we got a sentinel-style error, not a hang. - let _ = err; - } - - #[test] - fn append_sync_writer_reports_write_failed() { - // Writer encountered a write_all error; recv returns WriteFailed. - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - let recv = pool.try_send_append_sync(0, 1, Bytes::from_static(b"x")); - let ack = match rx0.try_recv() { - Ok(AofMessage::AppendSync { ack, .. }) => ack, - other => panic!("expected AppendSync, got {:?}", other.is_ok()), - }; - let _ = ack.send(AofAck::WriteFailed); - let result = recv.recv_blocking().expect("recv resolves"); - assert_eq!(result, AofAck::WriteFailed); - } - - #[test] - fn append_sync_writer_reports_fsync_failed() { - // Writer wrote the payload but fsync (sync_data) returned an error. - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - let recv = pool.try_send_append_sync(0, 1, Bytes::from_static(b"x")); - let ack = match rx0.try_recv() { - Ok(AofMessage::AppendSync { ack, .. }) => ack, - other => panic!("expected AppendSync, got {:?}", other.is_ok()), - }; - let _ = ack.send(AofAck::FsyncFailed); - let result = recv.recv_blocking().expect("recv resolves"); - assert_eq!(result, AofAck::FsyncFailed); - } - - /// Issue #140: an AppendSync drained during a BGREWRITEAOF must NOT be acked - /// `Synced` inside the drain — that reports the write durable before the - /// post-drain boundary fsync, so a crash in the window loses an entry the - /// client was told was safe. The drain must PARK the ack and only fulfil it - /// after the boundary `sync_data()`. (Framed / per-shard drain.) - #[test] - fn drain_framed_parks_appendsync_ack_until_boundary_fsync() { - let tmp = tempfile::tempdir().unwrap(); - let incr = tmp.path().join("incr.aof"); - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - let recv = pool.try_send_append_sync(0, 99, Bytes::from_static(b"SET k v")); - - let mut file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&incr) - .unwrap(); - let mut outcome = drain_pending_appends_framed(&rx0, &mut file).unwrap(); - - // CONTRACT: drained + parked, NOT yet acked. - assert_eq!(outcome.drained, 1, "the AppendSync was drained"); - assert_eq!( - outcome.pending_acks.len(), - 1, - "the ack must be parked until the boundary fsync, not sent during drain" - ); - - // Boundary fsync succeeds → fulfil Synced; only NOW is the client told durable. - file.sync_data().unwrap(); - outcome.fulfill_acks(true); - assert_eq!( - recv.recv_blocking().expect("ack resolves"), - AofAck::Synced, - "post-fsync the parked ack must resolve Synced" - ); - } - - /// Issue #140 failure path: if the rewrite-boundary fsync FAILS, a drained - /// AppendSync must resolve `FsyncFailed`, never `Synced`. Exercises the - /// non-framed `drain_pending_appends` — the DEFAULT `--shards 1` rewrite - /// path (`do_rewrite_single`), which is reachable under appendfsync=always. - #[cfg(feature = "runtime-monoio")] - #[test] - fn drain_single_fulfills_fsync_failure_as_fsync_failed() { - let tmp = tempfile::tempdir().unwrap(); - let incr = tmp.path().join("incr.aof"); - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard(vec![tx0, tx1]); - - let recv = pool.try_send_append_sync(0, 7, Bytes::from_static(b"SET a b")); - - let mut file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&incr) - .unwrap(); - let mut outcome = drain_pending_appends(&rx0, &mut file).unwrap(); - assert_eq!( - outcome.pending_acks.len(), - 1, - "ack parked, not sent in drain" - ); - - // Simulate a failed boundary fsync: the parked ack must report FsyncFailed. - outcome.fulfill_acks(false); - assert_eq!( - recv.recv_blocking().expect("ack resolves"), - AofAck::FsyncFailed, - "a failed boundary fsync must NOT be reported Synced to the client" - ); - } - - // F2 (design-for-failure): `appendfsync=always` must bound its fsync-ack - // await. A stalled writer must surface a hard error within the budget, - // never park the connection forever. Tokio-gated because it drives the - // runtime timer; the monoio path shares the proven `select! + sleep` - // shape from `cluster::failover`, exercised end-to-end by the crash tests. - #[cfg(feature = "runtime-tokio")] - #[tokio::test] - async fn always_fsync_times_out_when_writer_never_acks() { - // Writer channel is held (kept open) but never drained → the - // AppendSync sits buffered with its ack sender alive, so the receiver - // never resolves. The bounded await MUST elapse and report failure. - let (tx0, _rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard_with_policy( - vec![tx0, tx1], - FsyncPolicy::Always, - Duration::from_millis(50), - ); - - let start = Instant::now(); - let res = pool - .try_send_append_durable(0, 1, Bytes::from_static(b"x")) - .await; - let elapsed = start.elapsed(); - - assert_eq!( - res, - Err(AofAck::FsyncFailed), - "timed-out fsync must map to FsyncFailed (durability unconfirmed)" - ); - assert!( - elapsed < Duration::from_secs(2), - "must fail within the bound, not hang (took {:?})", - elapsed - ); - // Keep the receivers alive until here so the message stays buffered. - drop((_rx0, _rx1)); - } - - #[cfg(feature = "runtime-tokio")] - #[tokio::test] - async fn always_fsync_succeeds_when_writer_acks_in_time() { - // Happy path: a writer drains the AppendSync and acks `Synced` well - // within the bound → the durable append returns Ok(()). - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard_with_policy( - vec![tx0, tx1], - FsyncPolicy::Always, - Duration::from_millis(500), - ); - - tokio::spawn(async move { - if let Ok(AofMessage::AppendSync { ack, .. }) = rx0.recv_async().await { - let _ = ack.send(AofAck::Synced); - } - }); - - let res = pool - .try_send_append_durable(0, 1, Bytes::from_static(b"x")) - .await; - assert_eq!(res, Ok(()), "ack within the bound must succeed"); - drop(_rx1); - } - - /// Parse the PerShard incr framing `[u64 lsn LE][u32 len LE][len bytes]`, - /// stopping at a truncated tail (the crash/torn boundary) — exactly what - /// `replay_incr_framed` does. Returns the cleanly-replayable prefix. - #[cfg(feature = "runtime-tokio")] - fn parse_framed(buf: &[u8]) -> Vec<(u64, Vec)> { - let mut out = Vec::new(); - let mut i = 0usize; - while i + 12 <= buf.len() { - let lsn = u64::from_le_bytes(buf[i..i + 8].try_into().unwrap()); - let len = u32::from_le_bytes(buf[i + 8..i + 12].try_into().unwrap()) as usize; - if i + 12 + len > buf.len() { - break; // truncated tail → crash boundary, stop - } - out.push((lsn, buf[i + 12..i + 12 + len].to_vec())); - i += 12 + len; - } - out - } - - // Regression (PR #136 review, BUG #2): the tokio per-shard writer must carry - // a `write_error` latch like the single-file (~:1467) and monoio (~:2125) - // writers. A torn write (header lands, payload fails) must NOT be followed by - // more records — a lone orphaned header makes the framed replay misread the - // next record's bytes as the orphan's payload, corrupting everything after. - // The latch suppresses all writes after the tear and reports WriteFailed to - // AppendSync callers (so they error instead of ack'ing a corrupt write). - #[cfg(feature = "runtime-tokio")] - #[tokio::test] - async fn tokio_per_shard_writer_latches_after_torn_write() { - use crate::persistence::aof_manifest::AofManifest; - use std::sync::atomic::Ordering; - - let tmp = tempfile::tempdir().unwrap(); - let base_dir = tmp.path().to_path_buf(); - // PerShard layout, 2 shards (the per-shard pool needs >=2); drive shard 0. - let manifest = AofManifest::initialize_multi(&base_dir, 2).unwrap(); - let incr = manifest.shard_incr_path(0); - - // Inject: the 2nd Append tears (header written, payload "fails"). - TEST_FAIL_WRITE_AT.store(2, Ordering::SeqCst); - - let (tx, rx) = channel::mpsc_bounded::(16); - let cancel = CancellationToken::new(); - let writer = tokio::spawn(per_shard_aof_writer_task( - rx, - base_dir.clone(), - 0, - FsyncPolicy::Always, - cancel.clone(), - )); - - // 1: clean. 2: torn (header only). 3: must be suppressed by the latch. - tx.try_send(AofMessage::Append { - lsn: 1, - bytes: Bytes::from_static(b"AAAA"), - }) - .unwrap(); - tx.try_send(AofMessage::Append { - lsn: 2, - bytes: Bytes::from_static(b"BBBB"), - }) - .unwrap(); - tx.try_send(AofMessage::Append { - lsn: 3, - bytes: Bytes::from_static(b"CCCC"), - }) - .unwrap(); - - // Barrier + assertion: an AppendSync after the tear MUST come back - // WriteFailed (latched), never Synced. - let (ack_tx, ack_rx) = crate::runtime::channel::oneshot::(); - tx.try_send(AofMessage::AppendSync { - lsn: 4, - bytes: Bytes::from_static(b"DDDD"), - ack: ack_tx, - }) - .unwrap(); - - let ack = tokio::time::timeout(std::time::Duration::from_secs(5), ack_rx) - .await - .expect("writer must answer the AppendSync within 5s") - .expect("ack channel must not drop"); - assert_eq!( - ack, - AofAck::WriteFailed, - "after a torn write the latch must reject further writes (got {ack:?})" - ); - - cancel.cancel(); - TEST_FAIL_WRITE_AT.store(0, Ordering::SeqCst); - let _ = tokio::time::timeout(std::time::Duration::from_secs(5), writer).await; - - // On disk: exactly one replayable frame (lsn=1, "AAAA"). The orphaned - // lsn=2 header is a truncated tail (crash boundary); lsn 3 and 4 were - // never written (latch held) — no corruption. - let raw = std::fs::read(&incr).unwrap(); - let frames = parse_framed(&raw); - assert_eq!( - frames, - vec![(1u64, b"AAAA".to_vec())], - "only the pre-tear record may replay; orphaned headers / suppressed \ - records must not corrupt the stream" - ); - } - - #[test] - fn broadcast_shutdown_reaches_every_writer() { - let (tx0, rx0) = channel::mpsc_bounded::(2); - let (tx1, rx1) = channel::mpsc_bounded::(2); - let (tx2, rx2) = channel::mpsc_bounded::(2); - let pool = AofWriterPool::per_shard(vec![tx0, tx1, tx2]); - - pool.broadcast_shutdown(); - - for (i, rx) in [&rx0, &rx1, &rx2].iter().enumerate() { - assert!( - matches!(rx.try_recv(), Ok(AofMessage::Shutdown)), - "writer {} did not receive Shutdown", - i - ); - } - } - - /// FIX-W1-1 contract: `try_send_append_durable` under `Always` policy MUST - /// return `Err(AofAck::FsyncFailed)` when the writer reports failure. - /// handler_single.rs must await this BEFORE flushing responses to the client. - /// - /// Uses spawn_blocking to simulate the mock writer responding on the ack - /// channel concurrently, which allows the async rendezvous to complete. - #[cfg(feature = "runtime-tokio")] - #[tokio::test] - async fn always_policy_try_send_append_durable_returns_err_on_fsync_fail() { - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = std::sync::Arc::new(AofWriterPool::per_shard_with_policy( - vec![tx0, tx1], - FsyncPolicy::Always, - Duration::ZERO, // legacy unbounded await — disconnect/ack resolves it - )); - - // Spawn a mock writer that drains AppendSync and responds with FsyncFailed. - // Runs in a blocking thread (flume's blocking recv) so it doesn't block - // the async executor while waiting for the handler to enqueue the message. - let mock_writer = tokio::task::spawn_blocking(move || { - // flume::Receiver::recv() blocks until a message is available - let msg = rx0.recv().expect("mock writer got message"); - if let AofMessage::AppendSync { ack, .. } = msg { - let _ = ack.send(AofAck::FsyncFailed); - } else { - panic!("expected AppendSync under Always policy"); - } - }); - - // The handler MUST await this BEFORE flushing responses to the client - let result = pool - .try_send_append_durable(0, 1, Bytes::from_static(b"SET k v")) - .await; - mock_writer.await.expect("mock writer completed"); - - assert_eq!( - result, - Err(AofAck::FsyncFailed), - "Always policy MUST propagate fsync failure so caller can return an error frame" - ); - } - - /// FIX-W1-1 ordering contract: when `aof_entries` carries `(resp_idx, bytes)` - /// tuples, the handler can patch `responses[resp_idx]` on AOF failure BEFORE - /// flushing to the client. This test verifies the indexing is sound. - #[test] - fn aof_entries_indexed_by_response_slot_patches_correctly() { - use crate::protocol::Frame; - let mut responses: Vec = vec![ - Frame::SimpleString(bytes::Bytes::from_static(b"OK")), - Frame::SimpleString(bytes::Bytes::from_static(b"OK")), - Frame::SimpleString(bytes::Bytes::from_static(b"OK")), - ]; - // Simulate two write commands at response indices 0 and 2 (index 1 was a read) - let aof_entries: Vec<(usize, Bytes)> = vec![ - (0, Bytes::from_static(b"SET a 1")), - (2, Bytes::from_static(b"SET c 3")), - ]; - - // AOF write at index 2 fails; patch that response slot - for (resp_idx, _bytes) in &aof_entries { - if *resp_idx == 2 { - // Simulate Err(AofAck::FsyncFailed) from try_send_append_durable - responses[*resp_idx] = - Frame::Error(Bytes::from_static(b"WRITEFAIL aof fsync failed")); - } - } - - assert!( - matches!(&responses[0], Frame::SimpleString(_)), - "index 0 (successful fsync) should remain +OK" - ); - assert!( - matches!(&responses[1], Frame::SimpleString(_)), - "index 1 (read, no AOF) should remain +OK" - ); - assert!( - matches!(&responses[2], Frame::Error(_)), - "index 2 (failed fsync) must be patched to error" - ); - } - - // NOTE (FIX-W1-1 r3): The H1 ordering regression test was moved to - // `src/server/conn/handler_single.rs` (test module, fn - // `flush_with_aof_ack_ack_precedes_response`). The previous inline - // reproduction here was non-discriminating — it reproduced the ack-first - // loop IN THE TEST BODY rather than calling the real production fn, so it - // passed on both pre-fix and post-fix binaries. - // - // The new test calls `flush_with_aof_ack` directly (the fn the handler now - // delegates to), so inverting Phase 1/Phase 2 order in that fn causes a - // measurable timing failure (`elapsed_ms ≈ 0ms < 55ms`). - // - // End-to-end ordering is also covered by: - // tests/crash_matrix_per_shard_aof.rs (CRASH-01-LITE — AlwaysPolicy shards) - - // ----------------------------------------------------------------------- - // FIX-W2-5: channel-full returns AofAck::ChannelFull + increments counter - // ----------------------------------------------------------------------- - #[test] - fn try_send_append_sync_channel_full_returns_channel_full_ack() { - // Create a channel with capacity 1 and fill it so the next try_send - // hits TrySendError::Full. - let (tx0, rx0) = channel::mpsc_bounded::(1); - // Fill the channel by pre-loading one message. - tx0.try_send(AofMessage::Shutdown).expect("pre-fill"); - // rx0 intentionally not consumed — channel is now at capacity. - - let pool = AofWriterPool::top_level(tx0); - - let before = AOF_BACKPRESSURE_DROPPED.load(std::sync::atomic::Ordering::Relaxed); - let recv = pool.try_send_append_sync(0, 1, Bytes::from_static(b"SET k v")); - - // The channel was full — ChannelFull is returned immediately without - // a writer round-trip. - let result = recv.recv_blocking().expect("pre-filled oneshot resolves"); - assert_eq!( - result, - AofAck::ChannelFull, - "channel-full must yield ChannelFull, not {:?}", - result - ); - - let after = AOF_BACKPRESSURE_DROPPED.load(std::sync::atomic::Ordering::Relaxed); - assert_eq!( - after, - before + 1, - "backpressure counter must increment by 1" - ); - - // No AppendSync should have reached the (blocked) reader. - drop(rx0); // drain without consuming — just verify nothing snuck through - } - - // ----------------------------------------------------------------------- - // FIX-W2-9: try_send_append_durable must be used for SWAPDB-like mutations - // - // Red test: documents the contract that handler_single.rs SHOULD honour. - // When appendfsync=always, try_send_append_durable MUST return Err on - // writer failure so callers can abort the mutation safely. - // ----------------------------------------------------------------------- - #[test] - fn try_send_append_durable_always_writer_dead_returns_write_failed() { - // Create a pool with Always policy. The writer task is not running — - // we model that by draining the channel message and then dropping the - // ack sender, simulating a dead writer. - let (tx0, rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard_with_policy( - vec![tx0, tx1], - FsyncPolicy::Always, - Duration::ZERO, // legacy unbounded await — disconnect resolves it - ); - - // Spawn a thread that pulls the AppendSync off the channel but drops - // the ack without sending — simulating a writer crash mid-fsync. - let rx0_clone = rx0; - let handle = std::thread::spawn(move || { - match rx0_clone.recv() { - Ok(AofMessage::AppendSync { ack, .. }) => drop(ack), // writer crash - other => panic!("unexpected message: {:?}", other.is_ok()), - } - }); - - // try_send_append_durable for Always must await the ack. - // With the ack sender dropped, it should resolve to Err(WriteFailed). - let result = futures::executor::block_on(pool.try_send_append_durable( - 0, - 55, - Bytes::from_static(b"SWAPDB 0 1"), - )); - - handle.join().expect("ack dropper thread"); - - assert!( - result.is_err(), - "try_send_append_durable with dead writer must return Err, got Ok" - ); - assert_eq!( - result.unwrap_err(), - AofAck::WriteFailed, - "dead writer must resolve to WriteFailed" - ); - } - - #[test] - fn try_send_append_durable_everysec_is_fire_and_forget() { - // EverySec policy: try_send_append_durable always returns Ok — the - // durability policy doesn't block on fsync. handler_single.rs must - // use try_send_append_durable so the policy is respected. - let (tx0, _rx0) = channel::mpsc_bounded::(4); - let (tx1, _rx1) = channel::mpsc_bounded::(4); - let pool = AofWriterPool::per_shard_with_policy( - vec![tx0, tx1], - FsyncPolicy::EverySec, - Duration::ZERO, - ); - - let result = futures::executor::block_on(pool.try_send_append_durable( - 0, - 56, - Bytes::from_static(b"SWAPDB 0 1"), - )); - - assert!( - result.is_ok(), - "EverySec policy must be fire-and-forget (Ok), got {:?}", - result - ); - } -} - -/// Serialize a Frame into RESP wire format bytes. -pub fn serialize_command(frame: &Frame) -> Bytes { - let mut buf = BytesMut::with_capacity(64); - serialize::serialize(frame, &mut buf); - buf.freeze() -} - -/// Background AOF writer task. Receives commands via mpsc channel and appends them -/// to the AOF file. Handles fsync according to the configured policy. -pub async fn aof_writer_task( - rx: channel::MpscReceiver, - aof_path: PathBuf, - fsync: FsyncPolicy, - cancel: CancellationToken, -) { - #[cfg(feature = "runtime-tokio")] - use tokio::io::AsyncWriteExt; - - // Open file in append mode (create if not exists) - #[cfg(feature = "runtime-tokio")] - let file: tokio::fs::File = match tokio::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&aof_path) - .await - { - Ok(f) => f, - Err(e) => { - error!("Failed to open AOF file {}: {}", aof_path.display(), e); - return; - } - }; - - #[cfg(feature = "runtime-tokio")] - let mut writer = tokio::io::BufWriter::new(file); - #[cfg(feature = "runtime-tokio")] - let mut last_fsync = Instant::now(); - #[cfg(feature = "runtime-tokio")] - let mut interval = tokio::time::interval(std::time::Duration::from_secs(1)); - #[cfg(feature = "runtime-tokio")] - interval.tick().await; // consume first tick - - // Monoio path: multi-part AOF (base RDB + incremental RESP) with sync I/O. - // - // On startup, if appendonlydir/ exists with a manifest, open the current - // incr file for appending. Otherwise start fresh with seq 1. - // On BGREWRITEAOF: snapshot → write new base RDB → create new incr → advance manifest. - #[cfg(feature = "runtime-monoio")] - { - use crate::persistence::aof_manifest::AofManifest; - use std::io::Write; - - // Resolve the persistence base directory from aof_path's parent. - let base_dir = aof_path.parent().unwrap_or(Path::new(".")).to_path_buf(); - - // Load manifest — do NOT create one here if it doesn't exist. - // main.rs recovery runs concurrently and must finish before a manifest - // is created, to avoid racing against legacy single-file AOF detection. - // main.rs will create the manifest after recovery completes. - // - // A corrupt manifest is fatal — exit the writer so the server startup - // notices and fails loud rather than silently overwriting. - // - // Bounded wait: check the cancellation token each iteration and enforce - // a hard timeout so the writer doesn't spin forever if main.rs fails to - // create the manifest (e.g. disk full, permission error). - let manifest_wait_start = Instant::now(); - const MANIFEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); - let mut manifest = loop { - if cancel.is_cancelled() { - info!("AOF writer: cancelled while waiting for manifest"); - return; - } - if manifest_wait_start.elapsed() > MANIFEST_TIMEOUT { - error!( - "AOF writer: manifest not found at {} after {:?}. Writer exiting; check recovery logs.", - base_dir.display(), - MANIFEST_TIMEOUT, - ); - return; - } - match AofManifest::load(&base_dir) { - Ok(Some(m)) => break m, - Ok(None) => { - // main.rs recovery hasn't created the manifest yet — wait. - std::thread::sleep(std::time::Duration::from_millis(50)); - } - Err(e) => { - error!( - "AOF manifest corrupt at {}: {}. Writer exiting; persistence disabled.", - base_dir.display(), - e - ); - return; - } - } - }; - - // Open the current incremental file for appending - let incr_path = manifest.incr_path(); - let mut file = match std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&incr_path) - { - Ok(f) => f, - Err(e) => { - error!( - "Failed to open AOF incr file {}: {}", - incr_path.display(), - e - ); - return; - } - }; - info!( - "AOF writer: seq {}, incr={}", - manifest.seq, - incr_path.display() - ); - - let mut last_fsync = Instant::now(); - - let mut write_error = false; - - // Test-only fault injection: same env var as the PerShard writer. - // Read once at task startup; zero cost in production (var absent). - let fail_fsync_for_test = std::env::var("MOON_TEST_AOF_FSYNC_FAIL").as_deref() == Ok("1"); - - loop { - match rx.recv() { - // TopLevel writer: legacy v1 disk format is plain RESP. The - // LSN is ignored — TopLevel is single-shard so per-shard merge - // by LSN is moot. - Ok(AofMessage::Append { - bytes: data, - lsn: _, - }) => { - if write_error { - continue; // Drop appends after persistent I/O failure - } - if let Err(e) = file.write_all(&data) { - error!( - "AOF write failed (seq {}): {}. Persistence degraded.", - manifest.seq, e - ); - write_error = true; - continue; - } - match fsync { - FsyncPolicy::Always => { - let t = Instant::now(); - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!("AOF sync failed (seq {}, always): {}", manifest.seq, e); - write_error = true; - } else { - crate::admin::metrics_setup::record_aof_fsync( - t.elapsed().as_micros() as u64, - ); - } - } - FsyncPolicy::EverySec => { - if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { - let t = Instant::now(); - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!( - "AOF sync failed (seq {}, everysec): {}", - manifest.seq, e - ); - // Non-fatal for everysec: retry next interval - } else { - crate::admin::metrics_setup::record_aof_fsync( - t.elapsed().as_micros() as u64, - ); - last_fsync = Instant::now(); - } - } - } - FsyncPolicy::No => {} - } - } - // TopLevel writer (monoio): legacy v1 plain RESP, lsn ignored. - // AppendSync ALWAYS fsyncs and acks before returning, regardless - // of the configured policy — that's the durability contract the - // caller signed up for by choosing AppendSync. - Ok(AofMessage::AppendSync { - bytes: data, - lsn: _, - ack, - }) => { - if write_error { - let _ = ack.send(AofAck::WriteFailed); - continue; - } - // Test-only: return FsyncFailed immediately without touching disk. - if fail_fsync_for_test { - let _ = ack.send(AofAck::FsyncFailed); - continue; - } - if let Err(e) = file.write_all(&data) { - error!( - "AOF AppendSync write failed (seq {}): {}. Persistence degraded.", - manifest.seq, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; - } - let t = Instant::now(); - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!("AOF AppendSync sync failed (seq {}): {}", manifest.seq, e); - write_error = true; - let _ = ack.send(AofAck::FsyncFailed); - } else { - crate::admin::metrics_setup::record_aof_fsync( - t.elapsed().as_micros() as u64 - ); - let _ = ack.send(AofAck::Synced); - } - } - Ok(AofMessage::Shutdown) | Err(_) => { - if !write_error { - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!("AOF final sync failed (seq {}): {}", manifest.seq, e); - } - } - info!("AOF writer shutting down (monoio, seq {})", manifest.seq); - break; - } - Ok(AofMessage::Rewrite(db)) => { - if !write_error { - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!("AOF pre-rewrite sync failed (seq {}): {}", manifest.seq, e); - } - } - match do_rewrite_single(&db, &mut manifest, &mut file, &rx) { - Ok(()) => { - write_error = false; // Reset on successful rewrite - } - Err(e) => error!("AOF rewrite failed (seq {}): {}", manifest.seq, e), - } - crate::command::persistence::AOF_REWRITE_IN_PROGRESS - .store(false, std::sync::atomic::Ordering::SeqCst); - } - Ok(AofMessage::RewriteSharded(shard_dbs)) => { - if !write_error { - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!("AOF pre-rewrite sync failed (seq {}): {}", manifest.seq, e); - } - } - match do_rewrite_sharded(&shard_dbs, &mut manifest, &mut file, &rx) { - Ok(()) => { - write_error = false; - } - Err(e) => error!("AOF rewrite failed (seq {}): {}", manifest.seq, e), - } - crate::command::persistence::AOF_REWRITE_IN_PROGRESS - .store(false, std::sync::atomic::Ordering::SeqCst); - } - // [F6] A TopLevel writer never owns per-shard files; receiving - // RewritePerShard means a routing bug. Self-abort so the - // coordinator's countdown completes and the flag clears. - Ok(AofMessage::RewritePerShard { coord, .. }) => { - warn!("AOF TopLevel writer received RewritePerShard — routing bug; aborting"); - coord.mark_failed(); - coord.shard_done(); - } - } - } - return; - } - - loop { - #[cfg(feature = "runtime-tokio")] - tokio::select! { - msg = rx.recv_async() => { - match msg { - // TopLevel writer (tokio): legacy v1 plain RESP, lsn ignored. - Ok(AofMessage::Append { bytes: data, lsn: _ }) => { - if let Err(e) = writer.write_all(&data).await { - error!("AOF write error: {}", e); - continue; - } - match fsync { - FsyncPolicy::Always => { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - } - FsyncPolicy::EverySec | FsyncPolicy::No => { - // EverySec handled by interval tick below; No does nothing - } - } - } - // AppendSync: write + fsync + ack, regardless of policy. - Ok(AofMessage::AppendSync { bytes: data, lsn: _, ack }) => { - if let Err(e) = writer.write_all(&data).await { - error!("AOF AppendSync write error: {}", e); - let _ = ack.send(AofAck::WriteFailed); - continue; - } - if let Err(e) = writer.flush().await { - error!("AOF AppendSync flush error: {}", e); - let _ = ack.send(AofAck::FsyncFailed); - continue; - } - if let Err(e) = writer.get_ref().sync_data().await { - error!("AOF AppendSync sync_data error: {}", e); - let _ = ack.send(AofAck::FsyncFailed); - continue; - } - let _ = ack.send(AofAck::Synced); - } - Ok(AofMessage::Rewrite(db)) => { - // Flush current writer before rewrite - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - - if let Err(e) = rewrite_aof(db, &aof_path).await { - error!("AOF rewrite failed: {}", e); - } - crate::command::persistence::AOF_REWRITE_IN_PROGRESS - .store(false, std::sync::atomic::Ordering::SeqCst); - - // Reopen file after rewrite (it was replaced) - let reopen_result: Result = tokio::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&aof_path) - .await; - match reopen_result { - Ok(f) => { - writer = tokio::io::BufWriter::new(f); - } - Err(e) => { - error!("Failed to reopen AOF file after rewrite: {}", e); - return; - } - } - } - Ok(AofMessage::RewriteSharded(shard_dbs)) => { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - if let Err(e) = rewrite_aof_sharded_sync(&shard_dbs, &aof_path) { - error!("AOF rewrite (sharded) failed: {}", e); - } - crate::command::persistence::AOF_REWRITE_IN_PROGRESS - .store(false, std::sync::atomic::Ordering::SeqCst); - let reopen_result: Result = tokio::fs::OpenOptions::new() - .create(true).append(true).open(&aof_path).await; - match reopen_result { - Ok(f) => writer = tokio::io::BufWriter::new(f), - Err(e) => { error!("Failed to reopen AOF after rewrite: {}", e); return; } - } - } - // [F6] TopLevel writer never owns per-shard files — routing - // bug. Self-abort so the countdown completes + flag clears. - Ok(AofMessage::RewritePerShard { coord, .. }) => { - warn!("AOF TopLevel writer received RewritePerShard — routing bug; aborting"); - coord.mark_failed(); - coord.shard_done(); - } - Ok(AofMessage::Shutdown) | Err(_) => { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - info!("AOF writer shutting down"); - break; - } - } - } - _ = interval.tick(), if fsync == FsyncPolicy::EverySec => { - if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - last_fsync = Instant::now(); - } - } - _ = cancel.cancelled() => { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - info!("AOF writer cancelled"); - break; - } - } - } -} - -/// Background per-shard AOF writer task (Option B step 2b). -/// -/// One instance is spawned per shard in `PerShard` layout. Each instance owns -/// `appendonlydir/shard-{shard_id}/moon.aof.{seq}.incr.aof` exclusively — no -/// other writer touches that file, so there is no per-file locking. -/// -/// Differences from [`aof_writer_task`] (TopLevel): -/// - Opens `manifest.shard_incr_path(shard_id)` instead of `manifest.incr_path()`. -/// - `Rewrite`/`RewriteSharded` variants are rejected (logged + dropped). -/// The legacy single-writer rewrite enum has no meaning when each shard -/// owns its own files; per-shard BGREWRITEAOF lands in RFC step 6. -/// - Refuses to start if the loaded manifest's layout is `TopLevel` — the -/// spawn site (step 2f) must only invoke this task body for `PerShard` -/// layouts. Mismatch is a programmer error. -/// -/// Wait/timeout/corruption semantics for manifest loading match the existing -/// `aof_writer_task` (60s bounded wait, hard fail on corrupt manifest). -/// Test-only torn-write injection for `per_shard_aof_writer_task`: when set to a -/// nonzero `N`, the `N`-th `Append` received by a tokio per-shard writer writes -/// its header and then simulates a payload write failure, exercising the -/// `write_error` latch. `0` disables. Atomic (not an env var) because -/// `std::env::set_var` is `unsafe` under edition 2024. Compiled out of release. -#[cfg(all(test, feature = "runtime-tokio"))] -pub(crate) static TEST_FAIL_WRITE_AT: std::sync::atomic::AtomicUsize = - std::sync::atomic::AtomicUsize::new(0); - -pub async fn per_shard_aof_writer_task( - rx: channel::MpscReceiver, - base_dir: PathBuf, - shard_id: u16, - fsync: FsyncPolicy, - cancel: CancellationToken, -) { - #[cfg(feature = "runtime-tokio")] - { - use crate::persistence::aof_manifest::{AofLayout, AofManifest}; - use tokio::io::AsyncWriteExt; - - // Wait for main.rs recovery to create/load the manifest. - let manifest_wait_start = Instant::now(); - const MANIFEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); - let manifest = loop { - if cancel.is_cancelled() { - info!( - "AOF writer shard {}: cancelled while waiting for manifest", - shard_id - ); - return; - } - if manifest_wait_start.elapsed() > MANIFEST_TIMEOUT { - error!( - "AOF writer shard {}: manifest not found at {} after {:?}. Writer exiting.", - shard_id, - base_dir.display(), - MANIFEST_TIMEOUT, - ); - return; - } - match AofManifest::load(&base_dir) { - Ok(Some(m)) => break m, - Ok(None) => { - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - } - Err(e) => { - error!( - "AOF writer shard {}: manifest corrupt at {}: {}. Persistence disabled.", - shard_id, - base_dir.display(), - e - ); - return; - } - } - }; - - if manifest.layout != AofLayout::PerShard { - error!( - "AOF writer shard {}: layout is {:?}, expected PerShard. \ - per_shard_aof_writer_task should only be spawned for PerShard layouts. \ - Writer exiting.", - shard_id, manifest.layout - ); - return; - } - if (shard_id as usize) >= manifest.shards.len() { - error!( - "AOF writer shard {}: out of range for manifest with {} shards. Writer exiting.", - shard_id, - manifest.shards.len() - ); - return; - } - - let incr_path = manifest.shard_incr_path(shard_id); - // Ensure shard-{N}/ exists. The manifest constructor for PerShard - // already creates these, but be defensive — a manual deletion or - // a manifest written by an older binary could leave them missing. - if let Some(parent) = incr_path.parent() { - if let Err(e) = tokio::fs::create_dir_all(parent).await { - error!( - "AOF writer shard {}: failed to create dir {}: {}", - shard_id, - parent.display(), - e - ); - return; - } - } - let file: tokio::fs::File = match tokio::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&incr_path) - .await - { - Ok(f) => f, - Err(e) => { - error!( - "AOF writer shard {}: failed to open incr {}: {}", - shard_id, - incr_path.display(), - e - ); - return; - } - }; - info!( - "AOF writer shard {}: seq {}, incr={}", - shard_id, - manifest.seq, - incr_path.display() - ); - - let mut writer = tokio::io::BufWriter::new(file); - let mut last_fsync = Instant::now(); - // (No `interval` here: the EverySec flush deadline is enforced by the - // timeout-bounded recv in the loop below, which wakes at least every - // 200ms regardless of message traffic. A long-lived `interval.tick()` - // select arm is fairness-starvable under sustained writes and proved - // unreliable when idle on this dedicated current-thread writer runtime.) - - // Test-only fault injection: if MOON_TEST_AOF_FSYNC_FAIL=1 is set in - // the environment at writer task startup, every AppendSync ack resolves - // as FsyncFailed instead of Synced. This lets integration tests exercise - // the AOF_FSYNC_ERR response path without requiring a real disk error. - // The env var is read once here (not per-message) so it costs zero on the - // hot path in production deployments where the var is absent. - let fail_fsync_for_test = std::env::var("MOON_TEST_AOF_FSYNC_FAIL").as_deref() == Ok("1"); - - // Torn-write latch: once any write to this incr file fails partway - // (e.g. the header landed but the payload did not), we must NEVER write - // another record — a lone orphaned header makes the framed reader - // misinterpret the next record's bytes as the orphan's payload, - // corrupting every record after it on replay. Stay latched for the - // writer's lifetime; recovery replays the clean prefix and a rewrite - // starts a fresh file. This mirrors the single-file (line ~1467) and - // monoio per-shard (line ~2125) writers, which already carry the latch. - let mut write_error = false; - // Test-only fault injection (no env var: edition-2024 set_var is unsafe). - // When `TEST_FAIL_WRITE_AT` is the ordinal of an incoming Append, that - // append writes its header then simulates a payload failure, exercising - // the latch. Compiled out of production builds. - #[cfg(test)] - let mut test_append_ordinal: usize = 0; - - loop { - tokio::select! { - // Bounded recv (EverySec durability): wake at least every 200ms - // even when idle so the flush deadline after this select! is - // honored within its 1s bound. flume's recv future is drop-safe - // on the Elapsed branch (no message consumed on timeout); the - // Ok(Ok(msg)) path below captures the message with no loss. - r = tokio::time::timeout( - std::time::Duration::from_millis(200), - rx.recv_async(), - ) => { - // On Elapsed (timeout) `r` is Err: skip the match and fall - // through to the EverySec deadline check after this select!. - if let Ok(msg) = r { - match msg { - // PerShard writer (tokio): per RFC § 2 Rule 1 the on-disk - // format is `[u64 lsn LE][u32 len LE][RESP bytes]`. Header - // is written sequentially with the body — both calls land - // in the same BufWriter so this is one syscall under load. - Ok(AofMessage::Append { lsn, bytes: data }) => { - // Latch: stream already torn — drop silently (Append - // is fire-and-forget; no ack channel to notify). - if write_error { - continue; - } - #[cfg(test)] - { - test_append_ordinal += 1; - let fail_at = TEST_FAIL_WRITE_AT - .load(std::sync::atomic::Ordering::Relaxed); - if fail_at != 0 && fail_at == test_append_ordinal { - // Reproduce a torn write: header lands, payload - // "fails". The orphaned header is flushed so the - // on-disk effect matches the real I/O-error case. - let mut header = [0u8; 12]; - header[..8].copy_from_slice(&lsn.to_le_bytes()); - header[8..] - .copy_from_slice(&(data.len() as u32).to_le_bytes()); - let _ = writer.write_all(&header).await; - let _ = writer.flush().await; - error!( - "AOF shard {}: injected torn write after header (test)", - shard_id - ); - write_error = true; - continue; - } - } - let mut header = [0u8; 12]; - header[..8].copy_from_slice(&lsn.to_le_bytes()); - header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); - if let Err(e) = writer.write_all(&header).await { - error!("AOF header write error shard {}: {}", shard_id, e); - write_error = true; - continue; - } - if let Err(e) = writer.write_all(&data).await { - error!("AOF write error shard {}: {}", shard_id, e); - write_error = true; - continue; - } - if matches!(fsync, FsyncPolicy::Always) { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - } - } - // AppendSync (tokio + PerShard): framed write + fsync + ack. - Ok(AofMessage::AppendSync { lsn, bytes: data, ack }) => { - // Latch: stream already torn — refuse to write more and - // report failure so the caller does not hang to the F2 - // timeout and does not ack a write into a corrupt stream. - if write_error { - let _ = ack.send(AofAck::WriteFailed); - continue; - } - let mut header = [0u8; 12]; - header[..8].copy_from_slice(&lsn.to_le_bytes()); - header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); - if let Err(e) = writer.write_all(&header).await { - error!( - "AOF AppendSync header write error shard {}: {}", - shard_id, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; - } - if let Err(e) = writer.write_all(&data).await { - error!( - "AOF AppendSync write error shard {}: {}", - shard_id, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; - } - // Test-only: skip real fsync and return FsyncFailed - // immediately when the fault-injection env var is set. - if fail_fsync_for_test { - let _ = ack.send(AofAck::FsyncFailed); - continue; - } - if let Err(e) = writer.flush().await { - error!( - "AOF AppendSync flush error shard {}: {}", - shard_id, e - ); - let _ = ack.send(AofAck::FsyncFailed); - continue; - } - if let Err(e) = writer.get_ref().sync_data().await { - error!( - "AOF AppendSync sync_data error shard {}: {}", - shard_id, e - ); - let _ = ack.send(AofAck::FsyncFailed); - continue; - } - let _ = ack.send(AofAck::Synced); - } - Ok(AofMessage::Rewrite(_)) | Ok(AofMessage::RewriteSharded(_)) => { - warn!( - "AOF writer shard {}: received Rewrite/RewriteSharded — \ - not applicable in PerShard layout, dropped.", - shard_id - ); - } - // [F6] Per-shard rewrite (tokio): reuse the proven - // synchronous fold (`do_rewrite_per_shard`) verbatim, so - // the exactly-once invariant carries over unchanged. This - // writer runs on a DEDICATED std::thread (block_on_local, - // main.rs) — not a shared tokio worker — so executing the - // blocking fold here cannot starve the runtime. We flush - // the BufWriter (its `into_inner` does NOT flush) so any - // buffered appends are durable in the OLD incr, convert - // `tokio::fs::File` -> `std::fs::File` for the sync fold, - // then wrap the (reopened) file back into the BufWriter. - Ok(AofMessage::RewritePerShard { shard_dbs, coord }) => { - if let Err(e) = writer.flush().await { - error!( - "F6 tokio per-shard rewrite: shard {} pre-fold flush \ - failed: {}. Aborting; old generation stays authoritative.", - shard_id, e - ); - coord.mark_failed(); - coord.shard_done(); - } else { - // `into_std().await` waits for in-flight ops and is - // infallible; the buffer is already flushed above. - let mut sf = writer.into_inner().into_std().await; - let res = do_rewrite_per_shard( - shard_id, &shard_dbs, &mut sf, &rx, &coord, - ); - // `sf` is left pointing at the committed generation - // by the fold's internal barrier: NEW incr on - // success, OLD incr on abort (phase-8 rollback) or - // on a pre-reopen error. The fold's ShardDoneGuard - // already performed `shard_done` for every exit, so - // we MUST NOT decrement again here. Wrap `sf` back so - // the writer stays valid either way. - writer = - tokio::io::BufWriter::new(tokio::fs::File::from_std(sf)); - if let Err(e) = res { - error!( - "F6 tokio per-shard rewrite: shard {} fold failed: {}. \ - Rewrite aborted by the fold guard; old generation \ - stays authoritative.", - shard_id, e - ); - } - } - } - Ok(AofMessage::Shutdown) | Err(_) => { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - info!("AOF writer shard {} shutting down", shard_id); - break; - } - } - } - } - _ = cancel.cancelled() => { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - info!("AOF writer shard {} cancelled", shard_id); - break; - } - } - // EverySec deadline — checked after EVERY wake (message OR timeout), - // so it is NOT subject to select! fairness and holds the 1s bound - // under sustained writes as well as when idle. (The old long-lived - // `interval.tick()` arm could be starved by the always-ready recv - // arm under load, leaving >1s of writes buffered in the BufWriter - // and lost on SIGKILL — the COMPOSE crash-matrix failure.) - if fsync == FsyncPolicy::EverySec - && last_fsync.elapsed() >= std::time::Duration::from_secs(1) - { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - last_fsync = Instant::now(); - } - } - } - - #[cfg(feature = "runtime-monoio")] - { - use crate::persistence::aof_manifest::{AofLayout, AofManifest}; - use std::io::Write; - - let manifest_wait_start = Instant::now(); - const MANIFEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); - let manifest = loop { - if cancel.is_cancelled() { - info!( - "AOF writer shard {}: cancelled while waiting for manifest", - shard_id - ); - return; - } - if manifest_wait_start.elapsed() > MANIFEST_TIMEOUT { - error!( - "AOF writer shard {}: manifest not found at {} after {:?}. Writer exiting.", - shard_id, - base_dir.display(), - MANIFEST_TIMEOUT, - ); - return; - } - match AofManifest::load(&base_dir) { - Ok(Some(m)) => break m, - Ok(None) => { - std::thread::sleep(std::time::Duration::from_millis(50)); - } - Err(e) => { - error!( - "AOF writer shard {}: manifest corrupt at {}: {}. Persistence disabled.", - shard_id, - base_dir.display(), - e - ); - return; - } - } - }; - - if manifest.layout != AofLayout::PerShard { - error!( - "AOF writer shard {}: layout is {:?}, expected PerShard. Writer exiting.", - shard_id, manifest.layout - ); - return; - } - if (shard_id as usize) >= manifest.shards.len() { - error!( - "AOF writer shard {}: out of range for manifest with {} shards. Writer exiting.", - shard_id, - manifest.shards.len() - ); - return; - } - - let incr_path = manifest.shard_incr_path(shard_id); - if let Some(parent) = incr_path.parent() { - if let Err(e) = std::fs::create_dir_all(parent) { - error!( - "AOF writer shard {}: failed to create dir {}: {}", - shard_id, - parent.display(), - e - ); - return; - } - } - let mut file = match std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&incr_path) - { - Ok(f) => f, - Err(e) => { - error!( - "AOF writer shard {}: failed to open incr {}: {}", - shard_id, - incr_path.display(), - e - ); - return; - } - }; - info!( - "AOF writer shard {}: seq {}, incr={}", - shard_id, - manifest.seq, - incr_path.display() - ); - - let mut last_fsync = Instant::now(); - let mut write_error = false; - // Test-only fault injection: if MOON_TEST_AOF_FSYNC_FAIL=1 is set in - // the environment at writer task startup, every AppendSync ack resolves - // as FsyncFailed instead of Synced. Read once before the loop so there - // is zero cost in production deployments where the var is absent. - let fail_fsync_for_test = std::env::var("MOON_TEST_AOF_FSYNC_FAIL").as_deref() == Ok("1"); - - loop { - match rx.recv() { - // AppendSync (monoio + PerShard): framed write + fsync + ack. - Ok(AofMessage::AppendSync { - lsn, - bytes: data, - ack, - }) => { - if write_error { - let _ = ack.send(AofAck::WriteFailed); - continue; - } - let mut header = [0u8; 12]; - header[..8].copy_from_slice(&lsn.to_le_bytes()); - header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); - if let Err(e) = file.write_all(&header) { - error!( - "AOF AppendSync header write failed shard {} (seq {}): {}", - shard_id, manifest.seq, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; - } - if let Err(e) = file.write_all(&data) { - error!( - "AOF AppendSync write failed shard {} (seq {}): {}", - shard_id, manifest.seq, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; - } - // Test-only: skip real fsync and return FsyncFailed - // immediately when the fault-injection env var is set. - if fail_fsync_for_test { - let _ = ack.send(AofAck::FsyncFailed); - continue; - } - let t = Instant::now(); - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!( - "AOF AppendSync sync failed shard {} (seq {}): {}", - shard_id, manifest.seq, e - ); - write_error = true; - let _ = ack.send(AofAck::FsyncFailed); - } else { - crate::admin::metrics_setup::record_aof_fsync( - t.elapsed().as_micros() as u64 - ); - let _ = ack.send(AofAck::Synced); - } - } - // PerShard writer (monoio): framed `[u64 lsn LE][u32 len LE][RESP]`. - // See the tokio twin above for format rationale. - Ok(AofMessage::Append { lsn, bytes: data }) => { - if write_error { - continue; - } - let mut header = [0u8; 12]; - header[..8].copy_from_slice(&lsn.to_le_bytes()); - header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); - if let Err(e) = file.write_all(&header) { - error!( - "AOF header write failed shard {} (seq {}): {}. Persistence degraded.", - shard_id, manifest.seq, e - ); - write_error = true; - continue; - } - if let Err(e) = file.write_all(&data) { - error!( - "AOF write failed shard {} (seq {}): {}. Persistence degraded.", - shard_id, manifest.seq, e - ); - write_error = true; - continue; - } - match fsync { - FsyncPolicy::Always => { - let t = Instant::now(); - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!( - "AOF sync failed shard {} (seq {}, always): {}", - shard_id, manifest.seq, e - ); - write_error = true; - } else { - crate::admin::metrics_setup::record_aof_fsync( - t.elapsed().as_micros() as u64, - ); - } - } - FsyncPolicy::EverySec => { - if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { - let t = Instant::now(); - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!( - "AOF sync failed shard {} (seq {}, everysec): {}", - shard_id, manifest.seq, e - ); - } else { - crate::admin::metrics_setup::record_aof_fsync( - t.elapsed().as_micros() as u64, - ); - last_fsync = Instant::now(); - } - } - } - FsyncPolicy::No => {} - } - } - Ok(AofMessage::Rewrite(_)) | Ok(AofMessage::RewriteSharded(_)) => { - warn!( - "AOF writer shard {}: received Rewrite/RewriteSharded — \ - not applicable in PerShard layout (use per-shard \ - BGREWRITEAOF), dropped.", - shard_id - ); - } - // [F6] Per-shard rewrite fan-out (monoio). Fold THIS shard, - // then signal the coordinator; the last shard commits the - // manifest. On error the old generation stays authoritative - // (advance_shard did not commit the seq). - Ok(AofMessage::RewritePerShard { shard_dbs, coord }) => { - if let Err(e) = - do_rewrite_per_shard(shard_id, &shard_dbs, &mut file, &rx, &coord) - { - // The fold's ShardDoneGuard already marked the rewrite - // failed and decremented on this error exit (committing - // new_seq with a shard missing its new base would break - // recovery), so do NOT decrement again here. `file` is left - // on the OLD incr (error exits are pre-reopen). - error!( - "F6 per-shard rewrite: shard {} fold failed: {}. \ - Rewrite aborted by the fold guard; old generation \ - stays authoritative.", - shard_id, e - ); - } - } - Ok(AofMessage::Shutdown) | Err(_) => { - if !write_error { - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!( - "AOF final sync failed shard {} (seq {}): {}", - shard_id, manifest.seq, e - ); - } - } - info!( - "AOF writer shard {} shutting down (monoio, seq {})", - shard_id, manifest.seq - ); - break; - } - } - } - } -} - -/// Replay an AOF file by parsing RESP commands and dispatching them. -/// -/// Returns the number of commands successfully replayed. -/// -/// **Corruption recovery:** On mid-stream parse errors, logs a warning with the -/// byte offset, skips to the next RESP array marker (`*`), and continues replay. -/// At EOF, reports total corrupted entries skipped. Truncated tails are handled -/// gracefully (warn + stop). -pub fn replay_aof( - databases: &mut [Database], - path: &Path, - engine: &dyn CommandReplayEngine, -) -> Result { - let data = std::fs::read(path)?; - if data.is_empty() { - return Ok(0); - } - - // Detect RDB preamble: if the file starts with "MOON" magic, load the binary - // RDB section first, then replay any RESP commands appended after it. - let (rdb_keys, resp_start) = if data.starts_with(b"MOON") { - match crate::persistence::rdb::load_from_bytes(databases, &data) { - Ok((keys, consumed)) => { - info!( - "AOF RDB preamble loaded: {} keys ({} bytes)", - keys, consumed - ); - (keys, consumed) - } - Err(e) => { - // Data starts with MOON magic — it IS RDB format. - // Falling back to RESP would parse garbage. Propagate the error. - return Err(e); - } - } - } else { - (0, 0) - }; - - // If the entire file was RDB (no RESP tail), we're done - if resp_start >= data.len() { - return Ok(rdb_keys); - } - - let resp_data = &data[resp_start..]; - let total_len = resp_data.len(); - let mut buf = BytesMut::from(resp_data); - let config = ParseConfig::default(); - let mut selected_db: usize = 0; - let mut count: usize = 0; - let mut corruption_count: usize = 0; - - loop { - if buf.is_empty() { - break; - } - - match parse::parse(&mut buf, &config) { - Ok(Some(frame)) => { - // Extract command name and args, then dispatch - let (cmd, cmd_args) = match &frame { - Frame::Array(arr) if !arr.is_empty() => { - let name = match &arr[0] { - Frame::BulkString(s) => s.as_ref(), - Frame::SimpleString(s) => s.as_ref(), - _ => { - count += 1; - continue; - } - }; - (name as &[u8], &arr[1..]) - } - _ => { - count += 1; - continue; - } - }; - engine.replay_command(databases, cmd, cmd_args, &mut selected_db); - count += 1; - } - Ok(None) => { - // Incomplete frame at end of file - truncated AOF - if !buf.is_empty() { - let offset = total_len - buf.len(); - warn!( - "AOF truncated: {} unparseable bytes at offset {} (end of file)", - buf.len(), - offset - ); - } - break; - } - Err(e) => { - let error_offset = total_len - buf.len(); - warn!( - "AOF parse error at byte offset {} after {} commands: {}. Attempting skip.", - error_offset, count, e - ); - corruption_count += 1; - - // Skip past the corrupt byte(s) to the next RESP array marker ('*') - // Always discard at least 1 byte to guarantee forward progress. - let _ = buf.split_to(1); - if let Some(pos) = buf.iter().position(|&b| b == b'*') { - let _ = buf.split_to(pos); - } else if buf.is_empty() { - break; - } else { - // No more RESP array markers found; stop replay - warn!( - "AOF: no recoverable RESP frame found after offset {}; stopping", - error_offset - ); - break; - } - } - } - } - - if corruption_count > 0 { - warn!( - "AOF replay completed with {} corrupted entries skipped, {} commands replayed", - corruption_count, count - ); - } - - Ok(rdb_keys + count) -} - -/// Generate synthetic RESP commands from the current database state for AOF rewriting. -/// -/// Produces commands for all 5 data types plus PEXPIRE for keys with TTL. -#[allow(dead_code)] // Retained for RESP-only AOF rewrite fallback and testing -pub fn generate_rewrite_commands(databases: &[Database]) -> BytesMut { - let mut buf = BytesMut::new(); - let now_ms = current_time_ms(); - - for (db_idx, db) in databases.iter().enumerate() { - let base_ts = db.base_timestamp(); - let data = db.data(); - if data.is_empty() { - continue; - } - - // Generate SELECT if not db 0 - if db_idx > 0 { - let select_frame = Frame::Array(framevec![ - Frame::BulkString(Bytes::from_static(b"SELECT")), - Frame::BulkString(Bytes::from(db_idx.to_string())), - ]); - serialize::serialize(&select_frame, &mut buf); - } - - for (key, entry) in data { - // Skip expired entries - if entry.is_expired_at(base_ts, now_ms) { - continue; - } - - match entry.value.as_redis_value() { - RedisValueRef::String(val) => { - let frame = Frame::Array(framevec![ - Frame::BulkString(Bytes::from_static(b"SET")), - Frame::BulkString(key.to_bytes()), - Frame::BulkString(Bytes::copy_from_slice(val)), - ]); - serialize::serialize(&frame, &mut buf); - } - RedisValueRef::Hash(map) => { - if map.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"HSET")), - Frame::BulkString(key.to_bytes()), - ]; - for (field, val) in map.iter() { - args.push(Frame::BulkString(field.clone())); - args.push(Frame::BulkString(val.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - // Phase 200: for HashWithTtl we emit two RESP frames per key. - // 1. `HSET key f1 v1 f2 v2 ...` rebuilds the hash body. - // 2. `HPEXPIREAT key abs_ms FIELDS 1 field` for every entry - // in the TTL sidecar — one per TTL'd field for clarity - // (BGREWRITEAOF is rare; per-field framing keeps the - // replay shim simple, see `persistence::replay`). - RedisValueRef::HashWithTtl { fields, ttls, .. } => { - if fields.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"HSET")), - Frame::BulkString(key.to_bytes()), - ]; - for (field, val) in fields.iter() { - args.push(Frame::BulkString(field.clone())); - args.push(Frame::BulkString(val.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - - for (field, ttl_ms) in ttls.iter() { - let mut ttl_args = vec![ - Frame::BulkString(Bytes::from_static(b"HPEXPIREAT")), - Frame::BulkString(key.to_bytes()), - Frame::BulkString(Bytes::copy_from_slice( - ttl_ms.to_string().as_bytes(), - )), - Frame::BulkString(Bytes::from_static(b"FIELDS")), - Frame::BulkString(Bytes::from_static(b"1")), - Frame::BulkString(field.clone()), - ]; - ttl_args.shrink_to_fit(); - serialize::serialize(&Frame::Array(ttl_args.into()), &mut buf); - } - } - RedisValueRef::HashListpack(lp) => { - let map = lp.to_hash_map(); - if map.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"HSET")), - Frame::BulkString(key.to_bytes()), - ]; - for (field, val) in &map { - args.push(Frame::BulkString(field.clone())); - args.push(Frame::BulkString(val.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - RedisValueRef::List(list) => { - if list.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"RPUSH")), - Frame::BulkString(key.to_bytes()), - ]; - for elem in list.iter() { - args.push(Frame::BulkString(elem.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - RedisValueRef::ListListpack(lp) => { - let list = lp.to_vec_deque(); - if list.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"RPUSH")), - Frame::BulkString(key.to_bytes()), - ]; - for elem in &list { - args.push(Frame::BulkString(elem.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - RedisValueRef::Set(set) => { - if set.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"SADD")), - Frame::BulkString(key.to_bytes()), - ]; - for member in set.iter() { - args.push(Frame::BulkString(member.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - RedisValueRef::SetListpack(lp) => { - let set = lp.to_hash_set(); - if set.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"SADD")), - Frame::BulkString(key.to_bytes()), - ]; - for member in &set { - args.push(Frame::BulkString(member.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - RedisValueRef::SetIntset(is) => { - let set = is.to_hash_set(); - if set.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"SADD")), - Frame::BulkString(key.to_bytes()), - ]; - for member in &set { - args.push(Frame::BulkString(member.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - RedisValueRef::SortedSet { members, .. } - | RedisValueRef::SortedSetBPTree { members, .. } => { - if members.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"ZADD")), - Frame::BulkString(key.to_bytes()), - ]; - for (member, score) in members.iter() { - args.push(Frame::BulkString(Bytes::from(score.to_string()))); - args.push(Frame::BulkString(member.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - RedisValueRef::SortedSetListpack(lp) => { - let pairs: Vec<_> = lp.iter_pairs().collect(); - if pairs.is_empty() { - continue; - } - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"ZADD")), - Frame::BulkString(key.to_bytes()), - ]; - for (member_entry, score_entry) in &pairs { - let score_bytes = score_entry.as_bytes(); - args.push(Frame::BulkString(Bytes::from(score_bytes))); - args.push(Frame::BulkString(Bytes::from(member_entry.as_bytes()))); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - RedisValueRef::Stream(stream) => { - for (id, fields) in &stream.entries { - let mut args = vec![ - Frame::BulkString(Bytes::from_static(b"XADD")), - Frame::BulkString(key.to_bytes()), - Frame::BulkString(id.to_bytes()), - ]; - for (field, value) in fields { - args.push(Frame::BulkString(field.clone())); - args.push(Frame::BulkString(value.clone())); - } - serialize::serialize(&Frame::Array(args.into()), &mut buf); - } - } - } - - // Generate PEXPIRE for keys with TTL - if entry.has_expiry() { - let exp_ms = entry.expires_at_ms(base_ts); - if exp_ms > now_ms { - let remaining_ms = exp_ms - now_ms; - let pexpire_frame = Frame::Array(framevec![ - Frame::BulkString(Bytes::from_static(b"PEXPIRE")), - Frame::BulkString(key.to_bytes()), - Frame::BulkString(Bytes::from(remaining_ms.to_string())), - ]); - serialize::serialize(&pexpire_frame, &mut buf); - } - } - } - } - - buf -} - -/// Snapshot databases and generate compacted AOF commands. -/// -/// Shared by both the async (tokio) and sync (monoio) rewrite paths. -#[allow(dead_code)] -fn snapshot_and_generate(db: &SharedDatabases) -> BytesMut { - let snapshot: Vec<(Vec<(CompactKey, Entry)>, u32)> = db - .iter() - .map(|lock| { - let guard = lock.read(); - let base_ts = guard.base_timestamp(); - let entries = guard - .data() - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - (entries, base_ts) - }) - .collect(); - - let mut temp_dbs: Vec = Vec::with_capacity(snapshot.len()); - for (entries, _base_ts) in &snapshot { - let mut db = Database::new(); - for (key, entry) in entries { - db.set(key.to_bytes(), entry.clone()); - } - temp_dbs.push(db); - } - - generate_rewrite_commands(&temp_dbs) -} - -/// Drain any queued `AofMessage::Append` messages to the current incr file. -/// -/// Called during rewrite to catch in-flight appends that handlers sent before -/// the writer thread could enter the rewrite routine. Messages of other variants -/// are dropped silently (duplicate rewrites while a rewrite is in progress) or -/// returned via the flag for Shutdown (caller is responsible for honoring it -/// after the rewrite completes). -#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] -#[derive(Default)] -struct DrainOutcome { - drained: usize, - shutdown_requested: bool, - /// AppendSync ack senders for entries drained during a rewrite. Under - /// `appendfsync=always` the client must NOT be told `Synced` until the - /// post-drain boundary `sync_data()` makes those bytes durable, so the acks - /// are parked here and resolved by [`fulfill_acks`](Self::fulfill_acks) AFTER - /// the caller's fsync — `Synced` on success, `FsyncFailed` on failure. Acking - /// inside the drain (the old behaviour) reports a write durable before the - /// boundary fsync, so a crash in that window loses an entry the client was - /// told was safe. (Issue #140.) - pending_acks: Vec>, -} - -#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] -impl DrainOutcome { - /// Resolve every parked AppendSync ack after the rewrite-boundary fsync. - /// `synced=true` → `Synced`; `false` → `FsyncFailed`. A fresh `AofAck` is - /// built per sender so `AofAck` needs no `Copy`/`Clone`. Drains the vec so a - /// second call is a no-op. - fn fulfill_acks(&mut self, synced: bool) { - for tx in std::mem::take(&mut self.pending_acks) { - let _ = tx.send(if synced { - AofAck::Synced - } else { - AofAck::FsyncFailed - }); - } - } -} - -/// Fsync `file` at a rewrite drain boundary, then resolve the drained batch's -/// parked AppendSync acks against the result: `Synced` on success, or -/// `FsyncFailed` + propagate the IO error on failure. Centralizes the issue-#140 -/// durability ordering (ack strictly AFTER the bytes are durable) for every -/// `do_rewrite_*` drain site. -#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] -fn sync_and_fulfill_drain( - outcome: &mut DrainOutcome, - file: &mut std::fs::File, - incr_path: PathBuf, -) -> Result<(), MoonError> { - match file.sync_data() { - Ok(()) => { - outcome.fulfill_acks(true); - Ok(()) - } - Err(e) => { - // Boundary fsync failed: the drained writes are NOT durable. Tell the - // waiting clients FsyncFailed (never Synced) and propagate the error - // so the rewrite aborts. - outcome.fulfill_acks(false); - Err(AofError::Io { - path: incr_path, - source: e, - } - .into()) - } - } -} - -#[cfg(feature = "runtime-monoio")] -fn drain_pending_appends( - rx: &channel::MpscReceiver, - file: &mut std::fs::File, -) -> Result { - use std::io::Write; - let mut outcome = DrainOutcome::default(); - while let Ok(msg) = rx.try_recv() { - match msg { - // BGREWRITEAOF drain runs on the TopLevel writer (monoio) only; - // PerShard rewrite is RFC step 6. Legacy v1 disk format → ignore lsn. - AofMessage::Append { - bytes: data, - lsn: _, - } => { - file.write_all(&data).map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; - outcome.drained += 1; - } - // AppendSync during a rewrite drain: bytes are written and counted, - // but the ack is PARKED until the caller's post-drain boundary fsync - // (issue #140) — acking `Synced` here would report durability before - // the bytes are fsynced. If the write itself fails the `?` propagates - // the error and the parked ack is dropped with the outcome — the - // caller observes `RecvError`, which it treats as failure. - AofMessage::AppendSync { - bytes: data, - lsn: _, - ack, - } => { - file.write_all(&data).map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; - outcome.drained += 1; - outcome.pending_acks.push(ack); - } - AofMessage::Shutdown => { - outcome.shutdown_requested = true; - } - AofMessage::Rewrite(_) - | AofMessage::RewriteSharded(_) - | AofMessage::RewritePerShard { .. } => { - // Already rewriting — drop redundant request. - } - } - } - Ok(outcome) -} - -/// [F6] Drain a per-shard writer's queued appends into its OLD incr file using -/// the framed `[u64 lsn LE][u32 len LE][RESP bytes]` on-disk format that -/// per-shard recovery expects. -/// -/// This is the per-shard twin of [`drain_pending_appends`] (which writes the -/// legacy TopLevel raw-RESP format). Correctness depends on the framing -/// matching `replay_per_shard`'s reader — an unframed write here would make the -/// drained appends unparseable on restart. -#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] -fn drain_pending_appends_framed( - rx: &channel::MpscReceiver, - file: &mut std::fs::File, -) -> Result { - use std::io::Write; - let mut outcome = DrainOutcome::default(); - let write_framed = |file: &mut std::fs::File, lsn: u64, data: &[u8]| -> std::io::Result<()> { - let mut header = [0u8; 12]; - header[..8].copy_from_slice(&lsn.to_le_bytes()); - header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); - file.write_all(&header)?; - file.write_all(data) - }; - while let Ok(msg) = rx.try_recv() { - match msg { - AofMessage::Append { lsn, bytes: data } => { - write_framed(file, lsn, &data).map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; - outcome.drained += 1; - } - AofMessage::AppendSync { - lsn, - bytes: data, - ack, - } => { - write_framed(file, lsn, &data).map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; - outcome.drained += 1; - // Park the ack until the caller's post-drain boundary fsync - // (issue #140); resolved Synced/FsyncFailed by - // `sync_and_fulfill_drain`. Mirrors `drain_pending_appends`. - outcome.pending_acks.push(ack); - } - AofMessage::Shutdown => { - outcome.shutdown_requested = true; - } - AofMessage::Rewrite(_) - | AofMessage::RewriteSharded(_) - | AofMessage::RewritePerShard { .. } => { - // Already rewriting this shard — drop redundant request. - } - } - } - Ok(outcome) -} - -/// [F6] Per-shard rewrite fold (monoio). Run by a single per-shard writer for -/// ITS shard only; the manifest commit is coordinated across all shards by the -/// shared [`PerShardRewriteCoord`]. -/// -/// Correctness ordering (prevents double-apply of non-idempotent commands like -/// INCR after the rewrite) — identical discipline to [`do_rewrite_sharded`], -/// scoped to one shard: -/// -/// 1. Drain queued appends into the OLD incr (framed) and fsync. -/// 2. Acquire write locks on this shard's databases. -/// 3. Re-drain appends that arrived between phase 1 and the lock, into OLD -/// incr, and fsync. -/// 4. Snapshot this shard's databases under the locks. -/// 5. Release the locks before the expensive base-RDB write. -/// 6. Write the new base + new (empty) incr at `coord.new_seq` via -/// `advance_shard` (which does NOT bump `manifest.seq`), then reopen -/// `file` to the new incr. Subsequent appends land in the new generation. -/// 7. Signal completion to the coordinator; the last shard commits the -/// manifest (single seq flip) and prunes the old generation. -/// -/// Until step 7's commit, the on-disk manifest still resolves to the old seq, -/// so a crash anywhere in steps 1-6 recovers the intact old generation. -/// -/// # Cross-thread exactly-once invariant (load-bearing) -/// -/// This fold runs on the per-shard *writer* thread, which is distinct from the -/// shard event-loop thread that applies commands. Exactly-once across the -/// rewrite boundary depends on a single ordering fact: the live write path -/// enqueues each command's AOF append **inside** the same `RwLock` -/// write guard under which it mutated the db (see `spsc_handler.rs`: -/// `wal_append_and_fanout` is called before `drop(guard)`). Phase 2 here -/// acquires those *same* locks (`all_shard_dbs()[sidx]` is -/// `ShardDatabases::shards[sidx]`, the exact `RwLock`s `write_db` locks), so -/// RwLock mutual exclusion forces the order -/// `enqueue → guard-release → fold-acquire → mid-drain(phase 3)`. Hence every -/// INCR whose mutation lands in the phase-4 snapshot had its append drained -/// into the OLD incr (then pruned at commit) — never replayed on top of the -/// new base. Were the append enqueued *after* the guard drop, a snapshot would -/// capture the mutation while its append still raced toward the NEW incr → -/// double-apply. The in-guard append is therefore the invariant; do not move it. -/// -/// This also assumes the `RwLock`-backed `ShardDatabases` is the *live* store. -/// It is, because the thread-local `ShardSlice` fast path is dead code until -/// Phase 4 wires `init_shard` (`is_initialized()` is always false today). A -/// future Phase 4 that makes ShardSlice live MUST revisit this fold: the writer -/// thread cannot lock another thread's `!Send` `Rc>`, so the -/// per-shard rewrite would need a different snapshot-coordination mechanism. -/// -/// # Known limitation — channel saturation during the fold -/// -/// Exactly-once holds *absent append-channel saturation during the fold*. While -/// this function runs (phases 2-6, including the base-RDB serialize + write + -/// fsync of phase 6, which is hundreds of ms on a large shard) the writer is -/// NOT in its recv loop, so it is not draining the bounded -/// `mpsc_bounded::(10_000)` append channel. Post-snapshot appends -/// queue there for the new incr; the event loop enqueues them with -/// `try_send_append` (drop-on-full, return ignored — `spsc_handler.rs`). Under -/// *sustained concurrent* writes on a large dataset, > 10_000 appends can pile -/// up during the window and the overflow is silently dropped — lost even on a -/// clean restart (worse than the everysec contract, which only loses on crash). -/// The single-client crash matrix cannot surface this (serialized `redis-cli` -/// never pressures the channel). This window is *pre-existing*: the shipped -/// `do_rewrite_sharded` has the identical non-draining gap. Tracked as a -/// known limitation (F6 is behind `--experimental-per-shard-rewrite`); the fix -/// (keep draining during phase 6, or block-on-full for the rewrite's duration) -/// is a separate scoped task. See `tmp/F6-known-limitations.md`. -#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] -fn do_rewrite_per_shard( - shard_id: u16, - shard_dbs: &crate::shard::shared_databases::ShardDatabases, - file: &mut std::fs::File, - rx: &channel::MpscReceiver, - coord: &PerShardRewriteCoord, -) -> Result<(), MoonError> { - // Panic/early-error safety: guarantees `shard_done` runs on EVERY exit - // (success via `complete()`, `?`-error or panic-unwind via `Drop`). The - // phase-8 `await_outcome` barrier makes that a liveness requirement, so - // callers MUST NOT call `shard_done` after invoking this function — the - // guard owns the single decrement for all exits. See `ShardDoneGuard`. - let guard = ShardDoneGuard::new(coord); - - let sidx = shard_id as usize; - let all_shards = shard_dbs.all_shard_dbs(); - if sidx >= all_shards.len() { - return Err(AofError::RewriteFailed { - detail: format!( - "do_rewrite_per_shard: shard {} out of range ({} shards)", - sidx, - all_shards.len() - ), - } - .into()); - } - - // Phase 1: drain pre-rewrite queued appends into old incr (framed), fsync, - // then resolve their parked AppendSync acks (issue #140). - let mut pre_drain = drain_pending_appends_framed(rx, file)?; - sync_and_fulfill_drain(&mut pre_drain, file, PathBuf::from(""))?; - - // Phase 2: acquire write locks on this shard's db(s) (db_idx ascending). - let shard_locks = &all_shards[sidx]; - let guards: Vec<_> = shard_locks.iter().map(|lock| lock.write()).collect(); - - // Phase 3: drain appends that completed between phase 1 and phase 2, fsync, - // then resolve their parked AppendSync acks (issue #140). - let mut mid_drain = drain_pending_appends_framed(rx, file)?; - sync_and_fulfill_drain(&mut mid_drain, file, PathBuf::from(""))?; - - // Phase 4: snapshot this shard's databases under the locks. - let now_ms = current_time_ms(); - let mut snapshot: Vec<( - Vec<( - crate::storage::compact_key::CompactKey, - crate::storage::entry::Entry, - )>, - u32, - )> = Vec::with_capacity(guards.len()); - for guard in &guards { - let base_ts = guard.base_timestamp(); - let mut entries = Vec::new(); - for (key, entry) in guard.data().iter() { - if !entry.is_expired_at(base_ts, now_ms) { - entries.push((key.clone(), entry.clone())); - } - } - snapshot.push((entries, base_ts)); - } - - // Phase 5: release locks before the expensive disk write. - drop(guards); - - // Phase 6: write new base, advance THIS shard's manifest entry (no seq - // commit), reopen to the new incr. The manifest lock is held only for the - // brief, await-free advance_shard call. - let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot)?; - let new_incr = { - let mut m = coord.manifest.lock(); - m.advance_shard(shard_id, coord.new_seq, &rdb_bytes)? - }; - *file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&new_incr) - .map_err(|e| AofError::Io { - path: new_incr, - source: e, - })?; - - info!( - "F6 per-shard rewrite: shard {} folded (drained {}+{} appends), new seq {}", - shard_id, pre_drain.drained, mid_drain.drained, coord.new_seq - ); - if pre_drain.shutdown_requested || mid_drain.shutdown_requested { - warn!( - "F6 per-shard rewrite: shard {} saw shutdown during rewrite (honored after commit)", - shard_id - ); - } - - // Phase 7: signal completion; the last writer commits + prunes. `complete()` - // performs the single clean `shard_done` and disarms the guard's Drop. - guard.complete(); - - // Phase 8 (barrier-before-resume): block until the terminal writer publishes - // the committed generation, then make sure THIS writer's append file points - // at it. On the happy path committed == new_seq and *file already points at - // new_incr (phase 6) — nothing to do. On an abort/commit-failure the manifest - // kept old_seq and pruned our new_seq incr, so reopen *file onto old_seq's - // incr; otherwise we keep appending into a discarded generation that recovery - // ignores — silent data loss. Replaces the old "RESTART recommended" hazard. - let committed_seq = coord.await_outcome(); - if committed_seq != coord.new_seq { - let committed_incr = coord - .manifest - .lock() - .shard_incr_path_seq(shard_id, committed_seq); - *file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&committed_incr) - .map_err(|e| AofError::Io { - path: committed_incr, - source: e, - })?; - warn!( - "F6 per-shard rewrite ABORTED: shard {} rolled its append file back to \ - committed seq {} (no restart needed)", - shard_id, committed_seq - ); - } - Ok(()) -} - -/// Multi-part rewrite: snapshot single-shard databases → RDB base → advance manifest. -/// -/// Correctness ordering (prevents double-apply of non-idempotent commands like -/// INCR/LPUSH/SADD after rewrite): -/// -/// 1. Drain any queued appends into the OLD incr file and fsync. -/// 2. Acquire write locks on all databases in the shard. This blocks handlers -/// from applying new writes or queueing new appends for the locked dbs. -/// 3. Drain the channel once more — catches appends for writes that the -/// handler completed between step 1 and step 2. -/// 4. Snapshot every database under the write locks. Because no handler can -/// mutate the dbs while we hold the locks, the snapshot is atomic with -/// respect to the post-drain channel state. -/// 5. Release the write locks. New handler writes from here on queue in the -/// channel and will be processed into the NEW incr file after rotation. -/// 6. Write the new base RDB, advance the manifest, reopen the file handle. -/// -/// Invariant: any write captured in the new base is NOT in the new incr file -/// (handlers were blocked between drain and snapshot), and any write NOT in -/// the new base IS in the new incr file (queued after lock release). -#[cfg(feature = "runtime-monoio")] -fn do_rewrite_single( - db: &SharedDatabases, - manifest: &mut crate::persistence::aof_manifest::AofManifest, - file: &mut std::fs::File, - rx: &channel::MpscReceiver, -) -> Result<(), MoonError> { - // Phase 1: drain pre-rewrite queued appends into old incr, fsync, then - // resolve their parked AppendSync acks (issue #140). - let mut pre_drain = drain_pending_appends(rx, file)?; - sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; - - // Phase 2: acquire write locks on every database in the shard. - // Order is consistent (index-ascending) so concurrent callers would - // serialize without deadlock — but in practice only this thread - // acquires multi-db locks. - let guards: Vec<_> = db.iter().map(|lock| lock.write()).collect(); - - // Phase 3: drain any appends the handlers sent between phase 1 and phase 2, - // fsync, then resolve their parked AppendSync acks (issue #140). - let mut mid_drain = drain_pending_appends(rx, file)?; - sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; - - // Phase 4: snapshot under the write locks. No mutation is possible. - let now_ms = current_time_ms(); - let snapshot: Vec<( - Vec<( - crate::storage::compact_key::CompactKey, - crate::storage::entry::Entry, - )>, - u32, - )> = guards - .iter() - .map(|guard| { - let base_ts = guard.base_timestamp(); - let entries: Vec<_> = guard - .data() - .iter() - .filter(|(_, v)| !v.is_expired_at(base_ts, now_ms)) - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - (entries, base_ts) - }) - .collect(); - - // Phase 5: release locks. Handlers resume; new appends queue in the channel - // and will be processed into the new incr after step 6. - drop(guards); - - // Phase 6: write new base, advance manifest, reopen. - let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot)?; - let new_incr = manifest.advance(&rdb_bytes)?; - - *file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&new_incr) - .map_err(|e| AofError::Io { - path: new_incr, - source: e, - })?; - - info!( - "AOF rewrite complete (single): drained {}+{} pre-snapshot appends, seq={}", - pre_drain.drained, mid_drain.drained, manifest.seq - ); - if pre_drain.shutdown_requested || mid_drain.shutdown_requested { - // Caller doesn't currently observe this; logging is the escape hatch. - warn!("AOF writer: shutdown requested during rewrite (will honor on next recv)"); - } - Ok(()) -} - -/// Multi-part rewrite: snapshot all shards → merged RDB base → advance manifest. -/// -/// See [`do_rewrite_single`] for the ordering rationale. The multi-shard variant -/// holds write locks on every (shard, db) pair simultaneously for the duration -/// of the snapshot. This creates a brief global write pause, but it is the only -/// way to guarantee a torn-free snapshot without per-message sequence numbers. -#[cfg(feature = "runtime-monoio")] -fn do_rewrite_sharded( - shard_dbs: &crate::shard::shared_databases::ShardDatabases, - manifest: &mut crate::persistence::aof_manifest::AofManifest, - file: &mut std::fs::File, - rx: &channel::MpscReceiver, -) -> Result<(), MoonError> { - // Phase 1: drain pre-rewrite queued appends into old incr, fsync, then - // resolve their parked AppendSync acks (issue #140). - let mut pre_drain = drain_pending_appends(rx, file)?; - sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; - - // Phase 2: acquire write locks on ALL (shard, db) pairs simultaneously. - // Lock order is (shard_idx, db_idx) ascending — must match anywhere else - // that acquires multiple locks to prevent deadlock (currently no other - // call site does, but the ordering discipline is documented for future - // maintainers). - let all_shards = shard_dbs.all_shard_dbs(); - let mut guards: Vec> = Vec::with_capacity(all_shards.len()); - for shard_locks in all_shards { - let mut shard_guards = Vec::with_capacity(shard_locks.len()); - for lock in shard_locks { - shard_guards.push(lock.write()); - } - guards.push(shard_guards); - } - - // Phase 3: drain appends completed between phase 1 and phase 2, fsync, then - // resolve their parked AppendSync acks (issue #140). - let mut mid_drain = drain_pending_appends(rx, file)?; - sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; - - // Phase 4: snapshot under locks. - let db_count = shard_dbs.db_count(); - let mut merged: Vec<( - Vec<( - crate::storage::compact_key::CompactKey, - crate::storage::entry::Entry, - )>, - u32, - )> = (0..db_count).map(|_| (Vec::new(), 0u32)).collect(); - let now_ms = current_time_ms(); - for shard_guards in &guards { - for (db_idx, guard) in shard_guards.iter().enumerate() { - let base_ts = guard.base_timestamp(); - if merged[db_idx].0.is_empty() { - merged[db_idx].1 = base_ts; - } - for (key, entry) in guard.data().iter() { - if !entry.is_expired_at(base_ts, now_ms) { - merged[db_idx].0.push((key.clone(), entry.clone())); - } - } - } - } - - // Phase 5: release locks before the expensive disk write. - drop(guards); - - // Phase 6: write new base, advance manifest, reopen. - let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&merged)?; - let new_incr = manifest.advance(&rdb_bytes)?; - - *file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&new_incr) - .map_err(|e| AofError::Io { - path: new_incr, - source: e, - })?; - - info!( - "AOF rewrite complete (sharded): drained {}+{} pre-snapshot appends, seq={}", - pre_drain.drained, mid_drain.drained, manifest.seq - ); - if pre_drain.shutdown_requested || mid_drain.shutdown_requested { - warn!("AOF writer: shutdown requested during rewrite (will honor on next recv)"); - } - Ok(()) -} - -/// Rewrite the AOF file with RDB preamble (binary base + empty RESP incremental). -/// -/// Uses the same strategy as Redis 7+ `aof-use-rdb-preamble yes`: -/// the rewritten AOF starts with a full RDB snapshot (compact binary), -/// and new writes are appended as RESP after it. On startup, the loader -/// detects the RDB magic and reads the binary preamble, then switches -/// to RESP parsing for any incremental commands appended after. -#[allow(dead_code)] // Retained for legacy single-file and tokio path -fn rewrite_aof_sync(db: &SharedDatabases, aof_path: &Path) -> Result<(), MoonError> { - // Snapshot under read locks, build temp Database objects for RDB serialization - let snapshot: Vec = db - .iter() - .map(|lock| { - let guard = lock.read(); - let mut temp = Database::new(); - let now_ms = current_time_ms(); - for (k, v) in guard.data().iter() { - if !v.is_expired_at(guard.base_timestamp(), now_ms) { - temp.set(k.to_bytes(), v.clone()); - } - } - temp - }) - .collect(); - - let rdb_bytes = crate::persistence::rdb::save_to_bytes(&snapshot)?; - - let tmp_path = aof_path.with_extension("aof.tmp"); - std::fs::write(&tmp_path, &rdb_bytes).map_err(|e| AofError::Io { - path: tmp_path.clone(), - source: e, - })?; - std::fs::rename(&tmp_path, aof_path).map_err(|e| AofError::RewriteFailed { - detail: format!( - "rename {} -> {}: {}", - tmp_path.display(), - aof_path.display(), - e - ), - })?; - - info!( - "AOF rewrite complete (RDB preamble): {} bytes", - rdb_bytes.len() - ); - Ok(()) -} - -/// Rewrite the AOF in sharded mode with RDB preamble. -/// -/// Merges all shards' databases into a single RDB snapshot, writes it as -/// the AOF base file. New incremental writes are appended as RESP after. -#[allow(dead_code)] -fn rewrite_aof_sharded_sync( - shard_dbs: &crate::shard::shared_databases::ShardDatabases, - aof_path: &Path, -) -> Result<(), MoonError> { - let db_count = shard_dbs.db_count(); - let now_ms = current_time_ms(); - let mut merged_dbs: Vec = (0..db_count).map(|_| Database::new()).collect(); - - for shard_locks in shard_dbs.all_shard_dbs() { - for (db_idx, lock) in shard_locks.iter().enumerate() { - let guard = lock.read(); - for (key, entry) in guard.data().iter() { - if !entry.is_expired_at(guard.base_timestamp(), now_ms) { - merged_dbs[db_idx].set(key.to_bytes(), entry.clone()); - } - } - } - } - - let rdb_bytes = crate::persistence::rdb::save_to_bytes(&merged_dbs)?; - - let tmp_path = aof_path.with_extension("aof.tmp"); - std::fs::write(&tmp_path, &rdb_bytes).map_err(|e| AofError::Io { - path: tmp_path.clone(), - source: e, - })?; - std::fs::rename(&tmp_path, aof_path).map_err(|e| AofError::RewriteFailed { - detail: format!( - "rename {} -> {}: {}", - tmp_path.display(), - aof_path.display(), - e - ), - })?; - - info!( - "AOF rewrite (sharded, RDB preamble) complete: {} bytes", - rdb_bytes.len() - ); - Ok(()) -} - -/// Reopen AOF file in append mode after atomic rewrite replaced it. -#[allow(dead_code)] -fn reopen_aof_sync(aof_path: &Path) -> Result { - std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(aof_path) -} - -/// Rewrite the AOF file (tokio async wrapper). -/// -/// Delegates to `rewrite_aof_sync` — the actual I/O is synchronous (temp write + rename). -#[cfg(feature = "runtime-tokio")] -#[tracing::instrument(skip_all, level = "info")] -pub async fn rewrite_aof(db: SharedDatabases, aof_path: &Path) -> Result<(), MoonError> { - rewrite_aof_sync(&db, aof_path) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::persistence::replay::DispatchReplayEngine; - use ordered_float::OrderedFloat; - use tempfile::tempdir; - - fn make_command(parts: &[&[u8]]) -> Frame { - Frame::Array( - parts - .iter() - .map(|p| Frame::BulkString(Bytes::copy_from_slice(p))) - .collect(), - ) - } - - // --- serialize_command / generate_aof_command round-trip tests --- - - #[test] - fn test_generate_aof_command_produces_valid_resp_that_round_trips() { - let frame = make_command(&[b"SET", b"key", b"value"]); - let serialized = serialize_command(&frame); - - let mut buf = BytesMut::from(&serialized[..]); - let config = ParseConfig::default(); - let parsed = parse::parse(&mut buf, &config).unwrap().unwrap(); - assert_eq!(parsed, frame); - } - - #[test] - fn test_serialize_command_round_trip_hset() { - let frame = make_command(&[b"HSET", b"myhash", b"f1", b"v1"]); - let serialized = serialize_command(&frame); - let mut buf = BytesMut::from(&serialized[..]); - let parsed = parse::parse(&mut buf, &ParseConfig::default()) - .unwrap() - .unwrap(); - assert_eq!(parsed, frame); - } - - // --- AOF replay tests --- - - #[test] - fn test_aof_replay_set_commands_restores_string_keys() { - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("test.aof"); - - // Write SET commands in RESP format - let mut aof_data = BytesMut::new(); - serialize::serialize(&make_command(&[b"SET", b"k1", b"v1"]), &mut aof_data); - serialize::serialize(&make_command(&[b"SET", b"k2", b"v2"]), &mut aof_data); - std::fs::write(&aof_path, &aof_data).unwrap(); - - let mut dbs = vec![Database::new()]; - let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); - assert_eq!(count, 2); - - let entry = dbs[0].get(b"k1").unwrap(); - assert_eq!(entry.value.as_bytes().unwrap(), b"v1"); - let entry = dbs[0].get(b"k2").unwrap(); - assert_eq!(entry.value.as_bytes().unwrap(), b"v2"); - } - - #[test] - fn test_aof_replay_collection_types() { - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("test.aof"); - - let mut aof_data = BytesMut::new(); - // HSET - serialize::serialize( - &make_command(&[b"HSET", b"myhash", b"f1", b"v1"]), - &mut aof_data, - ); - // LPUSH - serialize::serialize( - &make_command(&[b"LPUSH", b"mylist", b"a", b"b"]), - &mut aof_data, - ); - // SADD - serialize::serialize( - &make_command(&[b"SADD", b"myset", b"x", b"y"]), - &mut aof_data, - ); - // ZADD - serialize::serialize( - &make_command(&[b"ZADD", b"myzset", b"1.5", b"alice"]), - &mut aof_data, - ); - std::fs::write(&aof_path, &aof_data).unwrap(); - - let mut dbs = vec![Database::new()]; - let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); - assert_eq!(count, 4); - - // Check hash - let hash = dbs[0].get_hash(b"myhash").unwrap().unwrap(); - assert_eq!( - hash.get(&Bytes::from_static(b"f1")).unwrap().as_ref(), - b"v1" - ); - - // Check list - let list = dbs[0].get_list(b"mylist").unwrap().unwrap(); - assert_eq!(list.len(), 2); - - // Check set - let set = dbs[0].get_set(b"myset").unwrap().unwrap(); - assert_eq!(set.len(), 2); - - // Check sorted set - let (members, _) = dbs[0].get_sorted_set(b"myzset").unwrap().unwrap(); - assert_eq!(*members.get(&Bytes::from_static(b"alice")).unwrap(), 1.5); - } - - #[test] - fn test_aof_replay_with_expire_preserves_ttls() { - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("test.aof"); - - let mut aof_data = BytesMut::new(); - serialize::serialize(&make_command(&[b"SET", b"mykey", b"myval"]), &mut aof_data); - serialize::serialize( - &make_command(&[b"PEXPIRE", b"mykey", b"60000"]), - &mut aof_data, - ); - std::fs::write(&aof_path, &aof_data).unwrap(); - - let mut dbs = vec![Database::new()]; - let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); - assert_eq!(count, 2); - - let base_ts = dbs[0].base_timestamp(); - let entry = dbs[0].get(b"mykey").unwrap(); - assert!(entry.has_expiry()); - let remaining_secs = (entry.expires_at_ms(base_ts) - current_time_ms()) / 1000; - assert!(remaining_secs >= 50); // Allow some tolerance - } - - #[test] - fn test_aof_replay_with_select_switches_databases() { - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("test.aof"); - - let mut aof_data = BytesMut::new(); - serialize::serialize(&make_command(&[b"SET", b"k0", b"v0"]), &mut aof_data); - serialize::serialize(&make_command(&[b"SELECT", b"1"]), &mut aof_data); - serialize::serialize(&make_command(&[b"SET", b"k1", b"v1"]), &mut aof_data); - std::fs::write(&aof_path, &aof_data).unwrap(); - - let mut dbs = vec![Database::new(), Database::new()]; - let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); - assert_eq!(count, 3); - - assert!(dbs[0].get(b"k0").is_some()); - assert!(dbs[1].get(b"k1").is_some()); - } - - #[test] - fn test_aof_replay_empty_file_produces_zero_keys() { - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("test.aof"); - std::fs::write(&aof_path, b"").unwrap(); - - let mut dbs = vec![Database::new()]; - let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); - assert_eq!(count, 0); - assert_eq!(dbs[0].len(), 0); - } - - #[test] - fn test_aof_replay_corrupt_truncated_logs_error_loads_what_it_can() { - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("test.aof"); - - let mut aof_data = BytesMut::new(); - serialize::serialize(&make_command(&[b"SET", b"k1", b"v1"]), &mut aof_data); - // Append corrupt data - aof_data.extend_from_slice(b"*3\r\n$3\r\nSET\r\n$2\r\nk2"); - std::fs::write(&aof_path, &aof_data).unwrap(); - - let mut dbs = vec![Database::new()]; - let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); - // Should have loaded the first command - assert_eq!(count, 1); - assert!(dbs[0].get(b"k1").is_some()); - } - - // --- FsyncPolicy tests --- - - #[test] - fn test_fsync_policy_from_str() { - assert_eq!(FsyncPolicy::from_str("always"), FsyncPolicy::Always); - assert_eq!(FsyncPolicy::from_str("everysec"), FsyncPolicy::EverySec); - assert_eq!(FsyncPolicy::from_str("no"), FsyncPolicy::No); - assert_eq!(FsyncPolicy::from_str("unknown"), FsyncPolicy::EverySec); - } - - // --- generate_rewrite_commands tests --- - - #[test] - fn test_generate_rewrite_commands_all_5_types() { - let mut dbs = vec![Database::new()]; - - // String - dbs[0].set_string(Bytes::from_static(b"str"), Bytes::from_static(b"val")); - // Hash - { - let map = dbs[0].get_or_create_hash(b"h").unwrap(); - map.insert(Bytes::from_static(b"f"), Bytes::from_static(b"v")); - } - // List - { - let list = dbs[0].get_or_create_list(b"l").unwrap(); - list.push_back(Bytes::from_static(b"item")); - } - // Set - { - let set = dbs[0].get_or_create_set(b"s").unwrap(); - set.insert(Bytes::from_static(b"m")); - } - // Sorted set - { - let (members, tree) = dbs[0].get_or_create_sorted_set(b"z").unwrap(); - members.insert(Bytes::from_static(b"a"), 1.0); - tree.insert(OrderedFloat(1.0), Bytes::from_static(b"a")); - } - - let commands = generate_rewrite_commands(&dbs); - assert!(!commands.is_empty()); - - // Replay and verify round-trip - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("rewrite.aof"); - std::fs::write(&aof_path, &commands).unwrap(); - - let mut loaded_dbs = vec![Database::new()]; - let count = replay_aof(&mut loaded_dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); - assert!(count >= 5, "Expected at least 5 commands, got {}", count); - - // Verify each type restored - assert_eq!( - loaded_dbs[0].get(b"str").unwrap().value.type_name(), - "string" - ); - assert!(loaded_dbs[0].get_hash(b"h").unwrap().is_some()); - assert!(loaded_dbs[0].get_list(b"l").unwrap().is_some()); - assert!(loaded_dbs[0].get_set(b"s").unwrap().is_some()); - assert!(loaded_dbs[0].get_sorted_set(b"z").unwrap().is_some()); - } - - #[test] - fn test_generate_rewrite_commands_with_ttl() { - let mut dbs = vec![Database::new()]; - let future_ms = current_time_ms() + 3_600_000; - dbs[0].set_string_with_expiry( - Bytes::from_static(b"key"), - Bytes::from_static(b"val"), - future_ms, - ); - - let commands = generate_rewrite_commands(&dbs); - - // Replay and check TTL is preserved - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("rewrite.aof"); - std::fs::write(&aof_path, &commands).unwrap(); - - let mut loaded_dbs = vec![Database::new()]; - let count = replay_aof(&mut loaded_dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); - assert_eq!(count, 2); // SET + PEXPIRE - - let base_ts = loaded_dbs[0].base_timestamp(); - let entry = loaded_dbs[0].get(b"key").unwrap(); - assert!(entry.has_expiry()); - let remaining_secs = (entry.expires_at_ms(base_ts) - current_time_ms()) / 1000; - assert!(remaining_secs > 3500); - } - - /// Helper: build a snapshot tuple from a Database slice — mirrors the - /// `merged` construction in `do_rewrite_sharded` (aof.rs:2070-2090) so - /// tests exercise the exact same production path. - fn db_slice_to_snapshot( - dbs: &[Database], - ) -> Vec<( - Vec<( - crate::storage::compact_key::CompactKey, - crate::storage::entry::Entry, - )>, - u32, - )> { - let now_ms = crate::storage::entry::current_time_ms(); - dbs.iter() - .map(|db| { - let base_ts = db.base_timestamp(); - let entries: Vec<_> = db - .data() - .iter() - .filter(|(_, e)| !e.is_expired_at(base_ts, now_ms)) - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - (entries, base_ts) - }) - .collect() - } - - /// FIX-W3-8: BGREWRITEAOF on a fresh empty database must produce a valid - /// RDB base and recover cleanly with 0 keys. - /// - /// Updated to call `save_snapshot_to_bytes` (the function `do_rewrite_sharded` - /// actually calls at aof.rs:2096) rather than `save_to_bytes` (the previous - /// test used the wrong function — tautological for empty input since both - /// produce an identical valid RDB, but would miss regressions on the - /// snapshot-tuple path). - #[test] - fn empty_database_rewrite_produces_valid_rdb_and_recovers() { - let dir = tempdir().unwrap(); - - // Build the snapshot tuple the same way do_rewrite_sharded does. - let empty_dbs: Vec = vec![Database::new()]; - let snapshot = db_slice_to_snapshot(&empty_dbs); - let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot) - .expect("save_snapshot_to_bytes must succeed for empty snapshot"); - - // Invariant 1: RDB is non-empty (has at least magic + version + EOF marker). - assert!( - !rdb_bytes.is_empty(), - "empty-database RDB must not be 0 bytes" - ); - - // Invariant 2: starts with valid MOON magic header. - assert!( - rdb_bytes.starts_with(b"MOON"), - "RDB bytes must start with MOON magic, got: {:?}", - &rdb_bytes[..rdb_bytes.len().min(8)] - ); - - // Invariant 3: recovery from this base succeeds with 0 keys loaded. - let base_path = dir.path().join("empty.rdb"); - std::fs::write(&base_path, &rdb_bytes).expect("write empty rdb"); - let mut recovery_dbs = vec![Database::new()]; - let loaded = - crate::persistence::rdb::load(&mut recovery_dbs, &base_path).expect("load empty rdb"); - assert_eq!( - loaded, 0, - "recovering from empty-database RDB yields 0 keys" - ); - assert_eq!( - recovery_dbs[0].len(), - 0, - "database must be empty after recovering from zero-key RDB" - ); - } - - /// FIX-W3-8: Genuine regression guard — save_snapshot_to_bytes preserves - /// a 1-key database through a full serialize→file→reload cycle. - /// - /// This is the substantive test the verifier asked for: verifies the - /// production code path (`save_snapshot_to_bytes` via the snapshot-tuple - /// form) against a non-trivial database, so a future regression that - /// swaps back to `save_to_bytes` or breaks TTL handling in the snapshot - /// path will be caught. - #[test] - fn snapshot_to_bytes_round_trips_one_key_database() { - let dir = tempdir().unwrap(); - - // Build a 1-key database with a string value. - let mut db = Database::new(); - db.set_string( - Bytes::from_static(b"rdb_key"), - Bytes::from_static(b"rdb_value"), - ); - let dbs = vec![db]; - - // Serialize via the production path (save_snapshot_to_bytes). - let snapshot = db_slice_to_snapshot(&dbs); - let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot) - .expect("save_snapshot_to_bytes must succeed for 1-key snapshot"); - - assert!(rdb_bytes.starts_with(b"MOON"), "must have MOON magic"); - - // Reload and assert the key survives. - let base_path = dir.path().join("one_key.rdb"); - std::fs::write(&base_path, &rdb_bytes).expect("write rdb"); - let mut recovery_dbs = vec![Database::new()]; - let loaded = crate::persistence::rdb::load(&mut recovery_dbs, &base_path) - .expect("load rdb must succeed"); - assert_eq!(loaded, 1, "exactly 1 key must be recovered"); - let val = recovery_dbs[0] - .get(b"rdb_key") - .expect("rdb_key must be present"); - assert_eq!( - val.value.as_bytes().expect("string value"), - b"rdb_value", - "recovered value must match written value" - ); - } - - #[test] - fn test_generate_rewrite_round_trip_preserves_state() { - let mut dbs = vec![Database::new()]; - dbs[0].set_string(Bytes::from_static(b"a"), Bytes::from_static(b"1")); - dbs[0].set_string(Bytes::from_static(b"b"), Bytes::from_static(b"2")); - { - let list = dbs[0].get_or_create_list(b"mylist").unwrap(); - list.push_back(Bytes::from_static(b"x")); - list.push_back(Bytes::from_static(b"y")); - list.push_back(Bytes::from_static(b"z")); - } - - let commands = generate_rewrite_commands(&dbs); - let dir = tempdir().unwrap(); - let aof_path = dir.path().join("rewrite.aof"); - std::fs::write(&aof_path, &commands).unwrap(); - - let mut loaded = vec![Database::new()]; - replay_aof(&mut loaded, &aof_path, &DispatchReplayEngine::new()).unwrap(); - - // Check strings - assert_eq!(loaded[0].get(b"a").unwrap().value.as_bytes().unwrap(), b"1"); - assert_eq!(loaded[0].get(b"b").unwrap().value.as_bytes().unwrap(), b"2"); - - // Check list order preserved - let list = loaded[0].get_list(b"mylist").unwrap().unwrap(); - assert_eq!(list.len(), 3); - assert_eq!(list[0].as_ref(), b"x"); - assert_eq!(list[1].as_ref(), b"y"); - assert_eq!(list[2].as_ref(), b"z"); - } - - // ----------------------------------------------------------------------- - // FIX-W2-4 r2: canonical AOF fsync error string - // - // Red criterion: AOF_FSYNC_ERR constant must exist and equal the canonical - // Redis-style ERR-prefixed string used by handler_monoio and handler_sharded. - // handler_single.rs previously used "WRITEFAIL aof fsync failed" which is - // both non-canonical (no ERR prefix, different verb) and inconsistent with - // the other two handlers. - // - // These tests compile-fail on the prior commit (constant absent) and pass - // once AOF_FSYNC_ERR is declared in this module with the correct value. - // ----------------------------------------------------------------------- - - #[test] - fn aof_fsync_err_constant_is_canonical() { - // The canonical error frame bytes sent to the client when an AOF - // fsync under appendfsync=always fails. Must match what - // handler_monoio/mod.rs and handler_sharded/mod.rs use. - assert_eq!( - AOF_FSYNC_ERR, b"ERR AOF fsync failed; write not durable", - "AOF_FSYNC_ERR must equal the canonical ERR-prefixed string" - ); - } - - #[test] - fn aof_fsync_err_has_err_prefix() { - // Redis convention: protocol-level errors must start with a word - // followed by a space, using `ERR` for generic errors. `WRITEFAIL` - // is not a standard Redis error prefix and confuses clients that - // pattern-match on error codes. - assert!( - AOF_FSYNC_ERR.starts_with(b"ERR "), - "AOF_FSYNC_ERR must start with 'ERR ' (got {:?})", - std::str::from_utf8(AOF_FSYNC_ERR).unwrap_or("") - ); - } -} diff --git a/src/persistence/aof/mod.rs b/src/persistence/aof/mod.rs new file mode 100644 index 00000000..90d1cde0 --- /dev/null +++ b/src/persistence/aof/mod.rs @@ -0,0 +1,1078 @@ +//! Append-Only File (AOF) persistence: logs every write command in RESP format +//! for crash recovery. Supports three fsync policies and AOF rewriting for compaction. +//! +//! ## Unwrap Classification +//! +//! | Context | Classification | Rationale | +//! |---------|---------------|-----------| +//! | `AofWriter::append` (hot path) | **fire-and-forget** | Channel send; no Result needed | +//! | `aof_writer_task` | **must-panic** | Writer task; errors logged inline | +//! | `replay_aof` | **should-recover** (`Result<_, MoonError>`) | Startup replay; log+skip on corruption | +//! | `rewrite_aof` | **should-recover** (`Result<_, MoonError>`) | Background rewrite; caller logs error | +//! | `#[cfg(test)]` code (55 unwraps) | **test-only** | Panics are appropriate in tests | +// Suppressions narrowed: only keep what's needed for conditional compilation +#![allow(unused_imports, unused_variables, unreachable_code, clippy::empty_loop)] + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use crate::runtime::cancel::CancellationToken; +use crate::runtime::channel; +use bytes::{Bytes, BytesMut}; +use tracing::{error, info, warn}; + +use crate::error::{AofError, MoonError}; +use crate::framevec; +use crate::persistence::replay::CommandReplayEngine; +use crate::protocol::{Frame, ParseConfig, parse, serialize}; +use crate::storage::compact_key::CompactKey; +use crate::storage::compact_value::RedisValueRef; +use crate::storage::db::Database; +use crate::storage::entry::{Entry, current_time_ms}; +/// Type alias for the per-database RwLock container. +type SharedDatabases = Arc>>; + +/// Canonical AOF fsync failure error string sent to the client as a +/// `Frame::Error` when `appendfsync=always` and the writer task does not +/// confirm durability before the response. +/// +/// All handler variants (handler_single, handler_monoio, handler_sharded) +/// MUST use this constant so operators see a consistent error regardless of +/// which connection path handles the request. +/// +/// Redis convention: errors begin with a single-word code (`ERR` for generic +/// failures) followed by a space and a human-readable message. +pub const AOF_FSYNC_ERR: &[u8] = b"ERR AOF fsync failed; write not durable"; + +/// High bit of the per-entry LSN reserved for `OrderedAcrossShards` +/// (RFC § 2 Rule 2). When set on a per-shard AOF entry, recovery treats +/// the entry as participating in a cross-shard atomic operation and +/// buffers it for the cross-shard merge replay after per-shard replay +/// completes. +/// +/// Practical LSN ceilings (even at 10 M writes/s sustained for a century) +/// sit near 2^58, so reserving bit 63 has no observable effect on normal +/// writes — the bit is always 0 in entries written by `try_send_append`. +/// Only `try_send_append_ordered` sets it. +pub const ORDERED_LSN_FLAG: u64 = 1u64 << 63; + +/// Outcome reported by the writer task back to an `AppendSync` caller +/// once the rendezvous completes. +/// +/// `Synced` is sent AFTER `sync_data()` returns successfully — the +/// caller may safely `+OK` the client. `WriteFailed`/`FsyncFailed` +/// surface the failure mode so the caller can return a specific error +/// frame; either way, durability was NOT achieved. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AofAck { + /// Bytes were written and fsynced. Durability guaranteed. + Synced, + /// `write_all()` returned an error. The entry may be partially on + /// disk; recovery handles partial-payload truncation as crash EOF. + WriteFailed, + /// `write_all()` succeeded but `sync_data()` returned an error. The + /// entry is in the kernel buffer but NOT on durable storage. + FsyncFailed, + /// The writer channel was full at the time of the send — the entry + /// was **not** enqueued. This is a backpressure signal: the writer + /// is unable to keep up with the current write rate. Callers MUST + /// treat this as a hard failure (same as `WriteFailed`) under + /// `appendfsync=always`; for `everysec`/`no` it is logged and counted. + ChannelFull, +} + +/// Global counter incremented each time an AOF `AppendSync` (or fire-and- +/// forget `Append`) is dropped because the writer channel was at capacity. +/// +/// Exposed under `# Persistence` in the INFO command as +/// `aof_backpressure_dropped`. A persistently non-zero value indicates the +/// writer is a bottleneck and the operator should investigate disk I/O or +/// switch to `appendfsync=everysec`. +pub static AOF_BACKPRESSURE_DROPPED: std::sync::atomic::AtomicU64 = + std::sync::atomic::AtomicU64::new(0); + +/// Result of awaiting an `AppendSync` ack under a bounded timeout (F2). +/// +/// Distinguishes the three terminal states the `Always` durability path +/// can reach so the caller can map each to the correct client-facing +/// outcome: +/// - `Ack(_)` — the writer reported back; inspect the `AofAck`. +/// - `Disconnected` — the writer task is gone / channel dropped (no ack +/// will ever arrive). Treated as `WriteFailed`. +/// - `TimedOut` — the fsync did not confirm within the configured +/// bound. Durability is unconfirmed; treated as `FsyncFailed`. The +/// entry may still reach disk later, so the caller must NOT report +/// success but also must not assume the write was rejected. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AckOutcome { + Ack(AofAck), + Disconnected, + TimedOut, +} + +/// AOF fsync policy controlling when data is flushed to disk. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum FsyncPolicy { + /// Fsync after every write command (safest, slowest). + Always, + /// Fsync once per second in the background (good balance). + EverySec, + /// Let the OS decide when to flush (fastest, least safe). + No, +} + +impl FsyncPolicy { + /// Parse a policy string (as from config). Defaults to EverySec for unknown values. + pub fn from_str(s: &str) -> Self { + match s { + "always" => FsyncPolicy::Always, + "no" => FsyncPolicy::No, + _ => FsyncPolicy::EverySec, + } + } +} + +/// Messages sent to the AOF writer task via mpsc channel. +pub enum AofMessage { + /// Append serialized RESP command bytes to the AOF file, tagged with the + /// LSN that was issued for this write (`ReplicationState::issue_lsn`). + /// + /// `lsn` semantics by writer task: + /// - **TopLevel** (`aof_writer_task`): `lsn` is **ignored**; the legacy + /// v1 disk format is plain RESP bytes with no per-entry framing. + /// - **PerShard** (`per_shard_aof_writer_task`): `lsn` is **written** as + /// a u64 header per RFC § 2 Rule 1. Disk format per entry: + /// `[u64 lsn LE][u32 len LE][RESP bytes of length len]`. + /// Recovery reads `(lsn, cmd)` pairs and merges cross-shard + /// `OrderedAcrossShards` writes by LSN (RFC § 2 Rule 2). + /// + /// Construction sites that issue a real LSN call + /// `ReplicationState::issue_lsn(shard_id, bytes.len() as u64)` and pass + /// the returned value. Sites with no replication state available pass 0 + /// (TopLevel ignores it; PerShard treats 0 as "no ordering hint"). + Append { lsn: u64, bytes: Bytes }, + /// Append + fsync + ack rendezvous (RFC § 4 — Fix 2 for the H1 + /// data-loss vector exposed by `appendfsync=always`). + /// + /// Same encoding as [`AofMessage::Append`], but the writer task ALWAYS + /// fsyncs after writing the payload and signals `ack` ONCE the + /// `sync_data()` syscall returns. The caller is expected to await + /// `ack` before responding `+OK` to the client so the durability + /// contract of `appendfsync=always` is honoured end-to-end. + /// + /// Failure semantics: on write or fsync error the writer drops `ack` + /// without sending — the caller's `OneshotReceiver` resolves with + /// `RecvError`, which it must treat as a hard failure (return an + /// error frame to the client, do NOT silently +OK). + /// + /// Production callers: none in step 7 — this commit ships the + /// mechanism plus tests. Per-handler integration (which sites use + /// AppendSync vs Append) is wired in step 9 before lifting the + /// `--unsafe-multishard-aof` gate. + AppendSync { + lsn: u64, + bytes: Bytes, + ack: crate::runtime::channel::OneshotSender, + }, + /// Trigger a full AOF rewrite (compaction) using current database state. + Rewrite(SharedDatabases), + /// Trigger AOF rewrite in sharded mode (all shards' databases). + RewriteSharded(Arc), + /// [F6] Trigger a per-shard AOF rewrite (compaction) in the PerShard + /// layout. Sent to EVERY per-shard writer at once. Each writer folds its + /// own shard (drain → lock → snapshot → write new base+incr at the + /// coordinator's `new_seq` → reopen), then decrements the shared + /// `PerShardRewriteCoord`; the last writer commits the manifest once + /// (single seq flip) and prunes the old generation. The synchronized seq + /// + single commit are what make multi-shard BGREWRITEAOF crash-safe. + RewritePerShard { + shard_dbs: Arc, + coord: Arc, + }, + /// Shut down the AOF writer task gracefully. + Shutdown, +} + +/// Coordinator shared by all per-shard writers participating in one +/// BGREWRITEAOF fan-out (F6). +/// +/// Crash-safety contract (mirrors `AofManifest::advance` ordering, but +/// distributed across N writer threads): +/// +/// 1. Each writer writes its new base+incr at `new_seq` via +/// `manifest.advance_shard(shard_id, new_seq, rdb)` — which does NOT bump +/// `manifest.seq` or rewrite the manifest. So until the final commit, the +/// on-disk manifest still resolves to `old_seq`; a crash here recovers the +/// intact old generation (no loss, no double-apply). +/// 2. The LAST writer to finish (countdown reaches zero) performs the single +/// durable commit: `manifest.seq = new_seq; write_manifest()`. This is the +/// atomic point at which recovery flips to the new generation. +/// 3. Only AFTER the commit are the old-generation files pruned. +/// +/// The manifest is shared via `Arc>` and locked ONLY for the brief, +/// await-free `advance_shard` and final-commit critical sections — never held +/// across a blocking disk write of the base RDB (that happens before the lock) +/// nor across `.await`. +pub struct PerShardRewriteCoord { + /// Writers still to finish. Starts at the shard count; the writer that + /// decrements it to zero performs the commit + prune. + remaining: std::sync::atomic::AtomicUsize, + /// Shared manifest, loaded fresh from disk by the BGREWRITEAOF handler at + /// rewrite time (normal appends never touch the manifest, and BGREWRITEAOF + /// is CAS-serialized, so a fresh load is the authoritative current state). + manifest: Arc>, + /// The generation every writer advances to. Computed once = old_seq + 1. + new_seq: u64, + /// The generation being retired; pruned only after the commit. + old_seq: u64, + /// Number of shards participating (= initial `remaining`). + n_shards: usize, + /// Set by any shard whose fold fails. The final writer checks this and + /// ABORTS the commit if set — committing `new_seq` while a shard never + /// wrote its new base would make recovery look for a missing base and + /// refuse to start. On abort the old generation (`old_seq`) stays the + /// committed state for all shards (crash-safe). + failed: std::sync::atomic::AtomicBool, + /// The committed generation, published exactly once by the terminal writer + /// (`new_seq` on success, `old_seq` on abort/commit-failure). Every folded + /// writer blocks on this in `await_outcome` after `shard_done` so it can + /// reopen its append file onto the COMMITTED generation before resuming + /// appends — without this barrier a writer that reopened onto `new_seq` in + /// phase 6 keeps appending there after an abort, into a generation recovery + /// ignores (silent data loss; the old "RESTART recommended" hazard). + outcome: parking_lot::Mutex>, + outcome_cv: parking_lot::Condvar, +} + +impl PerShardRewriteCoord { + /// Construct a coordinator for an `n_shards`-way rewrite advancing the + /// shared `manifest` from its current seq to `current_seq + 1`. + pub fn new( + manifest: Arc>, + current_seq: u64, + n_shards: usize, + ) -> Arc { + Arc::new(Self { + remaining: std::sync::atomic::AtomicUsize::new(n_shards), + manifest, + new_seq: current_seq + 1, + old_seq: current_seq, + n_shards, + failed: std::sync::atomic::AtomicBool::new(false), + outcome: parking_lot::Mutex::new(None), + outcome_cv: parking_lot::Condvar::new(), + }) + } + + /// Publish the committed generation to every writer blocked in + /// `await_outcome`. Called exactly once, by the terminal `shard_done`. + #[inline] + fn publish_outcome(&self, committed_seq: u64) { + let mut slot = self.outcome.lock(); + *slot = Some(committed_seq); + self.outcome_cv.notify_all(); + } + + /// Block until the terminal writer publishes the committed generation, then + /// return it (`new_seq` on success, `old_seq` on abort/commit-failure). Each + /// folded writer calls this AFTER `shard_done` so it can reopen its append + /// file onto the committed generation. Safe against missed wake-ups: the + /// outcome is set under the same lock the condvar waits on, so a writer that + /// arrives after the publish observes `Some` and skips the wait. + /// + /// Deadlock-free: all `n_shards` writers call `shard_done` exactly once + /// (success via phase 7, fold-error/send-failure via the caller), so the + /// countdown always reaches zero and the terminal branch always publishes. + /// Each per-shard writer runs on its own dedicated OS thread, so blocking + /// one here never starves the thread that must publish. + #[must_use] + fn await_outcome(&self) -> u64 { + let mut slot = self.outcome.lock(); + while slot.is_none() { + self.outcome_cv.wait(&mut slot); + } + slot.expect("outcome is Some after the wait loop") + } + + /// The generation writers advance to. + #[inline] + pub fn new_seq(&self) -> u64 { + self.new_seq + } + + /// Mark the whole rewrite as failed (called by a shard whose fold errored). + /// The final writer will abort the commit, leaving `old_seq` authoritative. + #[inline] + pub fn mark_failed(&self) { + self.failed + .store(true, std::sync::atomic::Ordering::Release); + } + + /// Called by each writer AFTER it has durably written its new base+incr at + /// `new_seq` and reopened its append file. Decrements the countdown; the + /// final caller commits the manifest (single seq flip) and prunes the old + /// generation, then clears the global in-progress flag. + /// + /// Crash-safety: the commit (`write_manifest`) is the atomic flip point; + /// pruning runs strictly after it, so a crash mid-prune only orphans + /// already-superseded files (recovery uses `new_seq`). + pub fn shard_done(&self) { + use std::sync::atomic::Ordering; + // AcqRel: the decrement-to-zero must observe all prior writers' + // advance_shard manifest mutations before committing. + if self.remaining.fetch_sub(1, Ordering::AcqRel) == 1 { + // Abort if any shard failed to fold: committing new_seq while a + // shard lacks its new base would break recovery. Keep old_seq. + if self.failed.load(Ordering::Acquire) { + let m = self.manifest.lock(); + // Best-effort: prune the orphaned new-seq files written by the + // shards that DID fold, so they don't linger. + for sid in 0..self.n_shards { + m.prune_shard_files(sid as u16, self.new_seq); + } + drop(m); + error!( + "F6 per-shard rewrite ABORTED: a shard failed to fold; seq stays {}. \ + Old generation remains authoritative (crash-safe). Folded writers \ + roll their append files back to the old generation at the \ + barrier — no restart required.", + self.old_seq + ); + // Publish BEFORE clearing the in-progress flag so every blocked + // writer wakes and reopens onto the (still-authoritative) old gen. + self.publish_outcome(self.old_seq); + crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); + return; + } + let mut m = self.manifest.lock(); + m.seq = self.new_seq; + if let Err(e) = m.write_manifest() { + error!( + "F6 per-shard rewrite: final manifest commit (seq {}) failed: {}. \ + Old generation remains authoritative; rewrite did not take effect.", + self.new_seq, e + ); + // Do NOT prune — old generation is still the committed state. + drop(m); + // Commit failed: old_seq is still authoritative, so folded + // writers must roll back to it (their phase-6 reopen targeted + // new_seq, which never committed). + self.publish_outcome(self.old_seq); + crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); + return; + } + for sid in 0..self.n_shards { + m.prune_shard_files(sid as u16, self.old_seq); + } + drop(m); + info!( + "F6 per-shard rewrite complete: committed seq {} across {} shards, pruned seq {}", + self.new_seq, self.n_shards, self.old_seq + ); + // Success: new_seq is committed. Folded writers already reopened onto + // new_seq in phase 6, so the barrier is a no-op for them — but it must + // still unblock them. + self.publish_outcome(self.new_seq); + crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst); + } + } +} + +/// Panic-safety guard for a per-shard fold's `shard_done` obligation. +/// +/// The `await_outcome` barrier in `do_rewrite_per_shard` turns "every shard +/// calls `shard_done` exactly once" from a tidiness rule into a LIVENESS +/// requirement: any folded writer that decrements then blocks in `await_outcome` +/// hangs forever if some other shard never decrements. The one path that would +/// skip `shard_done` is a panic mid-fold (e.g. `save_snapshot_to_bytes` +/// OOM-unwinding — issue #138), which under `panic = "unwind"` would otherwise +/// hang every healthy shard's AOF writer thread. +/// +/// This guard makes the decrement unconditional. The fold creates it on entry +/// and calls [`complete`](Self::complete) on the success path (clean +/// `shard_done`). On ANY other exit — `?` error or panic unwind — `Drop` fires +/// `mark_failed` + `shard_done`, so the countdown still reaches zero, the +/// terminal writer publishes `old_seq`, and every barrier waiter wakes and rolls +/// back. Because the guard owns the decrement for ALL exits, callers MUST NOT +/// call `shard_done` again after invoking `do_rewrite_per_shard`. +struct ShardDoneGuard<'a> { + coord: &'a PerShardRewriteCoord, + done: bool, +} + +impl<'a> ShardDoneGuard<'a> { + #[inline] + fn new(coord: &'a PerShardRewriteCoord) -> Self { + Self { coord, done: false } + } + + /// Success-path completion: a single clean `shard_done` (no `mark_failed`). + /// Consumes the guard so the subsequent `Drop` is a no-op. + #[inline] + fn complete(mut self) { + self.done = true; + self.coord.shard_done(); + } +} + +impl Drop for ShardDoneGuard<'_> { + fn drop(&mut self) { + if !self.done { + // Unwind or early-error before completion: abort the rewrite (keep + // old_seq) and still decrement so the barrier releases every waiter. + self.coord.mark_failed(); + self.coord.shard_done(); + } + } +} + +/// Reasons a pool send may be refused without queueing. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AofPoolSendError { + /// `Rewrite`/`RewriteSharded` sent to a `PerShard` pool. BGREWRITEAOF must + /// be issued per shard in the per-shard layout; the legacy single-writer + /// rewrite path is not applicable. + RewriteUnsupportedInPerShard, + /// Underlying channel send failed (writer task dead or channel full). + SendFailed, +} + +/// Bundle of per-shard AOF writer senders. +/// +/// The pool keeps the call-site API uniform regardless of layout: +/// - **TopLevel** (legacy v1, single-shard, also used for `--shards 1` v2): +/// exactly one writer thread; every `sender(shard_id)` returns the same +/// sender so all shards multiplex onto one file. +/// - **PerShard** (v2 multi-shard): one writer per shard; `sender(shard_id)` +/// returns the writer that owns `appendonlydir/shard-{shard_id}/`. +/// +/// Step 2a is additive — this type is defined here but no call site is wired +/// to it yet. Step 2c performs the type plumbing in `conn_state` and +/// `conn/core`; steps 2d/2e/2f update the call sites and spawn paths. +/// Default bound for the `appendfsync=always` fsync-ack await. Mirrors the +/// `--aof-fsync-timeout-ms` config default; used by constructors that don't +/// take an explicit timeout (non-production / test helpers). +pub const DEFAULT_AOF_FSYNC_TIMEOUT: Duration = Duration::from_millis(2000); + +// ── Submodule decomposition (refactor: aof.rs 4379 lines -> directory module) ── +// Codec (serialize_command/replay_aof) stays in this parent so children reach it +// via `use super::*`. AofWriterPool, writer tasks, and rewrite paths move out. +mod pool; +mod rewrite; +mod writer_task; + +pub use pool::AofWriterPool; +pub use rewrite::generate_rewrite_commands; +// `rewrite_aof` is defined only under the tokio runtime; gate its re-export to match. +#[cfg(feature = "runtime-tokio")] +pub use rewrite::rewrite_aof; +pub use writer_task::{aof_writer_task, per_shard_aof_writer_task}; +// Test-only torn-write injection hook, consumed by pool_tests; gated to match its +// definition in writer_task.rs. +#[cfg(all(test, feature = "runtime-tokio"))] +pub(crate) use writer_task::TEST_FAIL_WRITE_AT; + +/// Serialize a Frame into RESP wire format bytes. +pub fn serialize_command(frame: &Frame) -> Bytes { + let mut buf = BytesMut::with_capacity(64); + serialize::serialize(frame, &mut buf); + buf.freeze() +} + +/// Replay an AOF file by parsing RESP commands and dispatching them. +/// +/// Returns the number of commands successfully replayed. +/// +/// **Corruption recovery:** On mid-stream parse errors, logs a warning with the +/// byte offset, skips to the next RESP array marker (`*`), and continues replay. +/// At EOF, reports total corrupted entries skipped. Truncated tails are handled +/// gracefully (warn + stop). +pub fn replay_aof( + databases: &mut [Database], + path: &Path, + engine: &dyn CommandReplayEngine, +) -> Result { + let data = std::fs::read(path)?; + if data.is_empty() { + return Ok(0); + } + + // Detect RDB preamble: if the file starts with "MOON" magic, load the binary + // RDB section first, then replay any RESP commands appended after it. + let (rdb_keys, resp_start) = if data.starts_with(b"MOON") { + match crate::persistence::rdb::load_from_bytes(databases, &data) { + Ok((keys, consumed)) => { + info!( + "AOF RDB preamble loaded: {} keys ({} bytes)", + keys, consumed + ); + (keys, consumed) + } + Err(e) => { + // Data starts with MOON magic — it IS RDB format. + // Falling back to RESP would parse garbage. Propagate the error. + return Err(e); + } + } + } else { + (0, 0) + }; + + // If the entire file was RDB (no RESP tail), we're done + if resp_start >= data.len() { + return Ok(rdb_keys); + } + + let resp_data = &data[resp_start..]; + let total_len = resp_data.len(); + let mut buf = BytesMut::from(resp_data); + let config = ParseConfig::default(); + let mut selected_db: usize = 0; + let mut count: usize = 0; + let mut corruption_count: usize = 0; + + loop { + if buf.is_empty() { + break; + } + + match parse::parse(&mut buf, &config) { + Ok(Some(frame)) => { + // Extract command name and args, then dispatch + let (cmd, cmd_args) = match &frame { + Frame::Array(arr) if !arr.is_empty() => { + let name = match &arr[0] { + Frame::BulkString(s) => s.as_ref(), + Frame::SimpleString(s) => s.as_ref(), + _ => { + count += 1; + continue; + } + }; + (name as &[u8], &arr[1..]) + } + _ => { + count += 1; + continue; + } + }; + engine.replay_command(databases, cmd, cmd_args, &mut selected_db); + count += 1; + } + Ok(None) => { + // Incomplete frame at end of file - truncated AOF + if !buf.is_empty() { + let offset = total_len - buf.len(); + warn!( + "AOF truncated: {} unparseable bytes at offset {} (end of file)", + buf.len(), + offset + ); + } + break; + } + Err(e) => { + let error_offset = total_len - buf.len(); + warn!( + "AOF parse error at byte offset {} after {} commands: {}. Attempting skip.", + error_offset, count, e + ); + corruption_count += 1; + + // Skip past the corrupt byte(s) to the next RESP array marker ('*') + // Always discard at least 1 byte to guarantee forward progress. + let _ = buf.split_to(1); + if let Some(pos) = buf.iter().position(|&b| b == b'*') { + let _ = buf.split_to(pos); + } else if buf.is_empty() { + break; + } else { + // No more RESP array markers found; stop replay + warn!( + "AOF: no recoverable RESP frame found after offset {}; stopping", + error_offset + ); + break; + } + } + } + } + + if corruption_count > 0 { + warn!( + "AOF replay completed with {} corrupted entries skipped, {} commands replayed", + corruption_count, count + ); + } + + Ok(rdb_keys + count) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::persistence::replay::DispatchReplayEngine; + use ordered_float::OrderedFloat; + use tempfile::tempdir; + + fn make_command(parts: &[&[u8]]) -> Frame { + Frame::Array( + parts + .iter() + .map(|p| Frame::BulkString(Bytes::copy_from_slice(p))) + .collect(), + ) + } + + // --- serialize_command / generate_aof_command round-trip tests --- + + #[test] + fn test_generate_aof_command_produces_valid_resp_that_round_trips() { + let frame = make_command(&[b"SET", b"key", b"value"]); + let serialized = serialize_command(&frame); + + let mut buf = BytesMut::from(&serialized[..]); + let config = ParseConfig::default(); + let parsed = parse::parse(&mut buf, &config).unwrap().unwrap(); + assert_eq!(parsed, frame); + } + + #[test] + fn test_serialize_command_round_trip_hset() { + let frame = make_command(&[b"HSET", b"myhash", b"f1", b"v1"]); + let serialized = serialize_command(&frame); + let mut buf = BytesMut::from(&serialized[..]); + let parsed = parse::parse(&mut buf, &ParseConfig::default()) + .unwrap() + .unwrap(); + assert_eq!(parsed, frame); + } + + // --- AOF replay tests --- + + #[test] + fn test_aof_replay_set_commands_restores_string_keys() { + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("test.aof"); + + // Write SET commands in RESP format + let mut aof_data = BytesMut::new(); + serialize::serialize(&make_command(&[b"SET", b"k1", b"v1"]), &mut aof_data); + serialize::serialize(&make_command(&[b"SET", b"k2", b"v2"]), &mut aof_data); + std::fs::write(&aof_path, &aof_data).unwrap(); + + let mut dbs = vec![Database::new()]; + let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); + assert_eq!(count, 2); + + let entry = dbs[0].get(b"k1").unwrap(); + assert_eq!(entry.value.as_bytes().unwrap(), b"v1"); + let entry = dbs[0].get(b"k2").unwrap(); + assert_eq!(entry.value.as_bytes().unwrap(), b"v2"); + } + + #[test] + fn test_aof_replay_collection_types() { + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("test.aof"); + + let mut aof_data = BytesMut::new(); + // HSET + serialize::serialize( + &make_command(&[b"HSET", b"myhash", b"f1", b"v1"]), + &mut aof_data, + ); + // LPUSH + serialize::serialize( + &make_command(&[b"LPUSH", b"mylist", b"a", b"b"]), + &mut aof_data, + ); + // SADD + serialize::serialize( + &make_command(&[b"SADD", b"myset", b"x", b"y"]), + &mut aof_data, + ); + // ZADD + serialize::serialize( + &make_command(&[b"ZADD", b"myzset", b"1.5", b"alice"]), + &mut aof_data, + ); + std::fs::write(&aof_path, &aof_data).unwrap(); + + let mut dbs = vec![Database::new()]; + let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); + assert_eq!(count, 4); + + // Check hash + let hash = dbs[0].get_hash(b"myhash").unwrap().unwrap(); + assert_eq!( + hash.get(&Bytes::from_static(b"f1")).unwrap().as_ref(), + b"v1" + ); + + // Check list + let list = dbs[0].get_list(b"mylist").unwrap().unwrap(); + assert_eq!(list.len(), 2); + + // Check set + let set = dbs[0].get_set(b"myset").unwrap().unwrap(); + assert_eq!(set.len(), 2); + + // Check sorted set + let (members, _) = dbs[0].get_sorted_set(b"myzset").unwrap().unwrap(); + assert_eq!(*members.get(&Bytes::from_static(b"alice")).unwrap(), 1.5); + } + + #[test] + fn test_aof_replay_with_expire_preserves_ttls() { + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("test.aof"); + + let mut aof_data = BytesMut::new(); + serialize::serialize(&make_command(&[b"SET", b"mykey", b"myval"]), &mut aof_data); + serialize::serialize( + &make_command(&[b"PEXPIRE", b"mykey", b"60000"]), + &mut aof_data, + ); + std::fs::write(&aof_path, &aof_data).unwrap(); + + let mut dbs = vec![Database::new()]; + let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); + assert_eq!(count, 2); + + let base_ts = dbs[0].base_timestamp(); + let entry = dbs[0].get(b"mykey").unwrap(); + assert!(entry.has_expiry()); + let remaining_secs = (entry.expires_at_ms(base_ts) - current_time_ms()) / 1000; + assert!(remaining_secs >= 50); // Allow some tolerance + } + + #[test] + fn test_aof_replay_with_select_switches_databases() { + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("test.aof"); + + let mut aof_data = BytesMut::new(); + serialize::serialize(&make_command(&[b"SET", b"k0", b"v0"]), &mut aof_data); + serialize::serialize(&make_command(&[b"SELECT", b"1"]), &mut aof_data); + serialize::serialize(&make_command(&[b"SET", b"k1", b"v1"]), &mut aof_data); + std::fs::write(&aof_path, &aof_data).unwrap(); + + let mut dbs = vec![Database::new(), Database::new()]; + let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); + assert_eq!(count, 3); + + assert!(dbs[0].get(b"k0").is_some()); + assert!(dbs[1].get(b"k1").is_some()); + } + + #[test] + fn test_aof_replay_empty_file_produces_zero_keys() { + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("test.aof"); + std::fs::write(&aof_path, b"").unwrap(); + + let mut dbs = vec![Database::new()]; + let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); + assert_eq!(count, 0); + assert_eq!(dbs[0].len(), 0); + } + + #[test] + fn test_aof_replay_corrupt_truncated_logs_error_loads_what_it_can() { + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("test.aof"); + + let mut aof_data = BytesMut::new(); + serialize::serialize(&make_command(&[b"SET", b"k1", b"v1"]), &mut aof_data); + // Append corrupt data + aof_data.extend_from_slice(b"*3\r\n$3\r\nSET\r\n$2\r\nk2"); + std::fs::write(&aof_path, &aof_data).unwrap(); + + let mut dbs = vec![Database::new()]; + let count = replay_aof(&mut dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); + // Should have loaded the first command + assert_eq!(count, 1); + assert!(dbs[0].get(b"k1").is_some()); + } + + // --- FsyncPolicy tests --- + + #[test] + fn test_fsync_policy_from_str() { + assert_eq!(FsyncPolicy::from_str("always"), FsyncPolicy::Always); + assert_eq!(FsyncPolicy::from_str("everysec"), FsyncPolicy::EverySec); + assert_eq!(FsyncPolicy::from_str("no"), FsyncPolicy::No); + assert_eq!(FsyncPolicy::from_str("unknown"), FsyncPolicy::EverySec); + } + + // --- generate_rewrite_commands tests --- + + #[test] + fn test_generate_rewrite_commands_all_5_types() { + let mut dbs = vec![Database::new()]; + + // String + dbs[0].set_string(Bytes::from_static(b"str"), Bytes::from_static(b"val")); + // Hash + { + let map = dbs[0].get_or_create_hash(b"h").unwrap(); + map.insert(Bytes::from_static(b"f"), Bytes::from_static(b"v")); + } + // List + { + let list = dbs[0].get_or_create_list(b"l").unwrap(); + list.push_back(Bytes::from_static(b"item")); + } + // Set + { + let set = dbs[0].get_or_create_set(b"s").unwrap(); + set.insert(Bytes::from_static(b"m")); + } + // Sorted set + { + let (members, tree) = dbs[0].get_or_create_sorted_set(b"z").unwrap(); + members.insert(Bytes::from_static(b"a"), 1.0); + tree.insert(OrderedFloat(1.0), Bytes::from_static(b"a")); + } + + let commands = generate_rewrite_commands(&dbs); + assert!(!commands.is_empty()); + + // Replay and verify round-trip + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("rewrite.aof"); + std::fs::write(&aof_path, &commands).unwrap(); + + let mut loaded_dbs = vec![Database::new()]; + let count = replay_aof(&mut loaded_dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); + assert!(count >= 5, "Expected at least 5 commands, got {}", count); + + // Verify each type restored + assert_eq!( + loaded_dbs[0].get(b"str").unwrap().value.type_name(), + "string" + ); + assert!(loaded_dbs[0].get_hash(b"h").unwrap().is_some()); + assert!(loaded_dbs[0].get_list(b"l").unwrap().is_some()); + assert!(loaded_dbs[0].get_set(b"s").unwrap().is_some()); + assert!(loaded_dbs[0].get_sorted_set(b"z").unwrap().is_some()); + } + + #[test] + fn test_generate_rewrite_commands_with_ttl() { + let mut dbs = vec![Database::new()]; + let future_ms = current_time_ms() + 3_600_000; + dbs[0].set_string_with_expiry( + Bytes::from_static(b"key"), + Bytes::from_static(b"val"), + future_ms, + ); + + let commands = generate_rewrite_commands(&dbs); + + // Replay and check TTL is preserved + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("rewrite.aof"); + std::fs::write(&aof_path, &commands).unwrap(); + + let mut loaded_dbs = vec![Database::new()]; + let count = replay_aof(&mut loaded_dbs, &aof_path, &DispatchReplayEngine::new()).unwrap(); + assert_eq!(count, 2); // SET + PEXPIRE + + let base_ts = loaded_dbs[0].base_timestamp(); + let entry = loaded_dbs[0].get(b"key").unwrap(); + assert!(entry.has_expiry()); + let remaining_secs = (entry.expires_at_ms(base_ts) - current_time_ms()) / 1000; + assert!(remaining_secs > 3500); + } + + /// Helper: build a snapshot tuple from a Database slice — mirrors the + /// `merged` construction in `do_rewrite_sharded` (aof.rs:2070-2090) so + /// tests exercise the exact same production path. + fn db_slice_to_snapshot( + dbs: &[Database], + ) -> Vec<( + Vec<( + crate::storage::compact_key::CompactKey, + crate::storage::entry::Entry, + )>, + u32, + )> { + let now_ms = crate::storage::entry::current_time_ms(); + dbs.iter() + .map(|db| { + let base_ts = db.base_timestamp(); + let entries: Vec<_> = db + .data() + .iter() + .filter(|(_, e)| !e.is_expired_at(base_ts, now_ms)) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + (entries, base_ts) + }) + .collect() + } + + /// FIX-W3-8: BGREWRITEAOF on a fresh empty database must produce a valid + /// RDB base and recover cleanly with 0 keys. + /// + /// Updated to call `save_snapshot_to_bytes` (the function `do_rewrite_sharded` + /// actually calls at aof.rs:2096) rather than `save_to_bytes` (the previous + /// test used the wrong function — tautological for empty input since both + /// produce an identical valid RDB, but would miss regressions on the + /// snapshot-tuple path). + #[test] + fn empty_database_rewrite_produces_valid_rdb_and_recovers() { + let dir = tempdir().unwrap(); + + // Build the snapshot tuple the same way do_rewrite_sharded does. + let empty_dbs: Vec = vec![Database::new()]; + let snapshot = db_slice_to_snapshot(&empty_dbs); + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot) + .expect("save_snapshot_to_bytes must succeed for empty snapshot"); + + // Invariant 1: RDB is non-empty (has at least magic + version + EOF marker). + assert!( + !rdb_bytes.is_empty(), + "empty-database RDB must not be 0 bytes" + ); + + // Invariant 2: starts with valid MOON magic header. + assert!( + rdb_bytes.starts_with(b"MOON"), + "RDB bytes must start with MOON magic, got: {:?}", + &rdb_bytes[..rdb_bytes.len().min(8)] + ); + + // Invariant 3: recovery from this base succeeds with 0 keys loaded. + let base_path = dir.path().join("empty.rdb"); + std::fs::write(&base_path, &rdb_bytes).expect("write empty rdb"); + let mut recovery_dbs = vec![Database::new()]; + let loaded = + crate::persistence::rdb::load(&mut recovery_dbs, &base_path).expect("load empty rdb"); + assert_eq!( + loaded, 0, + "recovering from empty-database RDB yields 0 keys" + ); + assert_eq!( + recovery_dbs[0].len(), + 0, + "database must be empty after recovering from zero-key RDB" + ); + } + + /// FIX-W3-8: Genuine regression guard — save_snapshot_to_bytes preserves + /// a 1-key database through a full serialize→file→reload cycle. + /// + /// This is the substantive test the verifier asked for: verifies the + /// production code path (`save_snapshot_to_bytes` via the snapshot-tuple + /// form) against a non-trivial database, so a future regression that + /// swaps back to `save_to_bytes` or breaks TTL handling in the snapshot + /// path will be caught. + #[test] + fn snapshot_to_bytes_round_trips_one_key_database() { + let dir = tempdir().unwrap(); + + // Build a 1-key database with a string value. + let mut db = Database::new(); + db.set_string( + Bytes::from_static(b"rdb_key"), + Bytes::from_static(b"rdb_value"), + ); + let dbs = vec![db]; + + // Serialize via the production path (save_snapshot_to_bytes). + let snapshot = db_slice_to_snapshot(&dbs); + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot) + .expect("save_snapshot_to_bytes must succeed for 1-key snapshot"); + + assert!(rdb_bytes.starts_with(b"MOON"), "must have MOON magic"); + + // Reload and assert the key survives. + let base_path = dir.path().join("one_key.rdb"); + std::fs::write(&base_path, &rdb_bytes).expect("write rdb"); + let mut recovery_dbs = vec![Database::new()]; + let loaded = crate::persistence::rdb::load(&mut recovery_dbs, &base_path) + .expect("load rdb must succeed"); + assert_eq!(loaded, 1, "exactly 1 key must be recovered"); + let val = recovery_dbs[0] + .get(b"rdb_key") + .expect("rdb_key must be present"); + assert_eq!( + val.value.as_bytes().expect("string value"), + b"rdb_value", + "recovered value must match written value" + ); + } + + #[test] + fn test_generate_rewrite_round_trip_preserves_state() { + let mut dbs = vec![Database::new()]; + dbs[0].set_string(Bytes::from_static(b"a"), Bytes::from_static(b"1")); + dbs[0].set_string(Bytes::from_static(b"b"), Bytes::from_static(b"2")); + { + let list = dbs[0].get_or_create_list(b"mylist").unwrap(); + list.push_back(Bytes::from_static(b"x")); + list.push_back(Bytes::from_static(b"y")); + list.push_back(Bytes::from_static(b"z")); + } + + let commands = generate_rewrite_commands(&dbs); + let dir = tempdir().unwrap(); + let aof_path = dir.path().join("rewrite.aof"); + std::fs::write(&aof_path, &commands).unwrap(); + + let mut loaded = vec![Database::new()]; + replay_aof(&mut loaded, &aof_path, &DispatchReplayEngine::new()).unwrap(); + + // Check strings + assert_eq!(loaded[0].get(b"a").unwrap().value.as_bytes().unwrap(), b"1"); + assert_eq!(loaded[0].get(b"b").unwrap().value.as_bytes().unwrap(), b"2"); + + // Check list order preserved + let list = loaded[0].get_list(b"mylist").unwrap().unwrap(); + assert_eq!(list.len(), 3); + assert_eq!(list[0].as_ref(), b"x"); + assert_eq!(list[1].as_ref(), b"y"); + assert_eq!(list[2].as_ref(), b"z"); + } + + // ----------------------------------------------------------------------- + // FIX-W2-4 r2: canonical AOF fsync error string + // + // Red criterion: AOF_FSYNC_ERR constant must exist and equal the canonical + // Redis-style ERR-prefixed string used by handler_monoio and handler_sharded. + // handler_single.rs previously used "WRITEFAIL aof fsync failed" which is + // both non-canonical (no ERR prefix, different verb) and inconsistent with + // the other two handlers. + // + // These tests compile-fail on the prior commit (constant absent) and pass + // once AOF_FSYNC_ERR is declared in this module with the correct value. + // ----------------------------------------------------------------------- + + #[test] + fn aof_fsync_err_constant_is_canonical() { + // The canonical error frame bytes sent to the client when an AOF + // fsync under appendfsync=always fails. Must match what + // handler_monoio/mod.rs and handler_sharded/mod.rs use. + assert_eq!( + AOF_FSYNC_ERR, b"ERR AOF fsync failed; write not durable", + "AOF_FSYNC_ERR must equal the canonical ERR-prefixed string" + ); + } + + #[test] + fn aof_fsync_err_has_err_prefix() { + // Redis convention: protocol-level errors must start with a word + // followed by a space, using `ERR` for generic errors. `WRITEFAIL` + // is not a standard Redis error prefix and confuses clients that + // pattern-match on error codes. + assert!( + AOF_FSYNC_ERR.starts_with(b"ERR "), + "AOF_FSYNC_ERR must start with 'ERR ' (got {:?})", + std::str::from_utf8(AOF_FSYNC_ERR).unwrap_or("") + ); + } +} diff --git a/src/persistence/aof/pool.rs b/src/persistence/aof/pool.rs new file mode 100644 index 00000000..f97d34e5 --- /dev/null +++ b/src/persistence/aof/pool.rs @@ -0,0 +1,1382 @@ +//! AOF writer pool: per-shard / single writer-task handles and backpressure. +#![allow(unused_imports, unused_variables, unreachable_code, clippy::empty_loop)] + +use super::rewrite::{do_rewrite_per_shard, drain_pending_appends_framed}; +use super::*; +// `drain_pending_appends` (non-framed) exists only under the monoio runtime. +#[cfg(feature = "runtime-monoio")] +use super::rewrite::drain_pending_appends; + +#[derive(Clone)] +pub struct AofWriterPool { + senders: Vec>, + layout: crate::persistence::aof_manifest::AofLayout, + /// Fsync policy configured at writer-task construction. Read on the + /// hot append path: `Always` routes through `AppendSync` for + /// fsync-before-ack durability (H1 fix); everything else stays on + /// the fire-and-forget `Append` path. + fsync_policy: FsyncPolicy, + /// F2: max time `try_send_append_durable` waits for the `Always` fsync + /// ack before failing the write. `Duration::ZERO` means unbounded + /// (legacy behavior). Prevents a stalled disk from parking write + /// connections forever (design-for-failure). + fsync_timeout: Duration, + /// F6: persistence base dir (the parent of `appendonlydir/`), set only for + /// PerShard pools that may service a per-shard BGREWRITEAOF. Needed to load + /// the authoritative manifest fresh at rewrite time. `None` for TopLevel + /// pools and test pools that never rewrite. + base_dir: Option, +} + +impl AofWriterPool { + /// Build a TopLevel pool from a single existing writer sender. Used for + /// legacy v1 deployments and `--shards 1` v2 deployments where one writer + /// thread services every shard. + pub fn top_level(sender: channel::MpscSender) -> Arc { + Self::top_level_with_policy(sender, FsyncPolicy::EverySec, DEFAULT_AOF_FSYNC_TIMEOUT) + } + + /// Same as [`Self::top_level`] but with an explicit fsync policy. The + /// policy controls whether [`Self::try_send_append_durable`] takes the + /// fast (fire-and-forget) or rendezvous (`AppendSync`) path. + /// `fsync_timeout` bounds the `Always` ack await (F2); `Duration::ZERO` + /// = unbounded. + pub fn top_level_with_policy( + sender: channel::MpscSender, + fsync_policy: FsyncPolicy, + fsync_timeout: Duration, + ) -> Arc { + Arc::new(Self { + senders: vec![sender], + layout: crate::persistence::aof_manifest::AofLayout::TopLevel, + fsync_policy, + fsync_timeout, + base_dir: None, + }) + } + + /// Build a PerShard pool from N senders. `senders[i]` MUST be the writer + /// task that owns `appendonlydir/shard-{i}/`. The vector's length is the + /// shard count; passing a length-1 vector here is a bug — use + /// [`AofWriterPool::top_level`] instead. + pub fn per_shard(senders: Vec>) -> Arc { + Self::per_shard_with_policy(senders, FsyncPolicy::EverySec, DEFAULT_AOF_FSYNC_TIMEOUT) + } + + /// Same as [`Self::per_shard`] but with an explicit fsync policy. + /// `fsync_timeout` bounds the `Always` ack await (F2); `Duration::ZERO` + /// = unbounded. + pub fn per_shard_with_policy( + senders: Vec>, + fsync_policy: FsyncPolicy, + fsync_timeout: Duration, + ) -> Arc { + debug_assert!( + senders.len() >= 2, + "per_shard pool needs >=2 writers; use top_level for single-writer" + ); + Arc::new(Self { + senders, + layout: crate::persistence::aof_manifest::AofLayout::PerShard, + fsync_policy, + fsync_timeout, + base_dir: None, + }) + } + + /// F6: same as [`Self::per_shard_with_policy`] but records the persistence + /// `base_dir` so a per-shard BGREWRITEAOF can load the authoritative + /// manifest fresh at rewrite time. This is the production constructor used + /// by `main.rs` for the PerShard layout. + pub fn per_shard_with_base_dir( + senders: Vec>, + fsync_policy: FsyncPolicy, + fsync_timeout: Duration, + base_dir: PathBuf, + ) -> Arc { + debug_assert!( + senders.len() >= 2, + "per_shard pool needs >=2 writers; use top_level for single-writer" + ); + Arc::new(Self { + senders, + layout: crate::persistence::aof_manifest::AofLayout::PerShard, + fsync_policy, + fsync_timeout, + base_dir: Some(base_dir), + }) + } + + /// Returns the configured fsync policy. Hot-path callers read this to + /// decide between the fast (`try_send_append`) and durable + /// (`try_send_append_sync`) write paths. + #[inline] + pub fn fsync_policy(&self) -> FsyncPolicy { + self.fsync_policy + } + + /// Policy-aware AOF append. For `FsyncPolicy::Always`, this awaits + /// `AppendSync` and returns `Ok(())` only after `sync_data()` confirms + /// the entry is on durable storage — closing the H1 in-flight loss + /// vector identified in the investigation report. For `EverySec` and + /// `No`, it stays on the fire-and-forget path (zero new latency). + /// + /// Returns `Err(AofAck)` only on the Always path when the write or + /// fsync failed (or the writer task is gone). Callers MUST treat + /// `Err(_)` as a hard failure — return an error frame to the client, + /// do NOT respond `+OK`. + /// + /// Async because the Always branch awaits a oneshot receiver. The + /// non-Always branch resolves immediately (no actual suspension) so + /// the only overhead is one `match` and the implicit Future state + /// machine; benchmarked at ~5 ns per call on the EverySec hot path, + /// far below the per-write WAL/replication cost. + #[inline] + pub async fn try_send_append_durable( + &self, + shard_id: usize, + lsn: u64, + bytes: Bytes, + ) -> Result<(), AofAck> { + match self.fsync_policy { + FsyncPolicy::Always => { + let rx = self.try_send_append_sync(shard_id, lsn, bytes); + // F2 (design-for-failure): bound the wait so a stalled disk + // can't park this connection forever. On elapse the write is + // failed — the entry may still land on disk later, but + // durability is NOT confirmed, so the caller must not report + // success. `Duration::ZERO` keeps the legacy unbounded await. + match Self::await_ack(rx, self.fsync_timeout).await { + AckOutcome::Ack(AofAck::Synced) => Ok(()), + AckOutcome::Ack(other) => Err(other), + // Writer task gone / channel disconnected. + AckOutcome::Disconnected => Err(AofAck::WriteFailed), + // Fsync did not confirm within the bound. + AckOutcome::TimedOut => Err(AofAck::FsyncFailed), + } + } + FsyncPolicy::EverySec | FsyncPolicy::No => { + self.try_send_append(shard_id, lsn, bytes); + Ok(()) + } + } + } + + /// Await an `AppendSync` ack receiver under a bounded timeout (F2). + /// + /// `timeout == Duration::ZERO` preserves the legacy unbounded await + /// (used when the operator explicitly opts out via + /// `--aof-fsync-timeout-ms 0`). Otherwise the await is capped by the + /// runtime-appropriate timer; on elapse the in-flight fsync is + /// abandoned (the receiver is dropped) and `TimedOut` is returned. + /// + /// Runtime-agnostic: monoio uses `select! { rx, sleep }` (matching the + /// established `cluster::failover` pattern — monoio 0.2 has no + /// `time::timeout`); tokio uses `tokio::time::timeout`. Both resolve the + /// ack first if it arrives within the bound, otherwise `TimedOut`. + async fn await_ack( + rx: crate::runtime::channel::OneshotReceiver, + timeout: Duration, + ) -> AckOutcome { + if timeout.is_zero() { + return match rx.await { + Ok(ack) => AckOutcome::Ack(ack), + Err(_) => AckOutcome::Disconnected, + }; + } + + #[cfg(feature = "runtime-monoio")] + { + monoio::select! { + res = rx => match res { + Ok(ack) => AckOutcome::Ack(ack), + Err(_) => AckOutcome::Disconnected, + }, + _ = monoio::time::sleep(timeout) => AckOutcome::TimedOut, + } + } + #[cfg(all(feature = "runtime-tokio", not(feature = "runtime-monoio")))] + { + match tokio::time::timeout(timeout, rx).await { + Ok(Ok(ack)) => AckOutcome::Ack(ack), + Ok(Err(_)) => AckOutcome::Disconnected, + Err(_) => AckOutcome::TimedOut, + } + } + } + + /// Return the writer sender that owns the given shard's AOF file. + /// + /// For TopLevel pools, `shard_id` is ignored — all shards multiplex onto + /// the single sender. For PerShard pools, `shard_id` MUST be in range + /// `[0, num_writers())`; an out-of-range id is a programmer error and + /// panics in debug builds. + #[inline] + pub fn sender(&self, shard_id: usize) -> &channel::MpscSender { + use crate::persistence::aof_manifest::AofLayout; + match self.layout { + AofLayout::TopLevel => &self.senders[0], + AofLayout::PerShard => { + debug_assert!( + shard_id < self.senders.len(), + "shard_id {} out of range for per-shard pool of size {}", + shard_id, + self.senders.len() + ); + &self.senders[shard_id] + } + } + } + + /// Fire-and-forget append for the given shard, tagged with the LSN that + /// was issued for this write (see [`AofMessage::Append`] docs for LSN + /// semantics per layout). Call sites must source `lsn` from + /// `ReplicationState::issue_lsn(shard_id, bytes.len() as u64)` for writes + /// that participate in replication ordering; sites without a + /// replication-state handle pass 0. + #[inline] + pub fn try_send_append(&self, shard_id: usize, lsn: u64, bytes: Bytes) { + let _ = self + .sender(shard_id) + .try_send(AofMessage::Append { lsn, bytes }); + } + + /// Synchronous (fsync-before-ack) append for `appendfsync=always` + /// durability (RFC § 4 — Fix 2). Returns a receiver the caller MUST + /// await before responding to the client; `AofAck::Synced` means the + /// entry is on durable storage. + /// + /// **Failure handling:** if the write or fsync fails, the receiver + /// resolves with `AofAck::WriteFailed` / `AofAck::FsyncFailed`. If + /// the writer task is gone (shutdown / channel disconnect), the + /// receiver resolves with `Err(RecvError)`. In every failure mode the + /// caller MUST return an error frame to the client, NOT `+OK`. + /// + /// **Performance:** every call adds a writer round-trip plus an + /// fsync syscall on the critical path. This is the explicit Redis + /// contract for `appendfsync=always`; callers should gate on the + /// configured policy and prefer [`Self::try_send_append`] for + /// `everysec`/`no`. + /// + /// **`shard_id` semantics:** matches [`Self::try_send_append`] — for + /// TopLevel the parameter is ignored, for PerShard it routes to + /// `senders[shard_id]`. + pub fn try_send_append_sync( + &self, + shard_id: usize, + lsn: u64, + bytes: Bytes, + ) -> crate::runtime::channel::OneshotReceiver { + let (ack_tx, ack_rx) = crate::runtime::channel::oneshot::(); + match self.sender(shard_id).try_send(AofMessage::AppendSync { + lsn, + bytes, + ack: ack_tx, + }) { + Ok(()) => {} + Err(flume::TrySendError::Full(_)) => { + // Writer channel is at capacity — count the dropped entry and + // signal ChannelFull back to the caller via a pre-filled + // oneshot so the caller's `.await` resolves immediately to + // Err(AofAck::ChannelFull) without a writer round-trip. + AOF_BACKPRESSURE_DROPPED.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + warn!( + "AOF writer channel full (shard {}): AppendSync dropped; \ + backpressure_dropped={}", + shard_id, + AOF_BACKPRESSURE_DROPPED.load(std::sync::atomic::Ordering::Relaxed), + ); + // Pre-send ChannelFull into a fresh oneshot pair; the + // caller's `ack_rx` was already returned — we create a + // new pair and use its sender to pre-fill what the caller + // will receive. The original ack_tx (inside the dropped + // AppendSync) is dropped, causing its ack_rx to yield + // RecvError. We send ChannelFull via the *returned* ack_rx + // by using a second oneshot whose sender is immediately + // fulfilled, then return that receiver instead. + let (pre_tx, pre_rx) = crate::runtime::channel::oneshot::(); + let _ = pre_tx.send(AofAck::ChannelFull); + return pre_rx; + } + Err(flume::TrySendError::Disconnected(_)) => { + // Writer task is dead — let caller handle RecvError on ack_rx. + // ack_tx was dropped inside the Err value; ack_rx will + // resolve with RecvError, which try_send_append_durable maps + // to Err(AofAck::WriteFailed). + } + } + ack_rx + } + + /// Fire-and-forget append for a cross-shard atomic operation (RFC § 2 + /// Rule 2 — `OrderedAcrossShards` tagging). + /// + /// The high bit of `lsn` (`1 << 63`) is set before the entry is queued. + /// Recovery uses this bit to recognize cross-shard atomic entries, + /// buffer them per-shard, and replay them globally in LSN order after + /// per-shard replay completes — guaranteeing TXN/SCRIPT atomicity + /// survives a crash even when multiple shards participated. + /// + /// **Caller contract:** `lsn` MUST be < `1 << 63` (i.e. the high bit + /// MUST be clear when passed in). Practical LSN ceilings — even at + /// 10 M writes/s sustained for a century — sit around 2^58, so any + /// real LSN satisfies this. Debug builds assert; release builds mask + /// the input to keep the wire format well-formed rather than + /// corrupt-by-zero-extending. + /// + /// **Production callers today:** none. Step 5 ships the infrastructure + /// (writer, framing flag, recovery merge) so a future cross-shard TXN + /// or replicated SCRIPT command has a place to land. Until that + /// consumer exists, only test code emits ordered entries. + #[inline] + pub fn try_send_append_ordered(&self, shard_id: usize, lsn: u64, bytes: Bytes) { + debug_assert_eq!( + lsn & ORDERED_LSN_FLAG, + 0, + "try_send_append_ordered: lsn must not have the high bit set; got {:#x}", + lsn, + ); + let tagged_lsn = (lsn & !ORDERED_LSN_FLAG) | ORDERED_LSN_FLAG; + let _ = self.sender(shard_id).try_send(AofMessage::Append { + lsn: tagged_lsn, + bytes, + }); + } + + /// Issue an LSN for an AOF append at every call site that has the + /// `Option>>` shape. Wraps + /// `ReplicationState::issue_lsn` so handler call sites collapse to a + /// single line. + /// + /// Returns 0 when: + /// - `repl_state` is None (test fixtures or shutdown paths) + /// - the `RwLock` is poisoned (shouldn't happen in production — + /// ReplicationState is only `write()`-locked under known-safe paths) + /// + /// 0 is a sentinel meaning "no replication ordering for this write". + /// TopLevel writers ignore the LSN entirely so 0 is harmless there; + /// PerShard writers treat 0 the same as any other LSN (per-shard order + /// is preserved by write order, not by LSN value). The LSN only matters + /// for the cross-shard `OrderedAcrossShards` merge in RFC step 5. + #[inline] + pub fn issue_append_lsn( + repl_state: &Option>>, + shard_id: usize, + delta: usize, + ) -> u64 { + repl_state + .as_ref() + .and_then(|rs| rs.read().ok().map(|g| g.issue_lsn(shard_id, delta as u64))) + .unwrap_or(0) + } + + /// Submit a Rewrite/RewriteSharded message. Only legal for TopLevel pools; + /// PerShard rewrites are per-shard operations and must be initiated by + /// the BGREWRITEAOF code path in step 6, not via this enum variant. + pub fn try_send_rewrite(&self, msg: AofMessage) -> Result<(), AofPoolSendError> { + use crate::persistence::aof_manifest::AofLayout; + debug_assert!( + matches!(msg, AofMessage::Rewrite(_) | AofMessage::RewriteSharded(_)), + "try_send_rewrite called with a non-Rewrite variant", + ); + if self.layout == AofLayout::PerShard { + return Err(AofPoolSendError::RewriteUnsupportedInPerShard); + } + self.senders[0] + .try_send(msg) + .map_err(|_| AofPoolSendError::SendFailed) + } + + /// [F6] Initiate a per-shard BGREWRITEAOF across every writer in a + /// PerShard pool. + /// + /// Loads the authoritative manifest fresh from `base_dir` (normal appends + /// never mutate the manifest, and BGREWRITEAOF is CAS-serialized by + /// `AOF_REWRITE_IN_PROGRESS`, so a fresh load is the current committed + /// state), builds a shared [`PerShardRewriteCoord`] that advances the + /// generation by one, and hands every writer the same `coord` + a cheap + /// `Arc` clone of `shard_dbs`. + /// + /// **Reliable delivery (design-for-failure):** the fan-out uses the + /// *blocking* `send` rather than `try_send`. A dropped rewrite message + /// would leave the countdown unable to reach zero — folded writers would + /// have reopened to new-seq files that the manifest never commits, silently + /// losing their post-rewrite appends. The writers run on dedicated threads + /// draining continuously, so `send` blocks only until a channel slot frees + /// (sub-millisecond), which is acceptable for a rare admin command. + /// + /// Returns `SendFailed` if `base_dir` is unset, the manifest can't be + /// loaded, or a writer thread is gone (disconnected channel). On the last + /// case the rewrite aborts WITHOUT committing — the old generation stays + /// authoritative (crash-safe), but a dead writer already means that shard's + /// persistence was compromised before this call. + pub fn try_send_rewrite_per_shard( + &self, + shard_dbs: Arc, + ) -> Result<(), AofPoolSendError> { + use crate::persistence::aof_manifest::{AofLayout, AofManifest}; + if self.layout != AofLayout::PerShard { + // A TopLevel pool rewrites via try_send_rewrite; this entry point + // is PerShard-only. + return Err(AofPoolSendError::RewriteUnsupportedInPerShard); + } + let base_dir = self.base_dir.as_ref().ok_or(AofPoolSendError::SendFailed)?; + let manifest = match AofManifest::load(base_dir) { + Ok(Some(m)) if m.layout == AofLayout::PerShard => m, + Ok(_) => { + error!( + "F6 per-shard rewrite: manifest at {} missing or not PerShard; aborting", + base_dir.display() + ); + return Err(AofPoolSendError::SendFailed); + } + Err(e) => { + error!( + "F6 per-shard rewrite: failed to load manifest at {}: {}", + base_dir.display(), + e + ); + return Err(AofPoolSendError::SendFailed); + } + }; + let current_seq = manifest.seq; + let n_shards = self.senders.len(); + let shared_manifest = Arc::new(parking_lot::Mutex::new(manifest)); + let coord = PerShardRewriteCoord::new(shared_manifest, current_seq, n_shards); + for (idx, s) in self.senders.iter().enumerate() { + // Blocking send for guaranteed delivery — see the doc comment. + if s.send(AofMessage::RewritePerShard { + shard_dbs: shard_dbs.clone(), + coord: coord.clone(), + }) + .is_err() + { + error!( + "F6 per-shard rewrite: writer {} channel disconnected; \ + rewrite aborted (no manifest commit, old generation remains \ + authoritative). Inspect AOF writer threads.", + idx + ); + // Account for the shards that did NOT receive the message (this + // one plus the unsent tail). The writers that DID receive it will + // fold and call `shard_done`; decrementing for the rest here lets + // the countdown still reach zero, so the final `shard_done` runs + // the abort path (keeps `old_seq`, clears AOF_REWRITE_IN_PROGRESS) + // instead of wedging the rewrite flag forever. + coord.mark_failed(); + for _ in idx..n_shards { + coord.shard_done(); + } + return Err(AofPoolSendError::SendFailed); + } + } + info!( + "F6 per-shard rewrite dispatched: seq {} -> {} across {} shards", + current_seq, + current_seq + 1, + n_shards + ); + Ok(()) + } + + /// Broadcast `Shutdown` to every writer. Used by orchestrated shutdown + /// paths in `main.rs`/`embedded.rs`. Each writer drains its channel and + /// fsyncs before exiting. + pub fn broadcast_shutdown(&self) { + for s in &self.senders { + let _ = s.try_send(AofMessage::Shutdown); + } + } + + /// Number of underlying writer senders. 1 for TopLevel, num_shards for + /// PerShard. + #[inline] + pub fn num_writers(&self) -> usize { + self.senders.len() + } + + /// Reports the pool's layout. Useful for places that need to refuse + /// PerShard-incompatible legacy code paths with a clear error. + #[inline] + pub fn layout(&self) -> crate::persistence::aof_manifest::AofLayout { + self.layout + } +} + +#[cfg(test)] +mod pool_tests { + use super::*; + use crate::persistence::aof_manifest::AofLayout; + use crate::runtime::channel; + + /// A per-shard BGREWRITEAOF fan-out that fails partway (a writer channel is + /// disconnected) must NOT leave `AOF_REWRITE_IN_PROGRESS` stuck true. The + /// failed/unsent shards are accounted for (`mark_failed` + `shard_done`) so + /// the countdown still reaches zero and the abort path clears the flag, + /// leaving the old generation authoritative. + #[test] + fn rewrite_fan_out_partial_failure_clears_in_progress_flag() { + use crate::command::persistence::AOF_REWRITE_IN_PROGRESS; + use std::sync::atomic::Ordering; + + let tmp = tempfile::tempdir().unwrap(); + // PerShard manifest on disk so the rewrite reaches the fan-out loop. + crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); + + // Disconnect shard-0's writer so the FIRST send fails (before any + // successful send), making the abort accounting deterministic. + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + drop(rx0); + let pool = AofWriterPool::per_shard_with_base_dir( + vec![tx0, tx1], + FsyncPolicy::EverySec, + DEFAULT_AOF_FSYNC_TIMEOUT, + tmp.path().to_path_buf(), + ); + + let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![ + vec![crate::storage::Database::new()], + vec![crate::storage::Database::new()], + ]); + + // The command handler sets this before dispatching the rewrite. + AOF_REWRITE_IN_PROGRESS.store(true, Ordering::SeqCst); + let res = pool.try_send_rewrite_per_shard(shard_dbs); + assert!(res.is_err(), "disconnected writer must fail the fan-out"); + assert!( + !AOF_REWRITE_IN_PROGRESS.load(Ordering::SeqCst), + "partial fan-out failure must clear AOF_REWRITE_IN_PROGRESS, not wedge it" + ); + } + + /// A per-shard rewrite that ABORTS (another shard failed to fold) must roll + /// THIS shard's append file back onto the committed old generation. Phase 6 + /// reopens `*file` onto the new-seq incr; if the rewrite then aborts, the + /// manifest keeps `old_seq` and prunes the new-seq files — so without a + /// barrier-before-resume the writer keeps appending into a discarded incr + /// that recovery ignores (silent data loss). This test drives one shard's + /// real fold to completion with a second shard pre-marked failed, then proves + /// post-abort appends land in the COMMITTED old-gen incr. + #[test] + fn rewrite_abort_reopens_writer_onto_committed_old_generation() { + use crate::command::persistence::AOF_REWRITE_IN_PROGRESS; + use std::io::Write; + use std::sync::atomic::Ordering; + + let tmp = tempfile::tempdir().unwrap(); + // 2-shard manifest at seq=1 on disk; share it through the coord's Arc. + let manifest = + crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); + let old_seq = manifest.seq; + let old_incr_s0 = manifest.shard_incr_path_seq(0, old_seq); + assert!(old_incr_s0.exists(), "old-gen incr must exist pre-rewrite"); + let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); + + let coord = PerShardRewriteCoord::new(manifest.clone(), old_seq, 2); + let new_seq = coord.new_seq(); + assert_ne!(new_seq, old_seq); + + // Shard 1 "fails to fold": mark_failed + shard_done (countdown 2 -> 1). + // Shard 0's shard_done (inside do_rewrite_per_shard) will then be the + // terminal decrement and take the abort path. + coord.mark_failed(); + coord.shard_done(); + + // Shard 0's append file starts on the OLD incr (where the live writer + // appends before a rewrite). Empty databases keep the fold trivial; the + // reopen behaviour is independent of key count. + let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![ + vec![crate::storage::Database::new()], + vec![crate::storage::Database::new()], + ]); + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&old_incr_s0) + .unwrap(); + let (_tx, rx) = channel::mpsc_bounded::(4); + + AOF_REWRITE_IN_PROGRESS.store(true, Ordering::SeqCst); + // The fold itself succeeds; aborting is a coordinator decision, so this + // returns Ok even though the rewrite is discarded. + do_rewrite_per_shard(0, &shard_dbs, &mut file, &rx, &coord).unwrap(); + + // Abort kept the old generation committed and pruned the new-gen incr. + assert_eq!(manifest.lock().seq, old_seq, "abort must keep old_seq"); + let new_incr_s0 = manifest.lock().shard_incr_path_seq(0, new_seq); + assert!( + !new_incr_s0.exists(), + "aborted new-gen incr must be pruned (the dangling target)" + ); + + // The writer must have been rolled back onto the old-gen incr: appends + // through `*file` after the rewrite must land in the committed file, not + // the pruned/unlinked new-gen inode. + let marker = b"*1\r\n$4\r\nPING\r\n"; + file.write_all(marker).unwrap(); + file.sync_data().unwrap(); + drop(file); + let committed = std::fs::read(&old_incr_s0).unwrap(); + assert!( + committed.windows(marker.len()).any(|w| w == marker), + "post-abort appends must land in the committed old-gen incr (no silent loss)" + ); + } + + /// Cross-thread barrier wakeup: a folded writer that decrements then blocks + /// in `await_outcome` on one thread must be woken with the committed + /// generation by the terminal `shard_done` running on ANOTHER thread. The + /// single-threaded behavioural test above never actually blocks (its + /// `await_outcome` returns immediately), so this exercises the real condvar + /// wait/notify path across threads. + #[test] + fn rewrite_abort_wakes_waiter_cross_thread() { + let tmp = tempfile::tempdir().unwrap(); + let manifest = + crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); + let old_seq = manifest.seq; + let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); + let coord = PerShardRewriteCoord::new(manifest, old_seq, 2); + + // Shard A: decrement (countdown 2 -> 1, non-terminal) then block. + let c_a = coord.clone(); + let waiter = std::thread::spawn(move || { + c_a.shard_done(); + c_a.await_outcome() + }); + + // Shard B fails: mark_failed + the terminal decrement (-> 0). Whichever + // of the two `shard_done` calls is terminal sees `failed` set (B sets it + // first) and publishes old_seq. + coord.mark_failed(); + coord.shard_done(); + + let observed = waiter.join().expect("waiter must wake, not hang or panic"); + assert_eq!( + observed, old_seq, + "aborted rewrite must publish old_seq to the cross-thread waiter" + ); + } + + /// Panic-safety / liveness: a fold that PANICS mid-flight (the OOM-unwind + /// hazard of issue #138) must not hang the other shards' writers at the + /// barrier. The `ShardDoneGuard`'s `Drop` must fire `mark_failed` + + /// `shard_done` on unwind so the countdown still closes and every waiter + /// wakes. Without the guard this test hangs forever (the panicking shard + /// never decrements). + #[test] + fn rewrite_fold_panic_releases_barrier_for_other_writers() { + use std::sync::atomic::Ordering; + + let tmp = tempfile::tempdir().unwrap(); + let manifest = + crate::persistence::aof_manifest::AofManifest::initialize_multi(tmp.path(), 2).unwrap(); + let old_seq = manifest.seq; + let manifest = std::sync::Arc::new(parking_lot::Mutex::new(manifest)); + let coord = PerShardRewriteCoord::new(manifest, old_seq, 2); + + // Shard A folds, decrements, then blocks on the barrier in a thread. + let c_a = coord.clone(); + let waiter = std::thread::spawn(move || { + c_a.shard_done(); + c_a.await_outcome() + }); + + // Shard B "panics mid-fold": a ShardDoneGuard dropped during unwind. The + // guard's Drop must abort + decrement so A's barrier releases. + let c_b = coord.clone(); + let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _guard = ShardDoneGuard::new(&c_b); + panic!("simulated OOM unwind during fold"); + })); + assert!(panicked.is_err(), "the simulated fold must have panicked"); + + let observed = waiter + .join() + .expect("waiter must wake after the panicking shard's guard fires"); + assert_eq!( + observed, old_seq, + "panic-aborted rewrite must publish old_seq, not hang" + ); + assert!( + coord.failed.load(Ordering::Acquire), + "the dropped guard must have marked the rewrite failed" + ); + } + + #[test] + fn top_level_pool_routes_all_shards_to_writer_zero() { + let (tx, rx) = channel::mpsc_bounded::(8); + let pool = AofWriterPool::top_level(tx); + assert_eq!(pool.num_writers(), 1); + assert_eq!(pool.layout(), AofLayout::TopLevel); + + pool.try_send_append(0, 0, Bytes::from_static(b"a")); + pool.try_send_append(7, 0, Bytes::from_static(b"b")); + pool.try_send_append(42, 0, Bytes::from_static(b"c")); + + let mut seen = 0; + while rx.try_recv().is_ok() { + seen += 1; + } + assert_eq!(seen, 3, "all 3 appends should land on writer 0"); + } + + #[test] + fn per_shard_pool_routes_each_shard_to_its_own_writer() { + let (tx0, rx0) = channel::mpsc_bounded::(8); + let (tx1, rx1) = channel::mpsc_bounded::(8); + let (tx2, rx2) = channel::mpsc_bounded::(8); + let pool = AofWriterPool::per_shard(vec![tx0, tx1, tx2]); + assert_eq!(pool.num_writers(), 3); + assert_eq!(pool.layout(), AofLayout::PerShard); + + pool.try_send_append(0, 100, Bytes::from_static(b"shard0")); + pool.try_send_append(1, 200, Bytes::from_static(b"shard1a")); + pool.try_send_append(1, 300, Bytes::from_static(b"shard1b")); + pool.try_send_append(2, 400, Bytes::from_static(b"shard2")); + + let count = |rx: &channel::MpscReceiver| -> usize { + let mut n = 0; + while rx.try_recv().is_ok() { + n += 1; + } + n + }; + assert_eq!(count(&rx0), 1, "shard 0 writer should receive exactly 1"); + assert_eq!(count(&rx1), 2, "shard 1 writer should receive exactly 2"); + assert_eq!(count(&rx2), 1, "shard 2 writer should receive exactly 1"); + } + + #[test] + fn per_shard_pool_rejects_rewrite_with_explicit_error() { + let (tx0, _rx0) = channel::mpsc_bounded::(8); + let (tx1, _rx1) = channel::mpsc_bounded::(8); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let dummies: SharedDatabases = Arc::new(vec![]); + let err = pool + .try_send_rewrite(AofMessage::Rewrite(dummies)) + .unwrap_err(); + assert_eq!(err, AofPoolSendError::RewriteUnsupportedInPerShard); + } + + #[test] + fn top_level_pool_accepts_rewrite() { + let (tx, rx) = channel::mpsc_bounded::(8); + let pool = AofWriterPool::top_level(tx); + + let dummies: SharedDatabases = Arc::new(vec![]); + pool.try_send_rewrite(AofMessage::Rewrite(dummies)).unwrap(); + assert!(matches!(rx.try_recv(), Ok(AofMessage::Rewrite(_)))); + } + + #[test] + fn per_shard_pool_threads_lsn_field_to_each_writer() { + // Step 3 wire-format contract: try_send_append carries the issued LSN + // through to the writer task, which writes it as the per-entry header + // under PerShard layout. This unit test pins the channel-side contract + // (the disk-side framing is covered by writer-task integration). + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + pool.try_send_append(0, 42, Bytes::from_static(b"set foo 1")); + pool.try_send_append(1, 43, Bytes::from_static(b"set bar 2")); + pool.try_send_append(0, 44, Bytes::from_static(b"del foo")); + + // Shard 0 should see (42, "set foo 1") then (44, "del foo"). + match rx0.try_recv() { + Ok(AofMessage::Append { lsn, bytes }) => { + assert_eq!(lsn, 42, "shard 0 first entry lsn"); + assert_eq!(bytes.as_ref(), b"set foo 1"); + } + other => panic!( + "shard 0 first recv expected Append, got {:?}", + other.is_ok() + ), + } + match rx0.try_recv() { + Ok(AofMessage::Append { lsn, bytes }) => { + assert_eq!(lsn, 44, "shard 0 second entry lsn"); + assert_eq!(bytes.as_ref(), b"del foo"); + } + other => panic!( + "shard 0 second recv expected Append, got {:?}", + other.is_ok() + ), + } + // Shard 1 should see (43, "set bar 2") only. + match rx1.try_recv() { + Ok(AofMessage::Append { lsn, bytes }) => { + assert_eq!(lsn, 43, "shard 1 entry lsn"); + assert_eq!(bytes.as_ref(), b"set bar 2"); + } + other => panic!("shard 1 recv expected Append, got {:?}", other.is_ok()), + } + } + + #[test] + fn try_send_append_sync_queues_appendsync_with_ack() { + // Channel-level wiring contract for the H1 fix: `try_send_append_sync` + // queues `AofMessage::AppendSync { lsn, bytes, ack }`, and the + // returned receiver resolves to whatever value the (mocked) writer + // sends on `ack`. End-to-end durability is covered by step 8 + // (CRASH-01-LITE); this pins the API contract. + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let recv = pool.try_send_append_sync(0, 99, Bytes::from_static(b"SET k v")); + + // Drain the queue; the writer would normally do this. Capture the + // ack sender, do the (mock) durable write, then ack Synced. + let ack = match rx0.try_recv() { + Ok(AofMessage::AppendSync { lsn, bytes, ack }) => { + assert_eq!(lsn, 99, "lsn forwarded through the channel"); + assert_eq!(bytes.as_ref(), b"SET k v", "bytes forwarded"); + ack + } + other => panic!("expected AppendSync, got {:?}", other.is_ok()), + }; + + // Writer reports Synced — caller observes Synced. + let _ = ack.send(AofAck::Synced); + let result = recv.recv_blocking().expect("receiver resolves"); + assert_eq!(result, AofAck::Synced); + } + + #[test] + fn append_sync_writer_dropped_resolves_recv_error() { + // If the writer task is dead or the channel disconnects between + // queueing and the ack send, the receiver MUST resolve with an + // error rather than hang. Callers treat that as a hard failure + // (return an error frame, do not +OK). + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let recv = pool.try_send_append_sync(0, 7, Bytes::from_static(b"x")); + + // Drain the message but DROP the ack sender without sending. + match rx0.try_recv() { + Ok(AofMessage::AppendSync { ack, .. }) => drop(ack), + other => panic!("expected AppendSync, got {:?}", other.is_ok()), + } + + let err = recv.recv_blocking().expect_err("dropped ack -> RecvError"); + // Crash-safe: we got a sentinel-style error, not a hang. + let _ = err; + } + + #[test] + fn append_sync_writer_reports_write_failed() { + // Writer encountered a write_all error; recv returns WriteFailed. + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let recv = pool.try_send_append_sync(0, 1, Bytes::from_static(b"x")); + let ack = match rx0.try_recv() { + Ok(AofMessage::AppendSync { ack, .. }) => ack, + other => panic!("expected AppendSync, got {:?}", other.is_ok()), + }; + let _ = ack.send(AofAck::WriteFailed); + let result = recv.recv_blocking().expect("recv resolves"); + assert_eq!(result, AofAck::WriteFailed); + } + + #[test] + fn append_sync_writer_reports_fsync_failed() { + // Writer wrote the payload but fsync (sync_data) returned an error. + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let recv = pool.try_send_append_sync(0, 1, Bytes::from_static(b"x")); + let ack = match rx0.try_recv() { + Ok(AofMessage::AppendSync { ack, .. }) => ack, + other => panic!("expected AppendSync, got {:?}", other.is_ok()), + }; + let _ = ack.send(AofAck::FsyncFailed); + let result = recv.recv_blocking().expect("recv resolves"); + assert_eq!(result, AofAck::FsyncFailed); + } + + /// Issue #140: an AppendSync drained during a BGREWRITEAOF must NOT be acked + /// `Synced` inside the drain — that reports the write durable before the + /// post-drain boundary fsync, so a crash in the window loses an entry the + /// client was told was safe. The drain must PARK the ack and only fulfil it + /// after the boundary `sync_data()`. (Framed / per-shard drain.) + #[test] + fn drain_framed_parks_appendsync_ack_until_boundary_fsync() { + let tmp = tempfile::tempdir().unwrap(); + let incr = tmp.path().join("incr.aof"); + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let recv = pool.try_send_append_sync(0, 99, Bytes::from_static(b"SET k v")); + + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr) + .unwrap(); + let mut outcome = drain_pending_appends_framed(&rx0, &mut file).unwrap(); + + // CONTRACT: drained + parked, NOT yet acked. + assert_eq!(outcome.drained, 1, "the AppendSync was drained"); + assert_eq!( + outcome.pending_acks.len(), + 1, + "the ack must be parked until the boundary fsync, not sent during drain" + ); + + // Boundary fsync succeeds → fulfil Synced; only NOW is the client told durable. + file.sync_data().unwrap(); + outcome.fulfill_acks(true); + assert_eq!( + recv.recv_blocking().expect("ack resolves"), + AofAck::Synced, + "post-fsync the parked ack must resolve Synced" + ); + } + + /// Issue #140 failure path: if the rewrite-boundary fsync FAILS, a drained + /// AppendSync must resolve `FsyncFailed`, never `Synced`. Exercises the + /// non-framed `drain_pending_appends` — the DEFAULT `--shards 1` rewrite + /// path (`do_rewrite_single`), which is reachable under appendfsync=always. + #[cfg(feature = "runtime-monoio")] + #[test] + fn drain_single_fulfills_fsync_failure_as_fsync_failed() { + let tmp = tempfile::tempdir().unwrap(); + let incr = tmp.path().join("incr.aof"); + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + let recv = pool.try_send_append_sync(0, 7, Bytes::from_static(b"SET a b")); + + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr) + .unwrap(); + let mut outcome = drain_pending_appends(&rx0, &mut file).unwrap(); + assert_eq!( + outcome.pending_acks.len(), + 1, + "ack parked, not sent in drain" + ); + + // Simulate a failed boundary fsync: the parked ack must report FsyncFailed. + outcome.fulfill_acks(false); + assert_eq!( + recv.recv_blocking().expect("ack resolves"), + AofAck::FsyncFailed, + "a failed boundary fsync must NOT be reported Synced to the client" + ); + } + + // F2 (design-for-failure): `appendfsync=always` must bound its fsync-ack + // await. A stalled writer must surface a hard error within the budget, + // never park the connection forever. Tokio-gated because it drives the + // runtime timer; the monoio path shares the proven `select! + sleep` + // shape from `cluster::failover`, exercised end-to-end by the crash tests. + #[cfg(feature = "runtime-tokio")] + #[tokio::test] + async fn always_fsync_times_out_when_writer_never_acks() { + // Writer channel is held (kept open) but never drained → the + // AppendSync sits buffered with its ack sender alive, so the receiver + // never resolves. The bounded await MUST elapse and report failure. + let (tx0, _rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard_with_policy( + vec![tx0, tx1], + FsyncPolicy::Always, + Duration::from_millis(50), + ); + + let start = Instant::now(); + let res = pool + .try_send_append_durable(0, 1, Bytes::from_static(b"x")) + .await; + let elapsed = start.elapsed(); + + assert_eq!( + res, + Err(AofAck::FsyncFailed), + "timed-out fsync must map to FsyncFailed (durability unconfirmed)" + ); + assert!( + elapsed < Duration::from_secs(2), + "must fail within the bound, not hang (took {:?})", + elapsed + ); + // Keep the receivers alive until here so the message stays buffered. + drop((_rx0, _rx1)); + } + + #[cfg(feature = "runtime-tokio")] + #[tokio::test] + async fn always_fsync_succeeds_when_writer_acks_in_time() { + // Happy path: a writer drains the AppendSync and acks `Synced` well + // within the bound → the durable append returns Ok(()). + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard_with_policy( + vec![tx0, tx1], + FsyncPolicy::Always, + Duration::from_millis(500), + ); + + tokio::spawn(async move { + if let Ok(AofMessage::AppendSync { ack, .. }) = rx0.recv_async().await { + let _ = ack.send(AofAck::Synced); + } + }); + + let res = pool + .try_send_append_durable(0, 1, Bytes::from_static(b"x")) + .await; + assert_eq!(res, Ok(()), "ack within the bound must succeed"); + drop(_rx1); + } + + /// Parse the PerShard incr framing `[u64 lsn LE][u32 len LE][len bytes]`, + /// stopping at a truncated tail (the crash/torn boundary) — exactly what + /// `replay_incr_framed` does. Returns the cleanly-replayable prefix. + #[cfg(feature = "runtime-tokio")] + fn parse_framed(buf: &[u8]) -> Vec<(u64, Vec)> { + let mut out = Vec::new(); + let mut i = 0usize; + while i + 12 <= buf.len() { + let lsn = u64::from_le_bytes(buf[i..i + 8].try_into().unwrap()); + let len = u32::from_le_bytes(buf[i + 8..i + 12].try_into().unwrap()) as usize; + if i + 12 + len > buf.len() { + break; // truncated tail → crash boundary, stop + } + out.push((lsn, buf[i + 12..i + 12 + len].to_vec())); + i += 12 + len; + } + out + } + + // Regression (PR #136 review, BUG #2): the tokio per-shard writer must carry + // a `write_error` latch like the single-file (~:1467) and monoio (~:2125) + // writers. A torn write (header lands, payload fails) must NOT be followed by + // more records — a lone orphaned header makes the framed replay misread the + // next record's bytes as the orphan's payload, corrupting everything after. + // The latch suppresses all writes after the tear and reports WriteFailed to + // AppendSync callers (so they error instead of ack'ing a corrupt write). + #[cfg(feature = "runtime-tokio")] + #[tokio::test] + async fn tokio_per_shard_writer_latches_after_torn_write() { + use crate::persistence::aof_manifest::AofManifest; + use std::sync::atomic::Ordering; + + let tmp = tempfile::tempdir().unwrap(); + let base_dir = tmp.path().to_path_buf(); + // PerShard layout, 2 shards (the per-shard pool needs >=2); drive shard 0. + let manifest = AofManifest::initialize_multi(&base_dir, 2).unwrap(); + let incr = manifest.shard_incr_path(0); + + // Inject: the 2nd Append tears (header written, payload "fails"). + TEST_FAIL_WRITE_AT.store(2, Ordering::SeqCst); + + let (tx, rx) = channel::mpsc_bounded::(16); + let cancel = CancellationToken::new(); + let writer = tokio::spawn(per_shard_aof_writer_task( + rx, + base_dir.clone(), + 0, + FsyncPolicy::Always, + cancel.clone(), + )); + + // 1: clean. 2: torn (header only). 3: must be suppressed by the latch. + tx.try_send(AofMessage::Append { + lsn: 1, + bytes: Bytes::from_static(b"AAAA"), + }) + .unwrap(); + tx.try_send(AofMessage::Append { + lsn: 2, + bytes: Bytes::from_static(b"BBBB"), + }) + .unwrap(); + tx.try_send(AofMessage::Append { + lsn: 3, + bytes: Bytes::from_static(b"CCCC"), + }) + .unwrap(); + + // Barrier + assertion: an AppendSync after the tear MUST come back + // WriteFailed (latched), never Synced. + let (ack_tx, ack_rx) = crate::runtime::channel::oneshot::(); + tx.try_send(AofMessage::AppendSync { + lsn: 4, + bytes: Bytes::from_static(b"DDDD"), + ack: ack_tx, + }) + .unwrap(); + + let ack = tokio::time::timeout(std::time::Duration::from_secs(5), ack_rx) + .await + .expect("writer must answer the AppendSync within 5s") + .expect("ack channel must not drop"); + assert_eq!( + ack, + AofAck::WriteFailed, + "after a torn write the latch must reject further writes (got {ack:?})" + ); + + cancel.cancel(); + TEST_FAIL_WRITE_AT.store(0, Ordering::SeqCst); + let _ = tokio::time::timeout(std::time::Duration::from_secs(5), writer).await; + + // On disk: exactly one replayable frame (lsn=1, "AAAA"). The orphaned + // lsn=2 header is a truncated tail (crash boundary); lsn 3 and 4 were + // never written (latch held) — no corruption. + let raw = std::fs::read(&incr).unwrap(); + let frames = parse_framed(&raw); + assert_eq!( + frames, + vec![(1u64, b"AAAA".to_vec())], + "only the pre-tear record may replay; orphaned headers / suppressed \ + records must not corrupt the stream" + ); + } + + #[test] + fn broadcast_shutdown_reaches_every_writer() { + let (tx0, rx0) = channel::mpsc_bounded::(2); + let (tx1, rx1) = channel::mpsc_bounded::(2); + let (tx2, rx2) = channel::mpsc_bounded::(2); + let pool = AofWriterPool::per_shard(vec![tx0, tx1, tx2]); + + pool.broadcast_shutdown(); + + for (i, rx) in [&rx0, &rx1, &rx2].iter().enumerate() { + assert!( + matches!(rx.try_recv(), Ok(AofMessage::Shutdown)), + "writer {} did not receive Shutdown", + i + ); + } + } + + /// FIX-W1-1 contract: `try_send_append_durable` under `Always` policy MUST + /// return `Err(AofAck::FsyncFailed)` when the writer reports failure. + /// handler_single.rs must await this BEFORE flushing responses to the client. + /// + /// Uses spawn_blocking to simulate the mock writer responding on the ack + /// channel concurrently, which allows the async rendezvous to complete. + #[cfg(feature = "runtime-tokio")] + #[tokio::test] + async fn always_policy_try_send_append_durable_returns_err_on_fsync_fail() { + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = std::sync::Arc::new(AofWriterPool::per_shard_with_policy( + vec![tx0, tx1], + FsyncPolicy::Always, + Duration::ZERO, // legacy unbounded await — disconnect/ack resolves it + )); + + // Spawn a mock writer that drains AppendSync and responds with FsyncFailed. + // Runs in a blocking thread (flume's blocking recv) so it doesn't block + // the async executor while waiting for the handler to enqueue the message. + let mock_writer = tokio::task::spawn_blocking(move || { + // flume::Receiver::recv() blocks until a message is available + let msg = rx0.recv().expect("mock writer got message"); + if let AofMessage::AppendSync { ack, .. } = msg { + let _ = ack.send(AofAck::FsyncFailed); + } else { + panic!("expected AppendSync under Always policy"); + } + }); + + // The handler MUST await this BEFORE flushing responses to the client + let result = pool + .try_send_append_durable(0, 1, Bytes::from_static(b"SET k v")) + .await; + mock_writer.await.expect("mock writer completed"); + + assert_eq!( + result, + Err(AofAck::FsyncFailed), + "Always policy MUST propagate fsync failure so caller can return an error frame" + ); + } + + /// FIX-W1-1 ordering contract: when `aof_entries` carries `(resp_idx, bytes)` + /// tuples, the handler can patch `responses[resp_idx]` on AOF failure BEFORE + /// flushing to the client. This test verifies the indexing is sound. + #[test] + fn aof_entries_indexed_by_response_slot_patches_correctly() { + use crate::protocol::Frame; + let mut responses: Vec = vec![ + Frame::SimpleString(bytes::Bytes::from_static(b"OK")), + Frame::SimpleString(bytes::Bytes::from_static(b"OK")), + Frame::SimpleString(bytes::Bytes::from_static(b"OK")), + ]; + // Simulate two write commands at response indices 0 and 2 (index 1 was a read) + let aof_entries: Vec<(usize, Bytes)> = vec![ + (0, Bytes::from_static(b"SET a 1")), + (2, Bytes::from_static(b"SET c 3")), + ]; + + // AOF write at index 2 fails; patch that response slot + for (resp_idx, _bytes) in &aof_entries { + if *resp_idx == 2 { + // Simulate Err(AofAck::FsyncFailed) from try_send_append_durable + responses[*resp_idx] = + Frame::Error(Bytes::from_static(b"WRITEFAIL aof fsync failed")); + } + } + + assert!( + matches!(&responses[0], Frame::SimpleString(_)), + "index 0 (successful fsync) should remain +OK" + ); + assert!( + matches!(&responses[1], Frame::SimpleString(_)), + "index 1 (read, no AOF) should remain +OK" + ); + assert!( + matches!(&responses[2], Frame::Error(_)), + "index 2 (failed fsync) must be patched to error" + ); + } + + // NOTE (FIX-W1-1 r3): The H1 ordering regression test was moved to + // `src/server/conn/handler_single.rs` (test module, fn + // `flush_with_aof_ack_ack_precedes_response`). The previous inline + // reproduction here was non-discriminating — it reproduced the ack-first + // loop IN THE TEST BODY rather than calling the real production fn, so it + // passed on both pre-fix and post-fix binaries. + // + // The new test calls `flush_with_aof_ack` directly (the fn the handler now + // delegates to), so inverting Phase 1/Phase 2 order in that fn causes a + // measurable timing failure (`elapsed_ms ≈ 0ms < 55ms`). + // + // End-to-end ordering is also covered by: + // tests/crash_matrix_per_shard_aof.rs (CRASH-01-LITE — AlwaysPolicy shards) + + // ----------------------------------------------------------------------- + // FIX-W2-5: channel-full returns AofAck::ChannelFull + increments counter + // ----------------------------------------------------------------------- + #[test] + fn try_send_append_sync_channel_full_returns_channel_full_ack() { + // Create a channel with capacity 1 and fill it so the next try_send + // hits TrySendError::Full. + let (tx0, rx0) = channel::mpsc_bounded::(1); + // Fill the channel by pre-loading one message. + tx0.try_send(AofMessage::Shutdown).expect("pre-fill"); + // rx0 intentionally not consumed — channel is now at capacity. + + let pool = AofWriterPool::top_level(tx0); + + let before = AOF_BACKPRESSURE_DROPPED.load(std::sync::atomic::Ordering::Relaxed); + let recv = pool.try_send_append_sync(0, 1, Bytes::from_static(b"SET k v")); + + // The channel was full — ChannelFull is returned immediately without + // a writer round-trip. + let result = recv.recv_blocking().expect("pre-filled oneshot resolves"); + assert_eq!( + result, + AofAck::ChannelFull, + "channel-full must yield ChannelFull, not {:?}", + result + ); + + let after = AOF_BACKPRESSURE_DROPPED.load(std::sync::atomic::Ordering::Relaxed); + assert_eq!( + after, + before + 1, + "backpressure counter must increment by 1" + ); + + // No AppendSync should have reached the (blocked) reader. + drop(rx0); // drain without consuming — just verify nothing snuck through + } + + // ----------------------------------------------------------------------- + // FIX-W2-9: try_send_append_durable must be used for SWAPDB-like mutations + // + // Red test: documents the contract that handler_single.rs SHOULD honour. + // When appendfsync=always, try_send_append_durable MUST return Err on + // writer failure so callers can abort the mutation safely. + // ----------------------------------------------------------------------- + #[test] + fn try_send_append_durable_always_writer_dead_returns_write_failed() { + // Create a pool with Always policy. The writer task is not running — + // we model that by draining the channel message and then dropping the + // ack sender, simulating a dead writer. + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard_with_policy( + vec![tx0, tx1], + FsyncPolicy::Always, + Duration::ZERO, // legacy unbounded await — disconnect resolves it + ); + + // Spawn a thread that pulls the AppendSync off the channel but drops + // the ack without sending — simulating a writer crash mid-fsync. + let rx0_clone = rx0; + let handle = std::thread::spawn(move || { + match rx0_clone.recv() { + Ok(AofMessage::AppendSync { ack, .. }) => drop(ack), // writer crash + other => panic!("unexpected message: {:?}", other.is_ok()), + } + }); + + // try_send_append_durable for Always must await the ack. + // With the ack sender dropped, it should resolve to Err(WriteFailed). + let result = futures::executor::block_on(pool.try_send_append_durable( + 0, + 55, + Bytes::from_static(b"SWAPDB 0 1"), + )); + + handle.join().expect("ack dropper thread"); + + assert!( + result.is_err(), + "try_send_append_durable with dead writer must return Err, got Ok" + ); + assert_eq!( + result.unwrap_err(), + AofAck::WriteFailed, + "dead writer must resolve to WriteFailed" + ); + } + + #[test] + fn try_send_append_durable_everysec_is_fire_and_forget() { + // EverySec policy: try_send_append_durable always returns Ok — the + // durability policy doesn't block on fsync. handler_single.rs must + // use try_send_append_durable so the policy is respected. + let (tx0, _rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard_with_policy( + vec![tx0, tx1], + FsyncPolicy::EverySec, + Duration::ZERO, + ); + + let result = futures::executor::block_on(pool.try_send_append_durable( + 0, + 56, + Bytes::from_static(b"SWAPDB 0 1"), + )); + + assert!( + result.is_ok(), + "EverySec policy must be fire-and-forget (Ok), got {:?}", + result + ); + } +} diff --git a/src/persistence/aof/rewrite.rs b/src/persistence/aof/rewrite.rs new file mode 100644 index 00000000..c052cbd8 --- /dev/null +++ b/src/persistence/aof/rewrite.rs @@ -0,0 +1,941 @@ +//! AOF rewrite / compaction: snapshot generation, drain, and per-shard / +//! single / sharded rewrite paths. +#![allow(unused_imports, unused_variables, unreachable_code, clippy::empty_loop)] + +use super::*; + +/// Generate synthetic RESP commands from the current database state for AOF rewriting. +/// +/// Produces commands for all 5 data types plus PEXPIRE for keys with TTL. +#[allow(dead_code)] // Retained for RESP-only AOF rewrite fallback and testing +pub fn generate_rewrite_commands(databases: &[Database]) -> BytesMut { + let mut buf = BytesMut::new(); + let now_ms = current_time_ms(); + + for (db_idx, db) in databases.iter().enumerate() { + let base_ts = db.base_timestamp(); + let data = db.data(); + if data.is_empty() { + continue; + } + + // Generate SELECT if not db 0 + if db_idx > 0 { + let select_frame = Frame::Array(framevec![ + Frame::BulkString(Bytes::from_static(b"SELECT")), + Frame::BulkString(Bytes::from(db_idx.to_string())), + ]); + serialize::serialize(&select_frame, &mut buf); + } + + for (key, entry) in data { + // Skip expired entries + if entry.is_expired_at(base_ts, now_ms) { + continue; + } + + match entry.value.as_redis_value() { + RedisValueRef::String(val) => { + let frame = Frame::Array(framevec![ + Frame::BulkString(Bytes::from_static(b"SET")), + Frame::BulkString(key.to_bytes()), + Frame::BulkString(Bytes::copy_from_slice(val)), + ]); + serialize::serialize(&frame, &mut buf); + } + RedisValueRef::Hash(map) => { + if map.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"HSET")), + Frame::BulkString(key.to_bytes()), + ]; + for (field, val) in map.iter() { + args.push(Frame::BulkString(field.clone())); + args.push(Frame::BulkString(val.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + // Phase 200: for HashWithTtl we emit two RESP frames per key. + // 1. `HSET key f1 v1 f2 v2 ...` rebuilds the hash body. + // 2. `HPEXPIREAT key abs_ms FIELDS 1 field` for every entry + // in the TTL sidecar — one per TTL'd field for clarity + // (BGREWRITEAOF is rare; per-field framing keeps the + // replay shim simple, see `persistence::replay`). + RedisValueRef::HashWithTtl { fields, ttls, .. } => { + if fields.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"HSET")), + Frame::BulkString(key.to_bytes()), + ]; + for (field, val) in fields.iter() { + args.push(Frame::BulkString(field.clone())); + args.push(Frame::BulkString(val.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + + for (field, ttl_ms) in ttls.iter() { + let mut ttl_args = vec![ + Frame::BulkString(Bytes::from_static(b"HPEXPIREAT")), + Frame::BulkString(key.to_bytes()), + Frame::BulkString(Bytes::copy_from_slice( + ttl_ms.to_string().as_bytes(), + )), + Frame::BulkString(Bytes::from_static(b"FIELDS")), + Frame::BulkString(Bytes::from_static(b"1")), + Frame::BulkString(field.clone()), + ]; + ttl_args.shrink_to_fit(); + serialize::serialize(&Frame::Array(ttl_args.into()), &mut buf); + } + } + RedisValueRef::HashListpack(lp) => { + let map = lp.to_hash_map(); + if map.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"HSET")), + Frame::BulkString(key.to_bytes()), + ]; + for (field, val) in &map { + args.push(Frame::BulkString(field.clone())); + args.push(Frame::BulkString(val.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + RedisValueRef::List(list) => { + if list.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"RPUSH")), + Frame::BulkString(key.to_bytes()), + ]; + for elem in list.iter() { + args.push(Frame::BulkString(elem.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + RedisValueRef::ListListpack(lp) => { + let list = lp.to_vec_deque(); + if list.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"RPUSH")), + Frame::BulkString(key.to_bytes()), + ]; + for elem in &list { + args.push(Frame::BulkString(elem.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + RedisValueRef::Set(set) => { + if set.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"SADD")), + Frame::BulkString(key.to_bytes()), + ]; + for member in set.iter() { + args.push(Frame::BulkString(member.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + RedisValueRef::SetListpack(lp) => { + let set = lp.to_hash_set(); + if set.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"SADD")), + Frame::BulkString(key.to_bytes()), + ]; + for member in &set { + args.push(Frame::BulkString(member.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + RedisValueRef::SetIntset(is) => { + let set = is.to_hash_set(); + if set.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"SADD")), + Frame::BulkString(key.to_bytes()), + ]; + for member in &set { + args.push(Frame::BulkString(member.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + RedisValueRef::SortedSet { members, .. } + | RedisValueRef::SortedSetBPTree { members, .. } => { + if members.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"ZADD")), + Frame::BulkString(key.to_bytes()), + ]; + for (member, score) in members.iter() { + args.push(Frame::BulkString(Bytes::from(score.to_string()))); + args.push(Frame::BulkString(member.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + RedisValueRef::SortedSetListpack(lp) => { + let pairs: Vec<_> = lp.iter_pairs().collect(); + if pairs.is_empty() { + continue; + } + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"ZADD")), + Frame::BulkString(key.to_bytes()), + ]; + for (member_entry, score_entry) in &pairs { + let score_bytes = score_entry.as_bytes(); + args.push(Frame::BulkString(Bytes::from(score_bytes))); + args.push(Frame::BulkString(Bytes::from(member_entry.as_bytes()))); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + RedisValueRef::Stream(stream) => { + for (id, fields) in &stream.entries { + let mut args = vec![ + Frame::BulkString(Bytes::from_static(b"XADD")), + Frame::BulkString(key.to_bytes()), + Frame::BulkString(id.to_bytes()), + ]; + for (field, value) in fields { + args.push(Frame::BulkString(field.clone())); + args.push(Frame::BulkString(value.clone())); + } + serialize::serialize(&Frame::Array(args.into()), &mut buf); + } + } + } + + // Generate PEXPIRE for keys with TTL + if entry.has_expiry() { + let exp_ms = entry.expires_at_ms(base_ts); + if exp_ms > now_ms { + let remaining_ms = exp_ms - now_ms; + let pexpire_frame = Frame::Array(framevec![ + Frame::BulkString(Bytes::from_static(b"PEXPIRE")), + Frame::BulkString(key.to_bytes()), + Frame::BulkString(Bytes::from(remaining_ms.to_string())), + ]); + serialize::serialize(&pexpire_frame, &mut buf); + } + } + } + } + + buf +} + +/// Snapshot databases and generate compacted AOF commands. +/// +/// Shared by both the async (tokio) and sync (monoio) rewrite paths. +#[allow(dead_code)] +fn snapshot_and_generate(db: &SharedDatabases) -> BytesMut { + let snapshot: Vec<(Vec<(CompactKey, Entry)>, u32)> = db + .iter() + .map(|lock| { + let guard = lock.read(); + let base_ts = guard.base_timestamp(); + let entries = guard + .data() + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + (entries, base_ts) + }) + .collect(); + + let mut temp_dbs: Vec = Vec::with_capacity(snapshot.len()); + for (entries, _base_ts) in &snapshot { + let mut db = Database::new(); + for (key, entry) in entries { + db.set(key.to_bytes(), entry.clone()); + } + temp_dbs.push(db); + } + + generate_rewrite_commands(&temp_dbs) +} + +/// Drain any queued `AofMessage::Append` messages to the current incr file. +/// +/// Called during rewrite to catch in-flight appends that handlers sent before +/// the writer thread could enter the rewrite routine. Messages of other variants +/// are dropped silently (duplicate rewrites while a rewrite is in progress) or +/// returned via the flag for Shutdown (caller is responsible for honoring it +/// after the rewrite completes). +#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] +#[derive(Default)] +pub(crate) struct DrainOutcome { + pub(crate) drained: usize, + shutdown_requested: bool, + /// AppendSync ack senders for entries drained during a rewrite. Under + /// `appendfsync=always` the client must NOT be told `Synced` until the + /// post-drain boundary `sync_data()` makes those bytes durable, so the acks + /// are parked here and resolved by [`fulfill_acks`](Self::fulfill_acks) AFTER + /// the caller's fsync — `Synced` on success, `FsyncFailed` on failure. Acking + /// inside the drain (the old behaviour) reports a write durable before the + /// boundary fsync, so a crash in that window loses an entry the client was + /// told was safe. (Issue #140.) + pub(crate) pending_acks: Vec>, +} + +#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] +impl DrainOutcome { + /// Resolve every parked AppendSync ack after the rewrite-boundary fsync. + /// `synced=true` → `Synced`; `false` → `FsyncFailed`. A fresh `AofAck` is + /// built per sender so `AofAck` needs no `Copy`/`Clone`. Drains the vec so a + /// second call is a no-op. + pub(crate) fn fulfill_acks(&mut self, synced: bool) { + for tx in std::mem::take(&mut self.pending_acks) { + let _ = tx.send(if synced { + AofAck::Synced + } else { + AofAck::FsyncFailed + }); + } + } +} + +/// Fsync `file` at a rewrite drain boundary, then resolve the drained batch's +/// parked AppendSync acks against the result: `Synced` on success, or +/// `FsyncFailed` + propagate the IO error on failure. Centralizes the issue-#140 +/// durability ordering (ack strictly AFTER the bytes are durable) for every +/// `do_rewrite_*` drain site. +#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] +fn sync_and_fulfill_drain( + outcome: &mut DrainOutcome, + file: &mut std::fs::File, + incr_path: PathBuf, +) -> Result<(), MoonError> { + match file.sync_data() { + Ok(()) => { + outcome.fulfill_acks(true); + Ok(()) + } + Err(e) => { + // Boundary fsync failed: the drained writes are NOT durable. Tell the + // waiting clients FsyncFailed (never Synced) and propagate the error + // so the rewrite aborts. + outcome.fulfill_acks(false); + Err(AofError::Io { + path: incr_path, + source: e, + } + .into()) + } + } +} + +#[cfg(feature = "runtime-monoio")] +pub(crate) fn drain_pending_appends( + rx: &channel::MpscReceiver, + file: &mut std::fs::File, +) -> Result { + use std::io::Write; + let mut outcome = DrainOutcome::default(); + while let Ok(msg) = rx.try_recv() { + match msg { + // BGREWRITEAOF drain runs on the TopLevel writer (monoio) only; + // PerShard rewrite is RFC step 6. Legacy v1 disk format → ignore lsn. + AofMessage::Append { + bytes: data, + lsn: _, + } => { + file.write_all(&data).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + outcome.drained += 1; + } + // AppendSync during a rewrite drain: bytes are written and counted, + // but the ack is PARKED until the caller's post-drain boundary fsync + // (issue #140) — acking `Synced` here would report durability before + // the bytes are fsynced. If the write itself fails the `?` propagates + // the error and the parked ack is dropped with the outcome — the + // caller observes `RecvError`, which it treats as failure. + AofMessage::AppendSync { + bytes: data, + lsn: _, + ack, + } => { + file.write_all(&data).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + outcome.drained += 1; + outcome.pending_acks.push(ack); + } + AofMessage::Shutdown => { + outcome.shutdown_requested = true; + } + AofMessage::Rewrite(_) + | AofMessage::RewriteSharded(_) + | AofMessage::RewritePerShard { .. } => { + // Already rewriting — drop redundant request. + } + } + } + Ok(outcome) +} + +/// [F6] Drain a per-shard writer's queued appends into its OLD incr file using +/// the framed `[u64 lsn LE][u32 len LE][RESP bytes]` on-disk format that +/// per-shard recovery expects. +/// +/// This is the per-shard twin of [`drain_pending_appends`] (which writes the +/// legacy TopLevel raw-RESP format). Correctness depends on the framing +/// matching `replay_per_shard`'s reader — an unframed write here would make the +/// drained appends unparseable on restart. +#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] +pub(crate) fn drain_pending_appends_framed( + rx: &channel::MpscReceiver, + file: &mut std::fs::File, +) -> Result { + use std::io::Write; + let mut outcome = DrainOutcome::default(); + let write_framed = |file: &mut std::fs::File, lsn: u64, data: &[u8]| -> std::io::Result<()> { + let mut header = [0u8; 12]; + header[..8].copy_from_slice(&lsn.to_le_bytes()); + header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); + file.write_all(&header)?; + file.write_all(data) + }; + while let Ok(msg) = rx.try_recv() { + match msg { + AofMessage::Append { lsn, bytes: data } => { + write_framed(file, lsn, &data).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + outcome.drained += 1; + } + AofMessage::AppendSync { + lsn, + bytes: data, + ack, + } => { + write_framed(file, lsn, &data).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + outcome.drained += 1; + // Park the ack until the caller's post-drain boundary fsync + // (issue #140); resolved Synced/FsyncFailed by + // `sync_and_fulfill_drain`. Mirrors `drain_pending_appends`. + outcome.pending_acks.push(ack); + } + AofMessage::Shutdown => { + outcome.shutdown_requested = true; + } + AofMessage::Rewrite(_) + | AofMessage::RewriteSharded(_) + | AofMessage::RewritePerShard { .. } => { + // Already rewriting this shard — drop redundant request. + } + } + } + Ok(outcome) +} + +/// [F6] Per-shard rewrite fold (monoio). Run by a single per-shard writer for +/// ITS shard only; the manifest commit is coordinated across all shards by the +/// shared [`PerShardRewriteCoord`]. +/// +/// Correctness ordering (prevents double-apply of non-idempotent commands like +/// INCR after the rewrite) — identical discipline to [`do_rewrite_sharded`], +/// scoped to one shard: +/// +/// 1. Drain queued appends into the OLD incr (framed) and fsync. +/// 2. Acquire write locks on this shard's databases. +/// 3. Re-drain appends that arrived between phase 1 and the lock, into OLD +/// incr, and fsync. +/// 4. Snapshot this shard's databases under the locks. +/// 5. Release the locks before the expensive base-RDB write. +/// 6. Write the new base + new (empty) incr at `coord.new_seq` via +/// `advance_shard` (which does NOT bump `manifest.seq`), then reopen +/// `file` to the new incr. Subsequent appends land in the new generation. +/// 7. Signal completion to the coordinator; the last shard commits the +/// manifest (single seq flip) and prunes the old generation. +/// +/// Until step 7's commit, the on-disk manifest still resolves to the old seq, +/// so a crash anywhere in steps 1-6 recovers the intact old generation. +/// +/// # Cross-thread exactly-once invariant (load-bearing) +/// +/// This fold runs on the per-shard *writer* thread, which is distinct from the +/// shard event-loop thread that applies commands. Exactly-once across the +/// rewrite boundary depends on a single ordering fact: the live write path +/// enqueues each command's AOF append **inside** the same `RwLock` +/// write guard under which it mutated the db (see `spsc_handler.rs`: +/// `wal_append_and_fanout` is called before `drop(guard)`). Phase 2 here +/// acquires those *same* locks (`all_shard_dbs()[sidx]` is +/// `ShardDatabases::shards[sidx]`, the exact `RwLock`s `write_db` locks), so +/// RwLock mutual exclusion forces the order +/// `enqueue → guard-release → fold-acquire → mid-drain(phase 3)`. Hence every +/// INCR whose mutation lands in the phase-4 snapshot had its append drained +/// into the OLD incr (then pruned at commit) — never replayed on top of the +/// new base. Were the append enqueued *after* the guard drop, a snapshot would +/// capture the mutation while its append still raced toward the NEW incr → +/// double-apply. The in-guard append is therefore the invariant; do not move it. +/// +/// This also assumes the `RwLock`-backed `ShardDatabases` is the *live* store. +/// It is, because the thread-local `ShardSlice` fast path is dead code until +/// Phase 4 wires `init_shard` (`is_initialized()` is always false today). A +/// future Phase 4 that makes ShardSlice live MUST revisit this fold: the writer +/// thread cannot lock another thread's `!Send` `Rc>`, so the +/// per-shard rewrite would need a different snapshot-coordination mechanism. +/// +/// # Known limitation — channel saturation during the fold +/// +/// Exactly-once holds *absent append-channel saturation during the fold*. While +/// this function runs (phases 2-6, including the base-RDB serialize + write + +/// fsync of phase 6, which is hundreds of ms on a large shard) the writer is +/// NOT in its recv loop, so it is not draining the bounded +/// `mpsc_bounded::(10_000)` append channel. Post-snapshot appends +/// queue there for the new incr; the event loop enqueues them with +/// `try_send_append` (drop-on-full, return ignored — `spsc_handler.rs`). Under +/// *sustained concurrent* writes on a large dataset, > 10_000 appends can pile +/// up during the window and the overflow is silently dropped — lost even on a +/// clean restart (worse than the everysec contract, which only loses on crash). +/// The single-client crash matrix cannot surface this (serialized `redis-cli` +/// never pressures the channel). This window is *pre-existing*: the shipped +/// `do_rewrite_sharded` has the identical non-draining gap. Tracked as a +/// known limitation (F6 is behind `--experimental-per-shard-rewrite`); the fix +/// (keep draining during phase 6, or block-on-full for the rewrite's duration) +/// is a separate scoped task. See `tmp/F6-known-limitations.md`. +#[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] +pub(crate) fn do_rewrite_per_shard( + shard_id: u16, + shard_dbs: &crate::shard::shared_databases::ShardDatabases, + file: &mut std::fs::File, + rx: &channel::MpscReceiver, + coord: &PerShardRewriteCoord, +) -> Result<(), MoonError> { + // Panic/early-error safety: guarantees `shard_done` runs on EVERY exit + // (success via `complete()`, `?`-error or panic-unwind via `Drop`). The + // phase-8 `await_outcome` barrier makes that a liveness requirement, so + // callers MUST NOT call `shard_done` after invoking this function — the + // guard owns the single decrement for all exits. See `ShardDoneGuard`. + let guard = ShardDoneGuard::new(coord); + + let sidx = shard_id as usize; + let all_shards = shard_dbs.all_shard_dbs(); + if sidx >= all_shards.len() { + return Err(AofError::RewriteFailed { + detail: format!( + "do_rewrite_per_shard: shard {} out of range ({} shards)", + sidx, + all_shards.len() + ), + } + .into()); + } + + // Phase 1: drain pre-rewrite queued appends into old incr (framed), fsync, + // then resolve their parked AppendSync acks (issue #140). + let mut pre_drain = drain_pending_appends_framed(rx, file)?; + sync_and_fulfill_drain(&mut pre_drain, file, PathBuf::from(""))?; + + // Phase 2: acquire write locks on this shard's db(s) (db_idx ascending). + let shard_locks = &all_shards[sidx]; + let guards: Vec<_> = shard_locks.iter().map(|lock| lock.write()).collect(); + + // Phase 3: drain appends that completed between phase 1 and phase 2, fsync, + // then resolve their parked AppendSync acks (issue #140). + let mut mid_drain = drain_pending_appends_framed(rx, file)?; + sync_and_fulfill_drain(&mut mid_drain, file, PathBuf::from(""))?; + + // Phase 4: snapshot this shard's databases under the locks. + let now_ms = current_time_ms(); + let mut snapshot: Vec<( + Vec<( + crate::storage::compact_key::CompactKey, + crate::storage::entry::Entry, + )>, + u32, + )> = Vec::with_capacity(guards.len()); + for guard in &guards { + let base_ts = guard.base_timestamp(); + let mut entries = Vec::new(); + for (key, entry) in guard.data().iter() { + if !entry.is_expired_at(base_ts, now_ms) { + entries.push((key.clone(), entry.clone())); + } + } + snapshot.push((entries, base_ts)); + } + + // Phase 5: release locks before the expensive disk write. + drop(guards); + + // Phase 6: write new base, advance THIS shard's manifest entry (no seq + // commit), reopen to the new incr. The manifest lock is held only for the + // brief, await-free advance_shard call. + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot)?; + let new_incr = { + let mut m = coord.manifest.lock(); + m.advance_shard(shard_id, coord.new_seq, &rdb_bytes)? + }; + *file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&new_incr) + .map_err(|e| AofError::Io { + path: new_incr, + source: e, + })?; + + info!( + "F6 per-shard rewrite: shard {} folded (drained {}+{} appends), new seq {}", + shard_id, pre_drain.drained, mid_drain.drained, coord.new_seq + ); + if pre_drain.shutdown_requested || mid_drain.shutdown_requested { + warn!( + "F6 per-shard rewrite: shard {} saw shutdown during rewrite (honored after commit)", + shard_id + ); + } + + // Phase 7: signal completion; the last writer commits + prunes. `complete()` + // performs the single clean `shard_done` and disarms the guard's Drop. + guard.complete(); + + // Phase 8 (barrier-before-resume): block until the terminal writer publishes + // the committed generation, then make sure THIS writer's append file points + // at it. On the happy path committed == new_seq and *file already points at + // new_incr (phase 6) — nothing to do. On an abort/commit-failure the manifest + // kept old_seq and pruned our new_seq incr, so reopen *file onto old_seq's + // incr; otherwise we keep appending into a discarded generation that recovery + // ignores — silent data loss. Replaces the old "RESTART recommended" hazard. + let committed_seq = coord.await_outcome(); + if committed_seq != coord.new_seq { + let committed_incr = coord + .manifest + .lock() + .shard_incr_path_seq(shard_id, committed_seq); + *file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&committed_incr) + .map_err(|e| AofError::Io { + path: committed_incr, + source: e, + })?; + warn!( + "F6 per-shard rewrite ABORTED: shard {} rolled its append file back to \ + committed seq {} (no restart needed)", + shard_id, committed_seq + ); + } + Ok(()) +} + +/// Multi-part rewrite: snapshot single-shard databases → RDB base → advance manifest. +/// +/// Correctness ordering (prevents double-apply of non-idempotent commands like +/// INCR/LPUSH/SADD after rewrite): +/// +/// 1. Drain any queued appends into the OLD incr file and fsync. +/// 2. Acquire write locks on all databases in the shard. This blocks handlers +/// from applying new writes or queueing new appends for the locked dbs. +/// 3. Drain the channel once more — catches appends for writes that the +/// handler completed between step 1 and step 2. +/// 4. Snapshot every database under the write locks. Because no handler can +/// mutate the dbs while we hold the locks, the snapshot is atomic with +/// respect to the post-drain channel state. +/// 5. Release the write locks. New handler writes from here on queue in the +/// channel and will be processed into the NEW incr file after rotation. +/// 6. Write the new base RDB, advance the manifest, reopen the file handle. +/// +/// Invariant: any write captured in the new base is NOT in the new incr file +/// (handlers were blocked between drain and snapshot), and any write NOT in +/// the new base IS in the new incr file (queued after lock release). +#[cfg(feature = "runtime-monoio")] +pub(crate) fn do_rewrite_single( + db: &SharedDatabases, + manifest: &mut crate::persistence::aof_manifest::AofManifest, + file: &mut std::fs::File, + rx: &channel::MpscReceiver, +) -> Result<(), MoonError> { + // Phase 1: drain pre-rewrite queued appends into old incr, fsync, then + // resolve their parked AppendSync acks (issue #140). + let mut pre_drain = drain_pending_appends(rx, file)?; + sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; + + // Phase 2: acquire write locks on every database in the shard. + // Order is consistent (index-ascending) so concurrent callers would + // serialize without deadlock — but in practice only this thread + // acquires multi-db locks. + let guards: Vec<_> = db.iter().map(|lock| lock.write()).collect(); + + // Phase 3: drain any appends the handlers sent between phase 1 and phase 2, + // fsync, then resolve their parked AppendSync acks (issue #140). + let mut mid_drain = drain_pending_appends(rx, file)?; + sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; + + // Phase 4: snapshot under the write locks. No mutation is possible. + let now_ms = current_time_ms(); + let snapshot: Vec<( + Vec<( + crate::storage::compact_key::CompactKey, + crate::storage::entry::Entry, + )>, + u32, + )> = guards + .iter() + .map(|guard| { + let base_ts = guard.base_timestamp(); + let entries: Vec<_> = guard + .data() + .iter() + .filter(|(_, v)| !v.is_expired_at(base_ts, now_ms)) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + (entries, base_ts) + }) + .collect(); + + // Phase 5: release locks. Handlers resume; new appends queue in the channel + // and will be processed into the new incr after step 6. + drop(guards); + + // Phase 6: write new base, advance manifest, reopen. + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot)?; + let new_incr = manifest.advance(&rdb_bytes)?; + + *file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&new_incr) + .map_err(|e| AofError::Io { + path: new_incr, + source: e, + })?; + + info!( + "AOF rewrite complete (single): drained {}+{} pre-snapshot appends, seq={}", + pre_drain.drained, mid_drain.drained, manifest.seq + ); + if pre_drain.shutdown_requested || mid_drain.shutdown_requested { + // Caller doesn't currently observe this; logging is the escape hatch. + warn!("AOF writer: shutdown requested during rewrite (will honor on next recv)"); + } + Ok(()) +} + +/// Multi-part rewrite: snapshot all shards → merged RDB base → advance manifest. +/// +/// See [`do_rewrite_single`] for the ordering rationale. The multi-shard variant +/// holds write locks on every (shard, db) pair simultaneously for the duration +/// of the snapshot. This creates a brief global write pause, but it is the only +/// way to guarantee a torn-free snapshot without per-message sequence numbers. +#[cfg(feature = "runtime-monoio")] +pub(crate) fn do_rewrite_sharded( + shard_dbs: &crate::shard::shared_databases::ShardDatabases, + manifest: &mut crate::persistence::aof_manifest::AofManifest, + file: &mut std::fs::File, + rx: &channel::MpscReceiver, +) -> Result<(), MoonError> { + // Phase 1: drain pre-rewrite queued appends into old incr, fsync, then + // resolve their parked AppendSync acks (issue #140). + let mut pre_drain = drain_pending_appends(rx, file)?; + sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; + + // Phase 2: acquire write locks on ALL (shard, db) pairs simultaneously. + // Lock order is (shard_idx, db_idx) ascending — must match anywhere else + // that acquires multiple locks to prevent deadlock (currently no other + // call site does, but the ordering discipline is documented for future + // maintainers). + let all_shards = shard_dbs.all_shard_dbs(); + let mut guards: Vec> = Vec::with_capacity(all_shards.len()); + for shard_locks in all_shards { + let mut shard_guards = Vec::with_capacity(shard_locks.len()); + for lock in shard_locks { + shard_guards.push(lock.write()); + } + guards.push(shard_guards); + } + + // Phase 3: drain appends completed between phase 1 and phase 2, fsync, then + // resolve their parked AppendSync acks (issue #140). + let mut mid_drain = drain_pending_appends(rx, file)?; + sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; + + // Phase 4: snapshot under locks. + let db_count = shard_dbs.db_count(); + let mut merged: Vec<( + Vec<( + crate::storage::compact_key::CompactKey, + crate::storage::entry::Entry, + )>, + u32, + )> = (0..db_count).map(|_| (Vec::new(), 0u32)).collect(); + let now_ms = current_time_ms(); + for shard_guards in &guards { + for (db_idx, guard) in shard_guards.iter().enumerate() { + let base_ts = guard.base_timestamp(); + if merged[db_idx].0.is_empty() { + merged[db_idx].1 = base_ts; + } + for (key, entry) in guard.data().iter() { + if !entry.is_expired_at(base_ts, now_ms) { + merged[db_idx].0.push((key.clone(), entry.clone())); + } + } + } + } + + // Phase 5: release locks before the expensive disk write. + drop(guards); + + // Phase 6: write new base, advance manifest, reopen. + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&merged)?; + let new_incr = manifest.advance(&rdb_bytes)?; + + *file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&new_incr) + .map_err(|e| AofError::Io { + path: new_incr, + source: e, + })?; + + info!( + "AOF rewrite complete (sharded): drained {}+{} pre-snapshot appends, seq={}", + pre_drain.drained, mid_drain.drained, manifest.seq + ); + if pre_drain.shutdown_requested || mid_drain.shutdown_requested { + warn!("AOF writer: shutdown requested during rewrite (will honor on next recv)"); + } + Ok(()) +} + +/// Rewrite the AOF file with RDB preamble (binary base + empty RESP incremental). +/// +/// Uses the same strategy as Redis 7+ `aof-use-rdb-preamble yes`: +/// the rewritten AOF starts with a full RDB snapshot (compact binary), +/// and new writes are appended as RESP after it. On startup, the loader +/// detects the RDB magic and reads the binary preamble, then switches +/// to RESP parsing for any incremental commands appended after. +#[allow(dead_code)] // Retained for legacy single-file and tokio path +fn rewrite_aof_sync(db: &SharedDatabases, aof_path: &Path) -> Result<(), MoonError> { + // Snapshot under read locks, build temp Database objects for RDB serialization + let snapshot: Vec = db + .iter() + .map(|lock| { + let guard = lock.read(); + let mut temp = Database::new(); + let now_ms = current_time_ms(); + for (k, v) in guard.data().iter() { + if !v.is_expired_at(guard.base_timestamp(), now_ms) { + temp.set(k.to_bytes(), v.clone()); + } + } + temp + }) + .collect(); + + let rdb_bytes = crate::persistence::rdb::save_to_bytes(&snapshot)?; + + let tmp_path = aof_path.with_extension("aof.tmp"); + std::fs::write(&tmp_path, &rdb_bytes).map_err(|e| AofError::Io { + path: tmp_path.clone(), + source: e, + })?; + std::fs::rename(&tmp_path, aof_path).map_err(|e| AofError::RewriteFailed { + detail: format!( + "rename {} -> {}: {}", + tmp_path.display(), + aof_path.display(), + e + ), + })?; + + info!( + "AOF rewrite complete (RDB preamble): {} bytes", + rdb_bytes.len() + ); + Ok(()) +} + +/// Rewrite the AOF in sharded mode with RDB preamble. +/// +/// Merges all shards' databases into a single RDB snapshot, writes it as +/// the AOF base file. New incremental writes are appended as RESP after. +#[allow(dead_code)] +pub(crate) fn rewrite_aof_sharded_sync( + shard_dbs: &crate::shard::shared_databases::ShardDatabases, + aof_path: &Path, +) -> Result<(), MoonError> { + let db_count = shard_dbs.db_count(); + let now_ms = current_time_ms(); + let mut merged_dbs: Vec = (0..db_count).map(|_| Database::new()).collect(); + + for shard_locks in shard_dbs.all_shard_dbs() { + for (db_idx, lock) in shard_locks.iter().enumerate() { + let guard = lock.read(); + for (key, entry) in guard.data().iter() { + if !entry.is_expired_at(guard.base_timestamp(), now_ms) { + merged_dbs[db_idx].set(key.to_bytes(), entry.clone()); + } + } + } + } + + let rdb_bytes = crate::persistence::rdb::save_to_bytes(&merged_dbs)?; + + let tmp_path = aof_path.with_extension("aof.tmp"); + std::fs::write(&tmp_path, &rdb_bytes).map_err(|e| AofError::Io { + path: tmp_path.clone(), + source: e, + })?; + std::fs::rename(&tmp_path, aof_path).map_err(|e| AofError::RewriteFailed { + detail: format!( + "rename {} -> {}: {}", + tmp_path.display(), + aof_path.display(), + e + ), + })?; + + info!( + "AOF rewrite (sharded, RDB preamble) complete: {} bytes", + rdb_bytes.len() + ); + Ok(()) +} + +/// Reopen AOF file in append mode after atomic rewrite replaced it. +#[allow(dead_code)] +fn reopen_aof_sync(aof_path: &Path) -> Result { + std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(aof_path) +} + +/// Rewrite the AOF file (tokio async wrapper). +/// +/// Delegates to `rewrite_aof_sync` — the actual I/O is synchronous (temp write + rename). +#[cfg(feature = "runtime-tokio")] +#[tracing::instrument(skip_all, level = "info")] +pub async fn rewrite_aof(db: SharedDatabases, aof_path: &Path) -> Result<(), MoonError> { + rewrite_aof_sync(&db, aof_path) +} diff --git a/src/persistence/aof/writer_task.rs b/src/persistence/aof/writer_task.rs new file mode 100644 index 00000000..18b687ee --- /dev/null +++ b/src/persistence/aof/writer_task.rs @@ -0,0 +1,1017 @@ +//! AOF writer tasks: single-file and per-shard background append loops. +#![allow(unused_imports, unused_variables, unreachable_code, clippy::empty_loop)] + +use super::rewrite::{do_rewrite_per_shard, rewrite_aof_sharded_sync}; +use super::*; +// `do_rewrite_single` / `do_rewrite_sharded` exist only under the monoio runtime. +#[cfg(feature = "runtime-monoio")] +use super::rewrite::{do_rewrite_sharded, do_rewrite_single}; + +/// Background AOF writer task. Receives commands via mpsc channel and appends them +/// to the AOF file. Handles fsync according to the configured policy. +pub async fn aof_writer_task( + rx: channel::MpscReceiver, + aof_path: PathBuf, + fsync: FsyncPolicy, + cancel: CancellationToken, +) { + #[cfg(feature = "runtime-tokio")] + use tokio::io::AsyncWriteExt; + + // Open file in append mode (create if not exists) + #[cfg(feature = "runtime-tokio")] + let file: tokio::fs::File = match tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&aof_path) + .await + { + Ok(f) => f, + Err(e) => { + error!("Failed to open AOF file {}: {}", aof_path.display(), e); + return; + } + }; + + #[cfg(feature = "runtime-tokio")] + let mut writer = tokio::io::BufWriter::new(file); + #[cfg(feature = "runtime-tokio")] + let mut last_fsync = Instant::now(); + #[cfg(feature = "runtime-tokio")] + let mut interval = tokio::time::interval(std::time::Duration::from_secs(1)); + #[cfg(feature = "runtime-tokio")] + interval.tick().await; // consume first tick + + // Monoio path: multi-part AOF (base RDB + incremental RESP) with sync I/O. + // + // On startup, if appendonlydir/ exists with a manifest, open the current + // incr file for appending. Otherwise start fresh with seq 1. + // On BGREWRITEAOF: snapshot → write new base RDB → create new incr → advance manifest. + #[cfg(feature = "runtime-monoio")] + { + use crate::persistence::aof_manifest::AofManifest; + use std::io::Write; + + // Resolve the persistence base directory from aof_path's parent. + let base_dir = aof_path.parent().unwrap_or(Path::new(".")).to_path_buf(); + + // Load manifest — do NOT create one here if it doesn't exist. + // main.rs recovery runs concurrently and must finish before a manifest + // is created, to avoid racing against legacy single-file AOF detection. + // main.rs will create the manifest after recovery completes. + // + // A corrupt manifest is fatal — exit the writer so the server startup + // notices and fails loud rather than silently overwriting. + // + // Bounded wait: check the cancellation token each iteration and enforce + // a hard timeout so the writer doesn't spin forever if main.rs fails to + // create the manifest (e.g. disk full, permission error). + let manifest_wait_start = Instant::now(); + const MANIFEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + let mut manifest = loop { + if cancel.is_cancelled() { + info!("AOF writer: cancelled while waiting for manifest"); + return; + } + if manifest_wait_start.elapsed() > MANIFEST_TIMEOUT { + error!( + "AOF writer: manifest not found at {} after {:?}. Writer exiting; check recovery logs.", + base_dir.display(), + MANIFEST_TIMEOUT, + ); + return; + } + match AofManifest::load(&base_dir) { + Ok(Some(m)) => break m, + Ok(None) => { + // main.rs recovery hasn't created the manifest yet — wait. + std::thread::sleep(std::time::Duration::from_millis(50)); + } + Err(e) => { + error!( + "AOF manifest corrupt at {}: {}. Writer exiting; persistence disabled.", + base_dir.display(), + e + ); + return; + } + } + }; + + // Open the current incremental file for appending + let incr_path = manifest.incr_path(); + let mut file = match std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr_path) + { + Ok(f) => f, + Err(e) => { + error!( + "Failed to open AOF incr file {}: {}", + incr_path.display(), + e + ); + return; + } + }; + info!( + "AOF writer: seq {}, incr={}", + manifest.seq, + incr_path.display() + ); + + let mut last_fsync = Instant::now(); + + let mut write_error = false; + + // Test-only fault injection: same env var as the PerShard writer. + // Read once at task startup; zero cost in production (var absent). + let fail_fsync_for_test = std::env::var("MOON_TEST_AOF_FSYNC_FAIL").as_deref() == Ok("1"); + + loop { + match rx.recv() { + // TopLevel writer: legacy v1 disk format is plain RESP. The + // LSN is ignored — TopLevel is single-shard so per-shard merge + // by LSN is moot. + Ok(AofMessage::Append { + bytes: data, + lsn: _, + }) => { + if write_error { + continue; // Drop appends after persistent I/O failure + } + if let Err(e) = file.write_all(&data) { + error!( + "AOF write failed (seq {}): {}. Persistence degraded.", + manifest.seq, e + ); + write_error = true; + continue; + } + match fsync { + FsyncPolicy::Always => { + let t = Instant::now(); + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF sync failed (seq {}, always): {}", manifest.seq, e); + write_error = true; + } else { + crate::admin::metrics_setup::record_aof_fsync( + t.elapsed().as_micros() as u64, + ); + } + } + FsyncPolicy::EverySec => { + if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { + let t = Instant::now(); + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!( + "AOF sync failed (seq {}, everysec): {}", + manifest.seq, e + ); + // Non-fatal for everysec: retry next interval + } else { + crate::admin::metrics_setup::record_aof_fsync( + t.elapsed().as_micros() as u64, + ); + last_fsync = Instant::now(); + } + } + } + FsyncPolicy::No => {} + } + } + // TopLevel writer (monoio): legacy v1 plain RESP, lsn ignored. + // AppendSync ALWAYS fsyncs and acks before returning, regardless + // of the configured policy — that's the durability contract the + // caller signed up for by choosing AppendSync. + Ok(AofMessage::AppendSync { + bytes: data, + lsn: _, + ack, + }) => { + if write_error { + let _ = ack.send(AofAck::WriteFailed); + continue; + } + // Test-only: return FsyncFailed immediately without touching disk. + if fail_fsync_for_test { + let _ = ack.send(AofAck::FsyncFailed); + continue; + } + if let Err(e) = file.write_all(&data) { + error!( + "AOF AppendSync write failed (seq {}): {}. Persistence degraded.", + manifest.seq, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } + let t = Instant::now(); + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF AppendSync sync failed (seq {}): {}", manifest.seq, e); + write_error = true; + let _ = ack.send(AofAck::FsyncFailed); + } else { + crate::admin::metrics_setup::record_aof_fsync( + t.elapsed().as_micros() as u64 + ); + let _ = ack.send(AofAck::Synced); + } + } + Ok(AofMessage::Shutdown) | Err(_) => { + if !write_error { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF final sync failed (seq {}): {}", manifest.seq, e); + } + } + info!("AOF writer shutting down (monoio, seq {})", manifest.seq); + break; + } + Ok(AofMessage::Rewrite(db)) => { + if !write_error { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF pre-rewrite sync failed (seq {}): {}", manifest.seq, e); + } + } + match do_rewrite_single(&db, &mut manifest, &mut file, &rx) { + Ok(()) => { + write_error = false; // Reset on successful rewrite + } + Err(e) => error!("AOF rewrite failed (seq {}): {}", manifest.seq, e), + } + crate::command::persistence::AOF_REWRITE_IN_PROGRESS + .store(false, std::sync::atomic::Ordering::SeqCst); + } + Ok(AofMessage::RewriteSharded(shard_dbs)) => { + if !write_error { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF pre-rewrite sync failed (seq {}): {}", manifest.seq, e); + } + } + match do_rewrite_sharded(&shard_dbs, &mut manifest, &mut file, &rx) { + Ok(()) => { + write_error = false; + } + Err(e) => error!("AOF rewrite failed (seq {}): {}", manifest.seq, e), + } + crate::command::persistence::AOF_REWRITE_IN_PROGRESS + .store(false, std::sync::atomic::Ordering::SeqCst); + } + // [F6] A TopLevel writer never owns per-shard files; receiving + // RewritePerShard means a routing bug. Self-abort so the + // coordinator's countdown completes and the flag clears. + Ok(AofMessage::RewritePerShard { coord, .. }) => { + warn!("AOF TopLevel writer received RewritePerShard — routing bug; aborting"); + coord.mark_failed(); + coord.shard_done(); + } + } + } + return; + } + + loop { + #[cfg(feature = "runtime-tokio")] + tokio::select! { + msg = rx.recv_async() => { + match msg { + // TopLevel writer (tokio): legacy v1 plain RESP, lsn ignored. + Ok(AofMessage::Append { bytes: data, lsn: _ }) => { + if let Err(e) = writer.write_all(&data).await { + error!("AOF write error: {}", e); + continue; + } + match fsync { + FsyncPolicy::Always => { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + } + FsyncPolicy::EverySec | FsyncPolicy::No => { + // EverySec handled by interval tick below; No does nothing + } + } + } + // AppendSync: write + fsync + ack, regardless of policy. + Ok(AofMessage::AppendSync { bytes: data, lsn: _, ack }) => { + if let Err(e) = writer.write_all(&data).await { + error!("AOF AppendSync write error: {}", e); + let _ = ack.send(AofAck::WriteFailed); + continue; + } + if let Err(e) = writer.flush().await { + error!("AOF AppendSync flush error: {}", e); + let _ = ack.send(AofAck::FsyncFailed); + continue; + } + if let Err(e) = writer.get_ref().sync_data().await { + error!("AOF AppendSync sync_data error: {}", e); + let _ = ack.send(AofAck::FsyncFailed); + continue; + } + let _ = ack.send(AofAck::Synced); + } + Ok(AofMessage::Rewrite(db)) => { + // Flush current writer before rewrite + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + + if let Err(e) = rewrite_aof(db, &aof_path).await { + error!("AOF rewrite failed: {}", e); + } + crate::command::persistence::AOF_REWRITE_IN_PROGRESS + .store(false, std::sync::atomic::Ordering::SeqCst); + + // Reopen file after rewrite (it was replaced) + let reopen_result: Result = tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&aof_path) + .await; + match reopen_result { + Ok(f) => { + writer = tokio::io::BufWriter::new(f); + } + Err(e) => { + error!("Failed to reopen AOF file after rewrite: {}", e); + return; + } + } + } + Ok(AofMessage::RewriteSharded(shard_dbs)) => { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + if let Err(e) = rewrite_aof_sharded_sync(&shard_dbs, &aof_path) { + error!("AOF rewrite (sharded) failed: {}", e); + } + crate::command::persistence::AOF_REWRITE_IN_PROGRESS + .store(false, std::sync::atomic::Ordering::SeqCst); + let reopen_result: Result = tokio::fs::OpenOptions::new() + .create(true).append(true).open(&aof_path).await; + match reopen_result { + Ok(f) => writer = tokio::io::BufWriter::new(f), + Err(e) => { error!("Failed to reopen AOF after rewrite: {}", e); return; } + } + } + // [F6] TopLevel writer never owns per-shard files — routing + // bug. Self-abort so the countdown completes + flag clears. + Ok(AofMessage::RewritePerShard { coord, .. }) => { + warn!("AOF TopLevel writer received RewritePerShard — routing bug; aborting"); + coord.mark_failed(); + coord.shard_done(); + } + Ok(AofMessage::Shutdown) | Err(_) => { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + info!("AOF writer shutting down"); + break; + } + } + } + _ = interval.tick(), if fsync == FsyncPolicy::EverySec => { + if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + last_fsync = Instant::now(); + } + } + _ = cancel.cancelled() => { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + info!("AOF writer cancelled"); + break; + } + } + } +} + +/// Background per-shard AOF writer task (Option B step 2b). +/// +/// One instance is spawned per shard in `PerShard` layout. Each instance owns +/// `appendonlydir/shard-{shard_id}/moon.aof.{seq}.incr.aof` exclusively — no +/// other writer touches that file, so there is no per-file locking. +/// +/// Differences from [`aof_writer_task`] (TopLevel): +/// - Opens `manifest.shard_incr_path(shard_id)` instead of `manifest.incr_path()`. +/// - `Rewrite`/`RewriteSharded` variants are rejected (logged + dropped). +/// The legacy single-writer rewrite enum has no meaning when each shard +/// owns its own files; per-shard BGREWRITEAOF lands in RFC step 6. +/// - Refuses to start if the loaded manifest's layout is `TopLevel` — the +/// spawn site (step 2f) must only invoke this task body for `PerShard` +/// layouts. Mismatch is a programmer error. +/// +/// Wait/timeout/corruption semantics for manifest loading match the existing +/// `aof_writer_task` (60s bounded wait, hard fail on corrupt manifest). +/// Test-only torn-write injection for `per_shard_aof_writer_task`: when set to a +/// nonzero `N`, the `N`-th `Append` received by a tokio per-shard writer writes +/// its header and then simulates a payload write failure, exercising the +/// `write_error` latch. `0` disables. Atomic (not an env var) because +/// `std::env::set_var` is `unsafe` under edition 2024. Compiled out of release. +#[cfg(all(test, feature = "runtime-tokio"))] +pub(crate) static TEST_FAIL_WRITE_AT: std::sync::atomic::AtomicUsize = + std::sync::atomic::AtomicUsize::new(0); + +pub async fn per_shard_aof_writer_task( + rx: channel::MpscReceiver, + base_dir: PathBuf, + shard_id: u16, + fsync: FsyncPolicy, + cancel: CancellationToken, +) { + #[cfg(feature = "runtime-tokio")] + { + use crate::persistence::aof_manifest::{AofLayout, AofManifest}; + use tokio::io::AsyncWriteExt; + + // Wait for main.rs recovery to create/load the manifest. + let manifest_wait_start = Instant::now(); + const MANIFEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + let manifest = loop { + if cancel.is_cancelled() { + info!( + "AOF writer shard {}: cancelled while waiting for manifest", + shard_id + ); + return; + } + if manifest_wait_start.elapsed() > MANIFEST_TIMEOUT { + error!( + "AOF writer shard {}: manifest not found at {} after {:?}. Writer exiting.", + shard_id, + base_dir.display(), + MANIFEST_TIMEOUT, + ); + return; + } + match AofManifest::load(&base_dir) { + Ok(Some(m)) => break m, + Ok(None) => { + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + Err(e) => { + error!( + "AOF writer shard {}: manifest corrupt at {}: {}. Persistence disabled.", + shard_id, + base_dir.display(), + e + ); + return; + } + } + }; + + if manifest.layout != AofLayout::PerShard { + error!( + "AOF writer shard {}: layout is {:?}, expected PerShard. \ + per_shard_aof_writer_task should only be spawned for PerShard layouts. \ + Writer exiting.", + shard_id, manifest.layout + ); + return; + } + if (shard_id as usize) >= manifest.shards.len() { + error!( + "AOF writer shard {}: out of range for manifest with {} shards. Writer exiting.", + shard_id, + manifest.shards.len() + ); + return; + } + + let incr_path = manifest.shard_incr_path(shard_id); + // Ensure shard-{N}/ exists. The manifest constructor for PerShard + // already creates these, but be defensive — a manual deletion or + // a manifest written by an older binary could leave them missing. + if let Some(parent) = incr_path.parent() { + if let Err(e) = tokio::fs::create_dir_all(parent).await { + error!( + "AOF writer shard {}: failed to create dir {}: {}", + shard_id, + parent.display(), + e + ); + return; + } + } + let file: tokio::fs::File = match tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr_path) + .await + { + Ok(f) => f, + Err(e) => { + error!( + "AOF writer shard {}: failed to open incr {}: {}", + shard_id, + incr_path.display(), + e + ); + return; + } + }; + info!( + "AOF writer shard {}: seq {}, incr={}", + shard_id, + manifest.seq, + incr_path.display() + ); + + let mut writer = tokio::io::BufWriter::new(file); + let mut last_fsync = Instant::now(); + // (No `interval` here: the EverySec flush deadline is enforced by the + // timeout-bounded recv in the loop below, which wakes at least every + // 200ms regardless of message traffic. A long-lived `interval.tick()` + // select arm is fairness-starvable under sustained writes and proved + // unreliable when idle on this dedicated current-thread writer runtime.) + + // Test-only fault injection: if MOON_TEST_AOF_FSYNC_FAIL=1 is set in + // the environment at writer task startup, every AppendSync ack resolves + // as FsyncFailed instead of Synced. This lets integration tests exercise + // the AOF_FSYNC_ERR response path without requiring a real disk error. + // The env var is read once here (not per-message) so it costs zero on the + // hot path in production deployments where the var is absent. + let fail_fsync_for_test = std::env::var("MOON_TEST_AOF_FSYNC_FAIL").as_deref() == Ok("1"); + + // Torn-write latch: once any write to this incr file fails partway + // (e.g. the header landed but the payload did not), we must NEVER write + // another record — a lone orphaned header makes the framed reader + // misinterpret the next record's bytes as the orphan's payload, + // corrupting every record after it on replay. Stay latched for the + // writer's lifetime; recovery replays the clean prefix and a rewrite + // starts a fresh file. This mirrors the single-file (line ~1467) and + // monoio per-shard (line ~2125) writers, which already carry the latch. + let mut write_error = false; + // Test-only fault injection (no env var: edition-2024 set_var is unsafe). + // When `TEST_FAIL_WRITE_AT` is the ordinal of an incoming Append, that + // append writes its header then simulates a payload failure, exercising + // the latch. Compiled out of production builds. + #[cfg(test)] + let mut test_append_ordinal: usize = 0; + + loop { + tokio::select! { + // Bounded recv (EverySec durability): wake at least every 200ms + // even when idle so the flush deadline after this select! is + // honored within its 1s bound. flume's recv future is drop-safe + // on the Elapsed branch (no message consumed on timeout); the + // Ok(Ok(msg)) path below captures the message with no loss. + r = tokio::time::timeout( + std::time::Duration::from_millis(200), + rx.recv_async(), + ) => { + // On Elapsed (timeout) `r` is Err: skip the match and fall + // through to the EverySec deadline check after this select!. + if let Ok(msg) = r { + match msg { + // PerShard writer (tokio): per RFC § 2 Rule 1 the on-disk + // format is `[u64 lsn LE][u32 len LE][RESP bytes]`. Header + // is written sequentially with the body — both calls land + // in the same BufWriter so this is one syscall under load. + Ok(AofMessage::Append { lsn, bytes: data }) => { + // Latch: stream already torn — drop silently (Append + // is fire-and-forget; no ack channel to notify). + if write_error { + continue; + } + #[cfg(test)] + { + test_append_ordinal += 1; + let fail_at = TEST_FAIL_WRITE_AT + .load(std::sync::atomic::Ordering::Relaxed); + if fail_at != 0 && fail_at == test_append_ordinal { + // Reproduce a torn write: header lands, payload + // "fails". The orphaned header is flushed so the + // on-disk effect matches the real I/O-error case. + let mut header = [0u8; 12]; + header[..8].copy_from_slice(&lsn.to_le_bytes()); + header[8..] + .copy_from_slice(&(data.len() as u32).to_le_bytes()); + let _ = writer.write_all(&header).await; + let _ = writer.flush().await; + error!( + "AOF shard {}: injected torn write after header (test)", + shard_id + ); + write_error = true; + continue; + } + } + let mut header = [0u8; 12]; + header[..8].copy_from_slice(&lsn.to_le_bytes()); + header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); + if let Err(e) = writer.write_all(&header).await { + error!("AOF header write error shard {}: {}", shard_id, e); + write_error = true; + continue; + } + if let Err(e) = writer.write_all(&data).await { + error!("AOF write error shard {}: {}", shard_id, e); + write_error = true; + continue; + } + if matches!(fsync, FsyncPolicy::Always) { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + } + } + // AppendSync (tokio + PerShard): framed write + fsync + ack. + Ok(AofMessage::AppendSync { lsn, bytes: data, ack }) => { + // Latch: stream already torn — refuse to write more and + // report failure so the caller does not hang to the F2 + // timeout and does not ack a write into a corrupt stream. + if write_error { + let _ = ack.send(AofAck::WriteFailed); + continue; + } + let mut header = [0u8; 12]; + header[..8].copy_from_slice(&lsn.to_le_bytes()); + header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); + if let Err(e) = writer.write_all(&header).await { + error!( + "AOF AppendSync header write error shard {}: {}", + shard_id, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } + if let Err(e) = writer.write_all(&data).await { + error!( + "AOF AppendSync write error shard {}: {}", + shard_id, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } + // Test-only: skip real fsync and return FsyncFailed + // immediately when the fault-injection env var is set. + if fail_fsync_for_test { + let _ = ack.send(AofAck::FsyncFailed); + continue; + } + if let Err(e) = writer.flush().await { + error!( + "AOF AppendSync flush error shard {}: {}", + shard_id, e + ); + let _ = ack.send(AofAck::FsyncFailed); + continue; + } + if let Err(e) = writer.get_ref().sync_data().await { + error!( + "AOF AppendSync sync_data error shard {}: {}", + shard_id, e + ); + let _ = ack.send(AofAck::FsyncFailed); + continue; + } + let _ = ack.send(AofAck::Synced); + } + Ok(AofMessage::Rewrite(_)) | Ok(AofMessage::RewriteSharded(_)) => { + warn!( + "AOF writer shard {}: received Rewrite/RewriteSharded — \ + not applicable in PerShard layout, dropped.", + shard_id + ); + } + // [F6] Per-shard rewrite (tokio): reuse the proven + // synchronous fold (`do_rewrite_per_shard`) verbatim, so + // the exactly-once invariant carries over unchanged. This + // writer runs on a DEDICATED std::thread (block_on_local, + // main.rs) — not a shared tokio worker — so executing the + // blocking fold here cannot starve the runtime. We flush + // the BufWriter (its `into_inner` does NOT flush) so any + // buffered appends are durable in the OLD incr, convert + // `tokio::fs::File` -> `std::fs::File` for the sync fold, + // then wrap the (reopened) file back into the BufWriter. + Ok(AofMessage::RewritePerShard { shard_dbs, coord }) => { + if let Err(e) = writer.flush().await { + error!( + "F6 tokio per-shard rewrite: shard {} pre-fold flush \ + failed: {}. Aborting; old generation stays authoritative.", + shard_id, e + ); + coord.mark_failed(); + coord.shard_done(); + } else { + // `into_std().await` waits for in-flight ops and is + // infallible; the buffer is already flushed above. + let mut sf = writer.into_inner().into_std().await; + let res = do_rewrite_per_shard( + shard_id, &shard_dbs, &mut sf, &rx, &coord, + ); + // `sf` is left pointing at the committed generation + // by the fold's internal barrier: NEW incr on + // success, OLD incr on abort (phase-8 rollback) or + // on a pre-reopen error. The fold's ShardDoneGuard + // already performed `shard_done` for every exit, so + // we MUST NOT decrement again here. Wrap `sf` back so + // the writer stays valid either way. + writer = + tokio::io::BufWriter::new(tokio::fs::File::from_std(sf)); + if let Err(e) = res { + error!( + "F6 tokio per-shard rewrite: shard {} fold failed: {}. \ + Rewrite aborted by the fold guard; old generation \ + stays authoritative.", + shard_id, e + ); + } + } + } + Ok(AofMessage::Shutdown) | Err(_) => { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + info!("AOF writer shard {} shutting down", shard_id); + break; + } + } + } + } + _ = cancel.cancelled() => { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + info!("AOF writer shard {} cancelled", shard_id); + break; + } + } + // EverySec deadline — checked after EVERY wake (message OR timeout), + // so it is NOT subject to select! fairness and holds the 1s bound + // under sustained writes as well as when idle. (The old long-lived + // `interval.tick()` arm could be starved by the always-ready recv + // arm under load, leaving >1s of writes buffered in the BufWriter + // and lost on SIGKILL — the COMPOSE crash-matrix failure.) + if fsync == FsyncPolicy::EverySec + && last_fsync.elapsed() >= std::time::Duration::from_secs(1) + { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + last_fsync = Instant::now(); + } + } + } + + #[cfg(feature = "runtime-monoio")] + { + use crate::persistence::aof_manifest::{AofLayout, AofManifest}; + use std::io::Write; + + let manifest_wait_start = Instant::now(); + const MANIFEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + let manifest = loop { + if cancel.is_cancelled() { + info!( + "AOF writer shard {}: cancelled while waiting for manifest", + shard_id + ); + return; + } + if manifest_wait_start.elapsed() > MANIFEST_TIMEOUT { + error!( + "AOF writer shard {}: manifest not found at {} after {:?}. Writer exiting.", + shard_id, + base_dir.display(), + MANIFEST_TIMEOUT, + ); + return; + } + match AofManifest::load(&base_dir) { + Ok(Some(m)) => break m, + Ok(None) => { + std::thread::sleep(std::time::Duration::from_millis(50)); + } + Err(e) => { + error!( + "AOF writer shard {}: manifest corrupt at {}: {}. Persistence disabled.", + shard_id, + base_dir.display(), + e + ); + return; + } + } + }; + + if manifest.layout != AofLayout::PerShard { + error!( + "AOF writer shard {}: layout is {:?}, expected PerShard. Writer exiting.", + shard_id, manifest.layout + ); + return; + } + if (shard_id as usize) >= manifest.shards.len() { + error!( + "AOF writer shard {}: out of range for manifest with {} shards. Writer exiting.", + shard_id, + manifest.shards.len() + ); + return; + } + + let incr_path = manifest.shard_incr_path(shard_id); + if let Some(parent) = incr_path.parent() { + if let Err(e) = std::fs::create_dir_all(parent) { + error!( + "AOF writer shard {}: failed to create dir {}: {}", + shard_id, + parent.display(), + e + ); + return; + } + } + let mut file = match std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr_path) + { + Ok(f) => f, + Err(e) => { + error!( + "AOF writer shard {}: failed to open incr {}: {}", + shard_id, + incr_path.display(), + e + ); + return; + } + }; + info!( + "AOF writer shard {}: seq {}, incr={}", + shard_id, + manifest.seq, + incr_path.display() + ); + + let mut last_fsync = Instant::now(); + let mut write_error = false; + // Test-only fault injection: if MOON_TEST_AOF_FSYNC_FAIL=1 is set in + // the environment at writer task startup, every AppendSync ack resolves + // as FsyncFailed instead of Synced. Read once before the loop so there + // is zero cost in production deployments where the var is absent. + let fail_fsync_for_test = std::env::var("MOON_TEST_AOF_FSYNC_FAIL").as_deref() == Ok("1"); + + loop { + match rx.recv() { + // AppendSync (monoio + PerShard): framed write + fsync + ack. + Ok(AofMessage::AppendSync { + lsn, + bytes: data, + ack, + }) => { + if write_error { + let _ = ack.send(AofAck::WriteFailed); + continue; + } + let mut header = [0u8; 12]; + header[..8].copy_from_slice(&lsn.to_le_bytes()); + header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); + if let Err(e) = file.write_all(&header) { + error!( + "AOF AppendSync header write failed shard {} (seq {}): {}", + shard_id, manifest.seq, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } + if let Err(e) = file.write_all(&data) { + error!( + "AOF AppendSync write failed shard {} (seq {}): {}", + shard_id, manifest.seq, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } + // Test-only: skip real fsync and return FsyncFailed + // immediately when the fault-injection env var is set. + if fail_fsync_for_test { + let _ = ack.send(AofAck::FsyncFailed); + continue; + } + let t = Instant::now(); + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!( + "AOF AppendSync sync failed shard {} (seq {}): {}", + shard_id, manifest.seq, e + ); + write_error = true; + let _ = ack.send(AofAck::FsyncFailed); + } else { + crate::admin::metrics_setup::record_aof_fsync( + t.elapsed().as_micros() as u64 + ); + let _ = ack.send(AofAck::Synced); + } + } + // PerShard writer (monoio): framed `[u64 lsn LE][u32 len LE][RESP]`. + // See the tokio twin above for format rationale. + Ok(AofMessage::Append { lsn, bytes: data }) => { + if write_error { + continue; + } + let mut header = [0u8; 12]; + header[..8].copy_from_slice(&lsn.to_le_bytes()); + header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); + if let Err(e) = file.write_all(&header) { + error!( + "AOF header write failed shard {} (seq {}): {}. Persistence degraded.", + shard_id, manifest.seq, e + ); + write_error = true; + continue; + } + if let Err(e) = file.write_all(&data) { + error!( + "AOF write failed shard {} (seq {}): {}. Persistence degraded.", + shard_id, manifest.seq, e + ); + write_error = true; + continue; + } + match fsync { + FsyncPolicy::Always => { + let t = Instant::now(); + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!( + "AOF sync failed shard {} (seq {}, always): {}", + shard_id, manifest.seq, e + ); + write_error = true; + } else { + crate::admin::metrics_setup::record_aof_fsync( + t.elapsed().as_micros() as u64, + ); + } + } + FsyncPolicy::EverySec => { + if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { + let t = Instant::now(); + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!( + "AOF sync failed shard {} (seq {}, everysec): {}", + shard_id, manifest.seq, e + ); + } else { + crate::admin::metrics_setup::record_aof_fsync( + t.elapsed().as_micros() as u64, + ); + last_fsync = Instant::now(); + } + } + } + FsyncPolicy::No => {} + } + } + Ok(AofMessage::Rewrite(_)) | Ok(AofMessage::RewriteSharded(_)) => { + warn!( + "AOF writer shard {}: received Rewrite/RewriteSharded — \ + not applicable in PerShard layout (use per-shard \ + BGREWRITEAOF), dropped.", + shard_id + ); + } + // [F6] Per-shard rewrite fan-out (monoio). Fold THIS shard, + // then signal the coordinator; the last shard commits the + // manifest. On error the old generation stays authoritative + // (advance_shard did not commit the seq). + Ok(AofMessage::RewritePerShard { shard_dbs, coord }) => { + if let Err(e) = + do_rewrite_per_shard(shard_id, &shard_dbs, &mut file, &rx, &coord) + { + // The fold's ShardDoneGuard already marked the rewrite + // failed and decremented on this error exit (committing + // new_seq with a shard missing its new base would break + // recovery), so do NOT decrement again here. `file` is left + // on the OLD incr (error exits are pre-reopen). + error!( + "F6 per-shard rewrite: shard {} fold failed: {}. \ + Rewrite aborted by the fold guard; old generation \ + stays authoritative.", + shard_id, e + ); + } + } + Ok(AofMessage::Shutdown) | Err(_) => { + if !write_error { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!( + "AOF final sync failed shard {} (seq {}): {}", + shard_id, manifest.seq, e + ); + } + } + info!( + "AOF writer shard {} shutting down (monoio, seq {})", + shard_id, manifest.seq + ); + break; + } + } + } + } +} From edffce330730982172c71c462359316bc2f50d92 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 18:14:45 +0700 Subject: [PATCH 13/15] test(vector): fix VECTOR_INDEXES counter race via RwLock test isolation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `command::vector_search::tests::test_vector_metrics_increment_decrement` flaked under the tokio full-suite run (2861/1, "FT.DROPINDEX should decrement VECTOR_INDEXES"), passing in isolation. Root cause is a test-isolation defect, not a production bug: `VECTOR_INDEXES` (src/vector/metrics.rs) is a process- global AtomicU64 shared by every test in the lib binary. 36 tests mutate it via ft_create/ft_dropindex; only 8 held `METRICS_LOCK` (a plain Mutex), which serialized those 8 against each other but did nothing about the ~28 lock-free mutators. The failing test reads the counter, drops an index, re-reads, and asserts `after_drop < before_drop` — a relative delta that a concurrent ft_create (+1) landing in the read-modify-read window cancels (N-1+1 = N). Only this assertion flaked because VECTOR_INDEXES is the one counter that moves both directions; the others assert `after > before` on monotonic-up counters. Fix — RwLock isolation (production code unchanged): - METRICS_LOCK: Mutex<()> -> RwLock<()>. - The 8 delta-reader tests take `.write()` (exclusive) — no mutator runs during their critical section. - The 28 mutator-only tests take `.read()` (shared) — they keep running in parallel with each other; they are excluded only during a reader's brief exclusive window. Zero added serialization between mutators. - Upgrade the reader assertions to exact deltas (+1 create, +1 search, -1 drop), now deterministic under the write guard — stronger than the old inequalities. Regression guard — metrics_write_guard_isolates_index_counter_from_concurrent_mutator: deterministic in BOTH directions. A spawned lock-respecting mutator blocks on `.read()` for a 20ms window while the test holds `.write()`, so the drop delta is isolated (GREEN). Flipping the test's `write()` to `read()` lets the mutator increment during the window and turns the assertion RED (verified: left:1 right:0 — the exact original failure mode). Verification: - Red/green proven on the regression test (write=GREEN, read=RED). - tokio full lib suite (the config that flaked): 2863 passed / 0 failed x2 rounds. - vector_search module stress: 20 rounds x 16 threads, 0 failures. - monoio metrics tests 11/0 (RwLock is runtime-agnostic). - CI-exact clippy -D warnings + fmt clean on both runtimes. author: Tin Dang --- src/command/vector_search/tests.rs | 130 ++++++++++++++++++++++++----- 1 file changed, 109 insertions(+), 21 deletions(-) diff --git a/src/command/vector_search/tests.rs b/src/command/vector_search/tests.rs index c0fba026..ca8f717a 100644 --- a/src/command/vector_search/tests.rs +++ b/src/command/vector_search/tests.rs @@ -1,9 +1,9 @@ use super::*; use smallvec::SmallVec; -use std::sync::Mutex; +use std::sync::RwLock; /// Serialize tests that touch global atomic metrics to avoid flaky interference. -static METRICS_LOCK: Mutex<()> = Mutex::new(()); +static METRICS_LOCK: RwLock<()> = RwLock::new(()); fn bulk(s: &[u8]) -> Frame { Frame::BulkString(Bytes::from(s.to_vec())) @@ -292,6 +292,7 @@ fn ft_create_args() -> Vec { #[test] fn test_ft_create_parse_full_syntax() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); let result = ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -309,6 +310,7 @@ fn test_ft_create_parse_full_syntax() { #[test] fn test_ft_create_missing_dim() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); // Remove DIM param pair: keep TYPE FLOAT32 and DISTANCE_METRIC L2 (4 params = 2 pairs) let args = vec![ @@ -337,6 +339,7 @@ fn test_ft_create_missing_dim() { #[test] fn test_ft_create_duplicate() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); let r1 = ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -356,6 +359,7 @@ fn test_ft_create_duplicate() { #[test] fn test_ft_dropindex() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -518,6 +522,7 @@ fn test_merge_search_results_empty() { #[test] fn test_ft_search_dimension_mismatch() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -544,6 +549,7 @@ fn test_ft_search_dimension_mismatch() { #[test] fn test_ft_search_empty_index() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -570,6 +576,7 @@ fn test_ft_search_empty_index() { #[test] fn test_ft_info() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -645,6 +652,7 @@ fn build_ft_create_args( #[test] fn test_end_to_end_create_insert_search() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); // Initialize distance functions (required before any search) crate::vector::distance::init(); @@ -736,6 +744,7 @@ fn test_end_to_end_create_insert_search() { #[test] fn test_ft_info_returns_correct_data() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = build_ft_create_args("testidx", "test:", "vec", 128, "COSINE"); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -889,6 +898,7 @@ fn test_parse_filter_clause_none() { #[test] fn test_ft_search_with_filter_no_regression() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); // Unfiltered FT.SEARCH still works identically crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -915,6 +925,7 @@ fn test_ft_search_with_filter_no_regression() { #[test] fn test_vector_index_has_payload_index() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -927,18 +938,21 @@ fn test_vector_index_has_payload_index() { fn test_vector_metrics_increment_decrement() { use std::sync::atomic::Ordering; - let _guard = METRICS_LOCK.lock().unwrap(); + let _guard = METRICS_LOCK.write().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); - // FT.CREATE should increment VECTOR_INDEXES + // FT.CREATE should increment VECTOR_INDEXES by exactly 1. The exclusive + // write guard excludes every lock-respecting mutator, so the delta is + // deterministic (no concurrent ft_create/ft_dropindex can perturb it). let before_create = crate::vector::metrics::VECTOR_INDEXES.load(Ordering::Relaxed); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); let after_create = crate::vector::metrics::VECTOR_INDEXES.load(Ordering::Relaxed); - assert!( - after_create > before_create, - "FT.CREATE should increment VECTOR_INDEXES" + assert_eq!( + after_create, + before_create + 1, + "FT.CREATE should increment VECTOR_INDEXES by exactly 1" ); // FT.SEARCH should increment VECTOR_SEARCH_TOTAL @@ -955,12 +969,15 @@ fn test_vector_metrics_increment_decrement() { ]; ft_search(&mut store, &search_args, None, None, 0); let after_search = crate::vector::metrics::VECTOR_SEARCH_TOTAL.load(Ordering::Relaxed); - assert!( - after_search > before_search, - "FT.SEARCH should increment VECTOR_SEARCH_TOTAL" + assert_eq!( + after_search, + before_search + 1, + "FT.SEARCH should increment VECTOR_SEARCH_TOTAL by exactly 1" ); - // FT.DROPINDEX should decrement VECTOR_INDEXES + // FT.DROPINDEX should decrement VECTOR_INDEXES by exactly 1. Deterministic + // under the write guard — this is the assertion that flaked when concurrent + // mutators ran lock-free (a stray ft_create could cancel the decrement). let before_drop = crate::vector::metrics::VECTOR_INDEXES.load(Ordering::Relaxed); ft_dropindex( &mut store, @@ -969,10 +986,64 @@ fn test_vector_metrics_increment_decrement() { &[bulk(b"myidx")], ); let after_drop = crate::vector::metrics::VECTOR_INDEXES.load(Ordering::Relaxed); - assert!( - after_drop < before_drop, - "FT.DROPINDEX should decrement VECTOR_INDEXES" + assert_eq!( + after_drop, + before_drop - 1, + "FT.DROPINDEX should decrement VECTOR_INDEXES by exactly 1" + ); +} + +/// Deterministic regression for the `VECTOR_INDEXES` parallel-test flake. +/// +/// `VECTOR_INDEXES` is a process-global counter shared by every test in this +/// binary. Before the RwLock fix, ~28 tests mutated it lock-free while the +/// delta-reader tests asserted on it, so a concurrent `ft_create` (+1) could +/// land inside a reader's read-modify-read window and cancel an observed +/// decrement — breaking `after_drop < before_drop` (the failure first seen on +/// the tokio full-suite run, never in isolation). +/// +/// The fix: delta-readers take `METRICS_LOCK.write()` (exclusive) and every +/// mutator takes `METRICS_LOCK.read()` (shared, keeps their parallelism). This +/// test proves the write guard excludes a lock-respecting mutator. It is +/// deterministic in BOTH directions: GREEN as written; flipping the `write()` +/// below to `read()` lets the mutator run during the sleep and turns the +/// assertion RED. +#[test] +fn metrics_write_guard_isolates_index_counter_from_concurrent_mutator() { + use std::sync::atomic::Ordering; + + let _exclusive = METRICS_LOCK.write().unwrap(); + + // Seed +1 so the decrement is observable even with no other live index. + crate::vector::metrics::increment_indexes(); + let before_drop = crate::vector::metrics::VECTOR_INDEXES.load(Ordering::Relaxed); + + // A concurrent "ft_create" that respects the lock: it blocks on read() + // until we drop the write guard, so it cannot mutate inside our window. + let mutator = std::thread::spawn(|| { + let _shared = METRICS_LOCK.read().unwrap(); + crate::vector::metrics::increment_indexes(); + }); + // Let the mutator reach (and park on) the read lock. Under the write guard + // it stays blocked; without it, it would increment here and corrupt the + // delta below — that is exactly the RED case. + std::thread::sleep(std::time::Duration::from_millis(20)); + + crate::vector::metrics::decrement_indexes(); + let after_drop = crate::vector::metrics::VECTOR_INDEXES.load(Ordering::Relaxed); + assert_eq!( + after_drop, + before_drop - 1, + "write guard must isolate the drop delta from the concurrent mutator" ); + + // Release; the mutator proceeds with its +1. + drop(_exclusive); + mutator.join().unwrap(); + + // Restore the global under a fresh exclusive guard (undo the mutator's +1). + let _cleanup = METRICS_LOCK.write().unwrap(); + crate::vector::metrics::decrement_indexes(); } #[test] @@ -1275,6 +1346,7 @@ fn test_parse_ft_search_args_without_limit() { #[test] fn test_ft_config_autocompact_on_off() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -1344,6 +1416,7 @@ fn test_ft_config_autocompact_on_off() { #[test] fn test_ft_config_unknown_param() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -1394,6 +1467,7 @@ fn test_ft_config_unknown_index() { #[test] fn test_ft_config_autocompact_guards_try_compact() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -1433,6 +1507,7 @@ fn test_ft_config_autocompact_guards_try_compact() { #[test] fn test_ft_config_autocompact_accepts_variants() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -1869,6 +1944,7 @@ mod cache_search_tests { #[test] fn test_ft_cachesearch_miss_on_empty_store() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let create_args = ft_create_args(); ft_create( @@ -2157,6 +2233,7 @@ fn ft_create_multi_field_args() -> Vec { #[test] fn test_ft_create_multi_field() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_multi_field_args(); let result = ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -2182,6 +2259,7 @@ fn test_ft_create_multi_field() { #[test] fn test_ft_create_duplicate_field_rejected() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = vec![ bulk(b"dupidx"), @@ -2226,6 +2304,7 @@ fn test_ft_create_duplicate_field_rejected() { #[test] fn test_ft_create_exceeds_max_fields() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let mut args = vec![ bulk(b"toomanyidx"), @@ -2263,6 +2342,7 @@ fn test_ft_create_exceeds_max_fields() { #[test] fn test_ft_info_multi_field() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_multi_field_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -2325,6 +2405,7 @@ fn test_ft_info_multi_field() { #[test] fn test_ft_search_field_targeting() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2430,6 +2511,7 @@ fn test_ft_search_field_targeting() { #[test] fn test_ft_search_default_field_compat() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2477,6 +2559,7 @@ fn test_ft_search_default_field_compat() { #[test] fn test_ft_search_unknown_field_error() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -2745,7 +2828,7 @@ fn insert_hybrid_doc( #[test] fn test_hybrid_search_basic() { - let _lock = METRICS_LOCK.lock().unwrap(); + let _lock = METRICS_LOCK.write().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2827,7 +2910,7 @@ fn test_hybrid_search_basic() { #[test] fn test_hybrid_search_sparse_only() { - let _lock = METRICS_LOCK.lock().unwrap(); + let _lock = METRICS_LOCK.write().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2873,7 +2956,7 @@ fn test_hybrid_search_sparse_only() { #[test] fn test_hybrid_search_dense_only_backward_compat() { - let _lock = METRICS_LOCK.lock().unwrap(); + let _lock = METRICS_LOCK.write().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2903,7 +2986,7 @@ fn test_hybrid_search_dense_only_backward_compat() { #[test] fn test_hybrid_search_hit_counts() { - let _lock = METRICS_LOCK.lock().unwrap(); + let _lock = METRICS_LOCK.write().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -3039,7 +3122,7 @@ fn test_parse_range_clause_case_insensitive() { #[test] fn test_range_filter_l2_search() { - let _lock = METRICS_LOCK.lock().unwrap(); + let _lock = METRICS_LOCK.write().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -3215,7 +3298,7 @@ fn test_recommend_unknown_index() { #[test] fn test_recommend_missing_key_vectors() { - let _lock = METRICS_LOCK.lock().unwrap(); + let _lock = METRICS_LOCK.write().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -3244,7 +3327,7 @@ fn test_recommend_missing_key_vectors() { #[test] fn test_recommend_basic_with_vectors() { - let _lock = METRICS_LOCK.lock().unwrap(); + let _lock = METRICS_LOCK.write().unwrap(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -3405,6 +3488,7 @@ mod ft_navigate_tests { #[test] fn test_ft_dropindex_dd_deletes_docs() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); use crate::storage::db::Database; // Create database and vector store @@ -3477,6 +3561,7 @@ fn test_ft_dropindex_dd_deletes_docs() { #[test] fn test_ft_dropindex_preserves_docs() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); use crate::storage::db::Database; let mut db = Database::new(); @@ -3531,6 +3616,7 @@ fn test_ft_dropindex_preserves_docs() { #[test] fn test_ft_dropindex_dd_case_insensitive() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); use crate::storage::db::Database; // Test lowercase 'dd' @@ -3636,6 +3722,7 @@ fn test_ft_dropindex_dd_case_insensitive() { #[test] fn test_ft_dropindex_dd_unknown_index() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); use crate::storage::db::Database; let mut db = Database::new(); @@ -3658,6 +3745,7 @@ fn test_ft_dropindex_dd_unknown_index() { #[test] fn test_ft_dropindex_extra_args_error() { + let _metrics_guard = METRICS_LOCK.read().unwrap(); let mut store = VectorStore::new(); let mut text_store = crate::text::store::TextStore::new(); From 839ba8c985c900ac5314b420c1015a25f1d9ded9 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 20:24:18 +0700 Subject: [PATCH 14/15] docs(changelog): add PR #144 durability + decomposition + test-isolation entry CI Lint gate requires a CHANGELOG.md entry (or the skip-changelog label) per PR. Document the PR #144 scope under [0.2.0-alpha] Unreleased: the 8 CodeRabbit PR #136 durability follow-ups, the two PR #144-review Majors (async-spill eviction regression test, cfg(unix) migration gate), the aof_manifest.rs and aof.rs decompositions, and the VECTOR_INDEXES RwLock test-isolation fix. author: Tin Dang --- CHANGELOG.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e2f0e53..391ecc42 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,34 @@ The v0.2 enterprise beachhead. Built additively on per-shard WAL v3 + the dual-root manifest; no changes to the KV hot path, MVCC, page format, or transaction layer. +### Fixed — CodeRabbit PR #136 durability follow-ups + decomposition + test isolation (PR #144) + +Closes the 8 CodeRabbit findings left open after PR #136, plus two PR #144-review +Majors, oversized-file decomposition, and a parallel-test flake. No production +hot-path behaviour change. + +- **Disk-offload spill (data-loss fixes):** block instead of dropping spill + completions; salvage inline-batch spill failures per-entry rather than + wholesale; preserve spill context across tokio connection migration; recover + the cold tier under `appendonly=no`. +- **Per-shard AOF rewrite robustness:** clear the rewrite flag when fan-out + fails partway; roll per-shard rewrite writers back to the committed generation + on abort (barrier-before-resume + panic-safe `ShardDoneGuard`); ack drained + `AppendSync` only after the boundary fsync (issue #140 ordering). +- **Async-spill eviction:** corrected a stale doc comment, removed the dead + remove-first eviction path, and added a regression test locking the fail-safe + send-before-remove ordering (a full spill channel keeps the victim resident — + no data loss). +- **Platform hygiene:** gate the migrated-connection spawn fns behind + `cfg(all(..., unix))` to match their `RawFd` usage. +- **File decomposition (1500-line cap):** split `aof_manifest.rs` (3058 → + mod/shard_replay/shard_rewrite) and `aof.rs` (4379 → mod/pool/writer_task/ + rewrite); pure code relocation, verified line-exact on both runtimes. +- **Test isolation:** fixed the `VECTOR_INDEXES` counter flake — the process- + global metrics counter is now guarded by an `RwLock` (delta-reader tests take + `write()`, mutator tests take `read()`), making the index-count delta + assertions deterministic under the parallel test harness. + ### Persistence — Per-shard AOF migration complete (PR #129) Closes the P0 multi-shard AOF data-loss bug (~50% loss on SIGKILL with From e41f0652ae1b2233d9fdefcce68dfc4f30c4b68a Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Thu, 4 Jun 2026 20:46:21 +0700 Subject: [PATCH 15/15] test(vector): use parking_lot::RwLock for METRICS_LOCK (no poison cascade) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Self-review of the VECTOR_INDEXES isolation fix surfaced two issues with the std::sync::RwLock guard: 1. CLAUDE.md mandates parking_lot locks over std::sync everywhere. 2. std::sync::RwLock poisons on panic. The isolation fix grew the lock-holder set from 8 to 36 tests, so a single failing assertion under the write guard would poison the lock and cascade PoisonError through all ~28 newly-added read-guard tests' .unwrap() calls — burying the real failing assertion under a wall of identical PoisonError panics. Switch METRICS_LOCK to parking_lot::RwLock (const new() works as a static, same as the existing GATE_TEST_LOCK in src/command/persistence.rs). parking_lot does not poison, so a failing test now reports only its own assertion. Its guards are infallible, so the 29 read + 10 write sites drop their `.unwrap()`. Write- exclusive / read-shared semantics are identical, so the deterministic regression test's red/green behaviour is unchanged. Verified: metrics tests 11/0 on both monoio and tokio; clippy -D warnings + fmt clean on both runtimes. author: Tin Dang --- src/command/vector_search/tests.rs | 80 +++++++++++++++--------------- 1 file changed, 40 insertions(+), 40 deletions(-) diff --git a/src/command/vector_search/tests.rs b/src/command/vector_search/tests.rs index ca8f717a..4a620ce5 100644 --- a/src/command/vector_search/tests.rs +++ b/src/command/vector_search/tests.rs @@ -1,6 +1,6 @@ use super::*; +use parking_lot::RwLock; use smallvec::SmallVec; -use std::sync::RwLock; /// Serialize tests that touch global atomic metrics to avoid flaky interference. static METRICS_LOCK: RwLock<()> = RwLock::new(()); @@ -292,7 +292,7 @@ fn ft_create_args() -> Vec { #[test] fn test_ft_create_parse_full_syntax() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); let result = ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -310,7 +310,7 @@ fn test_ft_create_parse_full_syntax() { #[test] fn test_ft_create_missing_dim() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); // Remove DIM param pair: keep TYPE FLOAT32 and DISTANCE_METRIC L2 (4 params = 2 pairs) let args = vec![ @@ -339,7 +339,7 @@ fn test_ft_create_missing_dim() { #[test] fn test_ft_create_duplicate() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); let r1 = ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -359,7 +359,7 @@ fn test_ft_create_duplicate() { #[test] fn test_ft_dropindex() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -522,7 +522,7 @@ fn test_merge_search_results_empty() { #[test] fn test_ft_search_dimension_mismatch() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -549,7 +549,7 @@ fn test_ft_search_dimension_mismatch() { #[test] fn test_ft_search_empty_index() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -576,7 +576,7 @@ fn test_ft_search_empty_index() { #[test] fn test_ft_info() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -652,7 +652,7 @@ fn build_ft_create_args( #[test] fn test_end_to_end_create_insert_search() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); // Initialize distance functions (required before any search) crate::vector::distance::init(); @@ -744,7 +744,7 @@ fn test_end_to_end_create_insert_search() { #[test] fn test_ft_info_returns_correct_data() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = build_ft_create_args("testidx", "test:", "vec", 128, "COSINE"); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -898,7 +898,7 @@ fn test_parse_filter_clause_none() { #[test] fn test_ft_search_with_filter_no_regression() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); // Unfiltered FT.SEARCH still works identically crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -925,7 +925,7 @@ fn test_ft_search_with_filter_no_regression() { #[test] fn test_vector_index_has_payload_index() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -938,7 +938,7 @@ fn test_vector_index_has_payload_index() { fn test_vector_metrics_increment_decrement() { use std::sync::atomic::Ordering; - let _guard = METRICS_LOCK.write().unwrap(); + let _guard = METRICS_LOCK.write(); let mut store = VectorStore::new(); let args = ft_create_args(); @@ -1012,7 +1012,7 @@ fn test_vector_metrics_increment_decrement() { fn metrics_write_guard_isolates_index_counter_from_concurrent_mutator() { use std::sync::atomic::Ordering; - let _exclusive = METRICS_LOCK.write().unwrap(); + let _exclusive = METRICS_LOCK.write(); // Seed +1 so the decrement is observable even with no other live index. crate::vector::metrics::increment_indexes(); @@ -1021,7 +1021,7 @@ fn metrics_write_guard_isolates_index_counter_from_concurrent_mutator() { // A concurrent "ft_create" that respects the lock: it blocks on read() // until we drop the write guard, so it cannot mutate inside our window. let mutator = std::thread::spawn(|| { - let _shared = METRICS_LOCK.read().unwrap(); + let _shared = METRICS_LOCK.read(); crate::vector::metrics::increment_indexes(); }); // Let the mutator reach (and park on) the read lock. Under the write guard @@ -1042,7 +1042,7 @@ fn metrics_write_guard_isolates_index_counter_from_concurrent_mutator() { mutator.join().unwrap(); // Restore the global under a fresh exclusive guard (undo the mutator's +1). - let _cleanup = METRICS_LOCK.write().unwrap(); + let _cleanup = METRICS_LOCK.write(); crate::vector::metrics::decrement_indexes(); } @@ -1346,7 +1346,7 @@ fn test_parse_ft_search_args_without_limit() { #[test] fn test_ft_config_autocompact_on_off() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -1416,7 +1416,7 @@ fn test_ft_config_autocompact_on_off() { #[test] fn test_ft_config_unknown_param() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -1467,7 +1467,7 @@ fn test_ft_config_unknown_index() { #[test] fn test_ft_config_autocompact_guards_try_compact() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -1507,7 +1507,7 @@ fn test_ft_config_autocompact_guards_try_compact() { #[test] fn test_ft_config_autocompact_accepts_variants() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -1944,7 +1944,7 @@ mod cache_search_tests { #[test] fn test_ft_cachesearch_miss_on_empty_store() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let create_args = ft_create_args(); ft_create( @@ -2233,7 +2233,7 @@ fn ft_create_multi_field_args() -> Vec { #[test] fn test_ft_create_multi_field() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_multi_field_args(); let result = ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -2259,7 +2259,7 @@ fn test_ft_create_multi_field() { #[test] fn test_ft_create_duplicate_field_rejected() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = vec![ bulk(b"dupidx"), @@ -2304,7 +2304,7 @@ fn test_ft_create_duplicate_field_rejected() { #[test] fn test_ft_create_exceeds_max_fields() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let mut args = vec![ bulk(b"toomanyidx"), @@ -2342,7 +2342,7 @@ fn test_ft_create_exceeds_max_fields() { #[test] fn test_ft_info_multi_field() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_multi_field_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -2405,7 +2405,7 @@ fn test_ft_info_multi_field() { #[test] fn test_ft_search_field_targeting() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2511,7 +2511,7 @@ fn test_ft_search_field_targeting() { #[test] fn test_ft_search_default_field_compat() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2559,7 +2559,7 @@ fn test_ft_search_default_field_compat() { #[test] fn test_ft_search_unknown_field_error() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let args = ft_create_args(); ft_create(&mut store, &mut crate::text::store::TextStore::new(), &args); @@ -2828,7 +2828,7 @@ fn insert_hybrid_doc( #[test] fn test_hybrid_search_basic() { - let _lock = METRICS_LOCK.write().unwrap(); + let _lock = METRICS_LOCK.write(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2910,7 +2910,7 @@ fn test_hybrid_search_basic() { #[test] fn test_hybrid_search_sparse_only() { - let _lock = METRICS_LOCK.write().unwrap(); + let _lock = METRICS_LOCK.write(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2956,7 +2956,7 @@ fn test_hybrid_search_sparse_only() { #[test] fn test_hybrid_search_dense_only_backward_compat() { - let _lock = METRICS_LOCK.write().unwrap(); + let _lock = METRICS_LOCK.write(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -2986,7 +2986,7 @@ fn test_hybrid_search_dense_only_backward_compat() { #[test] fn test_hybrid_search_hit_counts() { - let _lock = METRICS_LOCK.write().unwrap(); + let _lock = METRICS_LOCK.write(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -3122,7 +3122,7 @@ fn test_parse_range_clause_case_insensitive() { #[test] fn test_range_filter_l2_search() { - let _lock = METRICS_LOCK.write().unwrap(); + let _lock = METRICS_LOCK.write(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -3298,7 +3298,7 @@ fn test_recommend_unknown_index() { #[test] fn test_recommend_missing_key_vectors() { - let _lock = METRICS_LOCK.write().unwrap(); + let _lock = METRICS_LOCK.write(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -3327,7 +3327,7 @@ fn test_recommend_missing_key_vectors() { #[test] fn test_recommend_basic_with_vectors() { - let _lock = METRICS_LOCK.write().unwrap(); + let _lock = METRICS_LOCK.write(); crate::vector::distance::init(); let mut store = VectorStore::new(); @@ -3488,7 +3488,7 @@ mod ft_navigate_tests { #[test] fn test_ft_dropindex_dd_deletes_docs() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); use crate::storage::db::Database; // Create database and vector store @@ -3561,7 +3561,7 @@ fn test_ft_dropindex_dd_deletes_docs() { #[test] fn test_ft_dropindex_preserves_docs() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); use crate::storage::db::Database; let mut db = Database::new(); @@ -3616,7 +3616,7 @@ fn test_ft_dropindex_preserves_docs() { #[test] fn test_ft_dropindex_dd_case_insensitive() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); use crate::storage::db::Database; // Test lowercase 'dd' @@ -3722,7 +3722,7 @@ fn test_ft_dropindex_dd_case_insensitive() { #[test] fn test_ft_dropindex_dd_unknown_index() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); use crate::storage::db::Database; let mut db = Database::new(); @@ -3745,7 +3745,7 @@ fn test_ft_dropindex_dd_unknown_index() { #[test] fn test_ft_dropindex_extra_args_error() { - let _metrics_guard = METRICS_LOCK.read().unwrap(); + let _metrics_guard = METRICS_LOCK.read(); let mut store = VectorStore::new(); let mut text_store = crate::text::store::TextStore::new();