diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0e2f0e538..391ecc42d 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
diff --git a/src/command/vector_search/tests.rs b/src/command/vector_search/tests.rs
index c0fba0266..4a620ce59 100644
--- a/src/command/vector_search/tests.rs
+++ b/src/command/vector_search/tests.rs
@@ -1,9 +1,9 @@
use super::*;
+use parking_lot::RwLock;
use smallvec::SmallVec;
-use std::sync::Mutex;
/// 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();
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();
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();
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();
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();
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();
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();
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();
// 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();
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();
// 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();
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();
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();
+
+ // 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();
+ 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();
+ 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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
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();
let mut store = VectorStore::new();
let mut text_store = crate::text::store::TextStore::new();
diff --git a/src/persistence/aof.rs b/src/persistence/aof.rs
deleted file mode 100644
index 77b8d09b5..000000000
--- a/src/persistence/aof.rs
+++ /dev/null
@@ -1,3912 +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,
-}
-
-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),
- })
- }
-
- /// 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). A RESTART is \
- recommended so successful shards' writers stop appending to the \
- discarded new generation.",
- 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);
- 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
- );
- crate::command::persistence::AOF_REWRITE_IN_PROGRESS.store(false, Ordering::SeqCst);
- }
- }
-}
-
-/// 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 s in &self.senders {
- // 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: a writer channel is disconnected; \
- rewrite aborted (no manifest commit, old generation remains \
- authoritative). Inspect AOF writer threads."
- );
- 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;
-
- #[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);
- }
-
- // 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,
- );
- // 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.
- 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.",
- shard_id, e
- );
- coord.mark_failed();
- coord.shard_done();
- }
- }
- }
- 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)
- {
- error!(
- "F6 per-shard rewrite: shard {} fold failed: {}. \
- Aborting rewrite; 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(_) => {
- 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,
-}
-
-#[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;
- // 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.
- AofMessage::AppendSync {
- bytes: data,
- lsn: _,
- ack,
- } => {
- file.write_all(&data).map_err(|e| AofError::Io {
- path: PathBuf::from(""),
- source: e,
- })?;
- outcome.drained += 1;
- let _ = ack.send(AofAck::Synced);
- }
- 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;
- // Durability for these is covered by the post-drain fsync at
- // the rewrite boundary (mirrors drain_pending_appends).
- let _ = ack.send(AofAck::Synced);
- }
- 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> {
- 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).
- let pre_drain = drain_pending_appends_framed(rx, file)?;
- file.sync_data().map_err(|e| AofError::Io {
- path: PathBuf::from(""),
- source: e,
- })?;
-
- // 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 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.
- coord.shard_done();
- 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.
- let pre_drain = drain_pending_appends(rx, file)?;
- file.sync_data().map_err(|e| AofError::Io {
- path: manifest.incr_path(),
- source: e,
- })?;
-
- // 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.
- let mid_drain = drain_pending_appends(rx, file)?;
- file.sync_data().map_err(|e| AofError::Io {
- path: manifest.incr_path(),
- source: e,
- })?;
-
- // 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.
- let pre_drain = drain_pending_appends(rx, file)?;
- file.sync_data().map_err(|e| AofError::Io {
- path: manifest.incr_path(),
- source: e,
- })?;
-
- // 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.
- let mid_drain = drain_pending_appends(rx, file)?;
- file.sync_data().map_err(|e| AofError::Io {
- path: manifest.incr_path(),
- source: e,
- })?;
-
- // 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 000000000..90d1cde06
--- /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 000000000..f97d34e5e
--- /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 000000000..c052cbd8c
--- /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 000000000..18b687ee7
--- /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;
+ }
+ }
+ }
+ }
+}
diff --git a/src/persistence/aof_manifest.rs b/src/persistence/aof_manifest.rs
deleted file mode 100644
index 572e16e85..000000000
--- 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