diff --git a/CHANGELOG.md b/CHANGELOG.md index 5343b6de9..3d39f8439 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,72 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added — sharded MULTI/EXEC routes a single-owner-shard body to its owner (Phase B, PR #TBD) + +- **`src/shard/{dispatch,spsc_handler,coordinator}.rs`, + `src/server/conn/{handler_sharded,handler_monoio}/write.rs`**: with keys routed + per-key to their owning shard but MULTI/EXEC executing on the connection's + (random, SO_REUSEPORT) accept shard, a hash-tagged `{tag}` transaction only + worked from the ~1/N connections that happened to land on the owning shard — + Phase A rejected the rest with `CROSSSLOT`. Phase B adds a boxed + `ShardMessage::TxnExecute` (kept within the 64-byte cache-line cap like + `MqCommand`) and a `coordinator::execute_txn_on_owner` hop: when + `analyze_txn_locality` proves every key is owned by ONE shard that isn't the + accept shard, the whole body is routed there, executed atomically on the + owner's slice, and each write persisted to the OWNER's AOF/WAL via the same + `wal_append_and_fanout` path as normal cross-shard writes. PUBLISH fan-out is + deferred back to the originating connection (returned in the reply) so it keeps + the normal scatter path; under `appendfsync=always` the originator issues one + `fsync_barrier` to the owner before acking (H1-BARRIER parity), and owner-side + append backpressure surfaces `AOF_APPEND_LOST_ERR`. Genuinely multi-shard + bodies (`CrossShard`) stay rejected — a shared-nothing engine can't commit + across shards atomically. Red/green: `tests/sharded_multi_exec_routing.rs` + (hash-tagged txn succeeds from every connection + writes visible; routed writes + survive kill-9 + restart via the owner's AOF; multi-shard span still CROSSSLOT). + WATCH in sharded mode remains an unimplemented follow-up. + +### Fixed — sharded MULTI/EXEC writes are now persisted (were lost on restart) (PR #TBD) + +- **`src/server/conn/{shared,handler_sharded/write,handler_monoio/write}.rs`**: + `execute_transaction_sharded` — the MULTI/EXEC executor used by the monoio + handler at **every** shard count (including `--shards 1`) and by the tokio + sharded handler at `--shards ≥ 2` — appended **nothing** to the AOF. Every + transactional write was silently lost on restart under `appendonly=yes`, while + an identical write issued outside MULTI survived. The executor now returns the + serialized AOF bytes for each successful write (mirroring the single-shard + tokio `execute_transaction`), and a shared `persist_txn_aof` helper appends + them to the owning shard's writer via the normal group-commit path, issuing one + `fsync_barrier` under `appendfsync=always` before EXEC is acked. On a barrier + failure EXEC returns `AOF_FSYNC_ERR` and suppresses any queued PUBLISH fan-out, + rather than acking a durability it can't guarantee. `try_handle_multi_exec` is + now `async` in both sharded handlers. Red/green: + `tests/sharded_multi_exec_durability.rs` pins *EXEC-committed write survives + kill-9 + restart, exactly like a non-MULTI write* (fails on the pre-fix path: + the txn key returns nil after restart while the plain control key survives). + +### Security — sharded MULTI/EXEC no longer silently misplaces cross-shard writes (Phase A, PR #TBD) + +- **`src/server/conn/{shared,handler_sharded/write,handler_monoio/write}.rs`**: + `execute_transaction_sharded` runs the whole queued body on the connection's + OWN shard with no per-key routing, so at `--shards ≥ 2` a key owned by another + shard was silently written to / read from the wrong shard's table — EXEC + reported success while the data diverged (silent lost updates; other + connections routing to the true owner saw nothing). Hash tags did **not** help: + they co-locate keys with each other but not with the (random, SO_REUSEPORT) + execution shard, and migration is disabled during MULTI. A new + `analyze_txn_locality` classifies the queued body via the command-metadata key + specs + `key_to_shard`; EXEC now **rejects** with `CROSSSLOT` any transaction + whose keys aren't all owned by the executing shard, instead of corrupting. + Invariant restored: *EXEC-success ⇒ every write is visible to other + connections; CROSSSLOT ⇒ nothing was written.* No effect at `--shards 1`. + Red/green: `tests/sharded_multi_exec_locality.rs` (silent-divergence invariant, + multi-shard span rejection, single-shard unaffected) + `analyze_txn_locality` + unit tests. Phase B (above) now routes single-owner-shard bodies to their owner + instead of rejecting them, so only genuinely cross-shard bodies still get + CROSSSLOT. Note found in passing (unrelated, out of scope): a *solo* GET sent + mid-MULTI on the monoio single-shard handler executes immediately instead of + queueing. + ### Docs — tuning guide: vector bulk load & compaction (PR #TBD) - `docs/guides/tuning.md`: new "Vector bulk load and compaction" section — documents diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 58b340e85..39f0237c3 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -1006,7 +1006,9 @@ pub(crate) async fn handle_connection_sharded_monoio< ctx, &mut responses, &mut exec_publishes, - ) { + ) + .await + { // C2: PUBLISH queued inside MULTI fans out only now — after the // transaction body has been applied — and its placeholder in the // EXEC reply array is patched with the real receiver count. diff --git a/src/server/conn/handler_monoio/write.rs b/src/server/conn/handler_monoio/write.rs index 0ba39ac5c..4870de39f 100644 --- a/src/server/conn/handler_monoio/write.rs +++ b/src/server/conn/handler_monoio/write.rs @@ -562,7 +562,12 @@ async fn mq_hop_or_local( } /// Handle MULTI/EXEC/DISCARD commands. Returns `true` if consumed. -pub(super) fn try_handle_multi_exec( +/// +/// `async` because EXEC now persists the transaction body to the shard AOF via +/// the same group-commit path as normal writes (previously this path logged +/// nothing, so every transactional write was lost on restart — including at +/// `--shards 1` under the monoio TopLevel writer). +pub(super) async fn try_handle_multi_exec( cmd: &[u8], conn: &mut ConnectionState, ctx: &ConnectionContext, @@ -591,7 +596,115 @@ pub(super) fn try_handle_multi_exec( responses.push(Frame::Error(Bytes::from_static(b"ERR EXEC without MULTI"))); } else { conn.in_multi = false; - let result = execute_transaction_sharded( + // The body runs on THIS shard with no per-key routing, so a + // foreign-owned key would be silently misplaced. Classify locality: + // - CrossShard: genuinely spans shards — a shared-nothing engine + // can't commit it atomically, so reject with CROSSSLOT. + // - SingleShard(other): every key is owned by ONE remote shard — + // Phase B routes the whole body there for atomic execution + AOF + // on the owner (instead of the Phase-A CROSSSLOT rejection). + // - Keyless / SingleShard(self): fall through to local execution. + if ctx.num_shards > 1 { + match crate::server::conn::shared::analyze_txn_locality( + &conn.command_queue, + ctx.num_shards, + ) { + crate::server::conn::shared::TxnLocality::SingleShard(s) + if s != ctx.shard_id => + { + let commands: Vec = conn.command_queue.to_vec(); + conn.command_queue.clear(); + // Keep a copy of the body for CLIENT TRACKING + // invalidation ONLY when tracking is active — the routed + // reply carries results, not the command frames, and the + // owner does not invalidate (the tracking table is + // process-global; invalidation is issued originating-side). + let tracking_cmds = + crate::tracking::tracking_active().then(|| commands.clone()); + let reply = crate::shard::coordinator::execute_txn_on_owner( + s, + ctx.shard_id, + conn.selected_db, + commands, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + match reply { + Some(r) => { + if r.append_lost { + exec_publishes.clear(); + responses.push(Frame::Error(Bytes::from_static( + crate::shard::spsc_handler::AOF_APPEND_LOST_ERR, + ))); + return true; + } + // appendfsync=always: confirm the owner's queued + // appends are on disk before acking (H1-BARRIER + // parity for normal cross-shard writes). No-op + // under everysec/no. + if r.wrote { + if let Some(ref pool) = ctx.aof_pool { + if pool.fsync_barrier(s).await.is_err() { + exec_publishes.clear(); + responses.push(Frame::Error(Bytes::from_static( + crate::persistence::aof::AOF_FSYNC_ERR, + ))); + return true; + } + } + } + // CLIENT TRACKING: invalidate every key written by + // the routed body, same as the local EXEC path + // (which the early return would otherwise skip). + if let Some(cmds) = tracking_cmds.as_ref() { + if let Frame::Array(ref txn_results) = r.result { + for (i, cmd_frame) in cmds.iter().enumerate() { + if i >= txn_results.len() + || matches!(txn_results[i], Frame::Error(_)) + { + continue; + } + if let Some((c, a)) = + crate::server::conn::util::extract_command( + cmd_frame, + ) + { + crate::tracking::invalidation::invalidate_after_write( + &ctx.tracking_table, + c, + a, + conn.client_id, + ); + } + } + } + } + // Adopt the owner's deferred PUBLISH fan-out; the + // caller's post-EXEC loop patches placeholders + + // scatters from this (originating) shard. + exec_publishes.extend(r.exec_publishes); + responses.push(r.result); + } + None => { + responses.push(Frame::Error(Bytes::from_static( + b"CROSSSLOT MULTI/EXEC owner shard unavailable; retry", + ))); + } + } + return true; + } + crate::server::conn::shared::TxnLocality::CrossShard => { + conn.command_queue.clear(); + responses.push(Frame::Error(Bytes::from_static( + b"CROSSSLOT Keys in MULTI/EXEC don't hash to the same shard", + ))); + return true; + } + _ => {} + } + } + let (result, aof_entries) = execute_transaction_sharded( &ctx.shard_databases, ctx.shard_id, &conn.command_queue, @@ -599,6 +712,27 @@ pub(super) fn try_handle_multi_exec( &ctx.cached_clock, exec_publishes, ); + // DURABILITY: append every successful write in the body to THIS + // shard's AOF via the same group-commit path as normal writes, then + // issue ONE fsync barrier under appendfsync=always before acking. + // All keys are local here (Phase A rejected foreign-owned bodies), + // so ctx.shard_id is the correct AOF target. On barrier failure we + // surface AOF_FSYNC_ERR instead of a false EXEC success — parity + // with the normal write path. + if crate::server::conn::shared::persist_txn_aof(ctx, aof_entries) + .await + .is_err() + { + conn.command_queue.clear(); + // Durability could not be guaranteed: report the error and + // suppress any queued PUBLISH fan-out — the client sees EXEC + // fail, so it must not observe the txn's pub/sub side effects. + exec_publishes.clear(); + responses.push(Frame::Error(Bytes::from_static( + crate::persistence::aof::AOF_FSYNC_ERR, + ))); + return true; + } // CLIENT TRACKING: invalidate keys written inside the txn, same as // the normal write path (EXEC previously bypassed this). Self-gated // on tracking_active(); must run before command_queue is cleared. diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 7b688cfd5..77e990745 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -927,7 +927,7 @@ pub(crate) async fn handle_connection_sharded_inner< // --- MULTI / EXEC_CMD / DISCARD --- let mut exec_publishes: Vec<(usize, Bytes, Bytes)> = Vec::new(); - if write::try_handle_multi_exec(cmd, &mut conn, ctx, &mut responses, &mut exec_publishes) { + if write::try_handle_multi_exec(cmd, &mut conn, ctx, &mut responses, &mut exec_publishes).await { // C2: PUBLISH queued inside MULTI fans out only now — after the // transaction body has been applied — and its placeholder in the // EXEC reply array is patched with the real receiver count. diff --git a/src/server/conn/handler_sharded/write.rs b/src/server/conn/handler_sharded/write.rs index 5fd605613..7bb827f1c 100644 --- a/src/server/conn/handler_sharded/write.rs +++ b/src/server/conn/handler_sharded/write.rs @@ -513,7 +513,11 @@ pub(super) async fn try_handle_mq_command( } /// Handle MULTI/EXEC/DISCARD commands. Returns `true` if consumed. -pub(super) fn try_handle_multi_exec( +/// +/// `async` because EXEC now persists the transaction body to the shard AOF via +/// the same group-commit path as normal writes (previously this path logged +/// nothing, so every transactional write was lost on restart). +pub(super) async fn try_handle_multi_exec( cmd: &[u8], conn: &mut ConnectionState, ctx: &ConnectionContext, @@ -542,7 +546,115 @@ pub(super) fn try_handle_multi_exec( responses.push(Frame::Error(Bytes::from_static(b"ERR EXEC without MULTI"))); } else { conn.in_multi = false; - let result = execute_transaction_sharded( + // The body runs on THIS shard with no per-key routing, so a + // foreign-owned key would be silently misplaced. Classify locality: + // - CrossShard: genuinely spans shards — a shared-nothing engine + // can't commit it atomically, so reject with CROSSSLOT. + // - SingleShard(other): every key is owned by ONE remote shard — + // Phase B routes the whole body there for atomic execution + AOF + // on the owner (instead of the Phase-A CROSSSLOT rejection). + // - Keyless / SingleShard(self): fall through to local execution. + if ctx.num_shards > 1 { + match crate::server::conn::shared::analyze_txn_locality( + &conn.command_queue, + ctx.num_shards, + ) { + crate::server::conn::shared::TxnLocality::SingleShard(s) + if s != ctx.shard_id => + { + let commands: Vec = conn.command_queue.to_vec(); + conn.command_queue.clear(); + // Keep a copy of the body for CLIENT TRACKING + // invalidation ONLY when tracking is active — the routed + // reply carries results, not the command frames, and the + // owner does not invalidate (the tracking table is + // process-global; invalidation is issued originating-side). + let tracking_cmds = + crate::tracking::tracking_active().then(|| commands.clone()); + let reply = crate::shard::coordinator::execute_txn_on_owner( + s, + ctx.shard_id, + conn.selected_db, + commands, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + match reply { + Some(r) => { + if r.append_lost { + exec_publishes.clear(); + responses.push(Frame::Error(Bytes::from_static( + crate::shard::spsc_handler::AOF_APPEND_LOST_ERR, + ))); + return true; + } + // appendfsync=always: confirm the owner's queued + // appends are on disk before acking (H1-BARRIER + // parity for normal cross-shard writes). No-op + // under everysec/no. + if r.wrote { + if let Some(ref pool) = ctx.aof_pool { + if pool.fsync_barrier(s).await.is_err() { + exec_publishes.clear(); + responses.push(Frame::Error(Bytes::from_static( + crate::persistence::aof::AOF_FSYNC_ERR, + ))); + return true; + } + } + } + // CLIENT TRACKING: invalidate every key written by + // the routed body, same as the local EXEC path + // (which the early return would otherwise skip). + if let Some(cmds) = tracking_cmds.as_ref() { + if let Frame::Array(ref txn_results) = r.result { + for (i, cmd_frame) in cmds.iter().enumerate() { + if i >= txn_results.len() + || matches!(txn_results[i], Frame::Error(_)) + { + continue; + } + if let Some((c, a)) = + crate::server::conn::util::extract_command( + cmd_frame, + ) + { + crate::tracking::invalidation::invalidate_after_write( + &ctx.tracking_table, + c, + a, + conn.client_id, + ); + } + } + } + } + // Adopt the owner's deferred PUBLISH fan-out; the + // caller's post-EXEC loop patches placeholders + + // scatters from this (originating) shard. + exec_publishes.extend(r.exec_publishes); + responses.push(r.result); + } + None => { + responses.push(Frame::Error(Bytes::from_static( + b"CROSSSLOT MULTI/EXEC owner shard unavailable; retry", + ))); + } + } + return true; + } + crate::server::conn::shared::TxnLocality::CrossShard => { + conn.command_queue.clear(); + responses.push(Frame::Error(Bytes::from_static( + b"CROSSSLOT Keys in MULTI/EXEC don't hash to the same shard", + ))); + return true; + } + _ => {} + } + } + let (result, aof_entries) = execute_transaction_sharded( &ctx.shard_databases, ctx.shard_id, &conn.command_queue, @@ -550,6 +662,27 @@ pub(super) fn try_handle_multi_exec( &ctx.cached_clock, exec_publishes, ); + // DURABILITY: append every successful write in the body to THIS + // shard's AOF via the same group-commit path as normal writes, then + // issue ONE fsync barrier under appendfsync=always before acking. + // All keys are local here (Phase A rejected foreign-owned bodies), + // so ctx.shard_id is the correct AOF target. On barrier failure we + // surface AOF_FSYNC_ERR instead of a false EXEC success — parity + // with the normal write path. + if crate::server::conn::shared::persist_txn_aof(ctx, aof_entries) + .await + .is_err() + { + conn.command_queue.clear(); + // Durability could not be guaranteed: report the error and + // suppress any queued PUBLISH fan-out — the client sees EXEC + // fail, so it must not observe the txn's pub/sub side effects. + exec_publishes.clear(); + responses.push(Frame::Error(Bytes::from_static( + crate::persistence::aof::AOF_FSYNC_ERR, + ))); + return true; + } // CLIENT TRACKING: invalidate keys written inside the txn, same as // the normal write path (EXEC previously bypassed this). Self-gated // on tracking_active(); must run before command_queue is cleared. diff --git a/src/server/conn/shared.rs b/src/server/conn/shared.rs index 9a19c1862..61b1a85fd 100644 --- a/src/server/conn/shared.rs +++ b/src/server/conn/shared.rs @@ -194,8 +194,14 @@ pub(crate) fn execute_transaction( /// Execute a queued transaction on the local shard (sharded path). /// -/// Transactions in the shared-nothing architecture are restricted to local-shard -/// keys only. Cross-shard transactions require distributed coordination (future work). +/// The caller must ensure every key in `command_queue` is owned by THIS shard +/// (see `analyze_txn_locality`) — the body runs on the local slice with no +/// per-key routing. +/// +/// Returns the result Frame (an array of per-command responses) **and** the +/// serialized AOF bytes for each successful write, in order. The caller MUST +/// append those entries to the shard's AOF (previously this path did no +/// persistence, so MULTI/EXEC writes were silently lost on restart). pub(crate) fn execute_transaction_sharded( shard_databases: &std::sync::Arc, _shard_id: usize, @@ -203,10 +209,11 @@ pub(crate) fn execute_transaction_sharded( selected_db: usize, cached_clock: &CachedClock, exec_publishes: &mut Vec<(usize, Bytes, Bytes)>, -) -> Frame { +) -> (Frame, Vec) { let db_count = shard_databases.db_count(); let mut results = Vec::with_capacity(command_queue.len()); + let mut aof_entries: Vec = Vec::new(); let mut selected = selected_db; for cmd_frame in command_queue { @@ -228,6 +235,19 @@ pub(crate) fn execute_transaction_sharded( continue; } + // Serialize write commands for AOF *before* dispatch (matches the + // single-shard `execute_transaction` path). Without this the sharded + // MULTI/EXEC path logged nothing, so every transactional write was + // silently lost on restart. Fully-qualified paths because `metadata` + // is only `use`d under runtime-tokio but this fn compiles under both. + let aof_bytes = if crate::command::metadata::is_write(cmd) { + let mut buf = bytes::BytesMut::new(); + crate::protocol::serialize::serialize(cmd_frame, &mut buf); + Some(buf.freeze()) + } else { + None + }; + let result = crate::shard::slice::with_shard_db(selected, |db| { db.refresh_now_from_cache(cached_clock); dispatch(db, cmd, cmd_args, &mut selected, db_count) @@ -237,6 +257,14 @@ pub(crate) fn execute_transaction_sharded( DispatchResult::Quit(f) => f, }; + // Only log the write if it actually succeeded (parity with the + // single-shard path — an errored write must not reach the AOF). + if let Some(bytes) = aof_bytes { + if !matches!(&response, Frame::Error(_)) { + aof_entries.push(bytes); + } + } + // Auto-index: if HSET succeeded, check for vector index match if cmd.eq_ignore_ascii_case(b"HSET") && !matches!(response, Frame::Error(_)) { if let Some(Frame::BulkString(key_bytes)) = cmd_args.first() { @@ -284,7 +312,50 @@ pub(crate) fn execute_transaction_sharded( results.push(response); } - Frame::Array(results.into()) + (Frame::Array(results.into()), aof_entries) +} + +/// Persist the AOF entries of a just-executed sharded MULTI/EXEC body. +/// +/// Mirrors the normal write path's group-commit: each entry is enqueued +/// fire-and-forget on the owning shard's writer (`send_append_group`), and a +/// single `fsync_barrier` is issued at the end under `appendfsync=always` +/// (`send_append_group` returns `Ok(true)` when a barrier is owed). All keys in +/// the body are owned by `ctx.shard_id` (Phase A rejects foreign-owned bodies), +/// so that is the correct AOF target. +/// +/// Returns `Err(())` if any append or the barrier fails — the caller surfaces +/// `AOF_FSYNC_ERR` instead of acking a durability it can't guarantee. A no-op +/// (returns `Ok`) when AOF is disabled (`aof_pool` is `None`) or the body wrote +/// nothing. +pub(crate) async fn persist_txn_aof( + ctx: &crate::server::conn::core::ConnectionContext, + aof_entries: Vec, +) -> Result<(), ()> { + if aof_entries.is_empty() { + return Ok(()); + } + let Some(ref pool) = ctx.aof_pool else { + return Ok(()); + }; + let mut barrier_pending = false; + for bytes in aof_entries { + let lsn = crate::persistence::aof::AofWriterPool::issue_append_lsn( + &ctx.repl_state, + ctx.shard_id, + bytes.len(), + ); + match pool.send_append_group(ctx.shard_id, lsn, bytes).await { + Ok(true) => barrier_pending = true, + Ok(false) => {} + Err(_) => return Err(()), + } + } + // appendfsync=always: one barrier confirms the whole body is on disk. + if barrier_pending && pool.fsync_barrier(ctx.shard_id).await.is_err() { + return Err(()); + } + Ok(()) } /// Shared PUBLISH-inside-MULTI intercept for both transaction executors (C2). @@ -555,6 +626,51 @@ pub(crate) fn is_multi_key_command(cmd: &[u8], args: &[Frame]) -> bool { } } +/// Shard-locality of a queued MULTI/EXEC body. +/// +/// `execute_transaction_sharded` runs the whole body on ONE shard's slice with +/// no per-key routing, so a key owned by a different shard is silently written +/// to (or read from) the wrong table. This classifies a queued body so the +/// EXEC handler can either run it locally (all keys local), route it to the +/// owner shard (all keys on one remote shard — Phase B), or reject it +/// (`CrossShard`: a single-process shared-nothing engine can't atomically span +/// shards). Hash tags (`{tag}`) collapse a body to `SingleShard`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum TxnLocality { + /// No key-bearing commands — safe to run on any shard. + Keyless, + /// Every key in the body resolves to this single shard. + SingleShard(usize), + /// Keys span more than one shard — not atomically executable. + CrossShard, +} + +/// Classify a queued transaction body by the shard(s) its keys hash to. +/// +/// Uses the command-metadata key specs (`command_keys`) so multi-key commands +/// (MSET/DEL/…) contribute every key, and `key_to_shard` so `{hash tags}` are +/// 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; + 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, + _ => {} + } + } + } + match owner { + None => TxnLocality::Keyless, + Some(s) => TxnLocality::SingleShard(s), + } +} + #[cfg(test)] mod as_of_tests { //! Unit tests for `resolve_ft_search_as_of_lsn` (TEMP-04 + ACID-09). @@ -715,3 +831,83 @@ pub async fn resolve_local_leg_barrier( } idxs.clear(); } + +#[cfg(test)] +mod txn_locality_tests { + use super::{TxnLocality, analyze_txn_locality}; + use crate::protocol::Frame; + use bytes::Bytes; + + /// Build a queued command frame (name + args) as the wire form. + fn cmd(parts: &[&str]) -> Frame { + Frame::Array( + parts + .iter() + .map(|p| Frame::BulkString(Bytes::copy_from_slice(p.as_bytes()))) + .collect::>() + .into(), + ) + } + + #[test] + fn empty_and_keyless_are_keyless() { + assert_eq!(analyze_txn_locality(&[], 4), TxnLocality::Keyless); + let q = [cmd(&["PING"]), cmd(&["MULTI"])]; + assert_eq!(analyze_txn_locality(&q, 4), TxnLocality::Keyless); + } + + #[test] + fn single_shard_at_one_shard() { + // With one shard every key resolves to shard 0. + let q = [cmd(&["SET", "a", "1"]), cmd(&["SET", "b", "2"])]; + assert_eq!(analyze_txn_locality(&q, 1), TxnLocality::SingleShard(0)); + } + + #[test] + fn hash_tags_collapse_to_one_shard() { + // `{t}` forces co-location regardless of the surrounding key text. + let q = [ + cmd(&["SET", "user:{t}:name", "x"]), + cmd(&["INCR", "user:{t}:hits"]), + cmd(&["DEL", "user:{t}:tmp"]), + ]; + let owner = crate::shard::dispatch::key_to_shard(b"t", 8); + assert_eq!(analyze_txn_locality(&q, 8), TxnLocality::SingleShard(owner)); + } + + #[test] + fn spanning_keys_are_cross_shard() { + // Find two keys that hash to different shards, then confirm CrossShard. + let num = 8; + let base = crate::shard::dispatch::key_to_shard(b"k0", num); + let mut other = None; + for i in 1..1000 { + let k = format!("k{i}"); + if crate::shard::dispatch::key_to_shard(k.as_bytes(), num) != base { + other = Some(k); + break; + } + } + let other = other.expect("two keys on different shards must exist across 8 shards"); + let q = [cmd(&["SET", "k0", "1"]), cmd(&["SET", &other, "2"])]; + assert_eq!(analyze_txn_locality(&q, num), TxnLocality::CrossShard); + } + + #[test] + fn multi_key_command_contributes_every_key() { + // MSET's keys are checked individually; spanning keys ⇒ CrossShard. + let num = 8; + let base = crate::shard::dispatch::key_to_shard(b"m0", num); + let mut other = None; + for i in 1..1000 { + let k = format!("m{i}"); + if crate::shard::dispatch::key_to_shard(k.as_bytes(), num) != base { + other = Some(k); + break; + } + } + let other = other.expect("two keys on different shards must exist"); + let q = [cmd(&["MSET", "m0", "1", &other, "2"])]; + assert_eq!(analyze_txn_locality(&q, num), TxnLocality::CrossShard); + } +} diff --git a/src/shard/coordinator.rs b/src/shard/coordinator.rs index a89747bf2..ac8b0347c 100644 --- a/src/shard/coordinator.rs +++ b/src/shard/coordinator.rs @@ -202,6 +202,34 @@ async fn run_remote( } } +/// Route a whole MULTI/EXEC body to the shard that owns all of its keys and +/// await the executed reply (sharded MULTI/EXEC Phase B). +/// +/// The owner runs the body atomically on its slice and persists each write to +/// ITS own AOF/WAL; PUBLISH fan-out is deferred back to the caller in the reply +/// (`exec_publishes`) so the originator keeps the normal scatter path. Returns +/// `None` if the owner's reply channel closed (the owner shard died) — the +/// caller surfaces an error rather than a false success. +pub(crate) async fn execute_txn_on_owner( + owner: usize, + my_shard: usize, + db_index: usize, + commands: Vec, + dispatch_tx: &Rc>>>, + spsc_notifiers: &[Arc], +) -> Option { + debug_assert_ne!(owner, my_shard, "execute_txn_on_owner called for own shard"); + let (reply_tx, reply_rx) = channel::oneshot(); + let payload = crate::shard::dispatch::TxnExecutePayload { + db_index, + commands, + reply_tx, + }; + let msg = ShardMessage::TxnExecute(Box::new(payload)); + let _ = spsc_send(dispatch_tx, my_shard, owner, msg, spsc_notifiers).await; + reply_rx.recv().await.ok() +} + /// Run one full command on whichever shard owns `routing_key`. #[allow(clippy::too_many_arguments)] async fn run_on_owner( diff --git a/src/shard/dispatch.rs b/src/shard/dispatch.rs index 52310c3e2..7dcedaf65 100644 --- a/src/shard/dispatch.rs +++ b/src/shard/dispatch.rs @@ -658,6 +658,20 @@ pub enum ShardMessage { AofFold { reply_tx: channel::OneshotSender, }, + /// Execute a queued MULTI/EXEC body atomically on this (owning) shard + /// (sharded MULTI/EXEC Phase B). + /// + /// Sent by a connection handler when `analyze_txn_locality` proved every key + /// in the transaction is owned by ONE shard that is NOT the connection's + /// accept shard. The owner runs the whole body on its slice, persists each + /// write to ITS AOF/WAL (via `wal_append_and_fanout`), and replies with the + /// results array. PUBLISH fan-out is deferred back to the originating + /// connection (returned in `TxnExecReply.exec_publishes`) so it keeps using + /// the normal originator-side scatter path. + /// + /// Boxed (`Box`) to keep `ShardMessage` within the + /// 64-byte cache-line cap — mirrors `MqCommand`. + TxnExecute(Box), /// Graceful shutdown signal. Shutdown, } @@ -680,6 +694,38 @@ pub struct MqCommandPayload { pub reply_tx: channel::OneshotSender, } +/// Payload for [`ShardMessage::TxnExecute`] (sharded MULTI/EXEC Phase B). +/// +/// Boxed in the enum variant to keep `ShardMessage` within the 64-byte cap. +pub struct TxnExecutePayload { + /// Target database index (the connection's `selected_db`). + pub db_index: usize, + /// The queued transaction body — the exact `command_queue` frames the + /// connection accumulated between MULTI and EXEC. + pub commands: Vec, + /// Oneshot channel: the owner sends the executed result + deferred + /// PUBLISH fan-out back to the caller. + pub reply_tx: channel::OneshotSender, +} + +/// Reply for [`ShardMessage::TxnExecute`]. +/// +/// Carries the executed result frame (a `Frame::Array` of per-command +/// responses, with `Frame::Integer(0)` placeholders for any queued PUBLISH), +/// the deferred PUBLISH fan-out list `(result_index, channel, message)` the +/// originator patches after fanning out, and whether the body performed any +/// durable write (so the originator can issue one `fsync_barrier` to the owner +/// under `appendfsync=always`). +pub struct TxnExecReply { + pub result: crate::protocol::Frame, + pub exec_publishes: Vec<(usize, Bytes, Bytes)>, + pub wrote: bool, + /// `true` iff an AOF append could not be enqueued on the owner (bounded + /// backpressure exhausted) — the originator surfaces `AOF_APPEND_LOST_ERR` + /// instead of the (in-memory-applied but non-durable) result. + pub append_lost: bool, +} + /// Snapshot payload for [`ShardMessage::AofFold`]. /// /// Produced by the shard thread and consumed by the AOF rewrite writer thread. diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index 925d62562..fe41c691d 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -176,6 +176,7 @@ pub(crate) fn drain_spsc_shared( ShardMessage::Execute { .. } | ShardMessage::PipelineBatch { .. } | ShardMessage::MultiExecute { .. } + | ShardMessage::TxnExecute(_) | ShardMessage::ExecuteSlotted { .. } | ShardMessage::PipelineBatchSlotted { .. } | ShardMessage::MultiExecuteSlotted { .. } @@ -2317,6 +2318,62 @@ pub(crate) fn handle_shard_message_shared( let _ = reply_tx.send(snapshot); } + // ───────────────────────────────────────────────────────────────────── + ShardMessage::TxnExecute(payload) => { + // Sharded MULTI/EXEC Phase B: the originating connection proved via + // `analyze_txn_locality` that every key is owned by THIS shard (its + // own accept shard differs), so run the whole body on our slice and + // persist each write to OUR AOF/WAL. `execute_transaction_sharded` + // manages its own `with_shard`/`with_shard_db` visits, so it must be + // called at the top of this arm (never nested in a slice borrow). + let crate::shard::dispatch::TxnExecutePayload { + db_index, + commands, + reply_tx, + } = *payload; + let mut exec_publishes: Vec<(usize, bytes::Bytes, bytes::Bytes)> = Vec::new(); + let (result, aof_entries) = crate::server::conn::shared::execute_transaction_sharded( + shard_databases, + shard_id, + &commands, + db_index, + cached_clock, + &mut exec_publishes, + ); + // Persist via the SAME sync path as normal cross-shard writes (the + // MultiExecute arm): fire-and-forget append + WAL/replica fan-out, + // ONE shared backpressure budget for the whole body so a stall costs + // at most AOF_SPSC_BACKPRESSURE_BOUND, not that per command. The + // always-mode fsync barrier is issued by the ORIGINATOR after the + // reply (mirrors the H1-BARRIER for normal cross-shard writes). + let mut wrote = false; + let mut append_lost = false; + let mut aof_budget = crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; + for entry_bytes in &aof_entries { + wrote = true; + let ok = wal_append_and_fanout( + entry_bytes, + wal_writer, + repl_backlog, + replica_txs, + repl_state, + shard_id, + aof_pool, + wal_kv_log, + &mut aof_budget, + ); + if !ok { + append_lost = true; + } + } + let _ = reply_tx.send(crate::shard::dispatch::TxnExecReply { + result, + exec_publishes, + wrote, + append_lost, + }); + } + // ───────────────────────────────────────────────────────────────────── ShardMessage::Shutdown => { info!("Received shutdown via SPSC"); diff --git a/tests/client_tracking_invalidation.rs b/tests/client_tracking_invalidation.rs index e873adc35..b4caf25a1 100644 --- a/tests/client_tracking_invalidation.rs +++ b/tests/client_tracking_invalidation.rs @@ -298,3 +298,69 @@ fn mset_invalidates_every_second_arg_key() { String::from_utf8_lossy(&reader.buf) ); } + +/// A routed (Phase B) MULTI/EXEC — whose owning shard differs from the writer's +/// accept shard — must still invalidate tracking clients, exactly like a local +/// EXEC or a plain write. Before the fix the routed EXEC path returned before +/// reaching the tracking-invalidation block, so cached readers were never +/// notified. +/// +/// A single warmed writer connection has a fixed accept shard, so across 12 +/// distinct keys ~3/4 of the single-key bodies route to a remote owner. Delivery +/// to an idle cross-shard reader is inherently ~75% per attempt in this suite +/// (a documented 4-shard async-push timing property — it hits plain SET and +/// MULTI/EXEC equally), so this asserts a MAJORITY deliver rather than all: the +/// fix yields ~9/12, whereas the pre-fix routed-skip yields only ~2/12 (just the +/// occasional local iteration), which the threshold cleanly rejects. +#[test] +fn routed_multi_exec_invalidates_tracking_client() { + let Some(m) = spawn_moon_4shard() else { return }; + + // One warmed writer (fixed accept shard → most keys route) reused for seed + // + commit; one reader reused across keys. Low connection churn. + let mut txw = Resp::connect(m.port); + txw.cmd(&["PING"]); + assert!(txw.saw(b"PONG")); + + let total = 12; + let mut delivered = 0; + for i in 0..total { + let key = format!("ctxroute{i}"); + + txw.clear(); + txw.cmd(&["SET", &key, "v1"]); + assert!(txw.saw(b"+OK"), "seed SET must succeed"); + + let mut reader = tracking_client(m.port); + reader.cmd(&["GET", &key]); + assert!(reader.saw(b"v1"), "tracked GET must return the seed value"); + reader.clear(); + + txw.cmd(&["MULTI"]); + txw.cmd(&["SET", &key, "v2"]); + txw.clear(); + txw.cmd(&["EXEC"]); + + // Disambiguate EXEC failure from invalidation-delivery loss: a routed + // single-key EXEC that committed returns `*1\r\n+OK\r\n`. If the routed + // execution itself errored (e.g. owner-shard-unavailable → CROSSSLOT), + // fail loudly here rather than silently under-counting `delivered` and + // misreporting a routing bug as an invalidation regression. + assert!( + txw.saw(b"+OK"), + "routed EXEC must commit the queued SET (got: {:?})", + String::from_utf8_lossy(&txw.buf) + ); + + if reader.wait_for(INVALIDATE, Duration::from_secs(4)) && reader.saw(key.as_bytes()) { + delivered += 1; + } + } + + assert!( + delivered * 2 >= total, + "routed MULTI/EXEC invalidation regression: only {delivered}/{total} writes \ + invalidated the tracking client (pre-fix routed EXEC skips invalidation \ + entirely, delivering ~2/12; the fix restores ~9/12)" + ); +} diff --git a/tests/sharded_multi_exec_durability.rs b/tests/sharded_multi_exec_durability.rs new file mode 100644 index 000000000..fdc201fca --- /dev/null +++ b/tests/sharded_multi_exec_durability.rs @@ -0,0 +1,294 @@ +//! Sharded MULTI/EXEC AOF durability (2026-07 follow-up). +//! +//! `execute_transaction_sharded` — the MULTI/EXEC executor used by BOTH the +//! monoio handler (all shard counts, incl. `--shards 1`) and the tokio sharded +//! handler (`--shards >= 2`) — previously wrote nothing to the AOF. Every +//! transactional write was therefore silently lost on restart, while an +//! identical write issued outside MULTI survived. This is the worst durability +//! failure class: EXEC acked success, the data was gone after a restart. +//! +//! The invariant these tests pin: +//! +//! **A write committed by EXEC under `appendonly=yes` is present after the +//! server is killed and restarted from the same data dir — exactly like a +//! write issued outside MULTI.** +//! +//! `--shards 1` makes placement deterministic (the Phase-A locality guard is +//! `num_shards > 1`, so it never rejects), and the default (monoio) runtime +//! routes `--shards 1` MULTI/EXEC through `execute_transaction_sharded` — the +//! buggy path. `appendfsync=always` means EXEC only returns after the fsync +//! barrier, so a `kill -9` immediately afterward is a true durability probe. +//! +//! Skips gracefully when the moon binary is missing (MOON_BIN pin wins, then +//! target/release, then target/debug). + +use std::io::{Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +fn free_port() -> u16 { + TcpListener::bind("127.0.0.1:0") + .expect("bind") + .local_addr() + .unwrap() + .port() +} + +fn moon_binary() -> Option { + if let Ok(p) = std::env::var("MOON_BIN") { + return Some(std::path::PathBuf::from(p)); + } + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); + for rel in ["target/release/moon", "target/debug/moon"] { + let p = root.join(rel); + if p.exists() { + return Some(p); + } + } + None +} + +/// A running moon process bound to `port`, persisting into `dir`. Does NOT +/// remove `dir` on drop — the durability test restarts against the same dir. +struct Moon { + child: Child, + port: u16, +} + +impl Moon { + fn kill9(mut self) { + // Hard kill: exercise crash recovery, not graceful shutdown flush. + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +fn spawn_moon_persistent(port: u16, dir: &std::path::Path) -> Option { + let bin = moon_binary()?; + let child = Command::new(&bin) + .args([ + "--port", + &port.to_string(), + "--shards", + "1", + "--admin-port", + "0", + "--appendonly", + "yes", + "--appendfsync", + "always", + "--disk-free-min-pct", + "0", + "--dir", + dir.to_str().unwrap(), + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .ok()?; + let moon = Moon { child, port }; + let deadline = Instant::now() + Duration::from_secs(10); + while Instant::now() < deadline { + if let Ok(mut c) = TcpStream::connect(("127.0.0.1", moon.port)) { + let _ = c.set_read_timeout(Some(Duration::from_millis(500))); + if c.write_all(b"*1\r\n$4\r\nPING\r\n").is_ok() { + let mut buf = [0u8; 64]; + if let Ok(n) = c.read(&mut buf) { + if n > 0 && buf.starts_with(b"+PONG") { + return Some(moon); + } + } + } + } + std::thread::sleep(Duration::from_millis(100)); + } + eprintln!("skipping: moon did not become ready on port {port}"); + None +} + +/// One parsed RESP2 reply — only the shapes these tests need. +#[derive(Debug, Clone, PartialEq)] +enum Reply { + Simple(String), + Error(String), + Int(i64), + Bulk(Option), + Array(Vec), +} + +struct Client { + stream: TcpStream, + buf: Vec, + pos: usize, +} + +impl Client { + fn connect(port: u16) -> Self { + let stream = TcpStream::connect(("127.0.0.1", port)).expect("connect"); + stream + .set_read_timeout(Some(Duration::from_millis(3000))) + .unwrap(); + Self { + stream, + buf: Vec::new(), + pos: 0, + } + } + + fn send(&mut self, args: &[&str]) { + let mut out = format!("*{}\r\n", args.len()).into_bytes(); + for a in args { + out.extend_from_slice(format!("${}\r\n{a}\r\n", a.len()).as_bytes()); + } + self.stream.write_all(&out).expect("write"); + } + + fn fill(&mut self) { + let mut chunk = [0u8; 4096]; + match self.stream.read(&mut chunk) { + Ok(0) => panic!("connection closed by server"), + Ok(n) => self.buf.extend_from_slice(&chunk[..n]), + Err(e) => panic!("read error: {e}"), + } + } + + fn read_line(&mut self) -> String { + loop { + if let Some(rel) = self.buf[self.pos..].windows(2).position(|w| w == b"\r\n") { + let end = self.pos + rel; + let line = String::from_utf8_lossy(&self.buf[self.pos..end]).to_string(); + self.pos = end + 2; + return line; + } + self.fill(); + } + } + + fn read_exact_bytes(&mut self, n: usize) -> String { + while self.buf.len() - self.pos < n + 2 { + self.fill(); + } + let s = String::from_utf8_lossy(&self.buf[self.pos..self.pos + n]).to_string(); + self.pos += n + 2; + s + } + + fn read_reply(&mut self) -> Reply { + let line = self.read_line(); + let (tag, rest) = line.split_at(1); + match tag { + "+" => Reply::Simple(rest.to_string()), + "-" => Reply::Error(rest.to_string()), + ":" => Reply::Int(rest.parse().unwrap_or(0)), + "$" => { + let len: i64 = rest.parse().unwrap_or(-1); + if len < 0 { + Reply::Bulk(None) + } else { + Reply::Bulk(Some(self.read_exact_bytes(len as usize))) + } + } + "*" => { + let len: i64 = rest.parse().unwrap_or(-1); + if len < 0 { + Reply::Array(Vec::new()) + } else { + let mut items = Vec::with_capacity(len as usize); + for _ in 0..len { + items.push(self.read_reply()); + } + Reply::Array(items) + } + } + other => panic!("unexpected RESP tag {other:?} in line {line:?}"), + } + } + + fn cmd(&mut self, args: &[&str]) -> Reply { + self.send(args); + self.read_reply() + } +} + +/// GOLDEN INVARIANT: a write committed by EXEC must survive a kill-9 + restart, +/// exactly like a write issued outside MULTI. Before the fix the transactional +/// key vanished on restart while the plain control key survived. +#[test] +fn multi_exec_write_survives_restart() { + let port = free_port(); + let dir = std::env::temp_dir().join(format!("moon-txn-dur-{port}")); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::create_dir_all(&dir); + + let Some(m1) = spawn_moon_persistent(port, &dir) else { + let _ = std::fs::remove_dir_all(&dir); + return; + }; + + { + let mut c = Client::connect(port); + // Control: a plain (non-MULTI) write — the baseline that always survived. + assert_eq!( + c.cmd(&["SET", "plain:key", "plainval"]), + Reply::Simple("OK".into()) + ); + // The write under test: committed via MULTI/EXEC. + assert_eq!(c.cmd(&["MULTI"]), Reply::Simple("OK".into())); + assert_eq!( + c.cmd(&["SET", "txn:key", "txnval"]), + Reply::Simple("QUEUED".into()) + ); + assert_eq!( + c.cmd(&["INCRBY", "txn:counter", "5"]), + Reply::Simple("QUEUED".into()) + ); + match c.cmd(&["EXEC"]) { + Reply::Array(items) => { + assert_eq!(items.len(), 2, "EXEC returns one reply per queued write"); + assert_eq!(items[0], Reply::Simple("OK".into())); + assert_eq!(items[1], Reply::Int(5)); + } + other => panic!("EXEC must return the results array, got: {other:?}"), + } + // Visible immediately (in-memory) before the restart. + assert_eq!( + c.cmd(&["GET", "txn:key"]), + Reply::Bulk(Some("txnval".into())) + ); + } + + // Hard kill (no graceful flush) — appendfsync=always already put the EXEC + // body on disk, so recovery must replay it. + m1.kill9(); + + let Some(m2) = spawn_moon_persistent(port, &dir) else { + let _ = std::fs::remove_dir_all(&dir); + panic!("moon failed to restart from {dir:?}"); + }; + + let mut r = Client::connect(port); + let plain = r.cmd(&["GET", "plain:key"]); + let txn = r.cmd(&["GET", "txn:key"]); + let counter = r.cmd(&["GET", "txn:counter"]); + m2.kill9(); + let _ = std::fs::remove_dir_all(&dir); + + // Control must survive (it always did). + assert_eq!( + plain, + Reply::Bulk(Some("plainval".into())), + "plain (non-MULTI) write lost on restart — harness/persistence broken" + ); + // The regression: the EXEC-committed writes must survive too. + assert_eq!( + txn, + Reply::Bulk(Some("txnval".into())), + "MULTI/EXEC SET lost on restart — transactional writes were not persisted" + ); + assert_eq!( + counter, + Reply::Bulk(Some("5".into())), + "MULTI/EXEC INCRBY lost on restart — transactional writes were not persisted" + ); +} diff --git a/tests/sharded_multi_exec_locality.rs b/tests/sharded_multi_exec_locality.rs new file mode 100644 index 000000000..35e8a89d7 --- /dev/null +++ b/tests/sharded_multi_exec_locality.rs @@ -0,0 +1,330 @@ +//! Sharded MULTI/EXEC key-locality safety (2026-07 follow-up). +//! +//! `execute_transaction_sharded` runs the whole queued body on the connection's +//! OWN shard with no per-key routing. A key owned by a different shard was +//! therefore silently written to / read from the wrong shard's table — EXEC +//! reported success while the data diverged (and WATCH-less "lost updates"). +//! +//! Phase A closed the silent-corruption hole by rejecting such a transaction +//! with `CROSSSLOT`; Phase B goes further and ROUTES a single-owner-shard body +//! to its owner (so single-key and hash-tagged bodies now succeed from any +//! connection — see `sharded_multi_exec_routing.rs`). Bodies that genuinely +//! span shards stay rejected. Either way the golden invariant these tests pin +//! is unchanged: +//! +//! **If EXEC returns a successful array, every write in it is visible to +//! other connections; if EXEC returns CROSSSLOT, nothing was written.** +//! +//! Never "EXEC says OK but the data isn't there." Skips gracefully when the +//! moon binary is missing (MOON_BIN pin wins, then target/release, then debug). + +use std::io::{Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +fn free_port() -> u16 { + TcpListener::bind("127.0.0.1:0") + .expect("bind") + .local_addr() + .unwrap() + .port() +} + +fn moon_binary() -> Option { + if let Ok(p) = std::env::var("MOON_BIN") { + return Some(std::path::PathBuf::from(p)); + } + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); + for rel in ["target/release/moon", "target/debug/moon"] { + let p = root.join(rel); + if p.exists() { + return Some(p); + } + } + None +} + +struct Moon { + child: Child, + port: u16, + tmp_dir: std::path::PathBuf, +} + +impl Drop for Moon { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + let _ = std::fs::remove_dir_all(&self.tmp_dir); + } +} + +fn spawn_moon(shards: &str) -> Option { + let bin = moon_binary()?; + let port = free_port(); + let tmp_dir = std::env::temp_dir().join(format!("moon-txn-loc-{port}")); + let _ = std::fs::create_dir_all(&tmp_dir); + let child = Command::new(&bin) + .args([ + "--port", + &port.to_string(), + "--shards", + shards, + "--admin-port", + "0", + "--appendonly", + "no", + "--disk-free-min-pct", + "0", + "--dir", + tmp_dir.to_str().unwrap(), + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .ok()?; + let moon = Moon { + child, + port, + tmp_dir, + }; + let deadline = Instant::now() + Duration::from_secs(10); + while Instant::now() < deadline { + if let Ok(mut c) = TcpStream::connect(("127.0.0.1", moon.port)) { + let _ = c.set_read_timeout(Some(Duration::from_millis(500))); + if c.write_all(b"*1\r\n$4\r\nPING\r\n").is_ok() { + let mut buf = [0u8; 64]; + if let Ok(n) = c.read(&mut buf) { + if n > 0 && buf.starts_with(b"+PONG") { + return Some(moon); + } + } + } + } + std::thread::sleep(Duration::from_millis(100)); + } + eprintln!("skipping: moon did not become ready on port {port}"); + None +} + +/// One parsed RESP2 reply — only the shapes these tests need. +#[derive(Debug, Clone, PartialEq)] +enum Reply { + Simple(String), + Error(String), + Int(i64), + Bulk(Option), + Array(Vec), +} + +/// Blocking RESP client that parses exactly one reply per command. +struct Client { + stream: TcpStream, + buf: Vec, + pos: usize, +} + +impl Client { + fn connect(port: u16) -> Self { + let stream = TcpStream::connect(("127.0.0.1", port)).expect("connect"); + stream + .set_read_timeout(Some(Duration::from_millis(2000))) + .unwrap(); + Self { + stream, + buf: Vec::new(), + pos: 0, + } + } + + fn send(&mut self, args: &[&str]) { + let mut out = format!("*{}\r\n", args.len()).into_bytes(); + for a in args { + out.extend_from_slice(format!("${}\r\n{a}\r\n", a.len()).as_bytes()); + } + self.stream.write_all(&out).expect("write"); + } + + fn fill(&mut self) { + let mut chunk = [0u8; 4096]; + match self.stream.read(&mut chunk) { + Ok(0) => panic!("connection closed by server"), + Ok(n) => self.buf.extend_from_slice(&chunk[..n]), + Err(e) => panic!("read error: {e}"), + } + } + + fn read_line(&mut self) -> String { + loop { + if let Some(rel) = self.buf[self.pos..].windows(2).position(|w| w == b"\r\n") { + let end = self.pos + rel; + let line = String::from_utf8_lossy(&self.buf[self.pos..end]).to_string(); + self.pos = end + 2; + return line; + } + self.fill(); + } + } + + fn read_exact_bytes(&mut self, n: usize) -> String { + while self.buf.len() - self.pos < n + 2 { + self.fill(); + } + let s = String::from_utf8_lossy(&self.buf[self.pos..self.pos + n]).to_string(); + self.pos += n + 2; // skip trailing CRLF + s + } + + fn read_reply(&mut self) -> Reply { + let line = self.read_line(); + let (tag, rest) = line.split_at(1); + match tag { + "+" => Reply::Simple(rest.to_string()), + "-" => Reply::Error(rest.to_string()), + ":" => Reply::Int(rest.parse().unwrap_or(0)), + "$" => { + let len: i64 = rest.parse().unwrap_or(-1); + if len < 0 { + Reply::Bulk(None) + } else { + Reply::Bulk(Some(self.read_exact_bytes(len as usize))) + } + } + "*" => { + let len: i64 = rest.parse().unwrap_or(-1); + if len < 0 { + Reply::Array(Vec::new()) + } else { + let mut items = Vec::with_capacity(len as usize); + for _ in 0..len { + items.push(self.read_reply()); + } + Reply::Array(items) + } + } + other => panic!("unexpected RESP tag {other:?} in line {line:?}"), + } + } + + fn cmd(&mut self, args: &[&str]) -> Reply { + self.send(args); + self.read_reply() + } +} + +/// GOLDEN INVARIANT: a transaction that touches a key owned by another shard +/// must never "succeed silently and lose the write." Over many distinct keys +/// (so both the local-shard and remote-shard branches are exercised at random +/// accept placement), assert: EXEC-array success ⇒ the value is readable from a +/// fresh connection; EXEC CROSSSLOT ⇒ the key does not exist. +#[test] +fn no_silent_divergence_across_shards() { + let Some(m) = spawn_moon("4") else { return }; + + let mut saw_ok = 0usize; + let mut saw_crossslot = 0usize; + + for i in 0..40 { + let key = format!("txnloc:{i}"); + let val = format!("v{i}"); + + let mut writer = Client::connect(m.port); + assert_eq!(writer.cmd(&["MULTI"]), Reply::Simple("OK".into())); + assert_eq!( + writer.cmd(&["SET", &key, &val]), + Reply::Simple("QUEUED".into()) + ); + let exec = writer.cmd(&["EXEC"]); + + // A fresh connection reads the key (routes to the true owner shard). + let mut reader = Client::connect(m.port); + let got = reader.cmd(&["GET", &key]); + + match &exec { + Reply::Array(items) => { + saw_ok += 1; + // The queued SET succeeded, so the value MUST be visible. + assert_eq!( + got, + Reply::Bulk(Some(val.clone())), + "EXEC reported success {items:?} but GET {key} = {got:?} (silent misplacement)" + ); + } + Reply::Error(e) if e.starts_with("CROSSSLOT") => { + saw_crossslot += 1; + // Rejected — nothing should have been written. + assert_eq!( + got, + Reply::Bulk(None), + "EXEC was rejected ({e}) but GET {key} = {got:?} (partial write leaked)" + ); + } + other => panic!("unexpected EXEC reply for {key}: {other:?}"), + } + } + + // Diagnostic only — the invariant above is what matters. Under Phase B a + // single-key body is routed to its owner, so `saw_ok` should be ~all 40 and + // `saw_crossslot` ~0 (pre-fix these "ok" cases silently returned Bulk(None); + // under Phase A the remote ones were CROSSSLOT). + eprintln!("no_silent_divergence: {saw_ok} succeeded, {saw_crossslot} rejected"); +} + +/// A transaction whose keys demonstrably span multiple shards is always +/// rejected with CROSSSLOT (20 distinct keys → span is near-certain at 4 +/// shards), and nothing it queued is written. +#[test] +fn multi_shard_span_rejected() { + let Some(m) = spawn_moon("4") else { return }; + let mut c = Client::connect(m.port); + + assert_eq!(c.cmd(&["MULTI"]), Reply::Simple("OK".into())); + for i in 0..20 { + let key = format!("span:{i}"); + assert_eq!(c.cmd(&["SET", &key, "x"]), Reply::Simple("QUEUED".into())); + } + match c.cmd(&["EXEC"]) { + Reply::Error(e) => assert!( + e.starts_with("CROSSSLOT"), + "20-key span must be CROSSSLOT, got: {e}" + ), + other => panic!("20-key span must be rejected, got: {other:?}"), + } + + // None of the queued keys should exist. + let mut r = Client::connect(m.port); + for i in 0..20 { + let key = format!("span:{i}"); + assert_eq!( + r.cmd(&["GET", &key]), + Reply::Bulk(None), + "rejected txn must not have written {key}" + ); + } +} + +/// Single-shard servers are unaffected: the `num_shards > 1` guard means the +/// locality check never runs, MULTI/EXEC returns the results array (never +/// CROSSSLOT), and the write is applied. (Uses writes only — a solo GET sent +/// mid-MULTI hits an unrelated monoio inline-read quirk, out of scope here.) +#[test] +fn single_shard_unaffected() { + let Some(m) = spawn_moon("1") else { return }; + let mut c = Client::connect(m.port); + + assert_eq!(c.cmd(&["MULTI"]), Reply::Simple("OK".into())); + assert_eq!(c.cmd(&["SET", "s:a", "1"]), Reply::Simple("QUEUED".into())); + assert_eq!(c.cmd(&["SET", "s:b", "2"]), Reply::Simple("QUEUED".into())); + match c.cmd(&["EXEC"]) { + Reply::Array(items) => { + assert_eq!(items.len(), 2, "EXEC array has one reply per queued write"); + assert_eq!(items[0], Reply::Simple("OK".into())); + assert_eq!(items[1], Reply::Simple("OK".into())); + } + other => panic!("single-shard EXEC must return the results array, got: {other:?}"), + } + // Writes are visible afterward (no CROSSSLOT swallowed them). + let mut r = Client::connect(m.port); + assert_eq!(r.cmd(&["GET", "s:a"]), Reply::Bulk(Some("1".into()))); + assert_eq!(r.cmd(&["GET", "s:b"]), Reply::Bulk(Some("2".into()))); +} diff --git a/tests/sharded_multi_exec_routing.rs b/tests/sharded_multi_exec_routing.rs new file mode 100644 index 000000000..bc419714d --- /dev/null +++ b/tests/sharded_multi_exec_routing.rs @@ -0,0 +1,354 @@ +//! Sharded MULTI/EXEC Phase B: route a single-owner-shard body to its owner. +//! +//! Phase A rejected any transaction whose keys weren't owned by the connection's +//! accept shard (random via SO_REUSEPORT) with `CROSSSLOT` — so a hash-tagged +//! `{tag}` multi-key transaction worked only from the ~1/N connections that +//! happened to land on the owning shard. Phase B routes the whole body to the +//! owning shard (via `ShardMessage::TxnExecute`), executes it atomically there, +//! and persists writes to THAT shard's AOF. Genuinely multi-shard bodies stay +//! rejected — a shared-nothing engine can't commit across shards atomically. +//! +//! Invariants pinned here: +//! 1. A hash-tagged MULTI/EXEC succeeds (returns the results array, never +//! CROSSSLOT) from EVERY connection, and its writes are visible to others. +//! 2. Those writes survive a kill-9 + restart (AOF persisted on the OWNER). +//! 3. A body whose keys span shards is still rejected with CROSSSLOT. +//! +//! Skips gracefully when the moon binary is missing (MOON_BIN pin wins, then +//! target/release, then target/debug). + +use std::io::{Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +fn free_port() -> u16 { + TcpListener::bind("127.0.0.1:0") + .expect("bind") + .local_addr() + .unwrap() + .port() +} + +fn moon_binary() -> Option { + if let Ok(p) = std::env::var("MOON_BIN") { + return Some(std::path::PathBuf::from(p)); + } + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); + for rel in ["target/release/moon", "target/debug/moon"] { + let p = root.join(rel); + if p.exists() { + return Some(p); + } + } + None +} + +struct Moon { + child: Child, + port: u16, +} + +impl Moon { + fn kill9(mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +/// Spawn moon at `--shards 4`. `durable` selects `appendonly yes` + +/// `appendfsync always` (for the restart test) vs `appendonly no`. +fn spawn_moon(port: u16, dir: &std::path::Path, durable: bool) -> Option { + let bin = moon_binary()?; + let mut args: Vec = vec![ + "--port".into(), + port.to_string(), + "--shards".into(), + "4".into(), + "--admin-port".into(), + "0".into(), + "--disk-free-min-pct".into(), + "0".into(), + "--dir".into(), + dir.to_str().unwrap().into(), + ]; + if durable { + args.extend(["--appendonly".into(), "yes".into()]); + args.extend(["--appendfsync".into(), "always".into()]); + } else { + args.extend(["--appendonly".into(), "no".into()]); + } + let child = Command::new(&bin) + .args(&args) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .ok()?; + let moon = Moon { child, port }; + let deadline = Instant::now() + Duration::from_secs(10); + while Instant::now() < deadline { + if let Ok(mut c) = TcpStream::connect(("127.0.0.1", moon.port)) { + let _ = c.set_read_timeout(Some(Duration::from_millis(500))); + if c.write_all(b"*1\r\n$4\r\nPING\r\n").is_ok() { + let mut buf = [0u8; 64]; + if let Ok(n) = c.read(&mut buf) { + if n > 0 && buf.starts_with(b"+PONG") { + return Some(moon); + } + } + } + } + std::thread::sleep(Duration::from_millis(100)); + } + eprintln!("skipping: moon did not become ready on port {port}"); + None +} + +#[derive(Debug, Clone, PartialEq)] +enum Reply { + Simple(String), + Error(String), + Int(i64), + Bulk(Option), + Array(Vec), +} + +struct Client { + stream: TcpStream, + buf: Vec, + pos: usize, +} + +impl Client { + fn connect(port: u16) -> Self { + let stream = TcpStream::connect(("127.0.0.1", port)).expect("connect"); + stream + .set_read_timeout(Some(Duration::from_millis(3000))) + .unwrap(); + Self { + stream, + buf: Vec::new(), + pos: 0, + } + } + + fn send(&mut self, args: &[&str]) { + let mut out = format!("*{}\r\n", args.len()).into_bytes(); + for a in args { + out.extend_from_slice(format!("${}\r\n{a}\r\n", a.len()).as_bytes()); + } + self.stream.write_all(&out).expect("write"); + } + + fn fill(&mut self) { + let mut chunk = [0u8; 4096]; + match self.stream.read(&mut chunk) { + Ok(0) => panic!("connection closed by server"), + Ok(n) => self.buf.extend_from_slice(&chunk[..n]), + Err(e) => panic!("read error: {e}"), + } + } + + fn read_line(&mut self) -> String { + loop { + if let Some(rel) = self.buf[self.pos..].windows(2).position(|w| w == b"\r\n") { + let end = self.pos + rel; + let line = String::from_utf8_lossy(&self.buf[self.pos..end]).to_string(); + self.pos = end + 2; + return line; + } + self.fill(); + } + } + + fn read_exact_bytes(&mut self, n: usize) -> String { + while self.buf.len() - self.pos < n + 2 { + self.fill(); + } + let s = String::from_utf8_lossy(&self.buf[self.pos..self.pos + n]).to_string(); + self.pos += n + 2; + s + } + + fn read_reply(&mut self) -> Reply { + let line = self.read_line(); + let (tag, rest) = line.split_at(1); + match tag { + "+" => Reply::Simple(rest.to_string()), + "-" => Reply::Error(rest.to_string()), + ":" => Reply::Int(rest.parse().unwrap_or(0)), + "$" => { + let len: i64 = rest.parse().unwrap_or(-1); + if len < 0 { + Reply::Bulk(None) + } else { + Reply::Bulk(Some(self.read_exact_bytes(len as usize))) + } + } + "*" => { + let len: i64 = rest.parse().unwrap_or(-1); + if len < 0 { + Reply::Array(Vec::new()) + } else { + let mut items = Vec::with_capacity(len as usize); + for _ in 0..len { + items.push(self.read_reply()); + } + Reply::Array(items) + } + } + other => panic!("unexpected RESP tag {other:?} in line {line:?}"), + } + } + + fn cmd(&mut self, args: &[&str]) -> Reply { + self.send(args); + self.read_reply() + } +} + +/// INVARIANT 1: a hash-tagged multi-key MULTI/EXEC succeeds from EVERY +/// connection (never CROSSSLOT — Phase B routes it to the owning shard) and its +/// writes are visible from a fresh connection. 40 distinct tags exercise all 4 +/// shards as owners across random accept placement. +#[test] +fn routed_hashtag_txn_succeeds_from_any_connection() { + let port = free_port(); + let dir = std::env::temp_dir().join(format!("moon-txn-route-{port}")); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::create_dir_all(&dir); + let Some(m) = spawn_moon(port, &dir, false) else { + let _ = std::fs::remove_dir_all(&dir); + return; + }; + + for i in 0..40 { + let ka = format!("user:{{{i}}}:a"); + let kb = format!("user:{{{i}}}:b"); + let kc = format!("user:{{{i}}}:c"); + let va = format!("va{i}"); + let vb = format!("vb{i}"); + + // Fresh connection each iteration → random accept shard. + let mut w = Client::connect(port); + assert_eq!(w.cmd(&["MULTI"]), Reply::Simple("OK".into())); + assert_eq!(w.cmd(&["SET", &ka, &va]), Reply::Simple("QUEUED".into())); + assert_eq!(w.cmd(&["SET", &kb, &vb]), Reply::Simple("QUEUED".into())); + assert_eq!(w.cmd(&["INCRBY", &kc, "7"]), Reply::Simple("QUEUED".into())); + match w.cmd(&["EXEC"]) { + Reply::Array(items) => { + assert_eq!(items.len(), 3, "one reply per queued command"); + assert_eq!(items[0], Reply::Simple("OK".into())); + assert_eq!(items[1], Reply::Simple("OK".into())); + assert_eq!(items[2], Reply::Int(7)); + } + other => { + panic!("hash-tagged txn (tag {i}) must succeed via Phase B routing, got: {other:?}") + } + } + + // Visible from a fresh reader (routes each key to its true owner). + let mut r = Client::connect(port); + assert_eq!( + r.cmd(&["GET", &ka]), + Reply::Bulk(Some(va.clone())), + "routed txn wrote {ka} but it isn't visible" + ); + assert_eq!(r.cmd(&["GET", &kb]), Reply::Bulk(Some(vb.clone()))); + assert_eq!(r.cmd(&["GET", &kc]), Reply::Bulk(Some("7".into()))); + } + + m.kill9(); + let _ = std::fs::remove_dir_all(&dir); +} + +/// INVARIANT 2: routed (Phase B) transactional writes survive kill-9 + restart — +/// they must be persisted to the OWNER shard's AOF, not silently dropped. +#[test] +fn routed_hashtag_txn_survives_restart() { + let port = free_port(); + let dir = std::env::temp_dir().join(format!("moon-txn-route-dur-{port}")); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::create_dir_all(&dir); + + let Some(m1) = spawn_moon(port, &dir, true) else { + let _ = std::fs::remove_dir_all(&dir); + return; + }; + + // 12 distinct tags → spread across all 4 owner shards; fresh conn each time + // so most bodies route cross-shard. + for i in 0..12 { + let ka = format!("acct:{{{i}}}:name"); + let kb = format!("acct:{{{i}}}:bal"); + let mut w = Client::connect(port); + assert_eq!(w.cmd(&["MULTI"]), Reply::Simple("OK".into())); + assert_eq!( + w.cmd(&["SET", &ka, "alice"]), + Reply::Simple("QUEUED".into()) + ); + assert_eq!( + w.cmd(&["INCRBY", &kb, "100"]), + Reply::Simple("QUEUED".into()) + ); + match w.cmd(&["EXEC"]) { + Reply::Array(items) => assert_eq!(items.len(), 2), + other => panic!("routed txn (tag {i}) must succeed, got: {other:?}"), + } + } + + m1.kill9(); + + let Some(m2) = spawn_moon(port, &dir, true) else { + let _ = std::fs::remove_dir_all(&dir); + panic!("moon failed to restart from {dir:?}"); + }; + let mut r = Client::connect(port); + let mut missing = Vec::new(); + for i in 0..12 { + let ka = format!("acct:{{{i}}}:name"); + let kb = format!("acct:{{{i}}}:bal"); + if r.cmd(&["GET", &ka]) != Reply::Bulk(Some("alice".into())) { + missing.push(ka); + } + if r.cmd(&["GET", &kb]) != Reply::Bulk(Some("100".into())) { + missing.push(kb); + } + } + m2.kill9(); + let _ = std::fs::remove_dir_all(&dir); + assert!( + missing.is_empty(), + "routed MULTI/EXEC writes lost on restart (not persisted on owner shard): {missing:?}" + ); +} + +/// INVARIANT 3: a body whose keys demonstrably span shards is still rejected +/// with CROSSSLOT — Phase B only routes single-owner-shard bodies. +#[test] +fn multi_shard_span_still_rejected() { + let port = free_port(); + let dir = std::env::temp_dir().join(format!("moon-txn-span-{port}")); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::create_dir_all(&dir); + let Some(m) = spawn_moon(port, &dir, false) else { + let _ = std::fs::remove_dir_all(&dir); + return; + }; + let mut c = Client::connect(port); + assert_eq!(c.cmd(&["MULTI"]), Reply::Simple("OK".into())); + for i in 0..20 { + // Distinct untagged keys → near-certain to span >1 of 4 shards. + let key = format!("span:{i}"); + assert_eq!(c.cmd(&["SET", &key, "x"]), Reply::Simple("QUEUED".into())); + } + match c.cmd(&["EXEC"]) { + Reply::Error(e) => assert!( + e.starts_with("CROSSSLOT"), + "20-key span must be CROSSSLOT, got: {e}" + ), + other => panic!("20-key span must be rejected, got: {other:?}"), + } + m.kill9(); + let _ = std::fs::remove_dir_all(&dir); +}