feat(replication): R1 real WAIT/ACK + consistency/durability fixes (silent replica drop, SELECT leak, AOF db attribution)

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 <db>` 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 <numreplicas> <timeout>` 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 <offset>` — 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 <offset>` 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

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 <db>` 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 {

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));
             }

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 <db>` 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<AofAck>,
     },
@@ -503,6 +513,75 @@ pub fn serialize_command(frame: &Frame) -> Bytes {
     buf.freeze()
 }
 
+/// Serialized `SELECT <db>` 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<Bytes> {
+    if payload_is_empty || db == *last_db {
+        return None;
+    }
+    *last_db = db;
+    Some(serialize_select_record(db))
+}
+
+/// Insert a synthetic `SELECT <db>` [`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<AofMessage>, last_db: &mut usize) -> Vec<AofMessage> {
+    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