diff --git a/CHANGELOG.md b/CHANGELOG.md index eb96c884..b3a6c910 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,60 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed — consistency/durability defects caught by the R1 gates (task #35) + +Three pre-existing data-integrity bugs surfaced by the new load/kill-9 gates +run for the R1 WAIT/ACK work (all reproduced on `main` before the fix): + +- **Replication: silent record drop for lagging replicas.** The per-replica + fan-out used `try_send` and *skipped* the record when the bounded channel + was full — under one pipelined burst a replica received ~2k of 40k keys, + stayed `master_link_status:up`, and diverged forever. A replica whose + channel overflows is now KICKED (shared `kicked` flag → drain loop closes + the socket) so it reconnects and resyncs from the backlog — Redis's + output-buffer-limit policy. Channel capacity raised 1024 → 16384 records + so bursts kick rarely. New e2e + `replica_converges_under_interleaved_multidb_load` locks in exact + post-load parity. +- **Client `SELECT` commands leaked into the AOF and replication stream.** + SELECT is W-flagged (routing), so every handler persisted the literal + client SELECT while bare writes carried no db context — two interleaved + connections on different dbs corrupted BOTH planes' db attribution + (observed: all 40k keys recovered into db2 after kill-9). SELECT is now + connection-state only (`metadata::is_persisted_write`), never persisted or + replicated. +- **AOF had no per-record db attribution.** The AOF writer now threads the + executing db through `AofMessage` and tracks the stream's current db, + prepending a `SELECT ` record exactly when it changes (Redis's + `aof_selected_db`), including through BGREWRITEAOF fold drains (context + resets per fresh incr segment; ambiguous drains use a force-emit + sentinel). New e2e `aof_multidb_kill9` (TopLevel + PerShard) proves + interleaved multi-db writes recover into the correct databases after + kill-9 on both runtimes. + +### Added — R1: real WAIT/ACK plumbing (replica acknowledgements) + +- **`WAIT ` now works on the production (monoio) + runtime.** It previously answered `:0` unconditionally from the synchronous + dispatch table; a new connection-layer intercept awaits + `wait_for_replicas` (10ms poll, early exit; `timeout 0` = block until + satisfied, capped at one year). Wired on the tokio sharded path too. +- **Replicas acknowledge their applied offset**: a dedicated 1s ticker task + owns the write half of the (split) replication socket and sends + `REPLCONF ACK ` — Redis's replicationCron cadence, doubling as an + idle keepalive for master-side lag detection. A timeout-wrapped read was + rejected: cancelling an in-flight io_uring read whose completion already + landed DISCARDS those bytes (silent stream corruption). +- **The master reads ACKs off the hijacked PSYNC socket**: the inline drain + loop splits the stream; a same-thread reader task parses + `REPLCONF ACK ` frames (dedicated parser — the shared replication + drainer deliberately drops REPLCONF as chatter) and records them into the + replica's `ack_offsets`/`last_ack_time` via `fetch_max` (reordered or + duplicate ACKs can never regress the recorded offset). +- New e2e `wait_returns_acked_replica_count`: WAIT with no replicas → 0 + fast; WAIT 1 after a write → 1 within the ACK cadence; WAIT 2 with one + replica → times out reporting 1 (RED before this change). + ### Fixed — multi-db replication: master streams `SELECT`, replicas serve it (HIGH-2) - **Writes outside db 0 landed in db 0 on the replica**: the replica-side diff --git a/src/command/metadata.rs b/src/command/metadata.rs index af961ed5..2338ab12 100644 --- a/src/command/metadata.rs +++ b/src/command/metadata.rs @@ -656,6 +656,20 @@ pub fn is_write(cmd: &[u8]) -> bool { lookup(cmd).is_some_and(|m| m.flags.contains(CommandFlags::WRITE)) } +/// Write commands that must be PERSISTED (AOF) and REPLICATED. +/// +/// Excludes SELECT: it is W-flagged so it routes through the write dispatch +/// paths, but it only mutates CONNECTION state. Persisting the literal client +/// SELECT poisons the shared AOF/replication stream's db context for every +/// OTHER connection interleaved with it — `conn A (db0): SET x` after +/// `conn B: SELECT 2` replayed `x` into db 2 (task #35). Both planes prepend +/// their own `SELECT ` records from per-stream db tracking instead +/// (`ReplicationState::stream_db`, the AOF writer's `last_db`). +#[inline] +pub fn is_persisted_write(cmd: &[u8]) -> bool { + is_write(cmd) && !cmd.eq_ignore_ascii_case(b"SELECT") +} + /// Check if a command is read-only via the metadata registry. #[inline] pub fn is_read(cmd: &[u8]) -> bool { diff --git a/src/command/mod.rs b/src/command/mod.rs index dfca68fd..fce3c6f3 100644 --- a/src/command/mod.rs +++ b/src/command/mod.rs @@ -1470,10 +1470,11 @@ fn dispatch_read_inner(db: &Database, cmd: &[u8], args: &[Frame], now_ms: u64) - } } (4, b'w') => { - // WAIT: no replication — always 0 (mirrors handler_single.rs:833). - // WAIT is in extract_primary_key's keyless table, so it routes - // locally: this arm fires on the local read path only, never the - // cross-shard fast path. + // WAIT fallback: the REAL implementation is the connection-layer + // `try_handle_wait` intercept (R1 — it awaits replica ACKs, which + // this synchronous dispatch cannot). This arm only answers when a + // path lacks that intercept (no repl_state configured): 0 is then + // truthful — no replicas can be attached. if cmd.eq_ignore_ascii_case(b"WAIT") { return resp(Frame::Integer(0)); } diff --git a/src/persistence/aof/mod.rs b/src/persistence/aof/mod.rs index bf2195a5..65d09365 100644 --- a/src/persistence/aof/mod.rs +++ b/src/persistence/aof/mod.rs @@ -159,7 +159,16 @@ pub enum AofMessage { /// `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 }, + /// + /// `db` is the database index the command EXECUTED in (task #35 — AOF + /// db-aware writer). The writer tracks a running `last_db` and, before + /// writing a non-empty record whose `db` differs, first writes a + /// `SELECT ` record (see [`serialize_select_record`]) so recovery + /// replays every record into the db it was actually written in — client + /// `SELECT` commands are connection-state only and are never persisted + /// (see `metadata::is_persisted_write`). A zero-length payload (the + /// `fsync_barrier` H1-BARRIER) never triggers or observes a db switch. + Append { lsn: u64, db: usize, bytes: Bytes }, /// Append + fsync + ack rendezvous (RFC § 4 — Fix 2 for the H1 /// data-loss vector exposed by `appendfsync=always`). /// @@ -180,6 +189,7 @@ pub enum AofMessage { /// `--unsafe-multishard-aof` gate. AppendSync { lsn: u64, + db: usize, bytes: Bytes, ack: crate::runtime::channel::OneshotSender, }, @@ -503,6 +513,75 @@ pub fn serialize_command(frame: &Frame) -> Bytes { buf.freeze() } +/// Serialized `SELECT ` RESP record for AOF db-context injection +/// (task #35). Same wire form the replay engines already execute +/// (`replay_incr_resp` / `replay_incr_framed` / `DispatchReplayEngine`). +pub fn serialize_select_record(db: usize) -> Bytes { + let mut n = itoa::Buffer::new(); + let d = n.format(db); + let mut buf = Vec::with_capacity(32); + buf.extend_from_slice(b"*2\r\n$6\r\nSELECT\r\n$"); + let mut l = itoa::Buffer::new(); + buf.extend_from_slice(l.format(d.len()).as_bytes()); + buf.extend_from_slice(b"\r\n"); + buf.extend_from_slice(d.as_bytes()); + buf.extend_from_slice(b"\r\n"); + Bytes::from(buf) +} + +/// Returns `Some(serialize_select_record(db))` and advances `*last_db` when a +/// non-barrier record's execution db differs from the writer's running +/// context (task #35 db-aware writer). Returns `None` (no state change) for +/// a zero-length barrier payload (`payload_is_empty` — `pool::fsync_barrier` +/// writes no record and must never trigger nor observe a db switch) or when +/// `db == *last_db`. Shared by every AOF write site — the batched writer +/// loops (via [`inject_select_records`]) and the rewrite-fold per-message +/// inline drains (`rewrite.rs`) alike — so the db-switch rule has exactly one +/// definition. +pub(crate) fn select_prefix_if_needed( + db: usize, + payload_is_empty: bool, + last_db: &mut usize, +) -> Option { + if payload_is_empty || db == *last_db { + return None; + } + *last_db = db; + Some(serialize_select_record(db)) +} + +/// Insert a synthetic `SELECT ` [`AofMessage::Append`] (lsn=0) before +/// every message in `data` whose db differs from the writer's running +/// `last_db` context (task #35). Barriers (zero-length payload) never +/// trigger nor observe a switch. `last_db` is updated in place so +/// consecutive batches stay consistent across writer-loop iterations. +/// +/// Used by all four batch-write writer loci (TopLevel/PerShard x +/// monoio/tokio) immediately after `collect_group_commit_batch` returns — +/// the injected message then rides the SAME write path as any other record +/// (a framed `[u64 lsn=0][u32 len]` header for the PerShard loops, raw bytes +/// for TopLevel), so no per-loop special-casing of the SELECT bytes is +/// needed. +pub(crate) fn inject_select_records(data: Vec, last_db: &mut usize) -> Vec { + let mut out = Vec::with_capacity(data.len()); + for msg in data { + let (db, is_empty) = match &msg { + AofMessage::Append { db, bytes, .. } => (*db, bytes.is_empty()), + AofMessage::AppendSync { db, bytes, .. } => (*db, bytes.is_empty()), + _ => (0, true), + }; + if let Some(select_bytes) = select_prefix_if_needed(db, is_empty, last_db) { + out.push(AofMessage::Append { + lsn: 0, + db, + bytes: select_bytes, + }); + } + out.push(msg); + } + out +} + /// Whether legacy best-effort skip-and-resync on mid-stream AOF corruption is /// enabled (prod-hardening #12). Default OFF: a hard parse error stops replay /// (keeping the valid prefix) rather than skipping forward to the next `*` and diff --git a/src/persistence/aof/pool.rs b/src/persistence/aof/pool.rs index 18680b18..b0d5fde0 100644 --- a/src/persistence/aof/pool.rs +++ b/src/persistence/aof/pool.rs @@ -251,11 +251,12 @@ impl AofWriterPool { &self, shard_id: usize, lsn: u64, + db: usize, bytes: Bytes, ) -> Result<(), AofAck> { match self.fsync_policy { FsyncPolicy::Always => { - let rx = self.try_send_append_sync(shard_id, lsn, bytes); + let rx = self.try_send_append_sync(shard_id, lsn, db, 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 @@ -277,7 +278,8 @@ impl AofWriterPool { // enqueue (NOT fsync) under `fsync_timeout`, and a failed // enqueue surfaces as Err so the client never gets a false // success for a write the durability machinery never saw. - self.send_append_backpressure(shard_id, lsn, bytes).await + self.send_append_backpressure(shard_id, lsn, db, bytes) + .await } } } @@ -302,9 +304,11 @@ impl AofWriterPool { &self, shard_id: usize, lsn: u64, + db: usize, bytes: Bytes, ) -> Result { - self.send_append_backpressure(shard_id, lsn, bytes).await?; + self.send_append_backpressure(shard_id, lsn, db, bytes) + .await?; Ok(matches!(self.fsync_policy, FsyncPolicy::Always)) } @@ -348,7 +352,9 @@ impl AofWriterPool { FsyncPolicy::Always => { // Enqueue a zero-length AppendSync. The writer will fsync all // preceding Append messages (ordered channel) then ack Synced. - let rx = self.try_send_append_sync(shard_id, 0, Bytes::new()); + // db=0: a zero-length barrier writes no record, so it can + // never trigger (or need) a SELECT injection. + let rx = self.try_send_append_sync(shard_id, 0, 0, Bytes::new()); // F2 bounded await — same semantics as try_send_append_durable. match Self::await_ack(rx, self.fsync_timeout).await { AckOutcome::Ack(AofAck::Synced) => Ok(()), @@ -441,10 +447,10 @@ impl AofWriterPool { /// [`Self::send_append_bounded_blocking`] (sync contexts) or /// [`Self::try_send_append_durable`] (async contexts). #[inline] - pub fn try_send_append(&self, shard_id: usize, lsn: u64, bytes: Bytes) -> bool { + pub fn try_send_append(&self, shard_id: usize, lsn: u64, db: usize, bytes: Bytes) -> bool { match self .sender(shard_id) - .try_send(AofMessage::Append { lsn, bytes }) + .try_send(AofMessage::Append { lsn, db, bytes }) { Ok(()) => true, Err(e) => { @@ -493,10 +499,11 @@ impl AofWriterPool { &self, shard_id: usize, lsn: u64, + db: usize, bytes: Bytes, budget: &mut Duration, ) -> bool { - let mut msg = AofMessage::Append { lsn, bytes }; + let mut msg = AofMessage::Append { lsn, db, bytes }; match self.sender(shard_id).try_send(msg) { Ok(()) => return true, Err(flume::TrySendError::Disconnected(_)) => { @@ -561,9 +568,10 @@ impl AofWriterPool { &self, shard_id: usize, lsn: u64, + db: usize, bytes: Bytes, ) -> Result<(), AofAck> { - let mut msg = AofMessage::Append { lsn, bytes }; + let mut msg = AofMessage::Append { lsn, db, bytes }; match self.sender(shard_id).try_send(msg) { Ok(()) => return Ok(()), Err(flume::TrySendError::Disconnected(_)) => return Err(AofAck::WriteFailed), @@ -640,11 +648,13 @@ impl AofWriterPool { &self, shard_id: usize, lsn: u64, + db: usize, 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, + db, bytes, ack: ack_tx, }) { @@ -704,7 +714,7 @@ impl AofWriterPool { /// 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) { + pub fn try_send_append_ordered(&self, shard_id: usize, lsn: u64, db: usize, bytes: Bytes) { debug_assert_eq!( lsn & ORDERED_LSN_FLAG, 0, @@ -714,6 +724,7 @@ impl AofWriterPool { let tagged_lsn = (lsn & !ORDERED_LSN_FLAG) | ORDERED_LSN_FLAG; let _ = self.sender(shard_id).try_send(AofMessage::Append { lsn: tagged_lsn, + db, bytes, }); } @@ -1032,6 +1043,7 @@ mod pool_tests { // on the reply. In this unit test there's no shard event loop consuming the // SPSC ring, so the fold will error out. The test verifies the abort path, // which is triggered by the fold guard's error handling. + let mut last_db: usize = 0; let _ = do_rewrite_per_shard( 0, &shard_dbs, @@ -1040,6 +1052,7 @@ mod pool_tests { &coord, &fold_producer, &fold_notifier, + &mut last_db, ); // Abort kept the old generation committed and pruned the new-gen incr. @@ -1152,9 +1165,9 @@ mod pool_tests { 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")); + pool.try_send_append(0, 0, 0, Bytes::from_static(b"a")); + pool.try_send_append(7, 0, 0, Bytes::from_static(b"b")); + pool.try_send_append(42, 0, 0, Bytes::from_static(b"c")); let mut seen = 0; while rx.try_recv().is_ok() { @@ -1172,10 +1185,10 @@ mod pool_tests { 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")); + pool.try_send_append(0, 100, 0, Bytes::from_static(b"shard0")); + pool.try_send_append(1, 200, 0, Bytes::from_static(b"shard1a")); + pool.try_send_append(1, 300, 0, Bytes::from_static(b"shard1b")); + pool.try_send_append(2, 400, 0, Bytes::from_static(b"shard2")); let count = |rx: &channel::MpscReceiver| -> usize { let mut n = 0; @@ -1222,13 +1235,13 @@ mod pool_tests { 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")); + pool.try_send_append(0, 42, 0, Bytes::from_static(b"set foo 1")); + pool.try_send_append(1, 43, 0, Bytes::from_static(b"set bar 2")); + pool.try_send_append(0, 44, 0, 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 }) => { + Ok(AofMessage::Append { lsn, bytes, .. }) => { assert_eq!(lsn, 42, "shard 0 first entry lsn"); assert_eq!(bytes.as_ref(), b"set foo 1"); } @@ -1238,7 +1251,7 @@ mod pool_tests { ), } match rx0.try_recv() { - Ok(AofMessage::Append { lsn, bytes }) => { + Ok(AofMessage::Append { lsn, bytes, .. }) => { assert_eq!(lsn, 44, "shard 0 second entry lsn"); assert_eq!(bytes.as_ref(), b"del foo"); } @@ -1249,7 +1262,7 @@ mod pool_tests { } // Shard 1 should see (43, "set bar 2") only. match rx1.try_recv() { - Ok(AofMessage::Append { lsn, bytes }) => { + Ok(AofMessage::Append { lsn, bytes, .. }) => { assert_eq!(lsn, 43, "shard 1 entry lsn"); assert_eq!(bytes.as_ref(), b"set bar 2"); } @@ -1268,12 +1281,14 @@ mod pool_tests { 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 recv = pool.try_send_append_sync(0, 99, 0, 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 }) => { + 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 @@ -1297,7 +1312,7 @@ mod pool_tests { 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")); + let recv = pool.try_send_append_sync(0, 7, 0, Bytes::from_static(b"x")); // Drain the message but DROP the ack sender without sending. match rx0.try_recv() { @@ -1317,7 +1332,7 @@ mod pool_tests { 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 recv = pool.try_send_append_sync(0, 1, 0, Bytes::from_static(b"x")); let ack = match rx0.try_recv() { Ok(AofMessage::AppendSync { ack, .. }) => ack, other => panic!("expected AppendSync, got {:?}", other.is_ok()), @@ -1334,7 +1349,7 @@ mod pool_tests { 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 recv = pool.try_send_append_sync(0, 1, 0, Bytes::from_static(b"x")); let ack = match rx0.try_recv() { Ok(AofMessage::AppendSync { ack, .. }) => ack, other => panic!("expected AppendSync, got {:?}", other.is_ok()), @@ -1357,14 +1372,16 @@ mod pool_tests { 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 recv = pool.try_send_append_sync(0, 99, 0, 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, usize::MAX).unwrap(); + let mut last_db: usize = 0; + let mut outcome = + drain_pending_appends_framed(&rx0, &mut file, usize::MAX, &mut last_db).unwrap(); // CONTRACT: drained + parked, NOT yet acked. assert_eq!(outcome.drained, 1, "the AppendSync was drained"); @@ -1398,15 +1415,17 @@ mod pool_tests { let pool = AofWriterPool::per_shard(vec![tx0, tx1]); // A real append followed by a barrier (empty payload). - pool.try_send_append(0, 5, Bytes::from_static(b"*1\r\n$4\r\nPING\r\n")); - let barrier_recv = pool.try_send_append_sync(0, 0, Bytes::new()); + pool.try_send_append(0, 5, 0, Bytes::from_static(b"*1\r\n$4\r\nPING\r\n")); + let barrier_recv = pool.try_send_append_sync(0, 0, 0, Bytes::new()); let mut file = std::fs::OpenOptions::new() .create(true) .append(true) .open(&incr) .unwrap(); - let mut outcome = drain_pending_appends_framed(&rx0, &mut file, usize::MAX).unwrap(); + let mut last_db: usize = 0; + let mut outcome = + drain_pending_appends_framed(&rx0, &mut file, usize::MAX, &mut last_db).unwrap(); assert_eq!(outcome.drained, 2, "both messages count toward drained"); assert_eq!(outcome.pending_acks.len(), 1, "barrier ack parked"); @@ -1429,6 +1448,77 @@ mod pool_tests { ); } + /// task #35 (C2 AOF db-aware writer): a db switch mid-stream must inject a + /// `SELECT ` record BEFORE the differing record — the writer is the + /// single consumer of a shard's AOF stream, so it (not the client's + /// SELECT, which is never persisted since C1) is what re-establishes + /// `aof_selected_db` for replay. Drives the REAL production drain + /// (`drain_pending_appends_framed`, one of the exact write loci this fix + /// touches) against a real temp file — not a mock — asserting the framed + /// on-disk bytes contain exactly `SELECT 2` before the db=2 record and + /// `SELECT 0` before the following db=0 record, with no redundant SELECT + /// around the two same-db (0) neighbors. + #[test] + fn drain_framed_injects_select_on_db_change() { + let tmp = tempfile::tempdir().unwrap(); + let incr = tmp.path().join("incr.aof"); + let (tx0, rx0) = channel::mpsc_bounded::(8); + let (tx1, _rx1) = channel::mpsc_bounded::(8); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + // Append(db=0), Append(db=2), Append(db=0) — exactly the spec's + // red/green scenario. + pool.try_send_append(0, 1, 0, Bytes::from_static(b"AAAA")); + pool.try_send_append(0, 2, 2, Bytes::from_static(b"BBBB")); + pool.try_send_append(0, 3, 0, Bytes::from_static(b"CCCC")); + + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr) + .unwrap(); + let mut last_db: usize = 0; + let outcome = + drain_pending_appends_framed(&rx0, &mut file, usize::MAX, &mut last_db).unwrap(); + assert_eq!(outcome.drained, 3, "all three real appends were drained"); + assert_eq!( + last_db, 0, + "context tracker ends at the 3rd record's db (0)" + ); + + file.sync_data().unwrap(); + let raw = std::fs::read(&incr).unwrap(); + + // Parse the PerShard incr framing `[u64 lsn LE][u32 len LE][len bytes]` + // — the same shape `replay_incr_framed` reads on restart. Local parser + // (not the sibling tokio-gated `parse_framed`) so this test runs under + // the default runtime-monoio build, not just runtime-tokio. + let mut frames: Vec> = Vec::new(); + let mut i = 0usize; + while i + 12 <= raw.len() { + let len = u32::from_le_bytes(raw[i + 8..i + 12].try_into().unwrap()) as usize; + assert!(i + 12 + len <= raw.len(), "no truncated tail expected here"); + frames.push(raw[i + 12..i + 12 + len].to_vec()); + i += 12 + len; + } + + assert_eq!( + frames, + vec![ + b"AAAA".to_vec(), + serialize_select_record(2).to_vec(), + b"BBBB".to_vec(), + serialize_select_record(0).to_vec(), + b"CCCC".to_vec(), + ], + "expected AAAA, SELECT 2, BBBB, SELECT 0, CCCC — got {:?}", + frames + .iter() + .map(|f| String::from_utf8_lossy(f).into_owned()) + .collect::>() + ); + } + /// 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 @@ -1442,14 +1532,15 @@ mod pool_tests { 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 recv = pool.try_send_append_sync(0, 7, 0, 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(); + let mut last_db: usize = 0; + let mut outcome = drain_pending_appends(&rx0, &mut file, &mut last_db).unwrap(); assert_eq!( outcome.pending_acks.len(), 1, @@ -1486,7 +1577,7 @@ mod pool_tests { let start = Instant::now(); let res = pool - .try_send_append_durable(0, 1, Bytes::from_static(b"x")) + .try_send_append_durable(0, 1, 0, Bytes::from_static(b"x")) .await; let elapsed = start.elapsed(); @@ -1524,7 +1615,7 @@ mod pool_tests { }); let res = pool - .try_send_append_durable(0, 1, Bytes::from_static(b"x")) + .try_send_append_durable(0, 1, 0, Bytes::from_static(b"x")) .await; assert_eq!(res, Ok(()), "ack within the bound must succeed"); drop(_rx1); @@ -1584,16 +1675,19 @@ mod pool_tests { // 1: clean. 2: torn (header only). 3: must be suppressed by the latch. tx.try_send(AofMessage::Append { lsn: 1, + db: 0, bytes: Bytes::from_static(b"AAAA"), }) .unwrap(); tx.try_send(AofMessage::Append { lsn: 2, + db: 0, bytes: Bytes::from_static(b"BBBB"), }) .unwrap(); tx.try_send(AofMessage::Append { lsn: 3, + db: 0, bytes: Bytes::from_static(b"CCCC"), }) .unwrap(); @@ -1603,6 +1697,7 @@ mod pool_tests { let (ack_tx, ack_rx) = crate::runtime::channel::oneshot::(); tx.try_send(AofMessage::AppendSync { lsn: 4, + db: 0, bytes: Bytes::from_static(b"DDDD"), ack: ack_tx, }) @@ -1685,7 +1780,7 @@ mod pool_tests { // 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")) + .try_send_append_durable(0, 1, 0, Bytes::from_static(b"SET k v")) .await; mock_writer.await.expect("mock writer completed"); @@ -1765,7 +1860,7 @@ mod pool_tests { 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")); + let recv = pool.try_send_append_sync(0, 1, 0, Bytes::from_static(b"SET k v")); // The channel was full — ChannelFull is returned immediately without // a writer round-trip. @@ -1823,6 +1918,7 @@ mod pool_tests { let result = futures::executor::block_on(pool.try_send_append_durable( 0, 55, + 0, Bytes::from_static(b"SWAPDB 0 1"), )); @@ -1855,6 +1951,7 @@ mod pool_tests { let result = futures::executor::block_on(pool.try_send_append_durable( 0, 56, + 0, Bytes::from_static(b"SWAPDB 0 1"), )); @@ -1887,6 +1984,7 @@ mod pool_tests { let result = futures::executor::block_on(pool.send_append_group( 0, 77, + 0, Bytes::from_static(b"MSET k v"), )); @@ -1919,6 +2017,7 @@ mod pool_tests { let result = futures::executor::block_on(pool.send_append_group( 0, 78, + 0, Bytes::from_static(b"MSET k v"), )); assert_eq!( @@ -1941,6 +2040,7 @@ mod pool_tests { let result = futures::executor::block_on(pool.send_append_group( 0, 79, + 0, Bytes::from_static(b"MSET k v"), )); assert!( @@ -2069,6 +2169,7 @@ mod pool_tests { // Fill the only slot; keep _rx0 alive so the channel is Full, not Disconnected. tx0.try_send(AofMessage::Append { lsn: 0, + db: 0, bytes: Bytes::from_static(b"x"), }) .unwrap(); @@ -2080,7 +2181,7 @@ mod pool_tests { let before = AOF_BACKPRESSURE_DROPPED.load(Ordering::Relaxed); let res = pool - .try_send_append_durable(0, 1, Bytes::from_static(b"SET k v")) + .try_send_append_durable(0, 1, 0, Bytes::from_static(b"SET k v")) .await; assert_eq!( res, @@ -2102,6 +2203,7 @@ mod pool_tests { let (tx0, rx0) = channel::mpsc_bounded::(1); tx0.try_send(AofMessage::Append { lsn: 0, + db: 0, bytes: Bytes::from_static(b"x"), }) .unwrap(); @@ -2119,7 +2221,7 @@ mod pool_tests { }); let res = pool - .try_send_append_durable(0, 1, Bytes::from_static(b"SET k v")) + .try_send_append_durable(0, 1, 0, Bytes::from_static(b"SET k v")) .await; assert_eq!( res, @@ -2142,6 +2244,7 @@ mod pool_tests { let (tx0, _rx0) = channel::mpsc_bounded::(1); tx0.try_send(AofMessage::Append { lsn: 0, + db: 0, bytes: Bytes::from_static(b"x"), }) .unwrap(); @@ -2155,7 +2258,7 @@ mod pool_tests { let start = std::time::Instant::now(); let mut budget = Duration::from_millis(30); let ok = - pool.send_append_bounded_blocking(0, 7, Bytes::from_static(b"SET k v"), &mut budget); + pool.send_append_bounded_blocking(0, 7, 0, Bytes::from_static(b"SET k v"), &mut budget); assert!(!ok, "still-full channel after the bound must report loss"); assert!( start.elapsed() >= Duration::from_millis(30), @@ -2181,6 +2284,7 @@ mod pool_tests { let (tx0, _rx0) = channel::mpsc_bounded::(1); tx0.try_send(AofMessage::Append { lsn: 0, + db: 0, bytes: Bytes::from_static(b"x"), }) .unwrap(); @@ -2194,7 +2298,7 @@ mod pool_tests { let mut budget = Duration::ZERO; // batch already spent its bound let start = std::time::Instant::now(); let ok = - pool.send_append_bounded_blocking(0, 8, Bytes::from_static(b"SET k v"), &mut budget); + pool.send_append_bounded_blocking(0, 8, 0, Bytes::from_static(b"SET k v"), &mut budget); assert!(!ok, "exhausted budget must report loss"); assert!( start.elapsed() < Duration::from_millis(10), @@ -2220,6 +2324,7 @@ mod pool_tests { assert!(pool.send_append_bounded_blocking( 0, 9, + 0, Bytes::from_static(b"SET a 1"), &mut budget, )); diff --git a/src/persistence/aof/rewrite.rs b/src/persistence/aof/rewrite.rs index 2c3e3f5c..b772610c 100644 --- a/src/persistence/aof/rewrite.rs +++ b/src/persistence/aof/rewrite.rs @@ -349,11 +349,21 @@ pub(crate) fn sync_and_fulfill_drain( /// bound (C4-DRAIN-BOUND): draining more would consume post-snapshot appends /// that belong in the NEW incr, and those bytes would then be lost when /// `manifest.advance()` deletes the old incr at the end of the fold. +/// +/// `db_ctx` is the writer's running db-attribution context for THIS incr file +/// (task #35): the caller passes `&mut` the writer's `last_db` when draining +/// into the file the live writer was already appending to (so the drain +/// continues that stream's context), or a fresh `0` when draining into a +/// brand-new incr (replay always starts a segment at db 0). A non-empty +/// record whose db differs from `*db_ctx` gets a raw `SELECT ` record +/// written first (same rule as the live writer loops); `*db_ctx` is updated +/// in place so the caller can read back the final context after the drain. #[cfg(feature = "runtime-monoio")] pub(crate) fn drain_pending_appends_bounded( rx: &channel::MpscReceiver, file: &mut std::fs::File, max_drain: usize, + db_ctx: &mut usize, ) -> Result { use std::io::Write; let mut outcome = DrainOutcome::default(); @@ -363,7 +373,14 @@ pub(crate) fn drain_pending_appends_bounded( AofMessage::Append { bytes: data, lsn: _, + db, } => { + if let Some(sel) = select_prefix_if_needed(db, data.is_empty(), db_ctx) { + file.write_all(&sel).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + } file.write_all(&data).map_err(|e| AofError::Io { path: PathBuf::from(""), source: e, @@ -373,8 +390,15 @@ pub(crate) fn drain_pending_appends_bounded( AofMessage::AppendSync { bytes: data, lsn: _, + db, ack, } => { + if let Some(sel) = select_prefix_if_needed(db, data.is_empty(), db_ctx) { + file.write_all(&sel).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + } file.write_all(&data).map_err(|e| AofError::Io { path: PathBuf::from(""), source: e, @@ -402,8 +426,9 @@ pub(crate) fn drain_pending_appends_bounded( pub(crate) fn drain_pending_appends( rx: &channel::MpscReceiver, file: &mut std::fs::File, + db_ctx: &mut usize, ) -> Result { - drain_pending_appends_bounded(rx, file, usize::MAX) + drain_pending_appends_bounded(rx, file, usize::MAX, db_ctx) } /// [F6] Drain at most `max_drain` pending [`AofMessage::Append`] / @@ -422,11 +447,17 @@ pub(crate) fn drain_pending_appends( /// 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. +/// +/// `db_ctx` carries the writer's running db-attribution context exactly like +/// [`drain_pending_appends_bounded`] (task #35) — see its doc for the +/// caller contract. The synthetic `SELECT ` record is written framed +/// (`lsn=0`), matching the live per-shard writer loops. #[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] pub(crate) fn drain_pending_appends_framed( rx: &channel::MpscReceiver, file: &mut std::fs::File, max_drain: usize, + db_ctx: &mut usize, ) -> Result { use std::io::Write; let mut outcome = DrainOutcome::default(); @@ -440,7 +471,17 @@ pub(crate) fn drain_pending_appends_framed( while outcome.drained < max_drain { match rx.try_recv() { Ok(msg) => match msg { - AofMessage::Append { lsn, bytes: data } => { + AofMessage::Append { + lsn, + db, + bytes: data, + } => { + if let Some(sel) = select_prefix_if_needed(db, data.is_empty(), db_ctx) { + write_framed(file, 0, &sel).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + } write_framed(file, lsn, &data).map_err(|e| AofError::Io { path: PathBuf::from(""), source: e, @@ -449,6 +490,7 @@ pub(crate) fn drain_pending_appends_framed( } AofMessage::AppendSync { lsn, + db, bytes: data, ack, } => { @@ -458,6 +500,12 @@ pub(crate) fn drain_pending_appends_framed( // still counts toward `drained` (sender.len() counted it) // and its ack still parks for the boundary fsync. if !data.is_empty() { + if let Some(sel) = select_prefix_if_needed(db, false, db_ctx) { + write_framed(file, 0, &sel).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + } write_framed(file, lsn, &data).map_err(|e| AofError::Io { path: PathBuf::from(""), source: e, @@ -560,6 +608,7 @@ pub(crate) fn do_rewrite_per_shard( coord: &PerShardRewriteCoord, fold_producer: &parking_lot::Mutex>, fold_notifier: &std::sync::Arc, + last_db: &mut usize, ) -> Result<(), MoonError> { use ringbuf::traits::Producer; // Panic/early-error safety: guarantees `shard_done` runs on EVERY exit @@ -581,7 +630,10 @@ pub(crate) fn do_rewrite_per_shard( // Without this bound, under sustained high write load the drain loops forever // because the INCR producer keeps the channel perpetually non-empty. let pre_drain_bound = rx.len(); - let mut pre_drain = drain_pending_appends_framed(rx, file, pre_drain_bound)?; + // task #35: pre/mid drain write into the OLD incr `file` — the SAME + // stream the live writer was appending to before the fold — so they + // continue the writer's running db context (`last_db`), not a fresh 0. + let mut pre_drain = drain_pending_appends_framed(rx, file, pre_drain_bound, last_db)?; sync_and_fulfill_drain(&mut pre_drain, file, PathBuf::from(""))?; info!( "F6 shard {} phase1 done: drained {} appends ({:.1}ms)", @@ -692,7 +744,7 @@ pub(crate) fn do_rewrite_per_shard( // of pre-snapshot appends the shard reported before building its snapshot. This // prevents an infinite drain loop under sustained high write load where new // (post-snapshot) appends arrive faster than we can drain them. - let mut mid_drain = drain_pending_appends_framed(rx, file, pending_aof_count)?; + let mut mid_drain = drain_pending_appends_framed(rx, file, pending_aof_count, last_db)?; sync_and_fulfill_drain(&mut mid_drain, file, PathBuf::from(""))?; info!( "F6 shard {} phase3 done: drained {} mid-appends ({:.1}ms total)", @@ -723,6 +775,11 @@ pub(crate) fn do_rewrite_per_shard( path: new_incr, source: e, })?; + // task #35: save the OLD incr's final db context (in case phase 8 rolls + // back to it on abort) then reset for the fresh NEW incr — replay always + // starts a segment at db 0, so the writer's running context must match. + let old_incr_last_db = *last_db; + *last_db = 0; info!( "F6 per-shard rewrite: shard {} folded (drained {}+{} appends), new seq {}", @@ -760,6 +817,9 @@ pub(crate) fn do_rewrite_per_shard( path: committed_incr, source: e, })?; + // task #35: rolled back onto the OLD (still-committed) incr — restore + // its db context instead of the fresh-incr 0 set above. + *last_db = old_incr_last_db; warn!( "F6 per-shard rewrite ABORTED: shard {} rolled its append file back to \ committed seq {} (no restart needed)", @@ -795,10 +855,12 @@ pub(crate) fn do_rewrite_single( manifest: &mut crate::persistence::aof_manifest::AofManifest, file: &mut std::fs::File, rx: &channel::MpscReceiver, + last_db: &mut usize, ) -> 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)?; + // resolve their parked AppendSync acks (issue #140). task #35: continues + // the writer's running db context — this IS the writer's live `file`. + let mut pre_drain = drain_pending_appends(rx, file, last_db)?; sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; // Phase 2: acquire write locks on every database in the shard. @@ -809,7 +871,7 @@ pub(crate) fn do_rewrite_single( // 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)?; + let mut mid_drain = drain_pending_appends(rx, file, last_db)?; sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; // Phase 4: snapshot under the write locks. No mutation is possible. @@ -850,6 +912,8 @@ pub(crate) fn do_rewrite_single( path: new_incr, source: e, })?; + // task #35: fresh incr — replay always starts a segment at db 0. + *last_db = 0; info!( "AOF rewrite complete (single): drained {}+{} pre-snapshot appends, seq={}", @@ -895,6 +959,7 @@ pub(crate) fn do_rewrite_sharded( Arc>>, Arc, )>, + last_db: &mut usize, ) -> Result<(), MoonError> { use ringbuf::traits::Producer; @@ -916,7 +981,9 @@ pub(crate) fn do_rewrite_sharded( }; // Phase 1: drain pre-rewrite queued appends into old incr (RAW RESP), fsync. - let mut pre_drain = drain_pending_appends(rx, file)?; + // task #35: continues the writer's running db context (this IS the + // writer's live `file`). + let mut pre_drain = drain_pending_appends(rx, file, last_db)?; sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; info!( "TopLevel fold phase1 done: drained {} appends ({:.1}ms)", @@ -993,7 +1060,7 @@ pub(crate) fn do_rewrite_sharded( // would be silently lost when `manifest.advance()` deletes the old incr file // at the end of phase 4. This mirrors the identical bound used by // `drain_pending_appends_framed` in `do_rewrite_per_shard`. - let mut mid_drain = drain_pending_appends_bounded(rx, file, pending_aof_count)?; + let mut mid_drain = drain_pending_appends_bounded(rx, file, pending_aof_count, last_db)?; sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; info!( "TopLevel fold phase3 done: drained {} mid-appends ({:.1}ms)", @@ -1019,6 +1086,8 @@ pub(crate) fn do_rewrite_sharded( path: new_incr, source: e, })?; + // task #35: fresh incr — replay always starts a segment at db 0. + *last_db = 0; info!( "TopLevel AOF rewrite complete: drained {}+{} appends, seq={}", @@ -1111,6 +1180,7 @@ pub(crate) fn rewrite_aof_sharded_sync( Arc>>, Arc, )>, + last_db: &mut usize, ) -> Result<(), MoonError> { use ringbuf::traits::Producer; use std::io::Write as _; @@ -1140,15 +1210,32 @@ pub(crate) fn rewrite_aof_sharded_sync( let mut pre_acks: Vec> = Vec::new(); while let Ok(msg) = rx.try_recv() { match msg { - AofMessage::Append { bytes: data, .. } => { + AofMessage::Append { + db, bytes: data, .. + } => { + if let Some(sel) = select_prefix_if_needed(db, data.is_empty(), last_db) { + old_file.write_all(&sel).map_err(|e| AofError::Io { + path: aof_path.to_path_buf(), + source: e, + })?; + } old_file.write_all(&data).map_err(|e| AofError::Io { path: aof_path.to_path_buf(), source: e, })?; } AofMessage::AppendSync { - bytes: data, ack, .. + db, + bytes: data, + ack, + .. } => { + if let Some(sel) = select_prefix_if_needed(db, data.is_empty(), last_db) { + old_file.write_all(&sel).map_err(|e| AofError::Io { + path: aof_path.to_path_buf(), + source: e, + })?; + } old_file.write_all(&data).map_err(|e| AofError::Io { path: aof_path.to_path_buf(), source: e, @@ -1250,7 +1337,15 @@ pub(crate) fn rewrite_aof_sharded_sync( while drained < pending_aof_count { match rx.try_recv() { Ok(msg) => match msg { - AofMessage::Append { bytes: data, .. } => { + AofMessage::Append { + db, bytes: data, .. + } => { + if let Some(sel) = select_prefix_if_needed(db, data.is_empty(), last_db) { + old_file.write_all(&sel).map_err(|e| AofError::Io { + path: aof_path.to_path_buf(), + source: e, + })?; + } old_file.write_all(&data).map_err(|e| AofError::Io { path: aof_path.to_path_buf(), source: e, @@ -1258,8 +1353,17 @@ pub(crate) fn rewrite_aof_sharded_sync( drained += 1; } AofMessage::AppendSync { - bytes: data, ack, .. + db, + bytes: data, + ack, + .. } => { + if let Some(sel) = select_prefix_if_needed(db, data.is_empty(), last_db) { + old_file.write_all(&sel).map_err(|e| AofError::Io { + path: aof_path.to_path_buf(), + source: e, + })?; + } old_file.write_all(&data).map_err(|e| AofError::Io { path: aof_path.to_path_buf(), source: e, @@ -1323,6 +1427,9 @@ pub(crate) fn rewrite_aof_sharded_sync( e ), })?; + // task #35: aof_path now points at a brand-new file (RDB base only) — + // replay always starts a segment at db 0. + *last_db = 0; info!( "rewrite_aof_sharded_sync (tokio) complete: {} bytes ({:.1}ms)", diff --git a/src/persistence/aof/writer_task.rs b/src/persistence/aof/writer_task.rs index 1af32663..ca36fdf9 100644 --- a/src/persistence/aof/writer_task.rs +++ b/src/persistence/aof/writer_task.rs @@ -181,6 +181,7 @@ mod poll_recv_tests { assert!( tx.try_send(AofMessage::Append { lsn: 7, + db: 0, bytes: bytes::Bytes::from_static(b"x"), }) .is_ok() @@ -202,6 +203,7 @@ mod poll_recv_tests { assert!( tx.try_send(AofMessage::Append { lsn: 1, + db: 0, bytes: bytes::Bytes::from_static(b"y"), }) .is_ok() @@ -396,6 +398,11 @@ pub async fn aof_writer_task( // state survives across iterations. #[cfg(feature = "runtime-tokio")] let mut idle_wait = IdleWait::new(); + // task #35: AOF db-aware writer — see the monoio TopLevel loop above for + // the full rationale. Resets to 0 on every successful rewrite (fresh + // file: replay starts a segment at db 0). + #[cfg(feature = "runtime-tokio")] + let mut last_db: usize = 0; // Monoio path: multi-part AOF (base RDB + incremental RESP) with sync I/O. // @@ -495,6 +502,12 @@ pub async fn aof_writer_task( // The bounded recv + end-of-loop proactive fsync below restore the // ~1s EverySec bound exactly like the PerShard writers. let mut idle_wait = IdleWait::new(); + // task #35: AOF db-aware writer. Tracks the db the last-written + // record executed in; a non-empty record whose db differs gets a + // `SELECT ` record injected first (see `inject_select_records`). + // Resets to 0 whenever a NEW incr/base file becomes the append + // target (fresh manifest segment — replay starts each incr at db 0). + let mut last_db: usize = 0; loop { // Group commit: wait (bounded) for one message, then @@ -536,6 +549,10 @@ pub async fn aof_writer_task( AOF_GROUP_COMMIT_MAX_BATCH, AOF_GROUP_COMMIT_MAX_BYTES, ); + // task #35: inject SELECT records ahead of any db switch + // before committing — rides the same write path as any other + // record (raw bytes on this TopLevel format). + batch.data = inject_select_records(std::mem::take(&mut batch.data), &mut last_db); // -- commit the data batch (one fsync under Always; deadline under everysec) -- if !batch.data.is_empty() { @@ -591,7 +608,7 @@ pub async fn aof_writer_task( error!("AOF pre-rewrite sync failed (seq {}): {}", manifest.seq, e); } } - match do_rewrite_single(&db, &mut manifest, &mut file, &rx) { + match do_rewrite_single(&db, &mut manifest, &mut file, &rx, &mut last_db) { Ok(()) => { write_error = false; // Reset on successful rewrite } @@ -617,6 +634,7 @@ pub async fn aof_writer_task( &mut file, &rx, fold_channels.as_ref(), + &mut last_db, ) { Ok(()) => { write_error = false; @@ -721,6 +739,11 @@ pub async fn aof_writer_task( AOF_GROUP_COMMIT_MAX_BATCH, AOF_GROUP_COMMIT_MAX_BYTES, ); + // task #35: inject SELECT records ahead of any db + // switch before committing (see the monoio TopLevel loop + // above). + batch.data = + inject_select_records(std::mem::take(&mut batch.data), &mut last_db); // -- write the data batch inline (async), then ONE fsync -- if !batch.data.is_empty() { @@ -803,7 +826,13 @@ pub async fn aof_writer_task( match rewrite_aof(db, &aof_path).await { // Rewrite replaced the file with a clean one — the // torn-write latch resets (mirrors monoio TopLevel). - Ok(()) => write_error = false, + // task #35: fresh file — replay starts at db 0. On + // Err the old file is untouched, so last_db stays + // whatever it already was. + Ok(()) => { + write_error = false; + last_db = 0; + } Err(e) => error!("AOF rewrite failed: {}", e), } crate::command::persistence::AOF_REWRITE_IN_PROGRESS @@ -851,6 +880,7 @@ pub async fn aof_writer_task( &rx, &mut sf, fold_channels.as_ref(), + &mut last_db, ) { // Fold rewrote aof_path clean — the latch resets. Ok(()) => write_error = false, @@ -1062,6 +1092,9 @@ pub async fn per_shard_aof_writer_task( // Idle-adaptive channel-poll wake cadence (RSS/CPU wave 5, item B) — // see `IdleWait` docs near the top of this file. let mut idle_wait = IdleWait::new(); + // task #35: AOF db-aware writer — see the monoio TopLevel loop's docs + // near the top of this file for the full rationale. + let mut last_db: usize = 0; // (No `interval` here: the EverySec flush deadline is enforced by the // timeout-bounded recv in the loop below, which wakes at least every // `idle_wait.current()` (50ms floor, escalates to 1s while idle) @@ -1130,6 +1163,12 @@ pub async fn per_shard_aof_writer_task( AOF_GROUP_COMMIT_MAX_BATCH, AOF_GROUP_COMMIT_MAX_BYTES, ); + // task #35: inject SELECT records on db-context + // changes before this batch's framed writes go out. + batch.data = inject_select_records( + std::mem::take(&mut batch.data), + &mut last_db, + ); if !batch.data.is_empty() { if write_error { @@ -1145,7 +1184,7 @@ pub async fn per_shard_aof_writer_task( let mut write_failed = false; for msg in &batch.data { let (lsn, data) = match msg { - AofMessage::Append { lsn, bytes } + AofMessage::Append { lsn, bytes, .. } | AofMessage::AppendSync { lsn, bytes, .. } => { (*lsn, bytes) } @@ -1284,7 +1323,7 @@ pub async fn per_shard_aof_writer_task( let mut sf = writer.into_inner().into_std().await; let res = do_rewrite_per_shard( shard_id, &shard_dbs, &mut sf, &rx, &coord, - &fold_producer, &fold_notifier, + &fold_producer, &fold_notifier, &mut last_db, ); // `sf` is left on the committed generation by // the fold's internal barrier: NEW incr on @@ -1443,6 +1482,9 @@ pub async fn per_shard_aof_writer_task( // Idle-adaptive channel-poll wake cadence (RSS/CPU wave 5, item B) — // see `IdleWait` docs near the top of this file. let mut idle_wait = IdleWait::new(); + // task #35: AOF db-aware writer — see the monoio TopLevel loop's docs + // near the top of this file for the full rationale. + let mut last_db: usize = 0; // 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 @@ -1506,6 +1548,9 @@ pub async fn per_shard_aof_writer_task( AOF_GROUP_COMMIT_MAX_BATCH, AOF_GROUP_COMMIT_MAX_BYTES, ); + // task #35: inject SELECT records on db-context changes + // before this batch's framed writes go into batch_buf. + batch.data = inject_select_records(std::mem::take(&mut batch.data), &mut last_db); if !batch.data.is_empty() { if write_error { @@ -1524,7 +1569,7 @@ pub async fn per_shard_aof_writer_task( batch_buf.clear(); for msg in &batch.data { let (lsn, data) = match msg { - AofMessage::Append { lsn, bytes } + AofMessage::Append { lsn, bytes, .. } | AofMessage::AppendSync { lsn, bytes, .. } => (*lsn, bytes), _ => continue, }; @@ -1622,6 +1667,7 @@ pub async fn per_shard_aof_writer_task( &coord, &fold_producer, &fold_notifier, + &mut last_db, ) { // The fold's ShardDoneGuard already marked the rewrite // failed and decremented on this error exit (committing @@ -1648,9 +1694,12 @@ pub async fn per_shard_aof_writer_task( // Combined, the two fsyncs bound the post-fold EverySec window // to ≤150ms — well within the test's 1500ms kill margin. if !write_error { - if let Ok(mut post_drain) = - drain_pending_appends_framed(&rx, &mut file, usize::MAX) - { + if let Ok(mut post_drain) = drain_pending_appends_framed( + &rx, + &mut file, + usize::MAX, + &mut last_db, + ) { if let Err(e) = sync_and_fulfill_drain( &mut post_drain, &mut file, diff --git a/src/persistence/aof_manifest/shard_replay.rs b/src/persistence/aof_manifest/shard_replay.rs index a6ce0652..4a7b012e 100644 --- a/src/persistence/aof_manifest/shard_replay.rs +++ b/src/persistence/aof_manifest/shard_replay.rs @@ -1206,7 +1206,7 @@ mod tests { let pool = AofWriterPool::per_shard(vec![tx0, tx1]); // Raw lsn = 42; high bit must end up set on the receive side. - pool.try_send_append_ordered(0, 42, bytes::Bytes::from_static(b"x")); + pool.try_send_append_ordered(0, 42, 0, bytes::Bytes::from_static(b"x")); let msg = rx0.try_recv().expect("ordered append delivered"); match msg { AofMessage::Append { lsn, .. } => { diff --git a/src/replication/master.rs b/src/replication/master.rs index fd84716f..760a6ae5 100644 --- a/src/replication/master.rs +++ b/src/replication/master.rs @@ -441,6 +441,10 @@ async fn register_replica_with_shards( let msg = crate::shard::dispatch::ShardMessage::RegisterReplica { replica_id, tx, + // Legacy multi-shard drain loops do not poll the kick flag + // (superseded by the R2 redesign); overflow still stops + // queueing via the fan-out's retain. + kicked: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), backlog_capacity, // Fire-and-forget: the multi-shard register paths are superseded // by the R2 PrepareReplicaSync redesign; the offset-reply catch-up @@ -535,6 +539,10 @@ async fn register_replica_with_shards( let msg = crate::shard::dispatch::ShardMessage::RegisterReplica { replica_id, tx, + // Legacy multi-shard drain loops do not poll the kick flag + // (superseded by the R2 redesign); overflow still stops + // queueing via the fan-out's retain. + kicked: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), backlog_capacity, // Fire-and-forget: the multi-shard register paths are superseded // by the R2 PrepareReplicaSync redesign; the offset-reply catch-up @@ -843,6 +851,11 @@ struct InlineReplicaRegistration { tx: crate::runtime::channel::MpscSender, rx: crate::runtime::channel::MpscReceiver, reg_rx: crate::runtime::channel::MpscReceiver, + /// Overflow disconnect signal shared with the shard fan-out — set when + /// this replica's channel filled and a record could not be queued + /// (task #35). The drain loop polls it and closes the socket so the + /// replica resyncs instead of silently diverging. + kicked: std::sync::Arc, } /// Push `RegisterReplica` onto shard 0's SPSC so the event loop captures the @@ -860,8 +873,13 @@ fn push_register_replica_inline( static NEXT_REPLICA_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); let replica_id = NEXT_REPLICA_ID.fetch_add(1, Ordering::Relaxed); - let (tx, rx) = crate::runtime::channel::mpsc_bounded::(1024); + // 16384 records (task #35): 1024 overflowed within one pipelined burst on + // the same host — every overflow now KICKS the replica into a resync, so + // headroom directly reduces resync churn. Records are Bytes handles; + // 16k × ~50 B typical ≈ under 1 MB queued worst-case per replica. + let (tx, rx) = crate::runtime::channel::mpsc_bounded::(16384); let (reg_tx, reg_rx) = crate::runtime::channel::mpsc_bounded::(1); + let kicked = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); // `--repl-backlog-size`, carried in RegisterReplica for the lazy fallback-init. let backlog_capacity = repl_state @@ -889,6 +907,7 @@ fn push_register_replica_inline( crate::shard::self_msg::push(crate::shard::dispatch::ShardMessage::RegisterReplica { replica_id, tx: tx.clone(), + kicked: kicked.clone(), backlog_capacity, registered: Some(reg_tx), push_offset: Some(push_offset), @@ -898,6 +917,7 @@ fn push_register_replica_inline( tx, rx, reg_rx, + kicked, }) } @@ -919,6 +939,7 @@ async fn drain_replica_inline_single_shard( tx, rx, reg_rx: _, + kicked, } = reg; // Bookkeeping for WAIT/INFO. @@ -938,17 +959,57 @@ async fn drain_replica_inline_single_shard( rs.replicas.push(replica_info); } - // Drain the channel and write to the stream until the replica disconnects. - let stream = std::cell::RefCell::new(stream); - #[allow(clippy::await_holding_refcell_ref)] - while let Ok(data) = rx.recv_async().await { - let buf = data.to_vec(); - let (wr, _) = stream.borrow_mut().write_all(buf).await; - if wr.is_err() { - info!("Replica {} disconnected", replica_id); - break; + // R1 (task #19): the hijacked PSYNC socket is full-duplex — the replica + // sends `REPLCONF ACK ` back on it (1s cadence). Split the stream + // so a local reader task records ACKs into this replica's + // `ack_offsets`/`last_ack_time` (the data WAIT and INFO lag read) while + // the write loop below streams live fan-out bytes. Same-thread + // `monoio::spawn` — the task is !Send, which is fine here. + use monoio::io::Splitable as _; + let (rd, mut wr_half) = stream.into_split(); + let ack_reader = monoio::spawn({ + let repl_state = repl_state.clone(); + async move { ack_read_loop(rd, replica_id, repl_state).await } + }); + + // Drain the channel and write to the stream until the replica + // disconnects — or until the shard fan-out KICKS this replica (task #35: + // its channel overflowed, so at least one record is already missing from + // the stream; continuing would deliver a silently-corrupt sequence). The + // kick cannot arrive as an in-band message (the trigger IS a full + // channel), so the recv races a coarse poll timer. `ReplicaInfo.shard_txs` + // and this task both hold sender clones, which is why channel closure + // can't signal this either. + loop { + monoio::select! { + recv = rx.recv_async() => { + let Ok(data) = recv else { break }; + let buf = data.to_vec(); + let (wr, _) = wr_half.write_all(buf).await; + if wr.is_err() { + info!("Replica {} disconnected", replica_id); + break; + } + } + _ = monoio::time::sleep(std::time::Duration::from_millis(250)) => { + if kicked.load(std::sync::atomic::Ordering::Acquire) { + tracing::warn!( + replica_id, + "closing kicked replica connection (fan-out overflow) — \ + replica will reconnect and resync" + ); + break; + } + } } } + // A kicked replica may still have queued records; they are stale (the + // stream already has a gap) — drop them with the channel. + // Dropping the write half closes our outbound side; the reader task ends + // on EOF/error when the peer closes (its socket dies with the write half + // on a disconnect-driven exit, so it does not linger). + drop(wr_half); + drop(ack_reader); // Remove from ReplicationState; the event loop will drop its replica_txs // entry on the next failed send via its own UnregisterReplica path. if let Ok(mut rs) = repl_state.write() { @@ -961,6 +1022,107 @@ async fn drain_replica_inline_single_shard( Ok(()) } +/// Read `REPLCONF ACK ` frames off the replica's half of the hijacked +/// PSYNC socket and record them (R1, task #19). Runs as a same-thread task +/// beside the write-drain loop in `drain_replica_inline_single_shard`; exits +/// on EOF/read error. Anything other than a well-formed ACK is logged and +/// skipped — a replica cannot corrupt master state through this path. +#[cfg(feature = "runtime-monoio")] +async fn ack_read_loop( + mut rd: monoio::net::tcp::TcpOwnedReadHalf, + replica_id: u64, + repl_state: Arc>, +) { + use monoio::io::AsyncReadRent; + use std::sync::atomic::Ordering; + + let mut buf = bytes::BytesMut::with_capacity(4096); + loop { + let tmp = vec![0u8; 4096]; + let (res, tmp) = rd.read(tmp).await; + let n = match res { + Ok(0) | Err(_) => return, // replica closed its send half + Ok(n) => n, + }; + buf.extend_from_slice(&tmp[..n]); + // Parse complete RESP frames directly — the shared replication + // drainer (`drain_replicated_commands`) deliberately DROPS REPLCONF + // as chatter, which is exactly the frame this loop exists to read. + let acks = match drain_ack_offsets(&mut buf) { + Ok(acks) => acks, + Err(()) => { + tracing::warn!( + replica_id, + "unparseable bytes on replica ACK channel — closing" + ); + return; + } + }; + for offset in acks { + if let Ok(rs) = repl_state.read() { + if let Some(info) = rs.replicas.iter().find(|r| r.id == replica_id) { + // fetch_max: ACKs can only move forward — a reordered or + // duplicate ACK never regresses the recorded offset. + if let Some(slot) = info.ack_offsets.first() { + slot.fetch_max(offset, Ordering::Relaxed); + } + info.last_ack_time.store( + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + Ordering::Relaxed, + ); + } + } + } + } +} + +/// Drain every complete RESP frame from `buf` and return the offsets of all +/// well-formed `REPLCONF ACK ` frames (R1). Non-ACK frames are +/// skipped at debug level; a parse error returns `Err(())` — the unframed +/// stream cannot be resynced, so the caller must drop the connection. +#[cfg(feature = "runtime-monoio")] +fn drain_ack_offsets(buf: &mut bytes::BytesMut) -> Result, ()> { + use crate::protocol::{Frame, ParseConfig, parse}; + + let config = ParseConfig::default(); + let mut acks = Vec::new(); + loop { + if buf.is_empty() { + return Ok(acks); + } + let frame = match parse::parse(buf, &config) { + Ok(Some(frame)) => frame, + Ok(None) => return Ok(acks), // partial trailing frame — wait for more + Err(_) => return Err(()), + }; + let Frame::Array(items) = &frame else { + continue; // inline keepalive etc. — ignore + }; + let bulk = |f: &Frame| -> Option { + match f { + Frame::BulkString(b) | Frame::SimpleString(b) => Some(b.clone()), + _ => None, + } + }; + let is_ack = items.len() >= 3 + && bulk(&items[0]).is_some_and(|c| c.eq_ignore_ascii_case(b"REPLCONF")) + && bulk(&items[1]).is_some_and(|s| s.eq_ignore_ascii_case(b"ACK")); + if !is_ack { + tracing::debug!("ignoring non-ACK frame on replica channel"); + continue; + } + if let Some(offset) = bulk(&items[2]) + .and_then(|b| std::str::from_utf8(&b).ok().map(|s| s.to_owned())) + .and_then(|s| s.trim().parse::().ok()) + { + acks.push(offset); + } + } +} + /// WAIT command: block until N replicas acknowledge >= target_offset, or timeout expires. /// /// Returns the count of replicas that have acknowledged the offset. diff --git a/src/replication/replica.rs b/src/replication/replica.rs index 329b2e01..7a95ae14 100644 --- a/src/replication/replica.rs +++ b/src/replication/replica.rs @@ -249,7 +249,7 @@ async fn run_handshake_and_stream( // stream-db in the snapshot-capture stretch, so post-snapshot bytes // always re-establish it with an explicit SELECT (task #22). cfg.stream_db.store(0, Ordering::Relaxed); - stream_commands(&mut stream, cfg).await?; + stream_commands(stream, cfg).await?; } else if response.starts_with(b"+CONTINUE") { // Partial resync: stream from current offset if let Ok(mut rs) = cfg.repl_state.write() { @@ -257,7 +257,7 @@ async fn run_handshake_and_stream( *state = ReplicaHandshakeState::Streaming; } } - stream_commands(&mut stream, cfg).await?; + stream_commands(stream, cfg).await?; } else { anyhow::bail!("Unexpected PSYNC response: {:?}", response); } @@ -270,7 +270,48 @@ async fn run_handshake_and_stream( /// Master sends RESP-encoded commands in the same format as the WAL (RESP Array frames). /// We parse each frame and route it to the correct shard via key_to_shard. #[cfg(feature = "runtime-tokio")] -async fn stream_commands(stream: &mut TcpStream, cfg: &ReplicaTaskConfig) -> anyhow::Result<()> { +async fn stream_commands(stream: TcpStream, cfg: &ReplicaTaskConfig) -> anyhow::Result<()> { + use tokio::io::AsyncWriteExt as _; + + // R1 (task #19): the replica acknowledges its applied offset back to the + // master with `REPLCONF ACK ` — the input WAIT needs to resolve. + // The socket is split so a dedicated 1s ticker task owns the write half + // (Redis's replicationCron cadence; also serves as an idle keepalive for + // master-side lag detection), while this loop keeps the read half. + let (mut stream, wr) = stream.into_split(); + let ack_stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let ack_handle = tokio::spawn({ + let stop = ack_stop.clone(); + let repl_state = cfg.repl_state.clone(); + let mut wr = wr; + async move { + loop { + tokio::time::sleep(std::time::Duration::from_secs(1)).await; + if stop.load(Ordering::Relaxed) { + break; + } + let offset = match repl_state.read() { + Ok(rs) => rs.master_repl_offset.load(Ordering::Relaxed), + Err(_) => break, + }; + if wr.write_all(&encode_replconf_ack(offset)).await.is_err() { + break; // socket dead — the read loop is exiting too + } + } + } + }); + let result = stream_commands_read_loop(&mut stream, cfg).await; + ack_stop.store(true, Ordering::Relaxed); + ack_handle.abort(); + result +} + +/// The read/apply half of `stream_commands` (tokio). +#[cfg(feature = "runtime-tokio")] +async fn stream_commands_read_loop( + stream: &mut tokio::net::tcp::OwnedReadHalf, + cfg: &ReplicaTaskConfig, +) -> anyhow::Result<()> { let mut buf = BytesMut::with_capacity(65536); // Seeded from the task-level slot (NOT 0): a +CONTINUE resume must keep // the db context the stream was in when the link dropped — see @@ -520,14 +561,14 @@ async fn run_handshake_and_stream( // stream-db in the snapshot-capture stretch, so post-snapshot bytes // always re-establish it with an explicit SELECT (task #22). cfg.stream_db.store(0, Ordering::Relaxed); - stream_commands(&mut stream, cfg).await?; + stream_commands(stream, cfg).await?; } else if response.starts_with(b"+CONTINUE") { if let Ok(mut rs) = cfg.repl_state.write() { if let ReplicationRole::Replica { ref mut state, .. } = rs.role { *state = ReplicaHandshakeState::Streaming; } } - stream_commands(&mut stream, cfg).await?; + stream_commands(stream, cfg).await?; } else { anyhow::bail!("Unexpected PSYNC response: {:?}", response); } @@ -538,7 +579,53 @@ async fn run_handshake_and_stream( /// Stream incoming WAL bytes from master under monoio using ownership read. #[cfg(feature = "runtime-monoio")] async fn stream_commands( - stream: &mut monoio::net::TcpStream, + stream: monoio::net::TcpStream, + cfg: &ReplicaTaskConfig, +) -> anyhow::Result<()> { + // R1 (task #19): acknowledge the applied offset back to the master with + // `REPLCONF ACK ` so WAIT resolves. Split the socket: a 1s ticker + // task owns the write half (Redis's replicationCron cadence + idle + // keepalive for lag detection); this loop keeps the read half. A + // timeout-wrapped read was rejected instead: cancelling an in-flight + // io_uring read whose CQE already completed DISCARDS those bytes — + // silent stream corruption. + use monoio::io::Splitable as _; + let (mut rd, wr) = stream.into_split(); + let ack_stop = std::rc::Rc::new(std::cell::Cell::new(false)); + let ack_handle = monoio::spawn({ + let stop = ack_stop.clone(); + let repl_state = cfg.repl_state.clone(); + let mut wr = wr; + async move { + use monoio::io::AsyncWriteRentExt; + loop { + monoio::time::sleep(std::time::Duration::from_secs(1)).await; + if stop.get() { + break; + } + let offset = match repl_state.read() { + Ok(rs) => rs.master_repl_offset.load(Ordering::Relaxed), + Err(_) => break, + }; + let (res, _) = wr.write_all(encode_replconf_ack(offset)).await; + if res.is_err() { + break; // socket dead — the read loop is exiting too + } + } + } + }); + let result = stream_commands_read_loop(&mut rd, cfg).await; + // Stop the ticker (checked each tick; the write half drops with the task, + // closing our side of the socket within ~1s of the read loop ending). + ack_stop.set(true); + drop(ack_handle); + result +} + +/// The read/apply half of `stream_commands` (monoio). +#[cfg(feature = "runtime-monoio")] +async fn stream_commands_read_loop( + stream: &mut monoio::net::tcp::TcpOwnedReadHalf, cfg: &ReplicaTaskConfig, ) -> anyhow::Result<()> { use monoio::io::AsyncReadRent; @@ -697,3 +784,17 @@ async fn read_rdb_bulk(stream: &mut TcpStream) -> anyhow::Result> { stream.read_exact(&mut data).await?; Ok(data) } + +/// RESP-serialize `REPLCONF ACK ` for the replication link (R1). +fn encode_replconf_ack(offset: u64) -> Vec { + let mut n = itoa::Buffer::new(); + let off = n.format(offset); + let mut ln = itoa::Buffer::new(); + let mut buf = Vec::with_capacity(48); + buf.extend_from_slice(b"*3\r\n$8\r\nREPLCONF\r\n$3\r\nACK\r\n$"); + buf.extend_from_slice(ln.format(off.len()).as_bytes()); + buf.extend_from_slice(b"\r\n"); + buf.extend_from_slice(off.as_bytes()); + buf.extend_from_slice(b"\r\n"); + buf +} diff --git a/src/server/conn/blocking.rs b/src/server/conn/blocking.rs index b7b65cff..fafb5153 100644 --- a/src/server/conn/blocking.rs +++ b/src/server/conn/blocking.rs @@ -1447,7 +1447,7 @@ pub(crate) fn try_inline_dispatch( // full. On loss, fail-loud: the SET is applied in memory, but the // client must see an error, not `+OK` (review finding, PR #211). let mut aof_budget = crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; - if !pool.send_append_bounded_blocking(shard_id, lsn, frozen, &mut aof_budget) { + if !pool.send_append_bounded_blocking(shard_id, lsn, selected_db, frozen, &mut aof_budget) { write_buf.extend_from_slice( b"-MOONERR AOF backpressure: write applied in memory but not queued \ for persistence\r\n", diff --git a/src/server/conn/handler_monoio/dispatch.rs b/src/server/conn/handler_monoio/dispatch.rs index bd35899e..706e1e9a 100644 --- a/src/server/conn/handler_monoio/dispatch.rs +++ b/src/server/conn/handler_monoio/dispatch.rs @@ -668,6 +668,56 @@ pub(super) async fn try_handle_info( true } +/// Handle WAIT (R1, task #19): block until N replicas acknowledge the current +/// master offset or the timeout expires; reply with the acked count. Runs at +/// the connection layer because it awaits — the generic dispatch path is +/// synchronous (it used to hard-code `:0`). Returns `true` if consumed. +pub(super) async fn try_handle_wait( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"WAIT") { + return false; + } + let int_arg = |f: &Frame| -> Option { + match f { + Frame::BulkString(b) | Frame::SimpleString(b) => { + std::str::from_utf8(b).ok()?.trim().parse().ok() + } + Frame::Integer(n) if *n >= 0 => Some(*n as u64), + _ => None, + } + }; + let (Some(num_required), Some(timeout_ms)) = ( + cmd_args.first().and_then(int_arg), + cmd_args.get(1).and_then(int_arg), + ) else { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'wait' command", + ))); + return true; + }; + // Redis: timeout 0 = block until satisfied. The poll loop is cooperative + // (10ms ticks), so an effectively-unbounded deadline cannot starve the + // shard thread — cap at one year rather than a literal infinity. + let timeout_ms = if timeout_ms == 0 { + 31_536_000_000 + } else { + timeout_ms + }; + let count = match ctx.repl_state.as_ref() { + Some(rs) => { + crate::replication::master::wait_for_replicas(num_required as usize, timeout_ms, rs) + .await + } + None => 0, + }; + responses.push(Frame::Integer(count as i64)); + true +} + /// Handle READONLY enforcement: reject writes on replicas. /// Returns `true` if the command was blocked. /// diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 0802c6cf..f3f3f80a 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -907,6 +907,11 @@ pub(crate) async fn handle_connection_sharded_monoio< { continue; } + // WAIT blocks on replica ACKs (R1) — must run at the connection + // layer; generic dispatch is synchronous and used to answer :0. + if cmd_len == 4 && dispatch::try_handle_wait(cmd, cmd_args, ctx, &mut responses).await { + continue; + } if dispatch::try_enforce_readonly(cmd, cmd_args, ctx, &mut responses) { continue; } @@ -1286,7 +1291,10 @@ pub(crate) async fn handle_connection_sharded_monoio< || conn.tracking_state.enabled || crate::replication::state::fanout_hint_active() { - metadata::is_write(cmd) + // `is_persisted_write`: SELECT is W-flagged but conn-state + // only — persisting the literal client SELECT poisons the + // stream db context for interleaved connections (task #35). + metadata::is_persisted_write(cmd) } else { false }; @@ -1355,7 +1363,10 @@ pub(crate) async fn handle_connection_sharded_monoio< responses.push(response); continue; }; - match pool.send_append_group(ctx.shard_id, lsn, serialized).await { + match pool + .send_append_group(ctx.shard_id, lsn, conn.selected_db, serialized) + .await + { // Always: durability confirmed by ONE // fsync_barrier per batch (resolve_local_leg_barrier // before serialization) instead of an awaited @@ -1438,7 +1449,15 @@ pub(crate) async fn handle_connection_sharded_monoio< responses.push(response); continue; }; - match pool.send_append_group(ctx.shard_id, lsn, serialized).await { + match pool + .send_append_group( + ctx.shard_id, + lsn, + conn.selected_db, + serialized, + ) + .await + { // Same one-barrier-per-batch contract as MOVE. Ok(true) => local_leg_write_idxs.push(responses.len()), Ok(false) => {} @@ -1742,7 +1761,15 @@ pub(crate) async fn handle_connection_sharded_monoio< ) }; if let Some(ref pool) = ctx.aof_pool { - match pool.send_append_group(ctx.shard_id, lsn, serialized).await { + match pool + .send_append_group( + ctx.shard_id, + lsn, + conn.selected_db, + serialized, + ) + .await + { Ok(true) => aof_barrier_pending = true, Ok(false) => {} Err(_) => { @@ -1925,7 +1952,9 @@ pub(crate) async fn handle_connection_sharded_monoio< let resp_idx = responses.len(); responses.push(Frame::Null); // placeholder, filled after batch dispatch // Pre-compute AOF bytes before moving frame into Arc - let aof_bytes = if ctx.aof_pool.is_some() && metadata::is_write(cmd) { + // `is_persisted_write`: never AOF/replicate a literal client + // SELECT (task #35 — poisons the stream db context). + let aof_bytes = if ctx.aof_pool.is_some() && metadata::is_persisted_write(cmd) { Some(aof::serialize_command(&dispatch_frame)) } else { None diff --git a/src/server/conn/handler_monoio/write.rs b/src/server/conn/handler_monoio/write.rs index 37cbe7f6..2d06226f 100644 --- a/src/server/conn/handler_monoio/write.rs +++ b/src/server/conn/handler_monoio/write.rs @@ -765,8 +765,11 @@ pub(super) async fn try_handle_multi_exec( // latency vector for cross-shard traffic sharing this thread. // Each drain iteration is just a try_send per replica, so the // burst is cheap; revisit only if EXEC bodies grow unbounded. - for bytes in &aof_entries { - super::ft::record_local_write_db(ctx, conn.selected_db, bytes.clone()); + // PR #282 review: per-entry db — a SELECT queued inside the + // body redirects the commands after it, so the replication + // stream must bind each record to ITS execution db. + for (entry_db, bytes) in &aof_entries { + super::ft::record_local_write_db(ctx, *entry_db, bytes.clone()); } } // DURABILITY: append every successful write in the body to THIS diff --git a/src/server/conn/handler_sharded/dispatch.rs b/src/server/conn/handler_sharded/dispatch.rs index cf0a644b..593dc98b 100644 --- a/src/server/conn/handler_sharded/dispatch.rs +++ b/src/server/conn/handler_sharded/dispatch.rs @@ -564,6 +564,56 @@ pub(super) async fn try_handle_swapdb( true } +/// Handle WAIT (R1, task #19): block until N replicas acknowledge the current +/// master offset or the timeout expires; reply with the acked count. Runs at +/// the connection layer because it awaits — the generic dispatch path is +/// synchronous (it used to hard-code `:0`). Returns `true` if consumed. +pub(super) async fn try_handle_wait( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"WAIT") { + return false; + } + let int_arg = |f: &Frame| -> Option { + match f { + Frame::BulkString(b) | Frame::SimpleString(b) => { + std::str::from_utf8(b).ok()?.trim().parse().ok() + } + Frame::Integer(n) if *n >= 0 => Some(*n as u64), + _ => None, + } + }; + let (Some(num_required), Some(timeout_ms)) = ( + cmd_args.first().and_then(int_arg), + cmd_args.get(1).and_then(int_arg), + ) else { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'wait' command", + ))); + return true; + }; + // Redis: timeout 0 = block until satisfied. The poll loop is cooperative + // (10ms ticks), so an effectively-unbounded deadline cannot starve the + // shard thread — cap at one year rather than a literal infinity. + let timeout_ms = if timeout_ms == 0 { + 31_536_000_000 + } else { + timeout_ms + }; + let count = match ctx.repl_state.as_ref() { + Some(rs) => { + crate::replication::master::wait_for_replicas(num_required as usize, timeout_ms, rs) + .await + } + None => 0, + }; + responses.push(Frame::Integer(count as i64)); + true +} + /// Handle READONLY enforcement for replicas. /// Returns `true` if the command was blocked (caller should `continue`). /// diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 24b74717..13ad4537 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -871,6 +871,12 @@ pub(crate) async fn handle_connection_sharded_inner< continue; } + // --- WAIT (R1): blocks on replica ACKs; must run at the + // connection layer (generic dispatch is synchronous) --- + if dispatch::try_handle_wait(cmd, cmd_args, ctx, &mut responses).await { + continue; + } + // --- READONLY enforcement --- if dispatch::try_enforce_readonly(cmd, cmd_args, ctx, &mut responses) { continue; @@ -1238,7 +1244,9 @@ pub(crate) async fn handle_connection_sharded_inner< } let is_write = if ctx.aof_pool.is_some() || conn.tracking_state.enabled { metadata::is_write(cmd) } else { false }; - let aof_bytes = if is_write && ctx.aof_pool.is_some() { Some(aof::serialize_command(&frame)) } else { None }; + // `is_persisted_write`: never AOF a literal client SELECT + // (task #35 — poisons the stream db context). + let aof_bytes = if is_write && ctx.aof_pool.is_some() && metadata::is_persisted_write(cmd) { Some(aof::serialize_command(&frame)) } else { None }; if is_local { // LOCAL PATH: split into read/write to avoid exclusive lock on reads. @@ -1285,7 +1293,12 @@ pub(crate) async fn handle_connection_sharded_inner< if let Some(ref pool) = ctx.aof_pool { let lsn = aof::AofWriterPool::issue_append_lsn(&ctx.repl_state, ctx.shard_id, bytes.len()); match pool - .send_append_group(ctx.shard_id, lsn, bytes.clone()) + .send_append_group( + ctx.shard_id, + lsn, + conn.selected_db, + bytes.clone(), + ) .await { // Always: one fsync_barrier per batch @@ -1343,7 +1356,12 @@ pub(crate) async fn handle_connection_sharded_inner< if let Some(ref pool) = ctx.aof_pool { let lsn = aof::AofWriterPool::issue_append_lsn(&ctx.repl_state, ctx.shard_id, bytes.len()); match pool - .send_append_group(ctx.shard_id, lsn, bytes.clone()) + .send_append_group( + ctx.shard_id, + lsn, + conn.selected_db, + bytes.clone(), + ) .await { // Always: one fsync_barrier per batch @@ -1694,7 +1712,15 @@ pub(crate) async fn handle_connection_sharded_inner< if !matches!(response, Frame::Error(_)) { if let Some(ref pool) = ctx.aof_pool { let lsn = aof::AofWriterPool::issue_append_lsn(&ctx.repl_state, ctx.shard_id, bytes.len()); - match pool.send_append_group(ctx.shard_id, lsn, bytes).await { + match pool + .send_append_group( + ctx.shard_id, + lsn, + conn.selected_db, + bytes, + ) + .await + { Ok(true) => aof_barrier_pending = true, Ok(false) => {} Err(_) => { diff --git a/src/server/conn/handler_single.rs b/src/server/conn/handler_single.rs index eaa95d7f..7d65555b 100644 --- a/src/server/conn/handler_single.rs +++ b/src/server/conn/handler_single.rs @@ -52,14 +52,15 @@ use crate::server::codec::RespCodec; /// tests pass any `Sink` mock. /// - `responses` — per-command response slots; fsync failures patch the /// corresponding slot to `Frame::Error`. -/// - `aof_entries` — `(resp_idx, bytes)`: bytes to fsync, slot to patch on failure. +/// - `aof_entries` — `(resp_idx, db, bytes)`: bytes to fsync (executed in `db`), +/// slot to patch on failure. /// - `pool` — AOF writer pool (caller must ensure Always policy). /// - `repl_state` — replication state for LSN issuance (`&None` in tests). /// - `change_counter` — auto-save dirty counter (`&None` if not configured). pub(crate) async fn flush_with_aof_ack( sink: &mut S, mut responses: Vec, - aof_entries: Vec<(usize, Bytes)>, + aof_entries: Vec<(usize, usize, Bytes)>, pool: &crate::persistence::aof::AofWriterPool, repl_state: &Option>>, change_counter: &Option>, @@ -73,10 +74,10 @@ where // resolve_local_leg_barrier. On barrier failure every enqueued write in // the batch is unconfirmed, so every joined slot is patched. let mut barrier_idxs: Vec = Vec::new(); - for (resp_idx, bytes) in aof_entries { + for (resp_idx, db, bytes) in aof_entries { let lsn = crate::persistence::aof::AofWriterPool::issue_append_lsn(repl_state, 0, bytes.len()); - match pool.send_append_group(0, lsn, bytes).await { + match pool.send_append_group(0, lsn, db, bytes).await { Ok(true) => barrier_idxs.push(resp_idx), Ok(false) => {} Err(_) => { @@ -427,10 +428,13 @@ pub async fn handle_connection( // Phase 1: Handle connection-level intercepts, collect dispatchable frames // Phase 2: Acquire ONE write lock, execute ALL dispatchable frames let mut responses: Vec = Vec::with_capacity(batch.len()); - // Each entry carries (resp_idx, bytes) so the Always-policy flush path - // can patch responses[resp_idx] with WRITEFAIL when fsync fails, + // Each entry carries (resp_idx, db, bytes) so the Always-policy flush + // path can patch responses[resp_idx] with WRITEFAIL when fsync fails, // before any response is sent to the client (H1 fix — FIX-W1-1). - let mut aof_entries: Vec<(usize, Bytes)> = Vec::new(); + // task #35: `db` is the connection's selected db at the moment this + // entry was recorded (only SELECT changes it, and SELECT itself is + // never persisted, so it is stable for every persisted entry). + let mut aof_entries: Vec<(usize, usize, Bytes)> = Vec::new(); let mut should_quit = false; let mut break_outer = false; @@ -776,7 +780,11 @@ pub async fn handle_connection( ); // Single-shard mode — shard_id = 0. let lsn = crate::persistence::aof::AofWriterPool::issue_append_lsn(&repl_state, 0, serialized.len()); - pool.try_send_append_durable(0, lsn, serialized) + // task #35: SWAPDB affects both `a` and + // `b` — no single db context applies; + // pass 0 (writer may emit a harmless + // redundant SELECT 0). + pool.try_send_append_durable(0, lsn, 0, serialized) .await .is_ok() } else { @@ -1006,7 +1014,7 @@ pub async fn handle_connection( if !matches!(&response, Frame::Error(_)) { // Carry resp_idx so the Always-policy flush can // patch responses[resp_idx] on fsync failure. - aof_entries.push((resp_idx, bytes)); + aof_entries.push((resp_idx, conn.selected_db, bytes)); } } // Apply RESP3 response conversion if needed @@ -1037,16 +1045,16 @@ pub async fn handle_connection( // fire-and-forget (returns Ok immediately) so no latency // penalty. // - // Note: aof_entries carries (resp_idx, bytes) from FIX-W1-1 - // — resp_idx is unused here because the all-or-nothing - // failure mode discards the entire response buffer; future - // per-slot patching could use it. + // Note: aof_entries carries (resp_idx, db, bytes) from + // FIX-W1-1 — resp_idx is unused here because the + // all-or-nothing failure mode discards the entire response + // buffer; future per-slot patching could use it. let mut aof_write_failed = false; let mut aof_barrier_needed = false; - for (_resp_idx, bytes) in aof_entries.drain(..) { + for (_resp_idx, entry_db, bytes) in aof_entries.drain(..) { if let Some(ref pool) = aof_pool { let lsn = crate::persistence::aof::AofWriterPool::issue_append_lsn(&repl_state, 0, bytes.len()); - match pool.send_append_group(0, lsn, bytes).await { + match pool.send_append_group(0, lsn, entry_db, bytes).await { Ok(true) => aof_barrier_needed = true, Ok(false) => {} Err(_) => aof_write_failed = true, @@ -1243,6 +1251,15 @@ pub async fn handle_connection( } else { conn.in_multi = false; let mut exec_publishes: Vec<(usize, Bytes, Bytes)> = Vec::new(); + // PR #282 review: `execute_transaction` holds + // ONE guard on the db selected at EXEC time — + // every body write physically lands there, + // even when a queued SELECT mutates + // `conn.selected_db` mid-body. Capture the + // guard's db NOW; attributing entries to the + // post-EXEC `conn.selected_db` would mis-place + // them on recovery. + let txn_db = conn.selected_db; let (mut result, txn_aof_entries) = execute_transaction( &db, &conn.command_queue, @@ -1380,8 +1397,15 @@ pub async fn handle_connection( // any command's fsync fails. let exec_resp_idx = responses.len(); responses.push(result); + // task #35 + PR #282 review: the whole body + // ran against the guard taken on `txn_db` + // (captured before EXEC) — NOT the possibly + // SELECT-mutated post-EXEC `conn.selected_db`. + // A queued SELECT itself is never persisted. aof_entries.extend( - txn_aof_entries.into_iter().map(|b| (exec_resp_idx, b)), + txn_aof_entries + .into_iter() + .map(|b| (exec_resp_idx, txn_db, b)), ); } continue; @@ -1768,7 +1792,10 @@ pub async fn handle_connection( // everysec/no. let bytes = bytes::Bytes::from(record); let lsn = crate::persistence::aof::AofWriterPool::issue_append_lsn(&repl_state, 0, bytes.len()); - match pool.send_append_group(0, lsn, bytes).await { + match pool + .send_append_group(0, lsn, conn.selected_db, bytes) + .await + { Ok(true) => graph_barrier_needed = true, Ok(false) => {} Err(_) => graph_aof_failed = true, @@ -1801,8 +1828,10 @@ pub async fn handle_connection( let is_write = metadata::is_write(cmd); - // Serialize for AOF before dispatch - let aof_bytes = if is_write && aof_pool.is_some() { + // Serialize for AOF before dispatch. + // `is_persisted_write`: never AOF a literal client + // SELECT (task #35 — poisons the stream db context). + let aof_bytes = if metadata::is_persisted_write(cmd) && aof_pool.is_some() { let mut buf = BytesMut::new(); crate::protocol::serialize::serialize(&frame, &mut buf); Some(buf.freeze()) @@ -2417,7 +2446,11 @@ pub async fn handle_connection( }; if matches!(response, Frame::Integer(1)) { if let Some(bytes) = &aof_bytes { - aof_entries.push((resp_idx, bytes.clone())); + // task #35: MOVE persists against its + // SOURCE db (src_db == conn.selected_db; + // MOVE never changes the connection's + // selected db). + aof_entries.push((resp_idx, src_db, bytes.clone())); } } responses[resp_idx] = response; @@ -2446,7 +2479,9 @@ pub async fn handle_connection( }; if matches!(response, Frame::Integer(1)) { if let Some(bytes) = &aof_bytes { - aof_entries.push((resp_idx, bytes.clone())); + // task #35: COPY ... DB n persists + // against its SOURCE db, same as MOVE. + aof_entries.push((resp_idx, src_db, bytes.clone())); } } responses[resp_idx] = response; @@ -2571,7 +2606,7 @@ pub async fn handle_connection( ); } if let Some(bytes) = aof_bytes { - aof_entries.push((resp_idx, bytes.clone())); + aof_entries.push((resp_idx, conn.selected_db, bytes.clone())); } } // Apply RESP3 response conversion if needed @@ -2625,10 +2660,10 @@ pub async fn handle_connection( break; } } - for (_, bytes) in aof_entries { + for (_, entry_db, bytes) in aof_entries { if let Some(ref pool) = aof_pool { let lsn = crate::persistence::aof::AofWriterPool::issue_append_lsn(&repl_state, 0, bytes.len()); - let _ = pool.try_send_append_durable(0, lsn, bytes).await; + let _ = pool.try_send_append_durable(0, lsn, entry_db, bytes).await; } if let Some(ref counter) = change_counter { counter.fetch_add(1, Ordering::Relaxed); @@ -2773,7 +2808,7 @@ mod tests { let start = Instant::now(); let responses = vec![Frame::SimpleString(bytes::Bytes::from_static(b"OK"))]; - let aof_entries = vec![(0usize, bytes::Bytes::from_static(b"SET k v\r\n"))]; + let aof_entries = vec![(0usize, 0usize, bytes::Bytes::from_static(b"SET k v\r\n"))]; let mut sink = RecordingSink::new(); let broke = flush_with_aof_ack( diff --git a/src/server/conn/shared.rs b/src/server/conn/shared.rs index 5980832a..4b9d8a50 100644 --- a/src/server/conn/shared.rs +++ b/src/server/conn/shared.rs @@ -161,8 +161,13 @@ pub(crate) fn execute_transaction( continue; } - // Check if this is a write command for AOF logging - let is_write = metadata::is_write(cmd); + // Check if this is a write command for AOF logging. + // `is_persisted_write` (PR #282 review): a SELECT queued inside MULTI + // must not be persisted as a literal record — it would shift the AOF + // stream's db context under every OTHER record (task #35). All body + // writes execute against the guard taken on the ENTRY db above, so + // the caller attributes every entry to that one db. + let is_write = metadata::is_persisted_write(cmd); // Serialize for AOF before dispatch let aof_bytes = if is_write { @@ -209,11 +214,16 @@ pub(crate) fn execute_transaction_sharded( selected_db: usize, cached_clock: &CachedClock, exec_publishes: &mut Vec<(usize, Bytes, Bytes)>, -) -> (Frame, Vec) { +) -> (Frame, Vec<(usize, Bytes)>) { let db_count = shard_databases.db_count(); let mut results = Vec::with_capacity(command_queue.len()); - let mut aof_entries: Vec = Vec::new(); + // Per-entry db (PR #282 review): this executor re-dispatches each body + // command against the CURRENT `selected`, so a SELECT queued inside MULTI + // really does redirect the commands after it — collapsing every entry to + // one db mis-attributes recovery and replication whenever the body + // switches dbs. + let mut aof_entries: Vec<(usize, Bytes)> = Vec::new(); let mut selected = selected_db; for cmd_frame in command_queue { @@ -240,13 +250,20 @@ pub(crate) fn execute_transaction_sharded( // MULTI/EXEC path logged nothing, so every transactional write was // silently lost on restart. Fully-qualified paths because `metadata` // is only `use`d under runtime-tokio but this fn compiles under both. - let aof_bytes = if crate::command::metadata::is_write(cmd) { + // `is_persisted_write` (PR #282 review): a queued literal SELECT is + // connection/txn state only — persisting it shifts the stream's db + // context under other records (task #35). + let aof_bytes = if crate::command::metadata::is_persisted_write(cmd) { let mut buf = bytes::BytesMut::new(); crate::protocol::serialize::serialize(cmd_frame, &mut buf); Some(buf.freeze()) } else { None }; + // The db THIS command executes in — captured before dispatch (a + // queued SELECT mutates `selected` for the commands after it, never + // for itself, and no persisted write mutates it mid-dispatch). + let entry_db = selected; let result = crate::shard::slice::with_shard_db(selected, |db| { db.refresh_now_from_cache(cached_clock); @@ -261,7 +278,7 @@ pub(crate) fn execute_transaction_sharded( // single-shard path — an errored write must not reach the AOF). if let Some(bytes) = aof_bytes { if !matches!(&response, Frame::Error(_)) { - aof_entries.push(bytes); + aof_entries.push((entry_db, bytes)); } } @@ -349,7 +366,11 @@ pub(crate) fn execute_transaction_sharded( /// master fanout is not wired; monoio is the production replication runtime). pub(crate) async fn persist_txn_aof( ctx: &crate::server::conn::core::ConnectionContext, - aof_entries: Vec, + // task #35 + PR #282 review: each entry carries the db THAT command + // executed in — a SELECT queued inside MULTI redirects the commands + // after it (sharded executor), so one collapsed db mis-attributes the + // body on recovery. + aof_entries: Vec<(usize, Bytes)>, repl_recorded: bool, ) -> Result<(), ()> { if aof_entries.is_empty() { @@ -359,7 +380,7 @@ pub(crate) async fn persist_txn_aof( return Ok(()); }; let mut barrier_pending = false; - for bytes in aof_entries { + for (db, bytes) in aof_entries { let lsn = if repl_recorded { 0 } else { @@ -369,7 +390,7 @@ pub(crate) async fn persist_txn_aof( bytes.len(), ) }; - match pool.send_append_group(ctx.shard_id, lsn, bytes).await { + match pool.send_append_group(ctx.shard_id, lsn, db, bytes).await { Ok(true) => barrier_pending = true, Ok(false) => {} Err(_) => return Err(()), diff --git a/src/shard/coordinator.rs b/src/shard/coordinator.rs index 8e96b005..bcaf6e83 100644 --- a/src/shard/coordinator.rs +++ b/src/shard/coordinator.rs @@ -359,7 +359,7 @@ async fn run_on_owner_persist( let serialized = crate::persistence::aof::serialize_command(&Frame::Array( command_parts.to_vec().into(), )); - match persist_local_leg(aof_pool, repl_state, my_shard, serialized).await { + match persist_local_leg(aof_pool, repl_state, my_shard, db_index, serialized).await { Ok(needs_barrier) => *local_barrier_pending |= needs_barrier, Err(()) => { return Frame::Error(Bytes::from_static(crate::persistence::aof::AOF_FSYNC_ERR)); @@ -401,6 +401,10 @@ async fn persist_local_leg( aof_pool: Option<&Arc>, repl_state: ReplStateRef<'_>, my_shard: usize, + // task #35: the db this local leg executed in (`db_index` at every call + // site) — threaded into the AOF pool so the writer can inject a + // `SELECT ` record on a db-context change. + db: usize, serialized: Bytes, ) -> Result { let Some(pool) = aof_pool else { @@ -411,7 +415,7 @@ async fn persist_local_leg( my_shard, serialized.len(), ); - match pool.send_append_group(my_shard, lsn, serialized).await { + match pool.send_append_group(my_shard, lsn, db, serialized).await { Ok(needs_barrier) => Ok(needs_barrier), Err(_) => Err(()), } @@ -1207,7 +1211,7 @@ async fn coordinate_mset( // owned by my_shard — matching the local single-key write contract. if let Some(pairs) = groups.get(&my_shard) { let serialized = serialize_local_mset(pairs); - match persist_local_leg(aof_pool, repl_state, my_shard, serialized).await { + match persist_local_leg(aof_pool, repl_state, my_shard, db_index, serialized).await { Ok(needs_barrier) => *local_barrier_pending |= needs_barrier, Err(()) => { return Frame::Error(Bytes::from_static( @@ -1283,7 +1287,7 @@ async fn coordinate_mset( // write keys this shard doesn't own). if let Some(pairs) = groups.get(&my_shard) { let serialized = serialize_local_mset(pairs); - match persist_local_leg(aof_pool, repl_state, my_shard, serialized).await { + match persist_local_leg(aof_pool, repl_state, my_shard, db_index, serialized).await { Ok(needs_barrier) => *local_barrier_pending |= needs_barrier, Err(()) => { return Frame::Error(Bytes::from_static(crate::persistence::aof::AOF_FSYNC_ERR)); @@ -1366,7 +1370,7 @@ async fn coordinate_msetnx( if matches!(resp, Frame::Integer(1)) { let serialized = crate::persistence::aof::serialize_command(&Frame::Array(command_parts.into())); - match persist_local_leg(aof_pool, repl_state, my_shard, serialized).await { + match persist_local_leg(aof_pool, repl_state, my_shard, db_index, serialized).await { Ok(needs_barrier) => *local_barrier_pending |= needs_barrier, Err(()) => { return Frame::Error(Bytes::from_static( @@ -1447,7 +1451,7 @@ async fn coordinate_multi_del_or_exists( parts.extend_from_slice(args); let serialized = crate::persistence::aof::serialize_command(&Frame::Array(parts.into())); - match persist_local_leg(aof_pool, repl_state, my_shard, serialized).await { + match persist_local_leg(aof_pool, repl_state, my_shard, db_index, serialized).await { Ok(needs_barrier) => *local_barrier_pending |= needs_barrier, Err(()) => { return Frame::Error(Bytes::from_static( @@ -1480,7 +1484,9 @@ async fn coordinate_multi_del_or_exists( parts.extend_from_slice(key_args); let serialized = crate::persistence::aof::serialize_command(&Frame::Array(parts.into())); - match persist_local_leg(aof_pool, repl_state, my_shard, serialized).await { + match persist_local_leg(aof_pool, repl_state, my_shard, db_index, serialized) + .await + { Ok(needs_barrier) => *local_barrier_pending |= needs_barrier, Err(()) => { return Frame::Error(Bytes::from_static( diff --git a/src/shard/dispatch.rs b/src/shard/dispatch.rs index 2622a009..c2d67a3e 100644 --- a/src/shard/dispatch.rs +++ b/src/shard/dispatch.rs @@ -411,6 +411,17 @@ pub struct CdcSubscribePayload { pub from_lsn: u64, } +/// One live replica fan-out endpoint held by a shard thread: +/// `(replica_id, live channel sender, kicked flag)`. +/// +/// `kicked` is the overflow disconnect signal — see +/// [`ShardMessage::RegisterReplica::kicked`]. +pub type ReplicaFanout = ( + u64, + channel::MpscSender, + std::sync::Arc, +); + /// Messages sent to a shard via SPSC channels from the connection layer /// or from other shards for cross-shard operations. pub enum ShardMessage { @@ -455,12 +466,20 @@ pub enum ShardMessage { BlockCancel { wait_id: u64 }, /// Register a connected replica's per-shard sender channel with this shard. /// Called once per shard per replica when a new replica connection is established. - /// The shard adds `tx` to its replica_txs list for WAL fan-out. + /// The shard adds `(id, tx, kicked)` to its replica_txs list for WAL fan-out. /// `backlog_capacity` (`--repl-backlog-size`) sizes the lazy backlog /// fallback-init so it can't diverge from the handshake-path allocation. RegisterReplica { replica_id: u64, tx: channel::MpscSender, + /// Set by the shard fan-out when this replica's bounded channel is + /// FULL: a record that cannot be queued would otherwise be silently + /// dropped, permanently diverging the replica while + /// `master_link_status` stays "up" (task #35 — observed 2k of 40k + /// keys delivered). The drain task polls the flag and disconnects, + /// converting silent divergence into a loud PSYNC resync (Redis + /// parity: output-buffer-limit disconnects). + kicked: std::sync::Arc, backlog_capacity: usize, /// When set, the event loop replies with the shard's replication /// offset AT registration — the exact point where live fan-out to diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 61a796bf..7a2d0971 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -670,7 +670,7 @@ impl super::Shard { }, None => std::sync::Arc::new(parking_lot::Mutex::new(None)), }; - let mut replica_txs: Vec<(u64, channel::MpscSender)> = Vec::new(); + let mut replica_txs: Vec = Vec::new(); let repl_state: Option>> = repl_state_ext; // QW3 (2026-06 review): lock-free offset handle cloned ONCE at shard // startup. The SPSC drain's per-write offset advance goes through this diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index 81f0a176..f5d572d4 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -101,7 +101,7 @@ pub(crate) fn drain_spsc_shared( snapshot_state: &mut Option, wal_writer: &mut Option, repl_backlog: &crate::replication::backlog::SharedBacklog, - replica_txs: &mut Vec<(u64, channel::MpscSender)>, + replica_txs: &mut Vec, repl_state: &Option, shard_id: usize, script_cache: &Rc>, @@ -368,7 +368,7 @@ pub(crate) fn handle_shard_message_shared( snapshot_state: &mut Option, wal_writer: &mut Option, repl_backlog: &crate::replication::backlog::SharedBacklog, - replica_txs: &mut Vec<(u64, channel::MpscSender)>, + replica_txs: &mut Vec, repl_state: &Option, shard_id: usize, script_cache: &Rc>, @@ -615,6 +615,7 @@ pub(crate) fn handle_shard_message_shared( crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; if !wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -691,6 +692,7 @@ pub(crate) fn handle_shard_message_shared( crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -855,6 +857,7 @@ pub(crate) fn handle_shard_message_shared( let serialized = aof::serialize_command(cmd_frame); aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -917,6 +920,7 @@ pub(crate) fn handle_shard_message_shared( let serialized = aof::serialize_command(cmd_frame); aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -1038,6 +1042,7 @@ pub(crate) fn handle_shard_message_shared( let serialized = aof::serialize_command(cmd_frame); aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -1099,6 +1104,7 @@ pub(crate) fn handle_shard_message_shared( let serialized = aof::serialize_command(cmd_frame); aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -1262,6 +1268,7 @@ pub(crate) fn handle_shard_message_shared( let mut aof_budget = crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; if !wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -1330,6 +1337,7 @@ pub(crate) fn handle_shard_message_shared( crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -1453,6 +1461,7 @@ pub(crate) fn handle_shard_message_shared( let serialized = aof::serialize_command(cmd_frame); aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -1515,6 +1524,7 @@ pub(crate) fn handle_shard_message_shared( let serialized = aof::serialize_command(cmd_frame); aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -1637,6 +1647,7 @@ pub(crate) fn handle_shard_message_shared( let serialized = aof::serialize_command(cmd_frame); aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -1698,6 +1709,7 @@ pub(crate) fn handle_shard_message_shared( let serialized = aof::serialize_command(cmd_frame); aof_ok = wal_append_and_fanout( &serialized, + db_idx, wal_writer, repl_backlog, replica_txs, @@ -2303,6 +2315,10 @@ pub(crate) fn handle_shard_message_shared( let mut aof_budget = crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; let _ = wal_append_and_fanout( &serialized, + // task #35: SWAPDB affects both `a` and `b` — no single db + // context applies; pass 0 (writer may emit a harmless + // redundant SELECT 0 if last_db was already non-zero). + 0, wal_writer, repl_backlog, replica_txs, @@ -2479,10 +2495,14 @@ pub(crate) fn handle_shard_message_shared( let mut wrote = false; let mut append_lost = false; let mut aof_budget = crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; - for entry_bytes in &aof_entries { + // PR #282 review: each entry carries the db it EXECUTED in (a + // SELECT queued inside the body redirects the commands after it) + // — attribute per entry, not the body's entry db. + for (entry_db, entry_bytes) in &aof_entries { wrote = true; let ok = wal_append_and_fanout( entry_bytes, + *entry_db, wal_writer, repl_backlog, replica_txs, @@ -2511,6 +2531,7 @@ pub(crate) fn handle_shard_message_shared( ShardMessage::RegisterReplica { replica_id, tx, + kicked, backlog_capacity, registered, push_offset, @@ -2534,7 +2555,7 @@ pub(crate) fn handle_shard_message_shared( *guard = Some(ReplicationBacklog::new_at(backlog_capacity, offset)); } drop(guard); - replica_txs.push((replica_id, tx)); + replica_txs.push((replica_id, tx, kicked)); // Reply with the offset at which live fan-out begins. For // same-thread self-queue registrations this is `push_offset`, // captured AT PUSH TIME: local writes advance the offset @@ -2561,17 +2582,15 @@ pub(crate) fn handle_shard_message_shared( } } ShardMessage::UnregisterReplica { replica_id } => { - replica_txs.retain(|(id, _)| *id != replica_id); + replica_txs.retain(|(id, _, _)| *id != replica_id); } ShardMessage::ReplicaLiveFanout { bytes } => { // Live-delivery leg ONLY: backlog append + offset advance already // happened synchronously at write time on this same thread // (`record_local_write`) — doing either again here would double- - // count. Lagging replicas are skipped (try_send), same policy as - // `wal_append_and_fanout`'s fan-out leg. - for (_id, tx) in replica_txs.iter() { - let _ = tx.try_send(bytes.clone()); - } + // count. A replica whose channel is FULL is KICKED (task #35): + // skipping the record would silently and permanently diverge it. + fanout_send_or_kick(replica_txs, &bytes); } ShardMessage::MigrateConnection(_) => { // MigrateConnection is collected by drain_spsc_shared into pending_migrations, @@ -3365,13 +3384,46 @@ pub(crate) fn cow_intercept( #[inline] pub(crate) fn wal_fanout_has_work( wal_writer: &Option, - replica_txs: &[(u64, channel::MpscSender)], + replica_txs: &[crate::shard::dispatch::ReplicaFanout], aof_pool: Option<&std::sync::Arc>, wal_kv_log: bool, ) -> bool { (wal_kv_log && wal_writer.is_some()) || !replica_txs.is_empty() || aof_pool.is_some() } +/// Deliver one replication record to every registered replica, KICKING any +/// replica whose bounded channel is full (task #35). +/// +/// The old policy — `try_send` and skip on `Full` — silently dropped records +/// for a lagging replica, leaving a permanent gap in its stream while +/// `master_link_status` stayed "up" (observed: 2k of 40k keys delivered under +/// pipelined load, offsets diverged forever, WAIT correctly reported 0). +/// A replica that cannot keep up must instead be disconnected so it retries +/// PSYNC and resyncs from the backlog — Redis's output-buffer-limit policy. +/// The kick is two-stage because the drain task and `ReplicaInfo.shard_txs` +/// hold sender clones (dropping our entry alone cannot close the channel): +/// set the shared `kicked` flag (the drain task polls it and closes the +/// socket), then drop our entry so this shard stops queueing immediately. +pub(crate) fn fanout_send_or_kick( + replica_txs: &mut Vec, + bytes: &bytes::Bytes, +) { + replica_txs.retain(|(id, tx, kicked)| match tx.try_send(bytes.clone()) { + Ok(()) => true, + Err(flume::TrySendError::Full(_)) => { + kicked.store(true, std::sync::atomic::Ordering::Release); + tracing::warn!( + replica_id = id, + "replica live fan-out channel FULL — kicking replica to force a \ + resync (a skipped record would silently diverge it forever)" + ); + false + } + // Drain task already gone; just stop queueing. + Err(flume::TrySendError::Disconnected(_)) => false, + }); +} + /// Error frame substituted for a write's success frame when the command /// mutated memory but its AOF record could not be enqueued within the /// backpressure budget — fail-loud so the client knows durability was not @@ -3388,9 +3440,12 @@ pub(crate) const AOF_APPEND_LOST_ERR: &[u8] = /// per drain arm, not per command. pub(crate) fn wal_append_and_fanout( data: &[u8], + // task #35: db the command executed in — threaded into the AOF pool so + // the writer can inject a `SELECT ` record on a db-context change. + db: usize, wal_writer: &mut Option, repl_backlog: &crate::replication::backlog::SharedBacklog, - replica_txs: &[(u64, channel::MpscSender)], + replica_txs: &mut Vec, repl_state: &Option, shard_id: usize, aof_pool: Option<&std::sync::Arc>, @@ -3435,12 +3490,11 @@ pub(crate) fn wal_append_and_fanout( if let Some(offsets) = repl_state { offsets.increment_shard_offset(shard_id, data.len() as u64); } - // 4. Fan-out to replica sender tasks (non-blocking: lagging replicas are skipped) + // 4. Fan-out to replica sender tasks (non-blocking: a replica whose + // channel is FULL is kicked to resync — see `fanout_send_or_kick`). if !replica_txs.is_empty() { let bytes = bytes::Bytes::copy_from_slice(data); - for (_id, tx) in replica_txs { - let _ = tx.try_send(bytes.clone()); - } + fanout_send_or_kick(replica_txs, &bytes); } // 5. Per-shard AOF pool (FIX-W1-2): route to the owning shard's writer. // Bounded-blocking (`send_append_bounded_blocking`) because this function @@ -3456,6 +3510,7 @@ pub(crate) fn wal_append_and_fanout( return pool.send_append_bounded_blocking( shard_id, 0, + db, bytes::Bytes::copy_from_slice(data), aof_budget, ); @@ -3498,10 +3553,11 @@ mod wal_append_tests { wal_append_and_fanout( b"hello", + 0, // db &mut None, // no writer &backlog, - &[], // no replicas - &None, // no repl_state + &mut vec![], // no replicas + &None, // no repl_state 0, None, // no aof_pool true, // wal_kv_log @@ -3523,13 +3579,18 @@ mod wal_append_tests { let backlog: SharedBacklog = std::sync::Arc::new(parking_lot::Mutex::new(Some(ReplicationBacklog::new(1024)))); let (tx, _rx) = crate::runtime::channel::mpsc_unbounded::(); - let replica_txs = vec![(1u64, tx)]; + let mut replica_txs: Vec = vec![( + 1u64, + tx, + std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), + )]; wal_append_and_fanout( b"hello", + 0, // db &mut None, &backlog, - &replica_txs, + &mut replica_txs, &None, 0, None, // no aof_pool @@ -3565,9 +3626,10 @@ mod wal_append_tests { wal_append_and_fanout( b"world", + 0, // db &mut None, // no writer &backlog, - &[], // no replicas — S3.5b bypass triggered without pool guard + &mut vec![], // no replicas — S3.5b bypass triggered without pool guard &None, // no repl_state 0, // shard_id Some(&pool), // aof_pool provided — bypass must NOT fire @@ -3629,13 +3691,14 @@ mod wal_append_tests { // Pre-fix this was `aof_pool` (Some), which caused the double-write. wal_append_and_fanout( b"*3\r\n$3\r\nSET\r\n$1\r\na\r\n$1\r\n1\r\n", + 0, // db &mut None, // no writer &backlog, - &[], // no replicas - &None, // no repl_state - 0, // shard_id - None, // PipelineBatch fix: None prevents double-write - true, // wal_kv_log + &mut vec![], // no replicas + &None, // no repl_state + 0, // shard_id + None, // PipelineBatch fix: None prevents double-write + true, // wal_kv_log &mut std::time::Duration::from_millis(5), ); assert!( @@ -3654,9 +3717,10 @@ mod wal_append_tests { // cross-shard MSET/DEL/EXISTS commands). wal_append_and_fanout( b"*3\r\n$4\r\nMSET\r\n$1\r\nb\r\n$1\r\n2\r\n", + 0, // db &mut None, &backlog, - &[], + &mut vec![], &None, 0, Some(&pool), // MultiExecute: pool must receive this entry @@ -3701,9 +3765,10 @@ mod wal_append_tests { wal_append_and_fanout( cmd, + 0, // db &mut w3, &backlog, - &[], + &mut vec![], &None, 0, Some(&pool), @@ -3725,9 +3790,10 @@ mod wal_append_tests { // the record must land in the WAL as before. wal_append_and_fanout( cmd, + 0, // db &mut w3, &backlog, - &[], + &mut vec![], &None, 0, Some(&pool), diff --git a/tests/aof_multidb_kill9.rs b/tests/aof_multidb_kill9.rs new file mode 100644 index 00000000..77aaec70 --- /dev/null +++ b/tests/aof_multidb_kill9.rs @@ -0,0 +1,283 @@ +//! AOF multi-db attribution under interleaved connections (durability gate). +//! +//! Red condition this locks in: the live AOF used to persist the CLIENT's +//! literal `SELECT` commands (SELECT is W-flagged) while bare writes carried +//! no db context of their own. Two interleaved connections on different dbs +//! then corrupted recovery: `conn A (db0): SET a1` / `conn B: SELECT 2, SET +//! b1` / `conn A: SET a2` replays `a2` into db2 — the stream's last literal +//! SELECT wins, not the db the master actually executed in. Caught by the +//! v0.7 R1 consistency/durability gates (task #35); pre-existing since the +//! first multi-db AOF. +//! +//! The fix mirrors Redis: the AOF WRITER tracks the stream's current db and +//! prepends `SELECT ` whenever a record's execution db differs; client +//! SELECT commands are connection-state only and never persisted. +//! +//! Black-box over a real `moon` process (needs a prebuilt binary): +//! +//! ```text +//! MOON_BIN=./target/release/moon \ +//! cargo test --test aof_multidb_kill9 -- --ignored --nocapture +//! ``` + +#![cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] + +use std::io::{BufRead, BufReader, Write}; +use std::net::TcpStream; +use std::process::{Child, Command, Stdio}; +use std::time::Duration; + +fn moon_bin() -> String { + std::env::var("MOON_BIN").unwrap_or_else(|_| "./target/release/moon".to_string()) +} + +fn start_moon(port: u16, dir: &str, shards: usize) -> Child { + Command::new(moon_bin()) + .args([ + "--port", + &port.to_string(), + "--shards", + &shards.to_string(), + "--dir", + dir, + "--appendonly", + "yes", + "--disk-free-min-pct", + "0", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("Failed to start moon (set MOON_BIN to a built binary)") +} + +fn sigkill(child: &mut Child) { + #[cfg(unix)] + { + let pid = child.id() as i32; + // SAFETY: kill(2) with a pid we own from Child::id; SIGKILL is the + // whole point of the crash test. + unsafe { + libc::kill(pid, libc::SIGKILL); + } + let _ = child.wait(); + } + #[cfg(not(unix))] + { + let _ = child.kill(); + let _ = child.wait(); + } +} + +/// One persistent connection; sends inline commands and reads one reply each. +struct Conn { + reader: BufReader, +} + +impl Conn { + fn connect(addr: &str) -> Self { + let stream = TcpStream::connect(addr).expect("connect"); + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("timeout"); + Conn { + reader: BufReader::new(stream), + } + } + + fn cmd(&mut self, line: &str) -> String { + self.reader + .get_mut() + .write_all(format!("{}\r\n", line).as_bytes()) + .expect("write"); + let mut reply = String::new(); + self.reader.read_line(&mut reply).expect("read"); + // Consume the bulk payload line for $-replies. + if reply.starts_with('$') && !reply.starts_with("$-1") { + let mut payload = String::new(); + self.reader.read_line(&mut payload).expect("read payload"); + return payload.trim_end().to_string(); + } + reply.trim_end().to_string() + } +} + +fn wait_ready(addr: &str) { + for _ in 0..100 { + if let Ok(stream) = TcpStream::connect(addr) { + drop(stream); + let mut c = Conn::connect(addr); + if c.cmd("PING") == "+PONG" { + return; + } + } + std::thread::sleep(Duration::from_millis(200)); + } + panic!("server at {addr} never became ready"); +} + +fn unique_port() -> u16 { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind ephemeral"); + let port = listener.local_addr().expect("local_addr").port(); + drop(listener); + port +} + +/// Interleave db0/db2 writes from two connections, kill -9, restart, assert +/// every key recovered into the db it was written to. +fn run_case(shards: usize) { + let port = unique_port(); + let addr = format!("127.0.0.1:{port}"); + let dir = tempfile::tempdir().expect("tempdir"); + let dir_path = dir.path().to_str().expect("utf8 dir").to_string(); + + let mut server = start_moon(port, &dir_path, shards); + wait_ready(&addr); + + let mut conn_a = Conn::connect(&addr); // stays on db0 + let mut conn_b = Conn::connect(&addr); + assert_eq!(conn_b.cmd("SELECT 2"), "+OK"); + + // Interleave so the AOF stream alternates db context between records. + for i in 0..20 { + assert_eq!(conn_a.cmd(&format!("SET a{i} v{i}")), "+OK"); + assert_eq!(conn_b.cmd(&format!("SET b{i} v{i}")), "+OK"); + } + // The poison pattern from the gate: a db0 write AFTER db2 activity. + assert_eq!(conn_a.cmd("SET tail v-tail"), "+OK"); + + // everysec fsync window. + std::thread::sleep(Duration::from_millis(1600)); + sigkill(&mut server); + + let mut server = start_moon(port, &dir_path, shards); + wait_ready(&addr); + + let mut c0 = Conn::connect(&addr); + let mut c2 = Conn::connect(&addr); + assert_eq!(c2.cmd("SELECT 2"), "+OK"); + + let db0 = c0.cmd("DBSIZE"); + let db2 = c2.cmd("DBSIZE"); + // Spot keys BEFORE the sizes so a failure message shows placement. + assert_eq!( + c0.cmd("GET tail"), + "v-tail", + "db0 tail write recovered into the wrong db (shards={shards}; db0={db0}, db2={db2})" + ); + assert_eq!(c0.cmd("GET a0"), "v0", "a0 missing from db0"); + assert_eq!(c0.cmd("GET a19"), "v19", "a19 missing from db0"); + assert_eq!(c2.cmd("GET b0"), "v0", "b0 missing from db2"); + assert_eq!(c2.cmd("GET b19"), "v19", "b19 missing from db2"); + assert_eq!(c2.cmd("GET a5"), "$-1", "db0 key leaked into db2"); + assert_eq!(c0.cmd("GET b5"), "$-1", "db2 key leaked into db0"); + assert_eq!(db0, ":21", "db0 size after recovery (shards={shards})"); + assert_eq!(db2, ":20", "db2 size after recovery (shards={shards})"); + + sigkill(&mut server); +} + +/// TopLevel AOF writer (shards=1). +#[test] +#[ignore] +fn aof_multidb_kill9_recovers_toplevel() { + run_case(1); +} + +/// PerShard framed AOF writers (shards=2) — exercises the SPSC/wal_append +/// fanout leg's db threading as well. +#[test] +#[ignore] +fn aof_multidb_kill9_recovers_per_shard() { + run_case(2); +} + +/// PR #282 review (CodeRabbit): a SELECT queued INSIDE MULTI/EXEC. The txn +/// executors persisted the literal queued SELECT (they gated on `is_write`, +/// missed by the top-level `is_persisted_write` sweep) and tagged every txn +/// AOF entry with one collapsed db. Poison: the literal `SELECT 2` shifts the +/// replay stream to db2, but the next top-level db0 write is tagged 0 == the +/// writer's tracker 0, so no corrective SELECT is injected — post-EXEC db0 +/// writes recover into db2. Fix: txn entries carry the db each command +/// actually executed in; queued SELECT is never persisted. +fn run_txn_case(shards: usize) { + let port = unique_port(); + let addr = format!("127.0.0.1:{port}"); + let dir = tempfile::tempdir().expect("tempdir"); + let dir_path = dir.path().to_str().expect("utf8 dir").to_string(); + + let mut server = start_moon(port, &dir_path, shards); + wait_ready(&addr); + + let mut conn = Conn::connect(&addr); + assert_eq!(conn.cmd("MULTI"), "+OK"); + assert_eq!(conn.cmd("SET t0 v0"), "+QUEUED"); + assert_eq!(conn.cmd("SELECT 2"), "+QUEUED"); + assert_eq!(conn.cmd("SET t2 v2"), "+QUEUED"); + let exec = conn.cmd("EXEC"); + assert!(exec.starts_with('*'), "EXEC failed: {exec}"); + // Drain the EXEC array body (4 reply lines: +OK, +OK, +OK — and the + // queued SET replies; read until the connection would block is overkill — + // reconnect instead for the follow-up write). + drop(conn); + + // The poison probe: a plain db0 write AFTER the txn. + let mut conn = Conn::connect(&addr); + assert_eq!(conn.cmd("SET post v-post"), "+OK"); + + // Where did the txn writes land in MEMORY? Record it — recovery must + // match memory exactly, whatever the executor's SELECT semantics. + let mem_t2_db0 = conn.cmd("GET t2"); + let mut conn2 = Conn::connect(&addr); + assert_eq!(conn2.cmd("SELECT 2"), "+OK"); + let mem_t2_db2 = conn2.cmd("GET t2"); + // PR #282 review: pin the baseline — a silently-failed `SET t2` would + // leave `$-1` in BOTH dbs and every recovery assertion would then pass + // vacuously. + assert!( + (mem_t2_db0 == "v2" && mem_t2_db2 == "$-1") || (mem_t2_db0 == "$-1" && mem_t2_db2 == "v2"), + "t2 must exist in exactly one database: db0={mem_t2_db0}, db2={mem_t2_db2}" + ); + + std::thread::sleep(Duration::from_millis(1600)); + sigkill(&mut server); + + let mut server = start_moon(port, &dir_path, shards); + wait_ready(&addr); + + let mut c0 = Conn::connect(&addr); + let mut c2 = Conn::connect(&addr); + assert_eq!(c2.cmd("SELECT 2"), "+OK"); + + assert_eq!( + c0.cmd("GET post"), + "v-post", + "post-EXEC db0 write recovered into the wrong db (shards={shards})" + ); + assert_eq!(c2.cmd("GET post"), "$-1", "post leaked into db2"); + assert_eq!(c0.cmd("GET t0"), "v0", "txn db0 write lost/misplaced"); + assert_eq!( + c0.cmd("GET t2"), + mem_t2_db0, + "t2 recovery placement in db0 diverges from memory placement" + ); + assert_eq!( + c2.cmd("GET t2"), + mem_t2_db2, + "t2 recovery placement in db2 diverges from memory placement" + ); + + sigkill(&mut server); +} + +#[test] +#[ignore] +fn aof_txn_select_inside_multi_kill9_toplevel() { + run_txn_case(1); +} + +#[test] +#[ignore] +fn aof_txn_select_inside_multi_kill9_per_shard() { + run_txn_case(2); +} diff --git a/tests/replication_streaming.rs b/tests/replication_streaming.rs index a161a7aa..adb2b230 100644 --- a/tests/replication_streaming.rs +++ b/tests/replication_streaming.rs @@ -763,3 +763,208 @@ fn replica_applies_multi_db_stream() { "db-0 write leaked into db 2 on the replica" ); } + +/// R1 (task #19): real WAIT/ACK plumbing. The replica sends +/// `REPLCONF ACK ` on the replication link (after each applied batch +/// + a 1s idle tick); the master reads them off the hijacked PSYNC socket and +/// records them per replica; `WAIT ` blocks until n replicas +/// acknowledge the master's current offset. Previously WAIT returned 0 +/// unconditionally on the monoio runtime and no ACK was ever sent or read. +#[test] +#[ignore] +fn wait_returns_acked_replica_count() { + let master_dir = tempfile::tempdir().unwrap(); + let replica_dir = tempfile::tempdir().unwrap(); + + let master_addr = "127.0.0.1:16736"; + let replica_addr = "127.0.0.1:16737"; + + let _master = Killer(start_moon(16736, master_dir.path().to_str().unwrap())); + assert!( + wait_until(Duration::from_secs(5), || send_cmd(master_addr, "PING") + .starts_with("+PONG")), + "master never became ready" + ); + + // No replicas attached: WAIT 1 with a short timeout must return 0 fast. + let r = send_cmd(master_addr, "WAIT 1 200"); + assert_eq!(r.trim(), ":0", "WAIT with no replicas must return 0: {r}"); + + let _replica = Killer(start_moon(16737, replica_dir.path().to_str().unwrap())); + assert!( + wait_until(Duration::from_secs(5), || send_cmd(replica_addr, "PING") + .starts_with("+PONG")), + "replica never became ready" + ); + send_cmd(replica_addr, &format!("REPLICAOF 127.0.0.1 {}", 16736)); + + // Prove the stream is live. + send_cmd(master_addr, "SET w1 v1"); + assert!( + wait_until(Duration::from_secs(10), || { + get(replica_addr, "w1").as_deref() == Some("v1") + }), + "live stream not flowing — WAIT assertions would be meaningless" + ); + + // Write, then WAIT for 1 replica to acknowledge everything so far. The + // ACK cadence is drain-driven + a 1s idle tick, so a 3s budget is ample. + send_cmd(master_addr, "SET w2 v2"); + let r = send_cmd(master_addr, "WAIT 1 3000"); + assert_eq!( + r.trim(), + ":1", + "WAIT must observe the replica's ACK of the current offset: {r}" + ); + + // Asking for MORE replicas than exist must time out and report the real + // acked count (1), not hang or claim 2. + let r = send_cmd(master_addr, "WAIT 2 300"); + assert_eq!(r.trim(), ":1", "WAIT 2 with one replica must return 1: {r}"); +} + +/// Task #35 (v0.7 consistency gate): interleaved multi-connection multi-db +/// load must CONVERGE on the replica. Locks in two pre-existing defects the +/// R1 gates caught: +/// +/// 1. Client `SELECT` commands leaked into the replication stream as literal +/// records (SELECT is W-flagged), desyncing the master's `stream_db` +/// tracking — a db-0 write after another connection's SELECT 2 handshake +/// applied into the replica's db 2. +/// 2. `wal_append_and_fanout` / `ReplicaLiveFanout` silently DROPPED records +/// when a replica's bounded fan-out channel filled under pipelined load — +/// the replica stayed "up" but permanently diverged (observed 2k of 40k +/// keys, link_status:up). The fix kicks the lagging replica so it +/// reconnects and resyncs — divergence becomes a loud, self-healing +/// resync, never silence. +#[test] +#[ignore] +fn replica_converges_under_interleaved_multidb_load() { + let master_dir = tempfile::tempdir().unwrap(); + let replica_dir = tempfile::tempdir().unwrap(); + + let master_addr = "127.0.0.1:16738"; + let replica_addr = "127.0.0.1:16739"; + + let _master = Killer(start_moon(16738, master_dir.path().to_str().unwrap())); + assert!( + wait_until(Duration::from_secs(5), || send_cmd(master_addr, "PING") + .starts_with("+PONG")), + "master never became ready" + ); + let _replica = Killer(start_moon(16739, replica_dir.path().to_str().unwrap())); + assert!( + wait_until(Duration::from_secs(5), || send_cmd(replica_addr, "PING") + .starts_with("+PONG")), + "replica never became ready" + ); + send_cmd(replica_addr, "REPLICAOF 127.0.0.1 16738"); + + // Prove the stream is up before loading. + send_cmd(master_addr, "SET warm 1"); + assert!( + wait_until(Duration::from_secs(10), || { + get(replica_addr, "warm").as_deref() == Some("1") + }), + "live stream not flowing — load assertions would be meaningless" + ); + + // Two persistent connections: conn0 stays on db 0; conn2 SELECTs db 2 (a + // literal client SELECT — the record that used to poison the stream). + let mut conn0 = TcpStream::connect(master_addr).expect("conn0"); + let mut conn2 = TcpStream::connect(master_addr).expect("conn2"); + conn0 + .set_read_timeout(Some(Duration::from_secs(30))) + .unwrap(); + conn2 + .set_read_timeout(Some(Duration::from_secs(30))) + .unwrap(); + pipeline_burst(&mut conn2, &["SELECT 2".to_string()]); + + // Interleaved pipelined bursts, large enough to overflow a bounded + // per-replica fan-out channel repeatedly. + const BATCH: usize = 2000; + const BATCHES: usize = 5; + for b in 0..BATCHES { + let cmds0: Vec = (0..BATCH) + .map(|i| format!("SET k0:{} v{}", b * BATCH + i, b * BATCH + i)) + .collect(); + let cmds2: Vec = (0..BATCH) + .map(|i| format!("SET k2:{} v{}", b * BATCH + i, b * BATCH + i)) + .collect(); + pipeline_burst(&mut conn0, &cmds0); + pipeline_burst(&mut conn2, &cmds2); + } + drop(conn0); + drop(conn2); + + let total = (BATCH * BATCHES) as i64; + assert_eq!(dbsize(master_addr), total + 1, "master db0 load incomplete"); + + // Convergence: the replica may be kicked + resync mid-load; give it a + // generous settle window but require exact parity at the end. + let converged = wait_until(Duration::from_secs(60), || { + dbsize_in_db(replica_addr, 0) == total + 1 && dbsize_in_db(replica_addr, 2) == total + }); + let (r0, r2) = (dbsize_in_db(replica_addr, 0), dbsize_in_db(replica_addr, 2)); + assert!( + converged, + "replica diverged after interleaved multi-db load: db0={r0}/{} db2={r2}/{}", + total + 1, + total + ); + // Spot checks: correct placement, no cross-db leaks. + assert_eq!( + get_in_db(replica_addr, 0, "k0:9999").as_deref(), + Some("v9999"), + "db0 tail key wrong" + ); + assert_eq!( + get_in_db(replica_addr, 2, "k2:9999").as_deref(), + Some("v9999"), + "db2 tail key wrong" + ); + assert_eq!( + get_in_db(replica_addr, 2, "k0:42"), + None, + "db0 key leaked into db2 (SELECT poisoning)" + ); + assert_eq!( + get_in_db(replica_addr, 0, "k2:42"), + None, + "db2 key leaked into db0" + ); +} + +/// Write `cmds` as one pipelined burst on a persistent connection and read +/// exactly one reply line per command (inline protocol; replies here are all +/// +OK/+PONG-class single lines). +fn pipeline_burst(stream: &mut TcpStream, cmds: &[String]) { + let mut payload = String::with_capacity(cmds.len() * 24); + for c in cmds { + payload.push_str(c); + payload.push_str("\r\n"); + } + stream.write_all(payload.as_bytes()).expect("burst write"); + stream.flush().ok(); + let mut reader = BufReader::new(&*stream); + for i in 0..cmds.len() { + let mut line = String::new(); + reader.read_line(&mut line).expect("burst reply"); + assert!( + line.starts_with("+OK"), + "burst cmd {i} ({}) got: {line}", + cmds[i] + ); + } +} + +/// DBSIZE in a specific db over one connection (SELECT + DBSIZE). +fn dbsize_in_db(addr: &str, db: usize) -> i64 { + let reply = send_seq(addr, &[&format!("SELECT {db}"), "DBSIZE"]); + reply + .rsplit(':') + .next() + .and_then(|s| s.trim().parse::().ok()) + .unwrap_or(-1) +} diff --git a/tests/wal_group_commit.rs b/tests/wal_group_commit.rs index 89578ef8..fe39b541 100644 --- a/tests/wal_group_commit.rs +++ b/tests/wal_group_commit.rs @@ -34,6 +34,7 @@ use std::collections::VecDeque; fn append(data: &[u8]) -> AofMessage { AofMessage::Append { lsn: 0, + db: 0, bytes: Bytes::copy_from_slice(data), } } @@ -43,6 +44,7 @@ fn append_sync(data: &[u8]) -> (AofMessage, OneshotReceiver) { ( AofMessage::AppendSync { lsn: 0, + db: 0, bytes: Bytes::copy_from_slice(data), ack: tx, },