diff --git a/CHANGELOG.md b/CHANGELOG.md index cfceb01e..39fd904b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,211 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed — Conn-plane: monoio central listener joins the SO_REUSEPORT group (PR #250) + +- Under the monoio runtime, each shard binds `bind:port` via `SO_REUSEPORT`, but + the central listener bound the same port with a plain + `monoio::net::TcpListener::bind` — unlike the tokio sibling, which uses + `create_reuseport_socket`. Linux requires every socket sharing a REUSEPORT port + to set the option, so depending on startup scheduling the central bind either + lost the race (EADDRINUSE propagated out of `run_sharded`, silently killing the + TLS listener + `conn_rx` fallback) or won it (forcing every shard's REUSEPORT + bind to fail and collapsing all traffic onto the single central accept loop) — + both silent and nondeterministic run-to-run. The monoio central listener now + binds via `create_reuseport_socket` + `from_std` like the tokio path (plain-bind + fallback on non-unix / unparseable addr). (audit finding 16) + +### Fixed — Durability: checkpoint Finalize backs off on repeated failure instead of flooding the WAL (PR #250) + +- When a checkpoint's Finalize step failed (`wal.wait_durable`, `manifest.commit`, + `graph_save`, or the control-file write), the state machine stayed `Finalizing` + and the next 1ms persistence tick re-ran finalize from step 1 — re-appending a + WAL Checkpoint record every single millisecond with no backoff. Under a + sustained failure (slow/degraded disk) this floods the WAL, and since + `last_checkpoint_lsn` never advances, `recycle_segments_before` never fires — + WAL disk usage grows fastest exactly during the disk-pressure incident it + should be backing off from. `CheckpointManager` now arms an exponential backoff + (50ms→5s cap) on each failed finalize; the tick checks `finalize_ready(now)` + before doing any finalize I/O, so retries are bounded instead of per-tick, and + a successful `complete()` clears it. (audit finding 13) + +### Fixed — Durability: legacy AOF replay stops at mid-stream corruption instead of resyncing (PR #250) + +- The default startup path `replay_aof` (single-file `appendonly.aof`, unframed + RESP with no per-record length/CRC) treated ANY mid-stream parse error by + discarding one byte, scanning forward to the next `*`, and resuming dispatch — + but a `*` appears freely inside binary value blobs, so the resync could land + mid-value and execute misaligned garbage as live SET/DEL/etc. against the + recovering keyspace (silent data corruption, WARN-only). The per-shard + `shard_replay` sibling already rejects exactly this. `replay_aof` now STOPS at + the first genuine mid-stream corruption (clean tail-truncation is still handled + gracefully), keeping only the valid prefix already applied and never + dispatching past the corruption point; it logs a loud error with the byte + offset and points at `redis-check-aof`. Operators who want the old best-effort + behavior can opt in with `MOON_AOF_BEST_EFFORT_RESYNC=1`. (audit finding 12) + +### Fixed — Xshard: bound cross-shard reply wait so a wedged shard can't hang a client (PR #250) + +- Every cross-shard leg (MGET/MSET/DEL/UNLINK/EXISTS/BITOP/COPY/MSETNX remote + legs, MULTI-on-owner, and the flush/scatter-gather loops) pushed a message via + the backoff-retried `spsc_send` and then did `reply_rx.recv().await` with **no + timeout**. `spsc_send`'s retry budget only bounds getting the message INTO the + target ring buffer; if the push succeeds but the target shard then stalls while + executing (wedged-disk fsync during snapshot/AOF, uninterruptible D-state I/O, + or a dead shard), the awaiting connection parked forever with no response and + no way to cancel. Added a `recv_reply_bounded` helper (races the receiver + against a runtime-agnostic 30s `XSHARD_REPLY_TIMEOUT` via `race2`) and applied + it to all 11 reply-await sites in the coordinator — a genuinely wedged shard + now surfaces the existing cross-shard-reply error instead of hanging. (audit finding 11) + +### Fixed — Hardening: bound TLS handshake + cluster-bus body reads against slow-loris (PR #250) + +- **TLS handshake has no timeout (#17):** after `try_accept_connection` consumed + a `maxclients` slot, `acceptor.accept(tcp_stream).await` ran with no deadline + on both the monoio and tokio TLS paths. A peer that completed the TCP handshake + on the TLS port and then sent no (or a partial) ClientHello pinned the task, + its fd, and its `maxclients` slot indefinitely at zero CPU cost — a classic + slow-loris that could exhaust the connection limit. Wrapped both accepts in a + 10s `TLS_HANDSHAKE_TIMEOUT`, releasing the slot via `record_connection_closed()` + on expiry. +- **Cluster-bus gossip body read has no timeout (#10):** `handle_cluster_peer` + guarded only the 4-byte length-prefix read with a shutdown-cancel select; once + a valid length was read, the body `read_exact` had no timeout and no + shutdown arm. A client could send a valid length header and never send the + body, blocking the spawned task forever (immune to graceful shutdown, socket + never reclaimed). Both runtimes now read the body under a + `GOSSIP_BODY_READ_TIMEOUT` (10s) + shutdown-cancel select. + +### Fixed — Txn: MULTI locality analysis honors SORT/GEORADIUS STORE destination (PR #250) + +- In a multi-shard deployment, `MULTI; SORT src STORE dst; EXEC` (or + `GEORADIUS ... STORE/STOREDIST dst`) where `src` and `dst` hash to different + shards was misclassified as single-shard and routed entirely to `src`'s + shard — writing `dst` into the WRONG shard's dataset, invisible to a later + normally-routed `GET dst`. The STORE/STOREDIST destination is positional + (the arg after the token) and so was not covered by the fixed command-metadata + key specs `analyze_txn_locality` consulted. It now detects the STORE clause + and forces a CrossShard classification, which the caller rejects with + CROSSSLOT instead of silently misrouting. (audit finding 15) + +### Fixed — Vector: DEL tombstones secondary fields + FT.INFO num_docs accuracy (PR #250) + +- **Multi-vector-field deletion resurrection (#20):** `DEL`/`UNLINK`/`HDEL` on a + document only tombstoned the *default* vector field. Secondary VECTOR fields + live in a separate `field_segments` map that `tombstone_key_in_index` never + iterated, so a subsequent `FT.SEARCH idx '@field2:[...]'` still returned the + deleted document (under a synthetic `vec:` key, since the shared + key-hash→key map was cleared). Now every field's segments are tombstoned. +- **FT.INFO `num_docs` over/under-count (#28):** the mutable segment's `len()` + counts tombstoned entries in place until compaction, inflating `num_docs` by + up to `compact_threshold` under DEL/HSET churn — switched to a new `live_len()` + that excludes tombstones. Separately, DiskANN cold-tier segments (experimental, + gated by `MOON_VEC_COLD_TIER`) were never summed at all even though cold docs + ARE returned by FT.SEARCH; added a `snap.cold` counting loop. (audit findings 20, 28) + +### Fixed — Vector: GraphUnion merge recall gate no longer lets a total collapse through (PR #250) + +- The background/manual GraphUnion merge recall gate read + `recall < tolerance && recall > 0.0`, so a merge whose verified recall was + **exactly 0.0** (total collapse) slipped past the guard and committed. Since + `verify_merge_recall` returns 1.0 (not 0.0) for the too-few-vectors cases, + 0.0 only ever means a real measured collapse — the `&& recall > 0.0` clause + was removed so such a merge now aborts. (audit finding 21) + +### Fixed — Hardening: APPEND 512MB cap + reject FT.* inside MULTI on prod handlers (PR #250) + +- **APPEND** now enforces the 512MB max-string-size limit (matching + SETRANGE/SETBIT); repeated APPENDs previously grew a value past the documented + limit unbounded. (audit finding 22) +- **FT.\*** commands are now explicitly rejected inside MULTI/EXEC on the + production sharded + monoio handlers (they aren't wired through the txn + execution path), matching the existing handler_single guard — a clear + "not supported inside MULTI/EXEC" error instead of an incidental one. + (audit finding 25) + +### Changed — Vector storage: fence experimental DiskANN cold tier + drop dead vector WAL (PR #250) + +- The DiskANN cold tier (WARM→COLD segment demotion + on-disk Vamana beam + search) was running by default (`--disk-offload` defaults on, `--segment-cold-after` + defaults 86400) despite being incomplete — **no cold-segment deletion** + (DEL/HDEL never removes cold docs), blocking per-hop uring I/O holding the + segment mutex, `FT.INFO num_docs` under-counting cold segments, unfinished + restart reload of PQ codebooks, and a library-code `panic!` in + `DiskAnnSegment::new`. It is now **fenced behind the experimental + `MOON_VEC_COLD_TIER=1` env gate** (default off); the WARM→COLD transition is + a no-op otherwise. Warm-tier mmap + LRU eviction (`--vec-warm-mmap-budget`) + continues to handle out-of-RAM indexes. (audit findings 19, 24, 28) +- `DiskAnnSegment::new` now returns `io::Result` instead of panicking on a + file-open error, so a cold transition can never abort the shard. +- Removed the dead vector WAL module (`vector/persistence/wal_record.rs`): it + was never on the live recovery path (superseded by manifest + segment + + keymap + dedup rescan) and carried a documented double-apply footgun. + +### Fixed — Cluster: gossip PING task leak + non-blocking accept send (PR #250) + +- The gossip PING ticker (both tokio + monoio) spawned an unbounded + connect+write+read task every 100ms with **no timeout**; against a dead or + partitioned peer the tasks piled up forever (task/FD/memory leak), and the + "random peer" it documented was always the first non-self node. Now a single + in-flight guard bounds outstanding probes to one, a probe timeout + (`node_timeout/2`, ≥100ms) cancels a hung connect/read, and targets rotate + across peers. (audit finding 7) +- The monoio per-shard accept task used the **blocking** `flume::Sender::send()` + on a bounded channel; because the accept task shares the shard's single + monoio thread with the event-loop consumer, a full channel deadlocked the + loop that drains it under connection-storm churn. Fixed by `send_async().await` + (cooperative yield, keeps backpressure). (audit finding 8, Batch B) + +### Fixed — Security: bound client-controlled allocation counts (DoS class, PR #250) + +- Six wire-reachable command paths fed a client-supplied count straight into + `Vec::with_capacity` / `Vec::resize` with no upper bound. On a real host the + oversized request fails allocation and Rust's `handle_alloc_error` calls + `abort()` — uncatchable, crashing every shard and connection. A single + unauthenticated request was a full-process DoS. Fixed by clamping each count + before allocating: **FT.SEARCH HIGHLIGHT/SUMMARIZE FIELDS** (bound by + remaining tokens), **HSCAN COUNT** (`count.min(total)`), + **SRANDMEMBER / HRANDFIELD / ZRANDMEMBER** negative COUNT (absolute 2²⁰ cap + with a loud `ERR COUNT is out of range` beyond it — the first-cut relative + `len()*10` cap silently truncated legitimate requests on small collections, + breaking Redis's exact-|COUNT|-with-duplicates contract; caught in review and + fixed at all six sites including the pre-existing HRANDFIELD/ZRANDMEMBER + ones), **BITFIELD SET/INCRBY** offset (reject past the 512MB limit, matching + SETBIT; GET exempt), and **ZMPOP COUNT** (`pop_count.min(card)` — this last + site was missed by the audit finders and caught by the allocation sweep). + Each fix has a red/green test. From the production-hardening audit (Batch A). + +### Fixed — Review follow-ups on the hardening sweep (PR #250) + +- **Gossip PONG fragmentation (monoio):** the failure-detector's PING probe read + the 4-byte length and the PONG body with single `read()` calls, so a valid + frame arriving fragmented was silently dropped — leaving `pong_recv_ms` stale + and pushing a healthy peer toward PFAIL. Both reads now use the bus listener's + exact-read loop, and both runtimes' probes cap the peer-supplied PONG length + at the bus's 64KB frame bound before sizing the buffer (the tokio probe + previously allocated an unchecked `vec![0u8; pong_len]`). +- **AOF corruption offsets are file-absolute:** replay's truncation/corruption + logs reported offsets relative to the RESP section; with an RDB preamble the + operator-facing repair guidance (redis-check-aof, truncation point) pointed at + the wrong file location. Offsets now include the preamble length. +- **Cross-shard MSET no longer OKs an unconfirmed leg:** a timed-out, closed, or + errored remote `MultiExecute` ack was discarded and the client still got `OK`. + All legs are drained (each was already dispatched, and the local leg still + persists), but any failed leg now surfaces as an error instead of a false ack. +- **`recv_reply_bounded` extended to the handler plane:** the coordinator's 30s + bounded reply-await now also covers the 11 remaining direct cross-shard acks + (WS.DROP cleanup, routed MQ commands, routed graph commands, TXN.COMMIT MQ + materialization acks, and the txn-abort remote rollback ack, on both + handler_monoio and handler_sharded) — a wedged owner shard can no longer park + those connections forever. MQ/graph/TXN owner execution is synchronous on the + shard thread (no blocking-wait commands route through these paths), so the + bound cannot cut off a legitimate long wait. +- **Checkpoint finalize backoff arms from the failure time:** the backoff was + armed with a timestamp captured before the finalize I/O; when + `wait_durable`/`manifest.commit`/`graph_save`/control-write took longer than + the backoff window, the retry deadline was already in the past and the next + 1ms tick retried immediately — exactly the flood the backoff exists to stop. + ### Added — CLI/moon.conf defaults for vector + graph tuning knobs (PR #TBD) - **`--vector-ef-runtime` / `--vector-rerank-mult` / `--vector-exact-beam`** diff --git a/src/cluster/bus.rs b/src/cluster/bus.rs index 9f7405c9..3c775f4a 100644 --- a/src/cluster/bus.rs +++ b/src/cluster/bus.rs @@ -27,6 +27,21 @@ use crate::cluster::gossip::{ pub type SharedVoteTx = Arc>>>; +/// Maximum time to wait for a gossip message BODY after its 4-byte length +/// prefix has been read (prod-hardening #10). Without this bound, any TCP +/// client that can reach the cluster bus port can send a valid length header +/// and then never send the body — the spawned `handle_cluster_peer` task +/// blocks in `read_exact` forever, is immune to graceful shutdown (the +/// cancellation token was only checked on the length read), and never releases +/// its socket/fd. A handful of such connections exhaust the bus listener's +/// accept capacity. On expiry we return an error, dropping the connection. +const GOSSIP_BODY_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); + +/// Maximum accepted gossip frame length (alloc guard: the 4-byte length prefix +/// is attacker-controlled, so it must never size an allocation unchecked). +/// Shared with the PING/PONG probe path in `gossip.rs`. +pub(crate) const MAX_GOSSIP_FRAME_LEN: usize = 64 * 1024; + /// Run the cluster bus listener loop. /// /// Spawns a new task for each incoming peer connection. @@ -97,13 +112,25 @@ async fn handle_cluster_peer( _ = shutdown.cancelled() => return Ok(()), } let msg_len = u32::from_be_bytes(len_buf) as usize; - if msg_len > 64 * 1024 { + if msg_len > MAX_GOSSIP_FRAME_LEN { anyhow::bail!("gossip message too large: {} bytes", msg_len); } - // Read message body + // Read message body — #10: bound with a timeout + shutdown-cancel so a + // client that sends the length prefix but stalls the body cannot pin + // the task/socket forever or dodge graceful shutdown. let mut buf = vec![0u8; msg_len]; - stream.read_exact(&mut buf).await?; + tokio::select! { + result = stream.read_exact(&mut buf) => { result?; } + _ = shutdown.cancelled() => return Ok(()), + _ = tokio::time::sleep(GOSSIP_BODY_READ_TIMEOUT) => { + anyhow::bail!( + "gossip body read from {} timed out after {}s", + peer_addr, + GOSSIP_BODY_READ_TIMEOUT.as_secs() + ); + } + } // Deserialize let msg = match deserialize_gossip(&buf) { @@ -179,7 +206,7 @@ async fn handle_cluster_peer( /// /// Monoio has no `read_exact` — we loop on `stream.read()` accumulating bytes. #[cfg(feature = "runtime-monoio")] -async fn monoio_read_exact( +pub(crate) async fn monoio_read_exact( stream: &mut monoio::net::TcpStream, total: usize, ) -> std::io::Result> { @@ -267,14 +294,26 @@ async fn handle_cluster_peer( }; let len_buf: [u8; 4] = len_data[..4].try_into().unwrap(); let msg_len = u32::from_be_bytes(len_buf) as usize; - if msg_len > 64 * 1024 { + if msg_len > MAX_GOSSIP_FRAME_LEN { anyhow::bail!("gossip message too large: {} bytes", msg_len); } - // Read message body - let buf = monoio_read_exact(&mut stream, msg_len) - .await - .map_err(|e| anyhow::anyhow!(e))?; + // Read message body — #10: bound with a timeout + shutdown-cancel so a + // client that sends the length prefix but stalls the body cannot pin + // the task/socket forever or dodge graceful shutdown. + let buf = monoio::select! { + result = monoio_read_exact(&mut stream, msg_len) => { + result.map_err(|e| anyhow::anyhow!(e))? + } + _ = shutdown.cancelled() => return Ok(()), + _ = monoio::time::sleep(GOSSIP_BODY_READ_TIMEOUT) => { + anyhow::bail!( + "gossip body read from {} timed out after {}s", + peer_addr, + GOSSIP_BODY_READ_TIMEOUT.as_secs() + ); + } + }; // Deserialize let msg = match deserialize_gossip(&buf) { diff --git a/src/cluster/gossip.rs b/src/cluster/gossip.rs index fb0d58e1..cad6ee54 100644 --- a/src/cluster/gossip.rs +++ b/src/cluster/gossip.rs @@ -364,6 +364,13 @@ pub async fn run_gossip_ticker( ) { let mut tick = tokio::time::interval(Duration::from_millis(100)); let mut election_spawned = false; + // Bound outstanding PING probes: a dead/partitioned peer must not leak a + // new connect+read task every 100ms. `ping_in_flight` allows at most one + // outstanding probe at a time; `ping_rotation` spreads probes across peers + // (the "random peer" the doc promised — the old code always picked the + // first non-self node). + let ping_in_flight = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let mut ping_rotation: usize = 0; loop { tokio::select! { _ = tick.tick() => { @@ -408,36 +415,66 @@ pub async fn run_gossip_ticker( } } - let target = cs.nodes.values() + // Rotate across peers instead of always pinging the first. + let peer_count = cs + .nodes + .values() .filter(|n| n.node_id != cs.node_id) - .next() - .map(|n| SocketAddr::new(n.addr.ip(), n.bus_port)); + .count(); + let target = if peer_count == 0 { + None + } else { + let idx = ping_rotation % peer_count; + ping_rotation = ping_rotation.wrapping_add(1); + cs.nodes + .values() + .filter(|n| n.node_id != cs.node_id) + .nth(idx) + .map(|n| SocketAddr::new(n.addr.ip(), n.bus_port)) + }; let msg = build_message(&cs, self_addr, GossipMsgType::Ping); cs.messages_sent += 1; (target, msg) }; if let Some(target_addr) = target_addr { - let cs = cluster_state.clone(); - tokio::spawn(async move { - if let Ok(mut stream) = TcpStream::connect(target_addr).await { - let data = serialize_gossip(&ping_msg); - let len = (data.len() as u32).to_be_bytes(); - let _ = stream.write_all(&len).await; - let _ = stream.write_all(&data).await; - // Read PONG response - let mut len_buf = [0u8; 4]; - if stream.read_exact(&mut len_buf).await.is_ok() { - let pong_len = u32::from_be_bytes(len_buf) as usize; - let mut pong_buf = vec![0u8; pong_len]; - if stream.read_exact(&mut pong_buf).await.is_ok() { - if let Ok(pong) = deserialize_gossip(&pong_buf) { - let mut cs2 = cs.write().unwrap(); - merge_gossip_into_state(&mut cs2, &pong); + // Skip if the previous probe hasn't resolved yet — bounds + // outstanding tasks to one even against a dead peer. + if !ping_in_flight.swap(true, std::sync::atomic::Ordering::AcqRel) { + let cs = cluster_state.clone(); + let inflight = ping_in_flight.clone(); + let probe_timeout = + Duration::from_millis((node_timeout_ms / 2).max(100)); + tokio::spawn(async move { + let _ = tokio::time::timeout(probe_timeout, async move { + if let Ok(mut stream) = TcpStream::connect(target_addr).await { + let data = serialize_gossip(&ping_msg); + let len = (data.len() as u32).to_be_bytes(); + let _ = stream.write_all(&len).await; + let _ = stream.write_all(&data).await; + // Read PONG response + let mut len_buf = [0u8; 4]; + if stream.read_exact(&mut len_buf).await.is_ok() { + let pong_len = u32::from_be_bytes(len_buf) as usize; + // Alloc guard: the length prefix is + // peer-controlled — cap it like the bus + // listener before sizing the buffer. + if pong_len > crate::cluster::bus::MAX_GOSSIP_FRAME_LEN { + return; + } + let mut pong_buf = vec![0u8; pong_len]; + if stream.read_exact(&mut pong_buf).await.is_ok() { + if let Ok(pong) = deserialize_gossip(&pong_buf) { + let mut cs2 = cs.write().unwrap(); + merge_gossip_into_state(&mut cs2, &pong); + } + } } } - } - } - }); + }) + .await; + inflight.store(false, std::sync::atomic::Ordering::Release); + }); + } } } _ = shutdown.cancelled() => break, @@ -460,6 +497,10 @@ pub async fn run_gossip_ticker( repl_state: std::sync::Arc>, ) { let mut election_spawned = false; + // Bound outstanding PING probes (see the tokio variant): one probe at a + // time, rotating across peers, so a dead peer can't leak a task per tick. + let ping_in_flight = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let mut ping_rotation: usize = 0; loop { monoio::select! { _ = monoio::time::sleep(Duration::from_millis(100)) => { @@ -502,45 +543,80 @@ pub async fn run_gossip_ticker( } } - let target = cs.nodes.values() + // Rotate across peers instead of always pinging the first. + let peer_count = cs + .nodes + .values() .filter(|n| n.node_id != cs.node_id) - .next() - .map(|n| SocketAddr::new(n.addr.ip(), n.bus_port)); + .count(); + let target = if peer_count == 0 { + None + } else { + let idx = ping_rotation % peer_count; + ping_rotation = ping_rotation.wrapping_add(1); + cs.nodes + .values() + .filter(|n| n.node_id != cs.node_id) + .nth(idx) + .map(|n| SocketAddr::new(n.addr.ip(), n.bus_port)) + }; let msg = build_message(&cs, self_addr, GossipMsgType::Ping); cs.messages_sent += 1; (target, msg) }; if let Some(target_addr) = target_addr { - let cs = cluster_state.clone(); - monoio::spawn(async move { - if let Ok(mut stream) = monoio::net::TcpStream::connect(target_addr).await { - let data = serialize_gossip(&ping_msg); - let len = (data.len() as u32).to_be_bytes(); - let len_vec = len.to_vec(); - let (wr, _) = stream.write_all(len_vec).await; - if wr.is_err() { return; } - let (wr, _) = stream.write_all(data).await; - if wr.is_err() { return; } - // Read PONG response - let len_buf = vec![0u8; 4]; - let (res, len_buf) = stream.read(len_buf).await; - if let Ok(4) = res { - let pong_len = u32::from_be_bytes( - len_buf[..4].try_into().unwrap(), - ) as usize; - let pong_buf = vec![0u8; pong_len]; - let (res, pong_buf) = stream.read(pong_buf).await; - if let Ok(n) = res { - if n == pong_len { - if let Ok(pong) = deserialize_gossip(&pong_buf) { - let mut cs2 = cs.write().unwrap(); - merge_gossip_into_state(&mut cs2, &pong); + // Skip if the previous probe hasn't resolved yet — bounds + // outstanding tasks to one even against a dead peer, and a + // sleep-race timeout cancels a hung connect/read. + if !ping_in_flight.swap(true, std::sync::atomic::Ordering::AcqRel) { + let cs = cluster_state.clone(); + let inflight = ping_in_flight.clone(); + let probe_timeout = + Duration::from_millis((node_timeout_ms / 2).max(100)); + monoio::spawn(async move { + monoio::select! { + _ = async { + if let Ok(mut stream) = monoio::net::TcpStream::connect(target_addr).await { + let data = serialize_gossip(&ping_msg); + let len = (data.len() as u32).to_be_bytes(); + let len_vec = len.to_vec(); + let (wr, _) = stream.write_all(len_vec).await; + if wr.is_err() { return; } + let (wr, _) = stream.write_all(data).await; + if wr.is_err() { return; } + // Read PONG response. Exact-read loops: + // a single `read()` can return short, and + // dropping a fragmented-but-valid PONG + // leaves `pong_recv_ms` stale, pushing a + // healthy peer toward PFAIL. Length is + // capped like the bus listener (the prefix + // is peer-controlled and sizes an alloc). + if let Ok(len_buf) = + crate::cluster::bus::monoio_read_exact(&mut stream, 4).await + { + let pong_len = u32::from_be_bytes([ + len_buf[0], len_buf[1], len_buf[2], len_buf[3], + ]) as usize; + if pong_len <= crate::cluster::bus::MAX_GOSSIP_FRAME_LEN + && let Ok(pong_buf) = + crate::cluster::bus::monoio_read_exact( + &mut stream, + pong_len, + ) + .await + && let Ok(pong) = deserialize_gossip(&pong_buf) + { + let mut cs2 = cs.write().unwrap(); + merge_gossip_into_state(&mut cs2, &pong); + } } } - } + } => {} + _ = monoio::time::sleep(probe_timeout) => {} } - } - }); + inflight.store(false, std::sync::atomic::Ordering::Release); + }); + } } } _ = shutdown.cancelled() => break, diff --git a/src/command/hash/hash_read.rs b/src/command/hash/hash_read.rs index a617b9cb..b4309718 100644 --- a/src/command/hash/hash_read.rs +++ b/src/command/hash/hash_read.rs @@ -275,7 +275,9 @@ pub fn hscan(db: &mut Database, args: &[Frame]) -> Frame { }; let total = entries.len(); - let mut results = Vec::with_capacity(count * 2); + // DoS guard: bound the pre-size by the actual field count so a huge COUNT + // hint can't drive an unbounded Vec::with_capacity -> allocator abort. + let mut results = Vec::with_capacity(count.min(total).saturating_mul(2)); let mut pos = cursor; let mut checked = 0; @@ -665,8 +667,13 @@ pub fn hrandfield(db: &mut Database, args: &[Frame]) -> Frame { Frame::Array(result.into()) } } else { - // Negative count: allow duplicates. Cap to entries.len() to prevent OOM on i64::MIN. - let n = std::cmp::min(count.unsigned_abs() as usize, entries.len() * 10); + // Negative count: allow duplicates — exactly |COUNT| of them (Redis + // contract). The DoS guard refuses extreme counts (i64::MIN → OOM) + // loudly instead of silently truncating. + let n = count.unsigned_abs() as usize; + if n > crate::command::RAND_DUP_COUNT_MAX { + return Frame::Error(Bytes::from_static(crate::command::ERR_RAND_COUNT_RANGE)); + } if with_values { let mut result = Vec::with_capacity(n * 2); for _ in 0..n { @@ -775,8 +782,13 @@ pub fn hrandfield_readonly(db: &Database, args: &[Frame], now_ms: u64) -> Frame Frame::Array(result.into()) } } else { - // Negative count: allow duplicates. Cap to prevent OOM on extreme values. - let n = std::cmp::min(count.unsigned_abs() as usize, entries.len() * 10); + // Negative count: allow duplicates — exactly |COUNT| of them (Redis + // contract). The DoS guard refuses extreme counts loudly instead of + // silently truncating. + let n = count.unsigned_abs() as usize; + if n > crate::command::RAND_DUP_COUNT_MAX { + return Frame::Error(Bytes::from_static(crate::command::ERR_RAND_COUNT_RANGE)); + } if with_values { let mut result = Vec::with_capacity(n * 2); for _ in 0..n { diff --git a/src/command/hash/mod.rs b/src/command/hash/mod.rs index 9b7c4afa..a0e91cf3 100644 --- a/src/command/hash/mod.rs +++ b/src/command/hash/mod.rs @@ -30,6 +30,27 @@ mod tests { assert_eq!(result, Frame::Integer(2)); } + #[test] + fn test_hscan_huge_count_no_abort() { + // A huge COUNT hint must not drive an unbounded Vec::with_capacity -> + // allocator abort (Batch A DoS guard); the scan still returns the + // real fields. + let mut db = Database::new(); + hset(&mut db, &make_args(&[b"h", b"f1", b"v1", b"f2", b"v2"])); + match hscan(&mut db, &make_args(&[b"h", b"0", b"COUNT", b"5000000000"])) { + Frame::Array(top) => { + assert_eq!(top.len(), 2, "HSCAN returns [cursor, entries]"); + match &top[1] { + Frame::Array(kv) => { + assert_eq!(kv.len(), 4, "expected 2 field/value pairs") + } + other => panic!("expected inner array, got {other:?}"), + } + } + other => panic!("expected array, got {other:?}"), + } + } + #[test] fn test_hset_update_existing() { let mut db = Database::new(); diff --git a/src/command/mod.rs b/src/command/mod.rs index cd48c8a2..dfca68fd 100644 --- a/src/command/mod.rs +++ b/src/command/mod.rs @@ -43,6 +43,18 @@ pub enum DispatchResult { Quit(Frame), } +/// Upper bound for the with-duplicates (negative COUNT) arms of +/// SRANDMEMBER / HRANDFIELD / ZRANDMEMBER. Redis returns exactly |COUNT| +/// elements; Moon honors that up to this cap and returns an error beyond it — +/// never a silently short reply. An unchecked |COUNT| would let one command +/// drive an arbitrarily large `Vec::with_capacity` (i64::MIN → allocator +/// abort), and the previous relative cap (`len() * 10`) silently truncated +/// legitimate requests on small collections. +pub(crate) const RAND_DUP_COUNT_MAX: usize = 1 << 20; + +/// Error returned when a negative COUNT exceeds [`RAND_DUP_COUNT_MAX`]. +pub(crate) const ERR_RAND_COUNT_RANGE: &[u8] = b"ERR COUNT is out of range"; + /// Dispatch a pre-extracted command to the appropriate handler. /// /// Uses two-level dispatch: match on `(command_length, first_byte_lowercase)` to narrow diff --git a/src/command/set/mod.rs b/src/command/set/mod.rs index 1f3486c4..3cad3700 100644 --- a/src/command/set/mod.rs +++ b/src/command/set/mod.rs @@ -151,6 +151,45 @@ mod tests { assert_eq!(result, Frame::Integer(1)); // only "d" is new } + #[test] + fn test_srandmember_negative_count_exact_redis_contract() { + // Redis contract: SRANDMEMBER key -N returns EXACTLY N elements + // (duplicates allowed) even when N exceeds the cardinality — a + // relative cap that silently truncates is a compliance break. + let mut db = Database::new(); + setup_set(&mut db, b"s", &[b"a", b"b", b"c"]); + match srandmember(&mut db, &[bs(b"s"), bs(b"-1000")]) { + Frame::Array(items) => { + assert_eq!(items.len(), 1000, "exactly |count| elements required") + } + other => panic!("expected array, got {other:?}"), + } + // The read-only twin shares the same contract. + match srandmember_readonly(&db, &[bs(b"s"), bs(b"-1000")], 0) { + Frame::Array(items) => assert_eq!(items.len(), 1000), + other => panic!("expected array, got {other:?}"), + } + } + + #[test] + fn test_srandmember_huge_negative_count_errors_no_abort() { + // Batch A DoS guard: a COUNT beyond RAND_DUP_COUNT_MAX must be refused + // loudly (never drive an unbounded Vec::with_capacity -> allocator + // abort, and never return a silently short reply). + let mut db = Database::new(); + setup_set(&mut db, b"s", &[b"a", b"b", b"c"]); + let huge = format!("-{}", crate::command::RAND_DUP_COUNT_MAX + 1); + match srandmember(&mut db, &[bs(b"s"), bs(huge.as_bytes())]) { + Frame::Error(e) => assert_eq!(&e[..], crate::command::ERR_RAND_COUNT_RANGE), + other => panic!("expected error, got {other:?}"), + } + // i64::MIN: unsigned_abs() of the most extreme input stays guarded. + match srandmember_readonly(&db, &[bs(b"s"), bs(i64::MIN.to_string().as_bytes())], 0) { + Frame::Error(e) => assert_eq!(&e[..], crate::command::ERR_RAND_COUNT_RANGE), + other => panic!("expected error, got {other:?}"), + } + } + #[test] fn test_srem_basic() { let mut db = Database::new(); diff --git a/src/command/set/set_read.rs b/src/command/set/set_read.rs index ab70d4f0..ad776e19 100644 --- a/src/command/set/set_read.rs +++ b/src/command/set/set_read.rs @@ -298,7 +298,13 @@ pub fn srandmember(db: &mut Database, args: &[Frame]) -> Frame { Frame::Array(chosen.into()) } else { // Allow duplicates + // DoS guard: refuse a huge negative COUNT loudly instead of letting it + // drive an unbounded Vec::with_capacity -> allocator abort. Within the + // cap, Redis semantics apply: exactly |COUNT| elements, duplicates ok. let n = count.unsigned_abs() as usize; + if n > crate::command::RAND_DUP_COUNT_MAX { + return Frame::Error(Bytes::from_static(crate::command::ERR_RAND_COUNT_RANGE)); + } let mut result = Vec::with_capacity(n); for _ in 0..n { let Some(chosen) = members.choose(&mut rng) else { @@ -658,7 +664,13 @@ pub fn srandmember_readonly(db: &Database, args: &[Frame], now_ms: u64) -> Frame .collect(); Frame::Array(chosen.into()) } else { + // DoS guard: refuse a huge negative COUNT loudly instead of letting it + // drive an unbounded Vec::with_capacity -> allocator abort. Within the + // cap, Redis semantics apply: exactly |COUNT| elements, duplicates ok. let n = count.unsigned_abs() as usize; + if n > crate::command::RAND_DUP_COUNT_MAX { + return Frame::Error(Bytes::from_static(crate::command::ERR_RAND_COUNT_RANGE)); + } let mut result = Vec::with_capacity(n); for _ in 0..n { let Some(chosen) = members.choose(&mut rng) else { diff --git a/src/command/sorted_set/mod.rs b/src/command/sorted_set/mod.rs index 2b4f702e..4c3e108b 100644 --- a/src/command/sorted_set/mod.rs +++ b/src/command/sorted_set/mod.rs @@ -597,6 +597,37 @@ mod tests { zcard(db, &frames) } + #[test] + fn test_zmpop_huge_count_no_abort() { + // A huge COUNT must not drive an unbounded Vec::with_capacity -> + // allocator abort. This site was MISSED by the audit finders and + // caught by the Batch A allocation sweep (unlike LMPOP it lacked a + // .min(card) cap). + let mut db = Database::new(); + run_zadd(&mut db, &[b"z", b"1", b"a", b"2", b"b"]); + match zmpop( + &mut db, + &[ + bulk(b"1"), + bulk(b"z"), + bulk(b"MIN"), + bulk(b"COUNT"), + bulk(b"5000000000"), + ], + ) { + Frame::Array(items) => { + assert_eq!(items.len(), 2, "ZMPOP returns [key, popped]"); + match &items[1] { + Frame::Array(popped) => { + assert_eq!(popped.len(), 2, "only 2 members exist to pop") + } + other => panic!("expected popped array, got {other:?}"), + } + } + other => panic!("expected array, got {other:?}"), + } + } + fn run_zincrby(db: &mut Database, args: &[&[u8]]) -> Frame { let frames: Vec = args.iter().map(|a| bulk(a)).collect(); zincrby(db, &frames) diff --git a/src/command/sorted_set/sorted_set_read.rs b/src/command/sorted_set/sorted_set_read.rs index a5548f93..6bb06301 100644 --- a/src/command/sorted_set/sorted_set_read.rs +++ b/src/command/sorted_set/sorted_set_read.rs @@ -1751,8 +1751,13 @@ pub fn zrandmember(db: &mut Database, args: &[Frame]) -> Frame { } Frame::Array(result.into()) } else { - // Negative count: allow duplicates. Cap to prevent OOM on extreme values. - let n = std::cmp::min(count.unsigned_abs() as usize, entries.len() * 10); + // Negative count: allow duplicates — exactly |COUNT| of them (Redis + // contract). The DoS guard refuses extreme counts loudly instead of + // silently truncating. + let n = count.unsigned_abs() as usize; + if n > crate::command::RAND_DUP_COUNT_MAX { + return Frame::Error(Bytes::from_static(crate::command::ERR_RAND_COUNT_RANGE)); + } let cap = if withscores { n * 2 } else { n }; let mut result = Vec::with_capacity(cap); for _ in 0..n { @@ -2033,7 +2038,13 @@ pub fn zrandmember_readonly(db: &Database, args: &[Frame], now_ms: u64) -> Frame } Frame::Array(result.into()) } else { - let n = std::cmp::min(count.unsigned_abs() as usize, entries.len() * 10); + // Negative count: allow duplicates — exactly |COUNT| of them (Redis + // contract). The DoS guard refuses extreme counts loudly instead of + // silently truncating. + let n = count.unsigned_abs() as usize; + if n > crate::command::RAND_DUP_COUNT_MAX { + return Frame::Error(Bytes::from_static(crate::command::ERR_RAND_COUNT_RANGE)); + } let cap = if withscores { n * 2 } else { n }; let mut result = Vec::with_capacity(cap); for _ in 0..n { diff --git a/src/command/sorted_set/sorted_set_write.rs b/src/command/sorted_set/sorted_set_write.rs index 4fc6507c..23f2aa8d 100644 --- a/src/command/sorted_set/sorted_set_write.rs +++ b/src/command/sorted_set/sorted_set_write.rs @@ -807,7 +807,10 @@ pub fn zmpop(db: &mut Database, args: &[Frame]) -> Frame { Err(e) => return e, }; - let mut popped = Vec::with_capacity(pop_count); + // DoS guard: bound the pre-size by the set's cardinality so a huge + // COUNT can't drive an unbounded Vec::with_capacity -> allocator abort + // (matches LMPOP's count.min(list_len); loop already breaks when empty). + let mut popped = Vec::with_capacity(pop_count.min(card)); let mut credit: usize = 0; for _ in 0..pop_count { let entry = if is_min { diff --git a/src/command/string/string_bit.rs b/src/command/string/string_bit.rs index 9a45b264..9edb207d 100644 --- a/src/command/string/string_bit.rs +++ b/src/command/string/string_bit.rs @@ -860,6 +860,18 @@ pub fn bitfield(db: &mut Database, args: &[Frame]) -> Frame { } }; + // DoS guard: SET/INCRBY grow the buffer to cover bit_offset; reject + // offsets past the 512MB string limit (matches SETBIT) so a client + // offset can't drive an unbounded Vec::resize -> allocator abort. GET + // is exempt (bf_get reads within data.len(), returning 0 past the end). + if !subcmd.eq_ignore_ascii_case(b"GET") + && bit_offset.saturating_add(bits as usize) > 512 * 1024 * 1024 * 8 + { + return Frame::Error(Bytes::from_static( + b"ERR bit offset is not an integer or out of range", + )); + } + if subcmd.eq_ignore_ascii_case(b"GET") { let val = bf_get(&buf, bit_offset, bits, signed); results.push(Frame::Integer(val)); @@ -1162,6 +1174,36 @@ mod tests { assert_eq!(result, Frame::Integer(0)); } + #[test] + fn test_bitfield_huge_offset_rejected_no_abort() { + // A bit offset past the 512MB string limit must be rejected, never + // drive an unbounded Vec::resize -> allocator abort (Batch A DoS guard). + let mut db = make_db(); + let r = bitfield( + &mut db, + &[ + bs(b"k"), + bs(b"SET"), + bs(b"u8"), + bs(b"#999999999999"), + bs(b"1"), + ], + ); + assert!( + matches!(r, Frame::Error(_)), + "huge BITFIELD SET offset must return an error, got {r:?}" + ); + // A normal offset still works (guard doesn't break legitimate use). + let ok = bitfield( + &mut db, + &[bs(b"k"), bs(b"SET"), bs(b"u8"), bs(b"0"), bs(b"200")], + ); + assert!( + matches!(ok, Frame::Array(_)), + "normal BITFIELD SET should succeed, got {ok:?}" + ); + } + #[test] fn test_getbit_out_of_range() { let mut db = make_db(); diff --git a/src/command/string/string_write.rs b/src/command/string/string_write.rs index 8f35c33b..e6719f8c 100644 --- a/src/command/string/string_write.rs +++ b/src/command/string/string_write.rs @@ -515,6 +515,15 @@ pub fn append(db: &mut Database, args: &[Frame]) -> Frame { None => (None, 0), }; + // Enforce the 512MB max string size (matches SETRANGE/SETBIT); APPEND + // previously let a string grow past the documented limit unbounded. + let combined_len = existing_data.as_ref().map_or(0, |v| v.len()) + append_val.len(); + if combined_len > 512 * 1024 * 1024 { + return Frame::Error(Bytes::from_static( + b"ERR string exceeds maximum allowed size (512MB)", + )); + } + let new_val = match existing_data { Some(existing) => { let mut combined = Vec::with_capacity(existing.len() + append_val.len()); diff --git a/src/command/vector_search/ft_info.rs b/src/command/vector_search/ft_info.rs index 9140bd23..9f08407e 100644 --- a/src/command/vector_search/ft_info.rs +++ b/src/command/vector_search/ft_info.rs @@ -45,7 +45,10 @@ pub fn ft_info( // fully searchable (see FT.INFO doc comment above re: summing across // segments). let snap = idx.segments.load(); - let mut num_docs = snap.mutable.len(); + // Prod-hardening #28: use live_len(), NOT len() — the mutable segment + // keeps tombstoned entries in place until compaction, so len() would + // over-count deleted docs by up to compact_threshold under DEL/HSET churn. + let mut num_docs = snap.mutable.live_len(); for imm in snap.immutable.iter() { num_docs += imm.live_count() as usize; } @@ -57,6 +60,13 @@ pub fn ft_info( for stub in snap.unloaded.iter() { num_docs += stub.live_count() as usize; } + // Prod-hardening #28: DiskANN cold-tier segments (experimental, gated by + // MOON_VEC_COLD_TIER) were never summed — cold docs ARE returned by + // FT.SEARCH, so omitting them under-reported num_docs once any segment + // aged into the cold tier. + for cold in snap.cold.iter() { + num_docs += cold.num_docs(); + } // HQ-1 observability (persistence-review R5): exact-rerank coverage. // A segment without the f16 sidecar silently answers with quantized ADC // distances only — surfacing the count makes a dropped sidecar (e.g. an @@ -146,7 +156,8 @@ pub fn ft_info( let (field_num_docs, field_mutable, field_immutable_count) = if i == 0 { // Default field: use top-level segments let s = idx.segments.load(); - let mut docs = s.mutable.len(); + // #28: live_len() excludes tombstoned mutable entries. + let mut docs = s.mutable.live_len(); let imm_count = s.immutable.len(); for imm in s.immutable.iter() { docs += imm.live_count() as usize; @@ -157,10 +168,13 @@ pub fn ft_info( for stub in s.unloaded.iter() { docs += stub.live_count() as usize; } - (docs, s.mutable.len(), imm_count) + for cold in s.cold.iter() { + docs += cold.num_docs(); + } + (docs, s.mutable.live_len(), imm_count) } else if let Some(fs) = idx.field_segments.get(&field_meta.field_name) { let s = fs.segments.load(); - let mut docs = s.mutable.len(); + let mut docs = s.mutable.live_len(); let imm_count = s.immutable.len(); for imm in s.immutable.iter() { docs += imm.live_count() as usize; @@ -171,7 +185,10 @@ pub fn ft_info( for stub in s.unloaded.iter() { docs += stub.live_count() as usize; } - (docs, s.mutable.len(), imm_count) + for cold in s.cold.iter() { + docs += cold.num_docs(); + } + (docs, s.mutable.live_len(), imm_count) } else { (0, 0, 0) }; diff --git a/src/command/vector_search/ft_text_search.rs b/src/command/vector_search/ft_text_search.rs index f2191e2f..90d3f619 100644 --- a/src/command/vector_search/ft_text_search.rs +++ b/src/command/vector_search/ft_text_search.rs @@ -118,6 +118,10 @@ pub fn parse_highlight_clause(args: &[Frame]) -> Option { _ => 0, }; i += 1; + // DoS guard: there can never be more field tokens than remaining + // args; clamp so a huge FIELDS count can't drive an unbounded + // Vec::with_capacity -> allocator abort. + let count = count.min(args.len().saturating_sub(i)); if count > 0 { let mut fields = Vec::with_capacity(count); for _ in 0..count { @@ -203,6 +207,10 @@ pub fn parse_summarize_clause(args: &[Frame]) -> Option { _ => 0, }; i += 1; + // DoS guard: there can never be more field tokens than remaining + // args; clamp so a huge FIELDS count can't drive an unbounded + // Vec::with_capacity -> allocator abort. + let count = count.min(args.len().saturating_sub(i)); if count > 0 { let mut fields = Vec::with_capacity(count); for _ in 0..count { @@ -2079,6 +2087,23 @@ mod tests { .collect() } + #[test] + fn parse_highlight_fields_huge_count_no_abort() { + // A huge FIELDS count must not drive an unbounded Vec::with_capacity -> + // allocator abort (Batch A DoS guard); fields are bounded by the real + // number of remaining tokens. + let args = args_of(&["HIGHLIGHT", "FIELDS", "999999999999", "title"]); + let opts = parse_highlight_clause(&args).expect("highlight opts"); + assert_eq!(opts.fields.as_deref().map(<[Bytes]>::len), Some(1)); + } + + #[test] + fn parse_summarize_fields_huge_count_no_abort() { + let args = args_of(&["SUMMARIZE", "FIELDS", "999999999999", "body"]); + let opts = parse_summarize_clause(&args).expect("summarize opts"); + assert_eq!(opts.fields.as_deref().map(<[Bytes]>::len), Some(1)); + } + #[test] fn has_sparse_clause_detects_standalone() { // R3: a standalone `SPARSE @field $param` clause is detected from the full args. diff --git a/src/config.rs b/src/config.rs index da2083e7..993eb556 100644 --- a/src/config.rs +++ b/src/config.rs @@ -438,14 +438,19 @@ pub struct ServerConfig { #[arg(long = "vec-warm-mmap-budget", default_value = "2gb")] pub vec_warm_mmap_budget: String, - // ── Cold-tier / DiskANN config stubs (not yet consumed) ───────── + // ── Cold-tier / DiskANN (EXPERIMENTAL — gated by MOON_VEC_COLD_TIER) ── + // The DiskANN cold tier is incomplete (no cold-segment deletion, ADC-only + // recall, restart reload of PQ codebooks unfinished). The WARM->COLD + // transition is a NO-OP unless an operator sets `MOON_VEC_COLD_TIER=1`; + // warm-tier mmap + LRU eviction handles out-of-RAM indexes by default. /// Seconds after last access before a WARM segment is promoted to COLD. - /// Not yet consumed — reserved for the WARM->COLD transition timer. + /// Consumed by the cold-transition timer ONLY when the experimental cold + /// tier is enabled (`MOON_VEC_COLD_TIER=1`); otherwise inert. #[arg(long = "segment-cold-after", default_value_t = 86_400)] pub segment_cold_after: u64, /// Minimum queries-per-second threshold; segments below this are COLD candidates. - /// Not yet consumed — reserved for the WARM->COLD transition heuristic. + /// Not yet consumed — reserved for the experimental cold-tier heuristic. #[arg(long = "segment-cold-min-qps", default_value_t = 0.1)] pub segment_cold_min_qps: f64, @@ -458,12 +463,14 @@ pub struct ServerConfig { pub memory_arenas_cap: u32, /// DiskANN beam width for disk-resident vector search. - /// Not yet consumed — reserved for the DiskANN search implementation. + /// Not yet consumed — reserved for the experimental cold-tier search path + /// (gated by MOON_VEC_COLD_TIER). #[arg(long = "vec-diskann-beam-width", default_value_t = 8)] pub vec_diskann_beam_width: u32, /// Number of HNSW upper levels cached in memory for DiskANN hybrid search. - /// Not yet consumed — reserved for the DiskANN cache layer. + /// Not yet consumed — reserved for the experimental cold-tier cache layer + /// (gated by MOON_VEC_COLD_TIER). #[arg(long = "vec-diskann-cache-levels", default_value_t = 3)] pub vec_diskann_cache_levels: u32, diff --git a/src/persistence/aof/mod.rs b/src/persistence/aof/mod.rs index 9fd1c8e3..bf2195a5 100644 --- a/src/persistence/aof/mod.rs +++ b/src/persistence/aof/mod.rs @@ -7,7 +7,7 @@ //! |---------|---------------|-----------| //! | `AofWriter::append` (hot path) | **fire-and-forget** | Channel send; no Result needed | //! | `aof_writer_task` | **must-panic** | Writer task; errors logged inline | -//! | `replay_aof` | **should-recover** (`Result<_, MoonError>`) | Startup replay; log+skip on corruption | +//! | `replay_aof` | **should-recover** (`Result<_, MoonError>`) | Startup replay; stop at mid-stream corruption (keep valid prefix) | //! | `rewrite_aof` | **should-recover** (`Result<_, MoonError>`) | Background rewrite; caller logs error | //! | `#[cfg(test)]` code (55 unwraps) | **test-only** | Panics are appropriate in tests | // Suppressions narrowed: only keep what's needed for conditional compilation @@ -503,18 +503,44 @@ pub fn serialize_command(frame: &Frame) -> Bytes { buf.freeze() } +/// 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 +/// risking dispatch of misaligned bytes as live commands. Operators who want +/// the old behavior opt in with `MOON_AOF_BEST_EFFORT_RESYNC=1`. +fn aof_best_effort_resync_enabled() -> bool { + std::env::var("MOON_AOF_BEST_EFFORT_RESYNC") + .map(|v| v == "1" || v.eq_ignore_ascii_case("enable")) + .unwrap_or(false) +} + /// Replay an AOF file by parsing RESP commands and dispatching them. /// /// Returns the number of commands successfully replayed. /// -/// **Corruption recovery:** On mid-stream parse errors, logs a warning with the -/// byte offset, skips to the next RESP array marker (`*`), and continues replay. -/// At EOF, reports total corrupted entries skipped. Truncated tails are handled -/// gracefully (warn + stop). +/// **Corruption recovery (#12):** A clean truncated tail is handled gracefully +/// (warn + stop). On genuine mid-stream corruption the default is to STOP +/// replay — keeping only the valid prefix already applied and never dispatching +/// bytes past the corruption point (which, in unframed RESP, could be a `*` +/// inside a binary value → misaligned garbage commands). Set +/// `MOON_AOF_BEST_EFFORT_RESYNC=1` to opt into the legacy skip-to-next-`*` +/// resync (loud warnings; may misapply corrupted data). pub fn replay_aof( databases: &mut [Database], path: &Path, engine: &dyn CommandReplayEngine, +) -> Result { + replay_aof_with_resync(databases, path, engine, aof_best_effort_resync_enabled()) +} + +/// [`replay_aof`] with the resync mode passed explicitly instead of read from +/// `MOON_AOF_BEST_EFFORT_RESYNC` — lets tests pin the branch deterministically +/// regardless of the ambient environment. +fn replay_aof_with_resync( + databases: &mut [Database], + path: &Path, + engine: &dyn CommandReplayEngine, + best_effort_resync: bool, ) -> Result { let data = std::fs::read(path)?; if data.is_empty() { @@ -586,7 +612,10 @@ pub fn replay_aof( Ok(None) => { // Incomplete frame at end of file - truncated AOF if !buf.is_empty() { - let offset = total_len - buf.len(); + // Absolute file offset: resp_data sits after any RDB + // preamble, and operator repair guidance must point at + // the real file location. + let offset = resp_start + (total_len - buf.len()); warn!( "AOF truncated: {} unparseable bytes at offset {} (end of file)", buf.len(), @@ -596,15 +625,43 @@ pub fn replay_aof( break; } Err(e) => { - let error_offset = total_len - buf.len(); + // Absolute file offset (includes any RDB preamble) — see above. + let error_offset = resp_start + (total_len - buf.len()); + corruption_count += 1; + + // Prod-hardening #12: DO NOT skip-and-resync by default. + // `appendonly.aof` is raw unframed RESP with no per-record + // length/CRC, and a `*` byte appears freely inside binary + // value blobs — resuming at "the next `*`" can land mid-value + // and dispatch misaligned garbage as live SET/DEL/etc. against + // the recovering keyspace (silent data corruption). Clean + // tail-truncation is already handled by the `Ok(None)` arm; + // reaching here means genuine mid-stream corruption. The safe + // default (matching the per-shard `shard_replay` sibling) + // stops here, keeping only the valid prefix already applied, + // and never dispatches anything past the corruption point. + if !best_effort_resync { + error!( + "AOF corruption at byte offset {} after {} commands: {}. \ + Stopping replay to avoid dispatching misaligned bytes as \ + live commands; the {} commands before this point are kept. \ + Run redis-check-aof (or set MOON_AOF_BEST_EFFORT_RESYNC=1 to \ + opt into the legacy skip-and-resync, which MAY misapply \ + corrupted data).", + error_offset, count, e, count + ); + break; + } + + // Opt-in legacy best-effort resync (loud): skip past the + // corrupt byte(s) to the next RESP array marker ('*'). Always + // discard at least 1 byte to guarantee forward progress. warn!( - "AOF parse error at byte offset {} after {} commands: {}. Attempting skip.", + "AOF parse error at byte offset {} after {} commands: {}. \ + MOON_AOF_BEST_EFFORT_RESYNC set — attempting skip (data may be \ + misapplied).", error_offset, count, e ); - corruption_count += 1; - - // Skip past the corrupt byte(s) to the next RESP array marker ('*') - // Always discard at least 1 byte to guarantee forward progress. let _ = buf.split_to(1); if let Some(pos) = buf.iter().position(|&b| b == b'*') { let _ = buf.split_to(pos); @@ -695,6 +752,67 @@ mod tests { assert_eq!(entry.value.as_bytes().unwrap(), b"v2"); } + /// Builds the corrupt AOF shared by the midstream-corruption tests: + /// valid prefix, a hard mid-stream parse error, then a valid command. + fn write_midstream_corrupt_aof(dir: &tempfile::TempDir) -> std::path::PathBuf { + let aof_path = dir.path().join("corrupt.aof"); + let mut aof_data = BytesMut::new(); + // Valid prefix. + serialize::serialize(&make_command(&[b"SET", b"before", b"1"]), &mut aof_data); + // Corruption: `*` marker with a non-integer count — complete (has CRLF) + // so it is a hard parse Err, not an Incomplete/tail-truncation. + aof_data.extend_from_slice(b"*Z\r\n"); + // A perfectly valid command AFTER the corruption. + serialize::serialize(&make_command(&[b"SET", b"after", b"2"]), &mut aof_data); + std::fs::write(&aof_path, &aof_data).unwrap(); + aof_path + } + + #[test] + fn test_aof_replay_stops_at_midstream_corruption_by_default() { + // Prod-hardening #12: a genuine mid-stream parse error must STOP replay + // (keeping the valid prefix), NOT skip forward to the next `*` and + // dispatch misaligned bytes as live commands. The resync mode is pinned + // explicitly so an ambient MOON_AOF_BEST_EFFORT_RESYNC in a dev shell + // or CI cannot flip this test onto the legacy branch. + let dir = tempdir().unwrap(); + let aof_path = write_midstream_corrupt_aof(&dir); + + let mut dbs = vec![Database::new()]; + let count = + replay_aof_with_resync(&mut dbs, &aof_path, &DispatchReplayEngine::new(), false) + .unwrap(); + + assert_eq!(count, 1, "only the pre-corruption command is replayed"); + assert!( + dbs[0].get(b"before").is_some(), + "valid prefix command must be applied" + ); + assert!( + dbs[0].get(b"after").is_none(), + "post-corruption command must NOT be dispatched (no skip-resync by default)" + ); + } + + #[test] + fn test_aof_replay_best_effort_resync_skips_past_corruption() { + // The opt-in legacy branch (MOON_AOF_BEST_EFFORT_RESYNC=1) skips to the + // next `*` marker and keeps replaying — pinned explicitly here. + let dir = tempdir().unwrap(); + let aof_path = write_midstream_corrupt_aof(&dir); + + let mut dbs = vec![Database::new()]; + let count = replay_aof_with_resync(&mut dbs, &aof_path, &DispatchReplayEngine::new(), true) + .unwrap(); + + assert_eq!(count, 2, "resync replays the post-corruption command too"); + assert!(dbs[0].get(b"before").is_some()); + assert!( + dbs[0].get(b"after").is_some(), + "resync mode must skip the corrupt bytes and replay the valid tail" + ); + } + #[test] fn test_aof_replay_collection_types() { let dir = tempdir().unwrap(); diff --git a/src/persistence/checkpoint.rs b/src/persistence/checkpoint.rs index 7b853083..7b0618b9 100644 --- a/src/persistence/checkpoint.rs +++ b/src/persistence/checkpoint.rs @@ -114,6 +114,15 @@ pub enum CheckpointAction { pub struct CheckpointManager { state: CheckpointState, trigger: CheckpointTrigger, + /// Backoff deadline for the Finalize step (prod-hardening #13). When a + /// finalize attempt fails, the state stays `Finalizing`, so without this + /// the next 1ms tick re-runs finalize (re-appending a WAL Checkpoint + /// record) immediately — flooding the WAL every millisecond under a + /// sustained failure. `None` = ready now; `Some(t)` = the next attempt is + /// gated until `t`. + finalize_retry_at: Option, + /// Consecutive failed finalize attempts, driving exponential backoff. + finalize_attempts: u32, } impl CheckpointManager { @@ -122,9 +131,39 @@ impl CheckpointManager { Self { state: CheckpointState::Idle, trigger, + finalize_retry_at: None, + finalize_attempts: 0, } } + /// Whether the Finalize step may be attempted at `now` (prod-hardening + /// #13). Returns `true` when no backoff is pending or the backoff window + /// has elapsed. The caller checks this BEFORE doing any finalize I/O + /// (notably the WAL checkpoint-record append) so a stuck finalize retries + /// on a bounded schedule instead of every tick. + #[inline] + pub fn finalize_ready(&self, now: Instant) -> bool { + self.finalize_retry_at.map_or(true, |t| now >= t) + } + + /// Record a failed finalize attempt and arm an exponential backoff + /// (50ms, 100ms, … capped at 5s) before the next attempt is allowed. + pub fn note_finalize_failed(&mut self, now: Instant) { + self.finalize_attempts = self.finalize_attempts.saturating_add(1); + // 50ms << (attempts-1), capped at 5s. + let shift = (self.finalize_attempts - 1).min(7); + let backoff = + std::time::Duration::from_millis(50u64 << shift).min(std::time::Duration::from_secs(5)); + self.finalize_retry_at = Some(now + backoff); + } + + /// Clear any finalize backoff (called on success / when leaving Finalizing). + #[inline] + fn reset_finalize_backoff(&mut self) { + self.finalize_retry_at = None; + self.finalize_attempts = 0; + } + /// Begin a new checkpoint. /// /// Records the REDO LSN and computes `pages_per_tick` based on the number @@ -202,6 +241,7 @@ impl CheckpointManager { pub fn complete(&mut self) { self.state = CheckpointState::Idle; self.trigger.reset(); + self.reset_finalize_backoff(); } /// Force-begin a checkpoint regardless of trigger conditions. @@ -235,6 +275,7 @@ impl CheckpointManager { #[cfg(test)] mod tests { use super::*; + use std::time::Duration; fn make_trigger(timeout_secs: u64, max_wal_bytes: u64, completion: f64) -> CheckpointTrigger { CheckpointTrigger::new(timeout_secs, max_wal_bytes, completion) @@ -320,6 +361,54 @@ mod tests { } } + #[test] + fn test_finalize_backoff_gates_retries() { + // Prod-hardening #13: after a failed finalize, finalize_ready must + // return false until the exponential backoff window elapses, so the + // caller does not re-append a WAL checkpoint record every tick. + let trigger = make_trigger(300, 256 * 1024 * 1024, 0.9); + let mut mgr = CheckpointManager::new(trigger); + + let t0 = Instant::now(); + // No failures yet → always ready. + assert!(mgr.finalize_ready(t0)); + + // First failure arms a 50ms backoff. + mgr.note_finalize_failed(t0); + assert!(!mgr.finalize_ready(t0), "immediately after failure: gated"); + assert!( + !mgr.finalize_ready(t0 + Duration::from_millis(49)), + "still within the 50ms window" + ); + assert!( + mgr.finalize_ready(t0 + Duration::from_millis(50)), + "ready once the 50ms window elapses" + ); + + // Second consecutive failure doubles the backoff to 100ms. + let t1 = t0 + Duration::from_millis(50); + mgr.note_finalize_failed(t1); + assert!(!mgr.finalize_ready(t1 + Duration::from_millis(99))); + assert!(mgr.finalize_ready(t1 + Duration::from_millis(100))); + + // complete() clears the backoff so the next checkpoint starts clean. + mgr.complete(); + assert!(mgr.finalize_ready(t1)); + } + + #[test] + fn test_finalize_backoff_caps_at_5s() { + let trigger = make_trigger(300, 256 * 1024 * 1024, 0.9); + let mut mgr = CheckpointManager::new(trigger); + let t = Instant::now(); + // Many consecutive failures must saturate at the 5s cap, not overflow. + for _ in 0..40 { + mgr.note_finalize_failed(t); + } + assert!(!mgr.finalize_ready(t + Duration::from_millis(4999))); + assert!(mgr.finalize_ready(t + Duration::from_secs(5))); + } + #[test] fn test_checkpoint_advance_flush_then_finalize() { let trigger = make_trigger(300, 256 * 1024 * 1024, 0.9); diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 243bc946..a56ed703 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -1104,6 +1104,15 @@ pub(crate) async fn handle_connection_sharded_monoio< // --- MULTI queue mode: queue commands when in transaction --- if conn.in_multi { + // FT.* vector commands aren't wired through the txn execution + // path; reject them explicitly inside MULTI (matches + // handler_single) instead of an incidental later error. + if cmd.len() > 3 && cmd[..3].eq_ignore_ascii_case(b"FT.") { + responses.push(Frame::Error(Bytes::from_static( + b"ERR FT.* commands are not supported inside MULTI/EXEC", + ))); + continue; + } conn.command_queue.push(frame); responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); continue; diff --git a/src/server/conn/handler_monoio/txn.rs b/src/server/conn/handler_monoio/txn.rs index 2ab9ad5b..7c6e1a93 100644 --- a/src/server/conn/handler_monoio/txn.rs +++ b/src/server/conn/handler_monoio/txn.rs @@ -177,8 +177,10 @@ pub(super) async fn try_handle_txn_commit( mq_lost.get_or_insert((owner, intent_count)); continue; } - // Await the ack before replying OK to the client. - let _ = reply_rx.recv().await; + // Await the ack before replying OK to the client + // (bounded: a stuck owner shard must not hang the + // connection forever). + let _ = crate::shard::coordinator::recv_reply_bounded(reply_rx).await; } } if let Some((owner, intent_count)) = mq_lost { @@ -305,12 +307,13 @@ pub(super) async fn try_handle_temporal_invalidate( &ctx.spsc_notifiers, ) .await; - let response = match reply_rx.recv().await { - Ok(f) => f, - Err(_) => Frame::Error(Bytes::from_static( - b"ERR cross-shard reply channel closed", - )), - }; + let response = + match crate::shard::coordinator::recv_reply_bounded(reply_rx).await { + Ok(f) => f, + Err(_) => Frame::Error(Bytes::from_static( + b"ERR cross-shard reply channel closed", + )), + }; responses.push(response); return true; } diff --git a/src/server/conn/handler_monoio/write.rs b/src/server/conn/handler_monoio/write.rs index 0ef94d06..7d1bc355 100644 --- a/src/server/conn/handler_monoio/write.rs +++ b/src/server/conn/handler_monoio/write.rs @@ -166,7 +166,8 @@ pub(super) async fn try_handle_ws_command( &ctx.spsc_notifiers, ) .await; - let _ = reply_rx.recv().await; + let _ = crate::shard::coordinator::recv_reply_bounded(reply_rx) + .await; } } responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); @@ -568,7 +569,7 @@ async fn mq_hop_or_local( &ctx.spsc_notifiers, ) .await; - match reply_rx.recv().await { + match crate::shard::coordinator::recv_reply_bounded(reply_rx).await { Ok(f) => f, Err(_) => crate::protocol::Frame::Error(bytes::Bytes::from_static( b"ERR cross-shard MQ reply channel closed", @@ -842,12 +843,13 @@ pub(super) async fn try_handle_graph_command( &ctx.spsc_notifiers, ) .await; - let mut response = match reply_rx.recv().await { - Ok(f) => f, - Err(_) => Frame::Error(bytes::Bytes::from_static( - b"ERR cross-shard reply channel closed", - )), - }; + let mut response = + match crate::shard::coordinator::recv_reply_bounded(reply_rx).await { + Ok(f) => f, + Err(_) => Frame::Error(bytes::Bytes::from_static( + b"ERR cross-shard reply channel closed", + )), + }; // Phase 166: explicit ADDNODE/ADDEDGE intents are captured // from the routed RESPONSE id, exactly like the local path; // the abort path routes the rollback back to the owner. diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index f488c97d..d246708b 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -1044,6 +1044,16 @@ pub(crate) async fn handle_connection_sharded_inner< // --- MULTI queue mode --- if conn.in_multi { + // FT.* vector commands aren't wired through the txn + // execution path; reject them explicitly inside MULTI + // (matches handler_single) instead of an incidental + // later error. + if cmd.len() > 3 && cmd[..3].eq_ignore_ascii_case(b"FT.") { + responses.push(Frame::Error(Bytes::from_static( + b"ERR FT.* commands are not supported inside MULTI/EXEC", + ))); + continue; + } conn.command_queue.push(frame); responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); continue; diff --git a/src/server/conn/handler_sharded/txn.rs b/src/server/conn/handler_sharded/txn.rs index cbdd2eeb..9775b196 100644 --- a/src/server/conn/handler_sharded/txn.rs +++ b/src/server/conn/handler_sharded/txn.rs @@ -193,7 +193,7 @@ pub(super) async fn try_handle_txn_commit( mq_lost.get_or_insert((owner, intent_count)); continue; } - match reply_rx.recv().await { + match crate::shard::coordinator::recv_reply_bounded(reply_rx).await { Ok(()) => {} Err(_) => { tracing::warn!( @@ -329,12 +329,13 @@ pub(super) async fn try_handle_temporal_invalidate( &ctx.spsc_notifiers, ) .await; - let response = match reply_rx.recv().await { - Ok(f) => f, - Err(_) => Frame::Error(Bytes::from_static( - b"ERR cross-shard reply channel closed", - )), - }; + let response = + match crate::shard::coordinator::recv_reply_bounded(reply_rx).await { + Ok(f) => f, + Err(_) => Frame::Error(Bytes::from_static( + b"ERR cross-shard reply channel closed", + )), + }; responses.push(response); return true; } diff --git a/src/server/conn/handler_sharded/write.rs b/src/server/conn/handler_sharded/write.rs index 0a6a5302..6ed78672 100644 --- a/src/server/conn/handler_sharded/write.rs +++ b/src/server/conn/handler_sharded/write.rs @@ -159,7 +159,7 @@ pub(super) async fn try_handle_ws_command( &ctx.spsc_notifiers, ) .await; - match reply_rx.recv().await { + match crate::shard::coordinator::recv_reply_bounded(reply_rx).await { Ok(count) => { tracing::debug!( "WS.DROP cleanup: deleted {} keys on shard {}", @@ -343,7 +343,7 @@ async fn mq_dispatch_to_owner( &ctx.spsc_notifiers, ) .await; - match reply_rx.recv().await { + match crate::shard::coordinator::recv_reply_bounded(reply_rx).await { Ok(f) => f, Err(_) => Frame::Error(bytes::Bytes::from_static( b"ERR cross-shard MQ reply channel closed", @@ -790,12 +790,13 @@ pub(super) async fn try_handle_graph_command( &ctx.spsc_notifiers, ) .await; - let mut response = match reply_rx.recv().await { - Ok(f) => f, - Err(_) => Frame::Error(bytes::Bytes::from_static( - b"ERR cross-shard reply channel closed", - )), - }; + let mut response = + match crate::shard::coordinator::recv_reply_bounded(reply_rx).await { + Ok(f) => f, + Err(_) => Frame::Error(bytes::Bytes::from_static( + b"ERR cross-shard reply channel closed", + )), + }; // Phase 166: explicit ADDNODE/ADDEDGE intents are captured // from the routed RESPONSE id, exactly like the local path; // the abort path routes the rollback back to the owner. diff --git a/src/server/conn/shared.rs b/src/server/conn/shared.rs index 75e8a4db..8230383b 100644 --- a/src/server/conn/shared.rs +++ b/src/server/conn/shared.rs @@ -664,16 +664,38 @@ pub(crate) enum TxnLocality { /// honored identically to the normal routing path. pub(crate) fn analyze_txn_locality(command_queue: &[Frame], num_shards: usize) -> TxnLocality { let mut owner: Option = None; + let visit = |key: &[u8], owner: &mut Option| -> bool { + let s = crate::shard::dispatch::key_to_shard(key, num_shards); + match *owner { + None => { + *owner = Some(s); + true + } + Some(existing) => existing == s, + } + }; for frame in command_queue { let Some((cmd, args)) = extract_command(frame) else { continue; }; for key in crate::tracking::invalidation::command_keys(cmd, args) { - let s = crate::shard::dispatch::key_to_shard(&key, num_shards); - match owner { - None => owner = Some(s), - Some(existing) if existing != s => return TxnLocality::CrossShard, - _ => {} + if !visit(&key, &mut owner) { + return TxnLocality::CrossShard; + } + } + // Prod-hardening #15: SORT/GEORADIUS/GEORADIUSBYMEMBER declare only + // their SOURCE key in command metadata (first_key==last_key==1); the + // optional `STORE`/`STOREDIST` destination is positional (the arg + // after the token) and is NOT covered by `command_keys`. Without + // this, `MULTI; SORT src STORE dst; EXEC` where src and dst hash to + // different shards is misclassified as SingleShard(owner-of-src) and + // the whole body routes to that one shard — `dst` gets written to the + // WRONG shard's dataset and is invisible to a later normally-routed + // `GET dst`. Detecting the STORE dest here forces CrossShard, which + // the caller rejects with CROSSSLOT instead of silently misrouting. + if let Some(dest) = store_clause_dest(cmd, args) { + if !visit(dest, &mut owner) { + return TxnLocality::CrossShard; } } } @@ -683,6 +705,36 @@ pub(crate) fn analyze_txn_locality(command_queue: &[Frame], num_shards: usize) - } } +/// Extract the `STORE`/`STOREDIST` destination key from a SORT or GEORADIUS +/// family command, if present. Returns `None` for commands that have no such +/// clause or when the token has no following argument. +/// +/// The destination is positional (the argument immediately after a `STORE` +/// or `STOREDIST` token in `args`, where `args` excludes the command name), +/// so it cannot be captured by the fixed `first_key/last_key/step` metadata +/// key specs the normal routing path uses. +fn store_clause_dest<'a>(cmd: &[u8], args: &'a [Frame]) -> Option<&'a [u8]> { + let is_store_cmd = cmd.eq_ignore_ascii_case(b"SORT") + || cmd.eq_ignore_ascii_case(b"GEORADIUS") + || cmd.eq_ignore_ascii_case(b"GEORADIUSBYMEMBER"); + if !is_store_cmd { + return None; + } + let mut i = 0; + while i < args.len() { + if let Frame::BulkString(tok) = &args[i] { + if tok.eq_ignore_ascii_case(b"STORE") || tok.eq_ignore_ascii_case(b"STOREDIST") { + if let Some(Frame::BulkString(dest)) = args.get(i + 1) { + return Some(dest.as_ref()); + } + return None; + } + } + i += 1; + } + None +} + #[cfg(test)] mod as_of_tests { //! Unit tests for `resolve_ft_search_as_of_lsn` (TEMP-04 + ACID-09). @@ -922,4 +974,60 @@ mod txn_locality_tests { let q = [cmd(&["MSET", "m0", "1", &other, "2"])]; assert_eq!(analyze_txn_locality(&q, num), TxnLocality::CrossShard); } + + #[test] + fn sort_store_dest_on_other_shard_is_cross_shard() { + // Prod-hardening #15: `SORT src STORE dst` where src and dst hash to + // different shards must classify CrossShard (→ CROSSSLOT reject), not + // silently SingleShard(owner-of-src) and misroute the STORE write. + let num = 8; + let base = crate::shard::dispatch::key_to_shard(b"src", num); + let mut dst = None; + for i in 0..1000 { + let k = format!("dst{i}"); + if crate::shard::dispatch::key_to_shard(k.as_bytes(), num) != base { + dst = Some(k); + break; + } + } + let dst = dst.expect("a dst on a different shard than src must exist"); + let q = [cmd(&["SORT", "src", "STORE", &dst])]; + assert_eq!(analyze_txn_locality(&q, num), TxnLocality::CrossShard); + } + + #[test] + fn sort_store_dest_same_shard_stays_single_shard() { + // Co-located source + STORE dest (hash tag) stays SingleShard — the + // STORE clause must not spuriously force CrossShard. + let q = [cmd(&["SORT", "{t}:src", "STORE", "{t}:dst"])]; + let owner = crate::shard::dispatch::key_to_shard(b"t", 8); + assert_eq!(analyze_txn_locality(&q, 8), TxnLocality::SingleShard(owner)); + } + + #[test] + fn georadius_storedist_dest_participates_in_locality() { + // GEORADIUS ... STOREDIST dst — the dest key must be considered too. + let num = 8; + let base = crate::shard::dispatch::key_to_shard(b"geo", num); + let mut dst = None; + for i in 0..1000 { + let k = format!("gd{i}"); + if crate::shard::dispatch::key_to_shard(k.as_bytes(), num) != base { + dst = Some(k); + break; + } + } + let dst = dst.expect("a dst on a different shard must exist"); + let q = [cmd(&[ + "GEORADIUS", + "geo", + "15", + "37", + "200", + "km", + "STOREDIST", + &dst, + ])]; + assert_eq!(analyze_txn_locality(&q, num), TxnLocality::CrossShard); + } } diff --git a/src/server/listener.rs b/src/server/listener.rs index 445342d5..6bfcdbae 100644 --- a/src/server/listener.rs +++ b/src/server/listener.rs @@ -492,6 +492,27 @@ pub async fn run_sharded( affinity_tracker: Arc>, ) -> anyhow::Result<()> { let addr = format!("{}:{}", config.bind, config.port); + // Prod-hardening #16: the per-shard monoio listeners bind this same + // `bind:port` via SO_REUSEPORT (`create_reuseport_socket`, see + // event_loop.rs). Linux requires EVERY socket sharing a REUSEPORT port — + // including this central one — to set the option; a plain `bind` here + // either loses the race (EADDRINUSE, propagated out of run_sharded) or + // wins it and forces every shard's REUSEPORT bind to fail, collapsing the + // per-shard accept fast-path onto this single loop. Join the group like + // the tokio sibling does. On non-unix (no socket2 REUSEPORT) fall back. + #[cfg(unix)] + let listener = match crate::shard::conn_accept::create_reuseport_socket(&addr) { + Ok(std_listener) => monoio::net::TcpListener::from_std(std_listener)?, + Err(e) => { + tracing::warn!( + "central SO_REUSEPORT bind failed ({}); falling back to plain bind on {}", + e, + addr + ); + monoio::net::TcpListener::bind(&addr)? + } + }; + #[cfg(not(unix))] let listener = monoio::net::TcpListener::bind(&addr)?; let num_shards = conn_txs.len(); info!("Listening on {} ({} shards, monoio)", addr, num_shards); diff --git a/src/shard/conn_accept.rs b/src/shard/conn_accept.rs index 5e749246..9a863437 100644 --- a/src/shard/conn_accept.rs +++ b/src/shard/conn_accept.rs @@ -31,6 +31,15 @@ use super::shared_databases::ShardDatabases; /// from parking_lot::RwLock (used for pubsub types). type StdRwLock = std::sync::RwLock; +/// Maximum time a client may take to complete the TLS handshake after the +/// TCP connection is accepted (prod-hardening #17). Without a bound, a peer +/// that completes the TCP 3-way handshake on the TLS port and then sends no +/// (or a partial) ClientHello parks the task, its fd, AND its already-consumed +/// `maxclients` slot indefinitely — a zero-CPU slow-loris that can exhaust the +/// connection limit for legitimate TLS clients. On expiry we drop the stream +/// and release the slot via `record_connection_closed()`. +const TLS_HANDSHAKE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); + /// Process-wide TCP listen backlog (S-5, `--tcp-backlog`). Set once at /// startup before any listener binds; 1024 is the historical hardcoded /// value. A process-wide static (rather than threading a param through the @@ -272,8 +281,10 @@ pub(crate) fn spawn_tokio_connection( #[cfg(not(unix))] let kill_fd = -1; let acceptor = tokio_rustls::TlsAcceptor::from(tls_cfg); - match acceptor.accept(tcp_stream).await { - Ok(tls_stream) => { + // #17: bound the handshake so a stalled ClientHello cannot pin the + // fd + maxclients slot forever. + match tokio::time::timeout(TLS_HANDSHAKE_TIMEOUT, acceptor.accept(tcp_stream)).await { + Ok(Ok(tls_stream)) => { let _ = handle_connection_sharded_inner( tls_stream, peer_addr, @@ -287,9 +298,16 @@ pub(crate) fn spawn_tokio_connection( ) .await; } - Err(e) => { + Ok(Err(e)) => { tracing::warn!("Shard {}: TLS handshake failed: {}", shard_id, e); } + Err(_) => { + tracing::warn!( + "Shard {}: TLS handshake timed out after {}s", + shard_id, + TLS_HANDSHAKE_TIMEOUT.as_secs() + ); + } } crate::admin::metrics_setup::record_connection_closed(); }); @@ -655,8 +673,12 @@ pub(crate) fn spawn_monoio_connection( return; } let acceptor = monoio_rustls::TlsAcceptor::from(tls_cfg); - match acceptor.accept(tcp_stream).await { - Ok(tls_stream) => { + // #17: bound the handshake so a stalled ClientHello cannot + // pin the fd + maxclients slot forever. + match monoio::time::timeout(TLS_HANDSHAKE_TIMEOUT, acceptor.accept(tcp_stream)) + .await + { + Ok(Ok(tls_stream)) => { let _ = handle_connection_sharded_monoio( tls_stream, peer_addr, @@ -671,13 +693,20 @@ pub(crate) fn spawn_monoio_connection( ) .await; } - Err(e) => { + Ok(Err(e)) => { tracing::warn!( "Shard {}: Monoio TLS handshake failed: {}", shard_id, e ); } + Err(_) => { + tracing::warn!( + "Shard {}: Monoio TLS handshake timed out after {}s", + shard_id, + TLS_HANDSHAKE_TIMEOUT.as_secs() + ); + } } crate::admin::metrics_setup::record_connection_closed(); }); diff --git a/src/shard/coordinator.rs b/src/shard/coordinator.rs index 30dd07ab..82222a1b 100644 --- a/src/shard/coordinator.rs +++ b/src/shard/coordinator.rs @@ -177,6 +177,42 @@ fn run_local( crate::shard::slice::with_shard_db(db_index, run) } +/// Maximum time to wait for a cross-shard reply after the request was +/// successfully pushed into the target shard's ring buffer (prod-hardening +/// #11). `spsc_send`'s retry budget only bounds getting the message INTO the +/// ring; once the push succeeds, if the target shard then stalls while +/// executing the command (a wedged-disk fsync during snapshot/AOF, an +/// uninterruptible D-state I/O stall, or a dead shard), `reply_rx.recv().await` +/// has no way to wake on its own and the awaiting client connection hangs +/// forever with no response and no cancel. This is a safety ceiling far above +/// any legitimate cross-shard command latency (including group-commit fsync +/// under load) — it exists only so a genuinely wedged shard surfaces an error +/// instead of an unbounded hang. +const XSHARD_REPLY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + +/// Await a cross-shard `reply_rx` with a bounded timeout (#11). +/// +/// Returns `Err(RecvError)` on EITHER a closed channel (the target shard +/// dropped the reply sender) OR expiry of [`XSHARD_REPLY_TIMEOUT`] — both mean +/// "no usable reply", so every call site's existing error handling applies +/// unchanged; the point is that neither case can hang the connection forever. +pub(crate) async fn recv_reply_bounded( + reply_rx: channel::OneshotReceiver, +) -> Result { + use crate::runtime::race::{Arm, race2}; + let recv = std::pin::pin!(reply_rx); + #[cfg(feature = "runtime-tokio")] + let sleep = std::pin::pin!(tokio::time::sleep(XSHARD_REPLY_TIMEOUT)); + #[cfg(feature = "runtime-monoio")] + let sleep = std::pin::pin!(monoio::time::sleep(XSHARD_REPLY_TIMEOUT)); + // race2 polls the recv arm first, so a ready reply always wins the tie and + // the timer future is dropped un-fired (cheap deregister on both runtimes). + match race2(recv, sleep).await { + Arm::First(r) => r, + Arm::Second(()) => Err(channel::RecvError), + } +} + /// Send one full command to a REMOTE shard and await its reply. async fn run_remote( target_shard: usize, @@ -196,7 +232,7 @@ async fn run_remote( reply_tx, }; let _ = spsc_send(dispatch_tx, my_shard, target_shard, msg, spsc_notifiers).await; - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(mut frames) if !frames.is_empty() => frames.swap_remove(0), _ => Frame::Error(Bytes::from_static(b"ERR cross-shard reply channel closed")), } @@ -227,7 +263,7 @@ pub(crate) async fn execute_txn_on_owner( }; let msg = ShardMessage::TxnExecute(Box::new(payload)); let _ = spsc_send(dispatch_tx, my_shard, owner, msg, spsc_notifiers).await; - reply_rx.recv().await.ok() + recv_reply_bounded(reply_rx).await.ok() } /// Run one full command on whichever shard owns `routing_key`. @@ -536,7 +572,7 @@ async fn coordinate_bitop( } } for (indices, reply_rx) in pending { - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(frames) => { for (idx, frame) in indices.into_iter().zip(frames) { match frame { @@ -977,7 +1013,7 @@ pub(crate) async fn coordinate_flush_broadcast( failed.get_or_insert(target); continue; } - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(frames) if frames.iter().all(|f| !matches!(f, Frame::Error(_))) => {} _ => { failed.get_or_insert(target); @@ -1088,7 +1124,7 @@ async fn coordinate_mget( // Await all remote results for (indices, reply_rx) in pending_shards { - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(frames) => { for (idx, frame) in indices.into_iter().zip(frames) { results[idx] = Some(frame); @@ -1216,8 +1252,28 @@ async fn coordinate_mset( } } + // Drain ALL remote acks even after a failure (every leg was already + // dispatched, and the local leg below must still persist what this shard + // applied) — but a timed-out, closed, or errored leg must NOT collapse + // into OK: that would acknowledge an unconfirmed distributed write. + let mut leg_err: Option = None; for reply_rx in pending_shards { - let _ = reply_rx.recv().await; + match recv_reply_bounded(reply_rx).await { + Ok(frames) => { + if leg_err.is_none() + && let Some(err) = frames.into_iter().find(|f| matches!(f, Frame::Error(_))) + { + leg_err = Some(err); + } + } + Err(_) => { + if leg_err.is_none() { + leg_err = Some(Frame::Error(Bytes::from_static( + b"ERR cross-shard MSET leg unconfirmed (timeout or closed reply channel); write may be partially applied", + ))); + } + } + } } // Local leg (review Finding 1): persist a synthesized MSET over ONLY the @@ -1235,6 +1291,9 @@ async fn coordinate_mset( } } + if let Some(err) = leg_err { + return err; + } Frame::SimpleString(Bytes::from_static(b"OK")) } @@ -1455,7 +1514,7 @@ async fn coordinate_multi_del_or_exists( } for reply_rx in pending_shards { - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(frames) => { for frame in frames { match frame { @@ -1536,7 +1595,7 @@ pub async fn coordinate_keys( // Collect remote results for reply_rx in pending_shards { - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(Frame::Array(keys)) => all_keys.extend(keys), Ok(_) => {} // Non-array response (e.g., error) — skip Err(_) => { @@ -1620,7 +1679,7 @@ pub async fn coordinate_scan( reply_tx, }; let _ = spsc_send(dispatch_tx, my_shard, target_shard_id, msg, spsc_notifiers).await; - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(frame) => frame, Err(_) => Frame::Error(Bytes::from_static( b"ERR cross-shard reply channel closed during SCAN", @@ -1706,7 +1765,7 @@ pub async fn coordinate_dbsize( } for reply_rx in pending_shards { - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(Frame::Integer(n)) => total += n, Ok(_) => {} // Non-integer response — skip Err(_) => { @@ -1766,7 +1825,7 @@ pub async fn coordinate_hotkeys( } for reply_rx in pending_shards { - match reply_rx.recv().await { + match recv_reply_bounded(reply_rx).await { Ok(Frame::Array(entries)) => { for entry in entries.iter() { if let Frame::Array(pair) = entry diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 01e5d77c..502210b8 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -36,6 +36,18 @@ use super::shared_databases::ShardDatabases; use super::uring_handler; use super::{conn_accept, persistence_tick, spsc_handler, timers}; +/// EXPERIMENTAL DiskANN cold-tier gate. The WARM->COLD transition + on-disk +/// Vamana beam search are incomplete (no cold-segment deletion, ADC-only +/// recall, restart reload of PQ codebooks unfinished — see the audit), so the +/// transition is a no-op unless an operator explicitly opts in via +/// `MOON_VEC_COLD_TIER=1`. Warm-tier mmap + LRU eviction (`--vec-warm-mmap-budget`) +/// handles out-of-RAM indexes by default. +fn cold_tier_experimental_enabled() -> bool { + std::env::var("MOON_VEC_COLD_TIER") + .map(|v| v == "1" || v.eq_ignore_ascii_case("enable")) + .unwrap_or(false) +} + impl super::Shard { /// Run the shard event loop on its dedicated current_thread runtime. /// @@ -330,7 +342,14 @@ impl super::Shard { // which relinquished ownership. We take sole ownership here. unsafe { std::net::TcpStream::from_raw_fd(fd) } }; - if tx.send(std_stream).is_err() { + // send_async (not the blocking send): the accept + // task shares this shard's single monoio thread + // with the event-loop consumer, so a blocking + // send on a full channel would stall the loop + // that drains it -> shard deadlock. Awaiting + // yields cooperatively and keeps backpressure + // (no dropped connections). + if tx.send_async(std_stream).await.is_err() { break; // receiver dropped, shard shutting down } } @@ -1568,7 +1587,10 @@ impl super::Shard { } // Cold tier transition check (60s, disk-offload only) _ = cold_check_interval.0.tick() => { - if server_config.disk_offload_enabled() && server_config.segment_cold_after > 0 { + if cold_tier_experimental_enabled() + && server_config.disk_offload_enabled() + && server_config.segment_cold_after > 0 + { if let Some(ref mut manifest) = shard_manifest { let shard_dir = server_config.effective_disk_offload_dir() .join(format!("shard-{}", shard_id)); @@ -2228,7 +2250,9 @@ impl super::Shard { } // Cold tier check: every cold_poll_secs * 1000 ticks if monoio_tick_counter % (cold_poll_secs as u64 * 1000) == 0 { - if server_config.disk_offload_enabled() && server_config.segment_cold_after > 0 + if cold_tier_experimental_enabled() + && server_config.disk_offload_enabled() + && server_config.segment_cold_after > 0 { if let Some(ref mut manifest) = shard_manifest { let shard_dir = server_config diff --git a/src/shard/persistence_tick.rs b/src/shard/persistence_tick.rs index 0e5b38cf..2150b3c7 100644 --- a/src/shard/persistence_tick.rs +++ b/src/shard/persistence_tick.rs @@ -1044,6 +1044,26 @@ pub(crate) fn handle_checkpoint_tick( false } CheckpointAction::Finalize { redo_lsn } => { + // Prod-hardening #13: a failed finalize leaves the state machine in + // `Finalizing`, so the next 1ms tick re-enters this arm and + // re-appends a WAL Checkpoint record (step 1) with no backoff. + // Under a sustained failure (slow/degraded disk causing repeated + // `wait_durable` timeouts, or a persistent `graph_save` failure) + // this floods the WAL with a Checkpoint marker every millisecond, + // and because `last_checkpoint_lsn` never advances, + // `recycle_segments_before` never fires — WAL usage grows fastest + // exactly during the disk-pressure incident. Gate re-attempts on an + // exponential backoff so a stuck finalize retries on a bounded + // schedule instead of hammering every tick. + // `now` gates readiness only; failure branches below re-stamp + // Instant::now() because wait_durable/commit/graph_save can take + // longer than the backoff — arming from this pre-work timestamp + // would put the retry deadline in the past (instant re-retry). + let now = std::time::Instant::now(); + if !checkpoint_mgr.finalize_ready(now) { + return false; + } + // 1. Write WAL checkpoint record with redo_lsn payload let mut payload = [0u8; 8]; payload.copy_from_slice(&redo_lsn.to_le_bytes()); @@ -1059,12 +1079,14 @@ pub(crate) fn handle_checkpoint_tick( crate::persistence::wal_v3::segment::WAIT_DURABLE_TIMEOUT, ) { tracing::error!("Checkpoint WAL flush failed: {}", e); + checkpoint_mgr.note_finalize_failed(std::time::Instant::now()); return false; } // 3. Commit manifest (atomic dual-root write) if let Err(e) = manifest.commit() { tracing::error!("Checkpoint manifest commit failed: {}", e); + checkpoint_mgr.note_finalize_failed(std::time::Instant::now()); return false; } @@ -1109,6 +1131,7 @@ pub(crate) fn handle_checkpoint_tick( // exactly on `current_lsn()` and must NOT be skipped). if !graph_save(wal.current_lsn().saturating_sub(1)) { tracing::error!("Checkpoint aborted: graph snapshot failed"); + checkpoint_mgr.note_finalize_failed(std::time::Instant::now()); return false; } @@ -1117,10 +1140,11 @@ pub(crate) fn handle_checkpoint_tick( control.last_checkpoint_epoch = manifest.epoch(); if let Err(e) = control.write(control_path) { tracing::error!("Checkpoint control file update failed: {}", e); + checkpoint_mgr.note_finalize_failed(std::time::Instant::now()); return false; } - // 5. Mark checkpoint complete + // 5. Mark checkpoint complete (also clears the finalize backoff). checkpoint_mgr.complete(); // 6. Recycle old WAL segments that are fully before redo_lsn diff --git a/src/storage/tiered/cold_tier.rs b/src/storage/tiered/cold_tier.rs index adb0915c..85e932b6 100644 --- a/src/storage/tiered/cold_tier.rs +++ b/src/storage/tiered/cold_tier.rs @@ -250,7 +250,7 @@ pub fn transition_to_cold( graph.entry_point(), graph.max_degree(), cold_file_id, - ); + )?; Ok(segment) } diff --git a/src/transaction/abort.rs b/src/transaction/abort.rs index 8ec9b44b..817b3358 100644 --- a/src/transaction/abort.rs +++ b/src/transaction/abort.rs @@ -508,7 +508,10 @@ pub async fn abort_cross_store_txn_routed( ); continue; } - if reply_rx.recv().await.is_err() { + if crate::shard::coordinator::recv_reply_bounded(reply_rx) + .await + .is_err() + { tracing::warn!( txn_id, owner, diff --git a/src/vector/diskann/segment.rs b/src/vector/diskann/segment.rs index 90e68f41..427412a2 100644 --- a/src/vector/diskann/segment.rs +++ b/src/vector/diskann/segment.rs @@ -72,15 +72,18 @@ impl DiskAnnSegment { entry_point: u32, max_degree: u32, file_id: u64, - ) -> Self { + ) -> std::io::Result { debug_assert_eq!( pq_codes.len(), num_vectors as usize * pq.m(), "pq_codes length must be num_vectors * m" ); + // Propagate the open error instead of panicking: this runs on the shard + // event loop during a WARM->COLD transition, and a library-code panic + // here would abort the whole shard. The caller skips + logs the + // transition on error. #[cfg(unix)] - let vamana_file = std::fs::File::open(&vamana_path) - .unwrap_or_else(|e| panic!("DiskAnnSegment: cannot open {:?}: {}", vamana_path, e)); + let vamana_file = std::fs::File::open(&vamana_path)?; // Try to open with O_DIRECT for io_uring beam search. Falls back // gracefully on filesystems that don't support O_DIRECT (e.g., tmpfs @@ -102,7 +105,7 @@ impl DiskAnnSegment { } }; - Self { + Ok(Self { pq_codes, pq, vamana_path, @@ -115,7 +118,7 @@ impl DiskAnnSegment { entry_point, max_degree, file_id, - } + }) } /// Load a DiskAnnSegment from on-disk files. @@ -183,6 +186,17 @@ impl DiskAnnSegment { }) } + /// Number of vectors stored in this cold-tier segment. + /// + /// DiskANN segments have no tombstone/MVCC layer yet (see prod-hardening + /// #19), so this is the raw stored count — used by FT.INFO `num_docs` + /// (#28) so cold-tier documents are counted at all rather than silently + /// dropped from the total once a segment ages into the cold tier. + #[inline] + pub fn num_docs(&self) -> usize { + self.num_vectors as usize + } + /// Approximate nearest neighbor search using PQ asymmetric distance /// and buffered Vamana beam traversal from disk. /// @@ -571,7 +585,8 @@ mod tests { graph.entry_point(), graph.max_degree(), 1, - ); + ) + .expect("build DiskAnnSegment"); (seg, vectors, tmp) } @@ -643,7 +658,8 @@ mod tests { 0, 4, 0, - ); + ) + .expect("build DiskAnnSegment"); let results = seg.search(&[0.0; 32], 5, 8); assert!(results.is_empty()); } @@ -733,7 +749,8 @@ mod tests { graph.entry_point(), graph.max_degree(), 1, - ); + ) + .expect("build DiskAnnSegment"); // If uring is None (tmpfs / O_DIRECT unsupported), skip gracefully. if !seg.has_uring() { diff --git a/src/vector/persistence/mod.rs b/src/vector/persistence/mod.rs index 2f02c010..15f58e11 100644 --- a/src/vector/persistence/mod.rs +++ b/src/vector/persistence/mod.rs @@ -9,6 +9,5 @@ pub mod recover_v2; pub mod sealed_mmap; pub mod segment_io; pub mod unloaded_segment; -pub mod wal_record; pub mod warm_search; pub mod warm_segment; diff --git a/src/vector/persistence/wal_record.rs b/src/vector/persistence/wal_record.rs deleted file mode 100644 index bec5c595..00000000 --- a/src/vector/persistence/wal_record.rs +++ /dev/null @@ -1,459 +0,0 @@ -//! Vector WAL record format with manual LE serialization and CRC32 framing. -//! -//! Frame format: -//! ```text -//! [u8: VECTOR_RECORD_TAG = 0x56] -- distinguishes from RESP block frames -//! [u32 LE: payload_len] -- length of record_type + payload bytes -//! [u8: record_type] -- 0=Upsert, 1=Delete, 2=TxnCommit, 3=TxnAbort, 4=Checkpoint -//! [payload bytes] -- record-specific fields, all LE -//! [u32 LE: crc32] -- CRC32 over record_type + payload -//! ``` -//! -//! **NOT wired on the live recovery path.** Vector-index recovery authority -//! is the B1/B2/B3 durability layout instead: `manifest.json` + -//! `keymap-.bin` + `segment-/` dirs -//! (`crate::vector::persistence::manifest`, `segment_io`), loaded by -//! `crate::vector::persistence::recover_v2` and reconciled by a dedup rescan -//! of the live keyspace. These record types are kept only because unit tests -//! in this module exercise their (de)serialization round-trip. Do not enable -//! WAL-based vector replay without first disabling the B3 rescan — the two -//! recovery mechanisms are not designed to run together (double-apply risk). - -/// Tag byte distinguishing vector WAL records from RESP block frames. -pub const VECTOR_RECORD_TAG: u8 = 0x56; // 'V' - -/// Error type for WAL record serialization/deserialization. -#[derive(Debug)] -pub enum WalRecordError { - Truncated, - InvalidTag(u8), - InvalidRecordType(u8), - CrcMismatch { expected: u32, actual: u32 }, - DeserializeFailed(String), -} - -impl std::fmt::Display for WalRecordError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Truncated => write!(f, "WAL record truncated"), - Self::InvalidTag(t) => write!(f, "invalid WAL record tag: 0x{t:02x}"), - Self::InvalidRecordType(t) => write!(f, "invalid WAL record type: {t}"), - Self::CrcMismatch { expected, actual } => { - write!( - f, - "CRC mismatch: expected 0x{expected:08x}, got 0x{actual:08x}" - ) - } - Self::DeserializeFailed(msg) => write!(f, "deserialize failed: {msg}"), - } - } -} - -/// Structured WAL record for vector operations. -/// -/// Each variant captures all fields needed to replay the operation during -/// crash recovery. Serialized with manual LE encoding (no serde/bincode) -/// for predictable format and zero overhead. -#[derive(Debug, Clone, PartialEq)] -pub enum VectorWalRecord { - VectorUpsert { - txn_id: u64, - collection_id: u64, - point_id: u64, - sq_vector: Vec, - tq_code: Vec, - norm: f32, - f32_vector: Vec, - }, - VectorDelete { - txn_id: u64, - collection_id: u64, - point_id: u64, - }, - TxnCommit { - txn_id: u64, - commit_lsn: u64, - }, - TxnAbort { - txn_id: u64, - }, - Checkpoint { - segment_id: u64, - last_lsn: u64, - }, -} - -impl VectorWalRecord { - /// Returns the record type discriminant (0-4). - fn record_type(&self) -> u8 { - match self { - Self::VectorUpsert { .. } => 0, - Self::VectorDelete { .. } => 1, - Self::TxnCommit { .. } => 2, - Self::TxnAbort { .. } => 3, - Self::Checkpoint { .. } => 4, - } - } - - /// Serialize record-specific fields to a byte buffer (all LE). - fn serialize_payload(&self, buf: &mut Vec) { - match self { - Self::VectorUpsert { - txn_id, - collection_id, - point_id, - sq_vector, - tq_code, - norm, - f32_vector, - } => { - buf.extend_from_slice(&txn_id.to_le_bytes()); - buf.extend_from_slice(&collection_id.to_le_bytes()); - buf.extend_from_slice(&point_id.to_le_bytes()); - // sq_vector: len:u32 + raw i8 bytes - buf.extend_from_slice(&(sq_vector.len() as u32).to_le_bytes()); - for &v in sq_vector { - buf.push(v as u8); - } - // tq_code: len:u32 + raw bytes - buf.extend_from_slice(&(tq_code.len() as u32).to_le_bytes()); - buf.extend_from_slice(tq_code); - // norm: f32 LE - buf.extend_from_slice(&norm.to_le_bytes()); - // f32_vector: len:u32 + f32 LE values - buf.extend_from_slice(&(f32_vector.len() as u32).to_le_bytes()); - for &v in f32_vector { - buf.extend_from_slice(&v.to_le_bytes()); - } - } - Self::VectorDelete { - txn_id, - collection_id, - point_id, - } => { - buf.extend_from_slice(&txn_id.to_le_bytes()); - buf.extend_from_slice(&collection_id.to_le_bytes()); - buf.extend_from_slice(&point_id.to_le_bytes()); - } - Self::TxnCommit { txn_id, commit_lsn } => { - buf.extend_from_slice(&txn_id.to_le_bytes()); - buf.extend_from_slice(&commit_lsn.to_le_bytes()); - } - Self::TxnAbort { txn_id } => { - buf.extend_from_slice(&txn_id.to_le_bytes()); - } - Self::Checkpoint { - segment_id, - last_lsn, - } => { - buf.extend_from_slice(&segment_id.to_le_bytes()); - buf.extend_from_slice(&last_lsn.to_le_bytes()); - } - } - } - - /// Deserialize record-specific fields from a byte slice. - fn deserialize_payload(record_type: u8, data: &[u8]) -> Result { - let mut pos = 0; - - let read_u32 = |pos: &mut usize| -> Result { - if *pos + 4 > data.len() { - return Err(WalRecordError::Truncated); - } - let val = - u32::from_le_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]); - *pos += 4; - Ok(val) - }; - - let read_u64 = |pos: &mut usize| -> Result { - if *pos + 8 > data.len() { - return Err(WalRecordError::Truncated); - } - let val = u64::from_le_bytes([ - data[*pos], - data[*pos + 1], - data[*pos + 2], - data[*pos + 3], - data[*pos + 4], - data[*pos + 5], - data[*pos + 6], - data[*pos + 7], - ]); - *pos += 8; - Ok(val) - }; - - let read_f32 = |pos: &mut usize| -> Result { - if *pos + 4 > data.len() { - return Err(WalRecordError::Truncated); - } - let val = - f32::from_le_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]); - *pos += 4; - Ok(val) - }; - - match record_type { - 0 => { - let txn_id = read_u64(&mut pos)?; - let collection_id = read_u64(&mut pos)?; - let point_id = read_u64(&mut pos)?; - // sq_vector - let sq_len = read_u32(&mut pos)? as usize; - if pos + sq_len > data.len() { - return Err(WalRecordError::Truncated); - } - let sq_vector: Vec = data[pos..pos + sq_len].iter().map(|&b| b as i8).collect(); - pos += sq_len; - // tq_code - let tq_len = read_u32(&mut pos)? as usize; - if pos + tq_len > data.len() { - return Err(WalRecordError::Truncated); - } - let tq_code = data[pos..pos + tq_len].to_vec(); - pos += tq_len; - // norm - let norm = read_f32(&mut pos)?; - // f32_vector - let f32_len = read_u32(&mut pos)? as usize; - if pos + f32_len * 4 > data.len() { - return Err(WalRecordError::Truncated); - } - let mut f32_vector = Vec::with_capacity(f32_len); - for _ in 0..f32_len { - f32_vector.push(read_f32(&mut pos)?); - } - Ok(Self::VectorUpsert { - txn_id, - collection_id, - point_id, - sq_vector, - tq_code, - norm, - f32_vector, - }) - } - 1 => { - let txn_id = read_u64(&mut pos)?; - let collection_id = read_u64(&mut pos)?; - let point_id = read_u64(&mut pos)?; - Ok(Self::VectorDelete { - txn_id, - collection_id, - point_id, - }) - } - 2 => { - let txn_id = read_u64(&mut pos)?; - let commit_lsn = read_u64(&mut pos)?; - Ok(Self::TxnCommit { txn_id, commit_lsn }) - } - 3 => { - let txn_id = read_u64(&mut pos)?; - Ok(Self::TxnAbort { txn_id }) - } - 4 => { - let segment_id = read_u64(&mut pos)?; - let last_lsn = read_u64(&mut pos)?; - Ok(Self::Checkpoint { - segment_id, - last_lsn, - }) - } - _ => Err(WalRecordError::InvalidRecordType(record_type)), - } - } - - /// Build a complete WAL frame: TAG + payload_len + record_type + payload + CRC32. - pub fn to_wal_frame(&self) -> Vec { - let mut payload = Vec::with_capacity(64); - payload.push(self.record_type()); - self.serialize_payload(&mut payload); - - let payload_len = payload.len() as u32; - - // CRC32 over record_type + payload (the entire payload vec) - let mut hasher = crc32fast::Hasher::new(); - hasher.update(&payload); - let crc = hasher.finalize(); - - // Frame: TAG(1) + payload_len(4) + payload(N) + crc32(4) - let frame_len = 1 + 4 + payload.len() + 4; - let mut frame = Vec::with_capacity(frame_len); - frame.push(VECTOR_RECORD_TAG); - frame.extend_from_slice(&payload_len.to_le_bytes()); - frame.extend_from_slice(&payload); - frame.extend_from_slice(&crc.to_le_bytes()); - frame - } - - /// Parse a WAL frame from a byte slice. - /// - /// Returns `(record, bytes_consumed)` on success. - /// Verifies CRC32. Returns `Err` on CRC mismatch, truncation, or invalid data. - pub fn from_wal_frame(data: &[u8]) -> Result<(Self, usize), WalRecordError> { - // Minimum frame: TAG(1) + payload_len(4) + record_type(1) + crc32(4) = 10 - if data.len() < 10 { - return Err(WalRecordError::Truncated); - } - - // Tag - if data[0] != VECTOR_RECORD_TAG { - return Err(WalRecordError::InvalidTag(data[0])); - } - - // Payload length - let payload_len = u32::from_le_bytes([data[1], data[2], data[3], data[4]]) as usize; - let frame_len = 1 + 4 + payload_len + 4; // TAG + len + payload + crc - - if data.len() < frame_len { - return Err(WalRecordError::Truncated); - } - - // Payload slice: starts at offset 5, length = payload_len - let payload = &data[5..5 + payload_len]; - - if payload.is_empty() { - return Err(WalRecordError::Truncated); - } - - // CRC32 check - let stored_crc = u32::from_le_bytes([ - data[5 + payload_len], - data[5 + payload_len + 1], - data[5 + payload_len + 2], - data[5 + payload_len + 3], - ]); - let mut hasher = crc32fast::Hasher::new(); - hasher.update(payload); - let computed_crc = hasher.finalize(); - - if stored_crc != computed_crc { - return Err(WalRecordError::CrcMismatch { - expected: stored_crc, - actual: computed_crc, - }); - } - - // Record type is first byte of payload - let record_type = payload[0]; - let record_data = &payload[1..]; - - let record = Self::deserialize_payload(record_type, record_data)?; - Ok((record, frame_len)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_upsert_roundtrip() { - let record = VectorWalRecord::VectorUpsert { - txn_id: 42, - collection_id: 7, - point_id: 100, - sq_vector: vec![1, -2, 3, -4], - tq_code: vec![0xAB, 0xCD, 0xEF], - norm: 1.5, - f32_vector: vec![0.1, 0.2, 0.3], - }; - let frame = record.to_wal_frame(); - let (decoded, consumed) = VectorWalRecord::from_wal_frame(&frame).unwrap(); - assert_eq!(consumed, frame.len()); - assert_eq!(decoded, record); - } - - #[test] - fn test_delete_roundtrip() { - let record = VectorWalRecord::VectorDelete { - txn_id: 10, - collection_id: 5, - point_id: 99, - }; - let frame = record.to_wal_frame(); - let (decoded, _) = VectorWalRecord::from_wal_frame(&frame).unwrap(); - assert_eq!(decoded, record); - } - - #[test] - fn test_txn_commit_roundtrip() { - let record = VectorWalRecord::TxnCommit { - txn_id: 123, - commit_lsn: 456, - }; - let frame = record.to_wal_frame(); - let (decoded, _) = VectorWalRecord::from_wal_frame(&frame).unwrap(); - assert_eq!(decoded, record); - } - - #[test] - fn test_txn_abort_roundtrip() { - let record = VectorWalRecord::TxnAbort { txn_id: 789 }; - let frame = record.to_wal_frame(); - let (decoded, _) = VectorWalRecord::from_wal_frame(&frame).unwrap(); - assert_eq!(decoded, record); - } - - #[test] - fn test_checkpoint_roundtrip() { - let record = VectorWalRecord::Checkpoint { - segment_id: 55, - last_lsn: 9999, - }; - let frame = record.to_wal_frame(); - let (decoded, _) = VectorWalRecord::from_wal_frame(&frame).unwrap(); - assert_eq!(decoded, record); - } - - #[test] - fn test_crc_mismatch_returns_error() { - let record = VectorWalRecord::VectorDelete { - txn_id: 1, - collection_id: 2, - point_id: 3, - }; - let mut frame = record.to_wal_frame(); - let len = frame.len(); - frame[len - 1] ^= 0xFF; - match VectorWalRecord::from_wal_frame(&frame) { - Err(WalRecordError::CrcMismatch { .. }) => {} - other => panic!("expected CrcMismatch, got {:?}", other), - } - } - - #[test] - fn test_truncated_frame_returns_error() { - let record = VectorWalRecord::TxnCommit { - txn_id: 1, - commit_lsn: 2, - }; - let frame = record.to_wal_frame(); - match VectorWalRecord::from_wal_frame(&frame[..3]) { - Err(WalRecordError::Truncated) => {} - other => panic!("expected Truncated, got {:?}", other), - } - } - - #[test] - fn test_to_wal_frame_has_tag_and_length() { - let record = VectorWalRecord::TxnAbort { txn_id: 1 }; - let frame = record.to_wal_frame(); - assert_eq!(frame[0], VECTOR_RECORD_TAG); - let payload_len = u32::from_le_bytes([frame[1], frame[2], frame[3], frame[4]]); - assert_eq!(frame.len(), 1 + 4 + payload_len as usize + 4); - } - - #[test] - fn test_from_wal_frame_rejects_bad_tag() { - let record = VectorWalRecord::TxnAbort { txn_id: 1 }; - let mut frame = record.to_wal_frame(); - frame[0] = 0x00; - match VectorWalRecord::from_wal_frame(&frame) { - Err(WalRecordError::InvalidTag(0x00)) => {} - other => panic!("expected InvalidTag, got {:?}", other), - } - } -} diff --git a/src/vector/segment/compaction.rs b/src/vector/segment/compaction.rs index 50b72b20..78cdef42 100644 --- a/src/vector/segment/compaction.rs +++ b/src/vector/segment/compaction.rs @@ -1502,7 +1502,11 @@ fn merge_graph_union( // search across the original segments. let recall = verify_merge_recall(&graph, &tq_bfs, segments, collection, dim, n, seed); - if recall < recall_tolerance && recall > 0.0 { + // NOTE: no `&& recall > 0.0` guard — recall == 0.0 is a genuine total + // collapse that MUST abort. verify_merge_recall returns 1.0 (not 0.0) for + // the "too few vectors to measure" cases, so 0.0 only ever means a real + // measured zero-recall merge. + if recall < recall_tolerance { tracing::warn!( "merge recall {recall:.4} < tolerance {recall_tolerance:.4}; aborting merge" ); diff --git a/src/vector/segment/mutable.rs b/src/vector/segment/mutable.rs index caa6e9f5..4a0595bf 100644 --- a/src/vector/segment/mutable.rs +++ b/src/vector/segment/mutable.rs @@ -1274,6 +1274,22 @@ impl MutableSegment { self.inner.read().entries.len() } + /// Returns the number of LIVE (non-tombstoned) entries. + /// + /// [`Self::len`] counts every slot including tombstoned ones — + /// `mark_deleted_*` only flips `delete_lsn` in place and never removes + /// the entry until the next compaction. FT.INFO's `num_docs` must use + /// this (prod-hardening #28): otherwise up to `compact_threshold` + /// DEL/HSET-update dead docs silently inflate the reported count. + pub fn live_len(&self) -> usize { + self.inner + .read() + .entries + .iter() + .filter(|e| e.delete_lsn == 0) + .count() + } + /// Returns true if no entries. #[allow(dead_code)] pub fn is_empty(&self) -> bool { @@ -2354,6 +2370,26 @@ mod tests { assert_eq!(frozen.entries[0].delete_lsn, 42); } + #[test] + fn test_live_len_excludes_tombstoned() { + // Prod-hardening #28: live_len() must NOT count tombstoned entries, + // while len() keeps counting them until compaction. + distance::init(); + let col = make_collection(128); + let seg = MutableSegment::new(128, col); + for i in 0..5u64 { + seg.append(1000 + i, &make_f32_vector(128, (i + 1) as u32), i + 1); + } + assert_eq!(seg.len(), 5); + assert_eq!(seg.live_len(), 5); + // Tombstone two entries by key_hash — they stay in the Vec (len + // unchanged) but must drop out of live_len. + assert_eq!(seg.mark_deleted_by_key_hash(1000, 100), 1); + assert_eq!(seg.mark_deleted_by_key_hash(1002, 100), 1); + assert_eq!(seg.len(), 5, "len still counts tombstoned entries"); + assert_eq!(seg.live_len(), 3, "live_len drops the two tombstoned"); + } + #[test] fn test_mvcc_backward_compat() { distance::init(); diff --git a/src/vector/store.rs b/src/vector/store.rs index ce6f2ed2..cbb05771 100644 --- a/src/vector/store.rs +++ b/src/vector/store.rs @@ -2225,11 +2225,12 @@ impl VectorStore { } } - fn tombstone_key_in_index(&mut self, idx_name: &Bytes, key_hash: u64) -> bool { - let Some(idx) = self.indexes.get_mut(idx_name) else { - return false; - }; - let snap = idx.segments.load(); + /// Tombstone `key_hash` across every tier of one [`SegmentHolder`]'s + /// current snapshot (mutable + immutable + warm + unloaded/cold-stub). + /// Shared by the default field and each secondary VECTOR field so a + /// deleted document cannot resurrect through any field's search path. + fn tombstone_key_in_holder(holder: &SegmentHolder, key_hash: u64) { + let snap = holder.load(); // Tombstone in mutable segment (always present). snap.mutable.mark_deleted_by_key_hash(key_hash, 1); // Also tombstone any already-compacted immutable segments that @@ -2248,7 +2249,24 @@ impl VectorStore { for stub in snap.unloaded.iter() { stub.mark_deleted_by_key_hash(key_hash); } - drop(snap); + } + + fn tombstone_key_in_index(&mut self, idx_name: &Bytes, key_hash: u64) -> bool { + let Some(idx) = self.indexes.get_mut(idx_name) else { + return false; + }; + // Default field (`vector_fields[0]`). + Self::tombstone_key_in_holder(&idx.segments, key_hash); + // Prod-hardening #20: secondary VECTOR fields live in a separate + // `field_segments` map that the default-field loop above never + // touches. Without this, `DEL`/`HDEL`/`UNLINK` on a doc leaves the + // vector alive in every non-default field's segments — a subsequent + // `FT.SEARCH idx '@field2:[...]'` resurrects the deleted document + // (under a synthetic `vec:` key, since the shared + // key_hash_to_key map below is cleared unconditionally). + for fs in idx.field_segments.values() { + Self::tombstone_key_in_holder(&fs.segments, key_hash); + } // QW7 (2026-06 review finding 6.3): prune the key-hash maps so // they track LIVE keys, not historical inserts — without this // they grow monotonically under key churn (~1GB / 24M deletes). @@ -4365,6 +4383,66 @@ mod bg_compact_tests { ); } + /// Prod-hardening #20: DEL of a doc must tombstone its vector in EVERY + /// vector field, not just the default field. A secondary VECTOR field + /// lives in `idx.field_segments`, which `tombstone_key_in_index` used to + /// skip entirely — leaving the deleted doc alive in that field's search. + #[test] + fn test_del_tombstones_secondary_vector_field() { + distance::init(); + let dim = 64u32; + let mut store = VectorStore::new(); + let mut meta = make_idx(dim); + // Two-field schema: vector_fields[0] is the default field, [1] is a + // secondary field that create_index materializes into field_segments. + meta.vector_fields = vec![ + VectorFieldMeta { + field_name: Bytes::from_static(b"vec"), + dimension: dim, + padded_dimension: padded_dimension(dim), + metric: DistanceMetric::L2, + quantization: QuantizationConfig::TurboQuant4, + build_mode: crate::vector::turbo_quant::collection::BuildMode::Light, + }, + VectorFieldMeta { + field_name: Bytes::from_static(b"vec2"), + dimension: dim, + padded_dimension: padded_dimension(dim), + metric: DistanceMetric::L2, + quantization: QuantizationConfig::TurboQuant4, + build_mode: crate::vector::turbo_quant::collection::BuildMode::Light, + }, + ]; + store.create_index(meta).unwrap(); + + let key = b"doc:1"; + let key_hash = xxhash_rust::xxh64::xxh64(key, 0); + // Default field insert (also registers key_hash_to_key). + insert(&mut store, key, random_vec(dim as usize, 1)); + // Secondary-field insert: mirror the production insert path by + // appending straight into vec2's mutable segment. + { + let idx = store.indexes.get_mut(b"idx".as_ref()).unwrap(); + let fs = idx.field_segments.get(b"vec2".as_ref()).unwrap(); + let snap = fs.segments.load(); + snap.mutable + .append(key_hash, &random_vec(dim as usize, 2), 1); + assert_eq!(snap.mutable.live_len(), 1, "vec2 has the live doc"); + } + + // DEL doc:1 → both fields must be tombstoned. + store.mark_deleted_for_key(key); + + let idx = store.indexes.get_mut(b"idx".as_ref()).unwrap(); + let fs = idx.field_segments.get(b"vec2".as_ref()).unwrap(); + let snap = fs.segments.load(); + assert_eq!( + snap.mutable.live_len(), + 0, + "secondary field vec2 must have no live docs after DEL (#20)" + ); + } + // ── Test 6: force_compact while in-flight ──────────────────────────────── /// force_compact_index() called while a background job is in-flight must