From 370f2f52c5442ccacf69d0b3381155da96df6486 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Tue, 25 Aug 2026 05:02:10 +0700 Subject: [PATCH 1/2] perf(pipeline): defer only for shards the batch has pending work on (#513) `must_wait_for_pending_remote`'s multi-key arm answered "wait" on the command NAME, without asking where its keys were. But `remote_groups` only ever holds FOREIGN shards -- the slotting branch is `else if let Some(target) = target_shard`, and `target_shard` is `None` for a local key -- so the moon#507 hazard (reading state a pending command is about to write, or writing state it then overwrites) requires the two to meet ON THE SAME SHARD. An MGET reading shards the batch has no pending work for was being cut for nothing. The cut is not free: it ends the batch pass, and the phase-2b drain then dispatches one PipelineBatchSlotted per target shard and awaits each reply slot in turn. The guard now takes a `pending` bitmask, maintained beside `remote_groups` at O(1) (set on insert, cleared with the map), and compares it against the shard mask of the command's keys. The mask comes from the shared key-position walker (moon#582) -- the same one ACL, cache invalidation and `cross_shard_multikey_rejection` use -- so layouts like ZUNIONSTORE are enumerated by the code that already knows them rather than a second copy. Only the multi-key arm is refined; the other two still always wait, because neither can be bounded by a key mask. An inline-intercepted command (EVAL, SWAPDB) executes against the LOCAL slice whatever keys it declares, and a keyless command (FLUSHALL, KEYS, SCAN) touches every shard. Every case the mask cannot enumerate -- `SORT ... BY w_*`, a key position holding a non-string, more shards than the mask has bits -- also waits: wrongly waiting costs a batch boundary, wrongly proceeding corrupts data. The two predicates were folded into one. The unmasked form had no callers left and keeping it would have left two answers to the same question. Measured on moon-dev (aarch64, 6 vCPU), --shards 4, 32 interleavings of SET,SET,MGET, six fresh server starts per side, interleaved: MGET reads shards the writes never touch: 41,600 -> 86,500 ops/s (2.08x), 64 deferrals -> 0 MGET reads the shards being written: 34,300 ops/s, 64 deferrals, both Fresh starts per measurement because SO_REUSEPORT decides which shard the connection lands on, and that changes the shape's cost as much as the code does. The deferral counts are placement-independent: 64/64 before and 0/0 after in every round. A co-located {tag} multi-key command still defers, and must: the coordinator executes it inline rather than slotting it, so skipping the wait would re-open moon#507. Routing a single-owner multi-key command into the slotted batch is tracked separately. Tests: `pco13` drives the disjoint shape and asserts 0 deferrals, with an overlap leg on the same harness that must stay non-zero -- without it a green disjoint leg could just mean the writes never went cross-shard. Both legs pin ONE key per shard rather than "the first n keys in this set", because two keys that both hashed to shard 0 made the overlap leg depend on where SO_REUSEPORT put the connection (caught as a real flake while benchmarking). Four unit tests cover the mask itself, including every fail-closed path. Verified by mutation: reverting the guard to name-only fails pco13 alone; making it never wait fails five of the moon#507 correctness tests. Refs #513, #507, #512 author: Tin Dang --- CHANGELOG.md | 36 +++++ src/server/conn/handler_monoio/mod.rs | 14 +- src/server/conn/handler_sharded/mod.rs | 14 +- src/server/conn/shared.rs | 198 ++++++++++++++++++++++++- tests/pipeline_cross_shard_ordering.rs | 140 +++++++++++++++++ 5 files changed, 394 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 578a12dd..8d4f5608 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,42 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- **A pipelined multi-key command no longer cuts the batch over shards it never touches** (#513). + + `must_wait_for_pending_remote`'s multi-key arm answered "wait" on the command NAME, without + asking where its keys were. But `remote_groups` only ever holds FOREIGN shards, so the #507 + hazard — reading state a pending command is about to write, or writing state it then + overwrites — requires the two to meet on the *same* shard. An `MGET` reading shards the batch + has no pending work for was being cut for nothing. + + The cut is not free: it ends the batch pass, and the phase-2b drain then dispatches one + `PipelineBatchSlotted` per target shard and awaits each reply slot in turn. + + Measured on moon-dev (aarch64, 6 vCPU), `--shards 4`, 32 interleavings of + `SET`,`SET`,`MGET`, six fresh server starts per side interleaved: + + | shape | before | after | + |---|---|---| + | `MGET` reads shards the writes never touch | 41,600 ops/s, 64 deferrals | 86,500 ops/s, **0** | + | `MGET` reads the shards being written | 34,300 ops/s, 64 deferrals | unchanged | + + Fresh starts per measurement because `SO_REUSEPORT` decides which shard the connection lands + on, and that changes the shape's cost as much as the code does — one start per side compares + placements as much as binaries. The deferral counts are placement-independent and were 64/64 + before and 0/0 after in every round. + + Only the multi-key arm is refined; the other two still always wait, because neither can be + bounded by a key mask. An inline-intercepted command (`EVAL`, `SWAPDB`, …) executes against + the local slice whatever keys it declares, and a keyless command (`FLUSHALL`, `KEYS`, `SCAN`) + touches every shard. A key layout the shared walker cannot enumerate — `SORT ... BY w_*`, a + key position holding a non-string, more shards than the mask has bits — also still waits: + wrongly waiting costs a batch boundary, wrongly proceeding corrupts data. + + A co-located `{tag}` multi-key command still defers, and must: the coordinator executes it + inline rather than slotting it, so skipping the wait would re-open #507. Routing a + single-owner multi-key command into the slotted batch is tracked separately. + ### Added - **`INFO stats` reports what the pipeline ordering guarantee costs** (`total_pipeline_remote_defer`, and the `moon_pipeline_remote_defer_total` Prometheus counter) — groundwork for #513. diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index d7c6815a..92a31810 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -1506,6 +1506,10 @@ pub(crate) async fn handle_connection_sharded_monoio< let mut should_quit = false; responses.clear(); remote_groups.clear(); + // moon#513: one bit per shard `remote_groups` holds an entry for. + // Maintained beside the map so the ordering guard can ask "does this + // command touch a shard with pending work?" without walking it. + let mut pending_mask: u64 = 0; local_leg_write_idxs.clear(); // The trailing bool marks a SHARDED publish. One batch map, split at flush: // the two namespaces share the fan-out plumbing but never the destination. @@ -1674,8 +1678,13 @@ pub(crate) async fn handle_connection_sharded_monoio< // phase 2 resolves the pending replies first, and the tail // re-parses at the top of the next batch with `remote_groups` // empty, so this cannot loop. - if !remote_groups.is_empty() - && crate::server::conn::shared::must_wait_for_pending_remote(cmd, cmd_args) + if pending_mask != 0 + && crate::server::conn::shared::must_wait_for_pending_remote( + cmd, + cmd_args, + ctx.num_shards, + pending_mask, + ) { frames[frame_idx - 1] = frame; crate::admin::metrics_setup::record_pipeline_remote_defer(); @@ -3331,6 +3340,7 @@ pub(crate) async fn handle_connection_sharded_monoio< track_keys, resp3_shape, )); + pending_mask |= 1u64 << (target % u64::BITS as usize); crate::admin::metrics_setup::record_dispatch_cross_spsc(); } } diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 3bbba278..398d4842 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -597,6 +597,10 @@ pub(crate) async fn handle_connection_sharded_inner< let mut local_leg_write_idxs: Vec = Vec::new(); let mut should_quit = false; let mut remote_groups: HashMap, Option, usize, Option, crate::protocol::resp3::Resp3Shape)>> = HashMap::with_capacity(ctx.num_shards); + // moon#513: one bit per shard `remote_groups` holds an entry for. + // Maintained beside the map so the ordering guard can ask "does + // this command touch a shard with pending work?" without walking it. + let mut pending_mask: u64 = 0; // Accumulate cross-shard PUBLISH pairs per target shard for batch dispatch // Key: target shard ID -> Vec of (response_index, channel, message) // Trailing bool marks a SHARDED publish. One batch map, split at flush: @@ -814,8 +818,13 @@ pub(crate) async fn handle_connection_sharded_inner< // pending replies first, and the tail re-parses at the top // of the next batch with `remote_groups` empty, so this // cannot loop. - if !remote_groups.is_empty() - && crate::server::conn::shared::must_wait_for_pending_remote(cmd, cmd_args) + if pending_mask != 0 + && crate::server::conn::shared::must_wait_for_pending_remote( + cmd, + cmd_args, + ctx.num_shards, + pending_mask, + ) { batch[frame_idx - 1] = frame; crate::admin::metrics_setup::record_pipeline_remote_defer(); @@ -2657,6 +2666,7 @@ pub(crate) async fn handle_connection_sharded_inner< let resp_idx = responses.len(); responses.push(Frame::Null); remote_groups.entry(target).or_default().push((resp_idx, std::sync::Arc::new(dispatch_frame), aof_bytes, conn.selected_db, track_keys, resp3_shape)); + pending_mask |= 1u64 << (target % u64::BITS as usize); cross_spsc_dispatches = cross_spsc_dispatches.saturating_add(1); } } diff --git a/src/server/conn/shared.rs b/src/server/conn/shared.rs index 72b7c1fa..cfcd7275 100644 --- a/src/server/conn/shared.rs +++ b/src/server/conn/shared.rs @@ -1597,10 +1597,88 @@ pub(crate) fn extract_primary_key<'a>(cmd: &[u8], args: &'a [Frame]) -> Option<& /// executed at the start of the next batch, which is always correct and costs /// one batch boundary. Wrongly calling something SAFE is the direction that /// corrupts data, so when in doubt, add it to the wait set. -pub(crate) fn must_wait_for_pending_remote(cmd: &[u8], args: &[Frame]) -> bool { - is_multi_key_command(cmd, args) - || is_inline_intercepted(cmd) - || extract_primary_key(cmd, args).is_none() +/// The shards a command's keys live on, as a bitmask — or `None` when the mask +/// cannot be trusted and the caller must fall back to waiting. +/// +/// `None` for: more shards than the mask has bits; a key position holding a +/// non-string (a malformed argv, which must reach the command so it earns its +/// own error rather than a routing decision); a layout the shared walker cannot +/// enumerate; and `AtPlusComputed` (`SORT ... BY w_*`), where a key nobody can +/// name would be missing from the mask — exactly the case that must keep +/// waiting. +/// +/// Uses the shared key-position walker (moon#582), the same one ACL, cache +/// invalidation and [`cross_shard_multikey_rejection`] use, so `ZUNIONSTORE dst +/// numkeys ...` is enumerated by the code that already knows that layout rather +/// than by a second, drifting copy. +#[must_use] +fn command_shard_mask(cmd: &[u8], args: &[Frame], num_shards: usize) -> Option { + if num_shards > u64::BITS as usize { + return None; + } + let idx = match crate::acl::keyspec::command_key_positions(cmd, args) { + crate::acl::keyspec::KeyPositions::At(idx) => idx, + _ => return None, + }; + let mut mask = 0u64; + for k in idx { + let key: &[u8] = match args.get(k.idx) { + Some(Frame::BulkString(b) | Frame::SimpleString(b)) => b.as_ref(), + _ => return None, + }; + mask |= 1u64 << crate::shard::dispatch::key_to_shard(key, num_shards); + } + // An empty mask would read as "touches nothing" and skip the wait. A + // multi-key command that named no key at all is malformed, so wait. + (mask != 0).then_some(mask) +} + +/// # moon#513: refined by WHERE the keys are +/// +/// `pending` is the set of shards this batch still has undispatched commands +/// for — one bit per shard, maintained alongside `remote_groups`. A zero mask +/// means the caller should not be asking at all. +/// +/// The hazard `must_wait_for_pending_remote` exists to stop is a command +/// reading state a pending command is about to write, or writing state a +/// pending command then overwrites. Both require the two to meet ON THE SAME +/// SHARD. `remote_groups` holds only FOREIGN shards — the slotting branch is +/// `else if let Some(target) = target_shard`, and `target_shard` is `None` for a +/// local key — so a command whose keys avoid every pending shard cannot be in +/// that hazard, and cutting the batch for it buys nothing. +/// +/// It buys nothing at a real price: a cut ends the batch pass, and the phase-2b +/// drain then dispatches one `PipelineBatchSlotted` per target shard and awaits +/// each reply slot in turn. Measured at `--shards 4`, a `{tag}`-co-located +/// `MGET` interleaved between writes — the co-location pattern the docs tell +/// users to adopt — cut 32 times in 32 interleavings, each one a full drain cycle. +/// +/// Only the multi-key arm is refined. The other two cannot be bounded by a key +/// mask and are unchanged: +/// +/// * an inline-intercepted command (EVAL, SWAPDB, …) executes against the LOCAL +/// slice whatever keys it declares, so its declared keys do not describe the +/// state it touches; +/// * a keyless command (FLUSHALL, KEYS, SCAN, DBSIZE) aggregates across every +/// shard, so its mask would be "all of them" anyway. +/// +/// The conservative direction is unchanged too: every `None` from +/// [`command_shard_mask`] waits. +#[must_use] +pub(crate) fn must_wait_for_pending_remote( + cmd: &[u8], + args: &[Frame], + num_shards: usize, + pending: u64, +) -> bool { + if is_inline_intercepted(cmd) || extract_primary_key(cmd, args).is_none() { + return true; + } + if !is_multi_key_command(cmd, args) { + // Routed by its own single key — the fast path, unchanged. + return false; + } + command_shard_mask(cmd, args, num_shards).is_none_or(|mask| mask & pending != 0) } /// Commands handled INLINE by a `try_handle_*` interceptor before the routing @@ -3689,3 +3767,115 @@ mod cross_shard_write_tests { ); } } + +#[cfg(test)] +mod pending_shard_mask_tests { + //! moon#513: the guard now asks WHERE a multi-key command's keys are. + //! + //! Shard membership is searched for with the routing hash the server + //! itself uses, never written as a literal — a hardcoded key that drifted + //! onto a different shard would make these pass for the wrong reason. + + use super::{command_shard_mask, must_wait_for_pending_remote}; + use crate::protocol::Frame; + use crate::shard::dispatch::key_to_shard; + use bytes::Bytes; + + const SHARDS: usize = 4; + + fn bulk(s: &str) -> Frame { + Frame::BulkString(Bytes::copy_from_slice(s.as_bytes())) + } + + /// Two keys owned by `want`, found by hashing rather than assumed. + fn keys_on(prefix: &str, want: usize) -> Vec { + let found: Vec = (0..10_000) + .map(|i| format!("{prefix}:{i}")) + .filter(|k| key_to_shard(k.as_bytes(), SHARDS) == want) + .take(2) + .map(|k| bulk(&k)) + .collect(); + assert_eq!(found.len(), 2, "no keys found for shard {want}"); + found + } + + fn mask_of(shards: &[usize]) -> u64 { + shards.iter().fold(0u64, |m, s| m | 1u64 << s) + } + + #[test] + fn mask_names_every_shard_the_keys_touch() { + let mut args = keys_on("psm_a", 0); + args.extend(keys_on("psm_b", 3)); + assert_eq!( + command_shard_mask(b"MGET", &args, SHARDS), + Some(mask_of(&[0, 3])), + "an MGET spanning shards 0 and 3 must report both" + ); + assert_eq!( + command_shard_mask(b"MGET", &keys_on("psm_c", 2), SHARDS), + Some(mask_of(&[2])), + "keys that all hash to one shard must report one bit" + ); + } + + #[test] + fn unenumerable_argv_fails_closed() { + // A key position holding a non-string is a malformed invocation: it must + // reach the command and earn its own error, not a routing decision. + assert_eq!( + command_shard_mask(b"MGET", &[bulk("k"), Frame::Integer(7)], SHARDS), + None + ); + // Nothing named at all — waiting is the only safe answer. + assert_eq!(command_shard_mask(b"MGET", &[], SHARDS), None); + // More shards than the mask has bits. + assert_eq!( + command_shard_mask(b"MGET", &keys_on("psm_d", 1), 65), + None, + "a shard id past bit 63 cannot be represented, so the mask must not \ + claim to know" + ); + } + + #[test] + fn waits_only_when_the_shard_sets_meet() { + let on2 = keys_on("psm_e", 2); + assert!( + must_wait_for_pending_remote(b"MGET", &on2, SHARDS, mask_of(&[2])), + "reading shard 2 while shard 2 has pending work is the moon#507 \ + hazard and must still wait" + ); + assert!( + !must_wait_for_pending_remote(b"MGET", &on2, SHARDS, mask_of(&[0, 1, 3])), + "reading shard 2 while every OTHER shard has pending work touches \ + nothing pending, so the batch must not be cut" + ); + assert!( + must_wait_for_pending_remote(b"MGET", &[bulk("k"), Frame::Integer(7)], SHARDS, 1), + "an unenumerable argv falls back to waiting" + ); + } + + #[test] + fn the_other_two_arms_are_unchanged() { + let on2 = keys_on("psm_f", 2); + // Inline-intercepted: runs against the LOCAL slice whatever keys it + // declares, so its declared keys do not describe the state it touches. + assert!( + must_wait_for_pending_remote(b"EVAL", &on2, SHARDS, mask_of(&[0])), + "EVAL must wait even when its declared keys avoid every pending shard" + ); + // Keyless: aggregates across every shard. + assert!( + must_wait_for_pending_remote(b"DBSIZE", &[], SHARDS, mask_of(&[0])), + "a keyless command touches every shard" + ); + // Routed by its own single key: the fast path, never cut. + assert!( + !must_wait_for_pending_remote(b"GET", &on2[..1], SHARDS, u64::MAX), + "a single-key command routes by its own key and is ordered by the \ + slotted batch itself" + ); + } +} diff --git a/tests/pipeline_cross_shard_ordering.rs b/tests/pipeline_cross_shard_ordering.rs index 9b7b382a..b091fac5 100644 --- a/tests/pipeline_cross_shard_ordering.rs +++ b/tests/pipeline_cross_shard_ordering.rs @@ -261,6 +261,146 @@ fn pco12_the_ordering_guard_reports_what_it_costs() { ); } +/// moon#513: a multi-key command must not cut the batch over shards it never +/// touches. +/// +/// `must_wait_for_pending_remote`'s multi-key arm answered "wait" without asking +/// WHERE the keys are. But `remote_groups` only ever holds FOREIGN shards — the +/// slotting branch is `else if let Some(target) = target_shard`, and +/// `target_shard` is `None` for a local key — so the moon#507 hazard (reading +/// state a pending command is about to write, or writing state it then +/// overwrites) requires the two to meet ON THE SAME SHARD. An `MGET` reading +/// shards the batch has no pending work for cannot be in that hazard, and the +/// cut buys nothing. +/// +/// It is not free: a cut ends the batch pass, and the phase-2b drain then +/// dispatches one `PipelineBatchSlotted` per target shard and awaits each reply +/// slot in turn. +/// +/// # Why the overlap leg is here +/// +/// The connection lands on whichever shard `SO_REUSEPORT` gives it, and this +/// test cannot choose. If every written key happened to be LOCAL, nothing would +/// enter `remote_groups`, the guard would never be consulted, and the disjoint +/// leg would pass while proving nothing. So the writes deliberately span two +/// shards — at `--shards 4` at least one of them is foreign whatever shard the +/// connection got — and the overlap leg asserts that an `MGET` reading THOSE +/// shards still defers. A green disjoint leg means something only because the +/// overlap leg is non-zero on the same harness. +/// +/// The co-located case (`{tag}` keys, so the `MGET` and the pending writes share +/// one shard) is deliberately NOT here: it must still defer, because the +/// coordinator executes a multi-key command inline rather than slotting it, so +/// skipping the wait would re-open moon#507. Routing a single-owner multi-key +/// command into the slotted batch is the separate follow-up. +#[test] +fn pco13_disjoint_shard_multikey_does_not_cut_the_batch() { + fn defers(port: u16) -> u64 { + let mut c = Conn::open(port); + let info = c.send(&["INFO", "stats"]); + info.split("\r\n") + .find_map(|l| l.strip_prefix("total_pipeline_remote_defer:")) + .and_then(|v| v.trim().parse().ok()) + .unwrap_or_else(|| { + panic!("INFO stats has no total_pipeline_remote_defer field: {info:?}") + }) + } + + let shards: usize = SHARDS.parse().expect("SHARDS is a number"); + assert!( + shards >= 4, + "this test needs >= 4 shards to build a disjoint pair of shard sets" + ); + let owner = |k: &String| moon::shard::dispatch::key_to_shard(k.as_bytes(), shards); + // Probed, never assumed: a key set that silently collapsed onto one shard + // would turn the disjoint leg vacuous in exactly the way this test exists + // to prevent. + // ONE key per named shard, never "the first n keys landing in this set" — + // two keys that both hashed to shard 0 would make the overlap leg depend on + // which shard SO_REUSEPORT gave the connection: with the writes' shard-0 leg + // local, `remote_groups` would hold only shard 1, the shard-0-only MGET + // would be disjoint from it, and the leg would measure 0 for a reason that + // has nothing to do with the code under test. + let on = |prefix: &str, want: &[usize]| -> Vec { + let keys: Vec = want + .iter() + .map(|&s| { + (0..) + .map(|i| format!("{prefix}{s}:{i}")) + .find(|k| owner(k) == s) + .expect("a key exists for every shard") + }) + .collect(); + let got: Vec = keys.iter().map(owner).collect(); + assert_eq!(got, want, "probe put {prefix} on the wrong shards"); + keys + }; + // Writes and the overlap read both cover shards 0 AND 1. Whatever shard the + // connection landed on, at most one of those is local, so `remote_groups` + // always holds one of them and the overlap read always meets it. + let w = on("pco13w", &[0, 1]); + let overlap = on("pco13o", &[0, 1]); + // Reads only shards 2 and 3, which the writes never touch. + let disjoint = on("pco13d", &[2, 3]); + + let m = spawn_moon(SHARDS); + + // Runs the same 32 interleavings, reading `read` between each write pair. + let run = |read: &[String]| -> (u64, Vec) { + let base = defers(m.port); + let mut c = Conn::open(m.port); + let mut wrong = Vec::new(); + for i in 0..32 { + let r = c.pipeline(&[ + &["SET", &w[0], "1"], + &["SET", &w[1], "2"], + &["MGET", &read[0], &read[1]], + ]); + if framed_len(r.as_bytes(), 3).is_none() { + wrong.push(format!(" i={i}: expected 3 framed replies, got {r:?}")); + } + } + drop(c); + (defers(m.port) - base, wrong) + }; + + let (overlap_defers, overlap_wrong) = run(&overlap); + assert!(overlap_wrong.is_empty(), "{}", overlap_wrong.join("\n")); + assert!( + overlap_defers > 0, + "the MGET reads the very shards the pending SETs are writing, so the \ + moon#507 guard MUST still cut the batch. Zero here means the writes \ + never reached `remote_groups` at all — the harness landed on a shard \ + that made them local — and the disjoint leg below would prove nothing" + ); + + let (disjoint_defers, disjoint_wrong) = run(&disjoint); + assert!(disjoint_wrong.is_empty(), "{}", disjoint_wrong.join("\n")); + assert_eq!( + disjoint_defers, 0, + "the MGET reads shards 2 and 3 while the pending SETs write shards 0 \ + and 1, so there is nothing for it to wait on — yet it cut the batch \ + {disjoint_defers} times, once per interleaving, each a full drain \ + cycle. The guard is still answering on the command NAME instead of on \ + where its keys are (the overlap leg above measured \ + {overlap_defers} on this same harness, so the guard is reachable)" + ); + + // The relaxation must not cost the guarantee the guard exists for: a + // co-located MGET reading the keys its own batch just wrote still sees them. + let mut c = Conn::open(m.port); + let r = c.pipeline(&[ + &["SET", &w[0], "11"], + &["SET", &w[1], "22"], + &["MGET", &w[0], &w[1]], + ]); + assert!( + r.contains("$2\r\n11\r\n") && r.contains("$2\r\n22\r\n"), + "an MGET reading the keys its own batch just wrote must observe them — \ + moon#507 is the whole reason the guard exists: {r:?}" + ); +} + #[test] fn pco1_mget_sees_writes_from_its_own_batch() { let m = spawn_moon(SHARDS); From b7632349a975577522c429565b922ab004acde05 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Tue, 25 Aug 2026 05:39:11 +0700 Subject: [PATCH 2/2] fix(pipeline): judge the ordering guard on the keys a workspace will actually use The shard mask added earlier in this branch hashes the keys visible AT the guard. In a workspace connection those are the RAW keys: `workspace_rewrite_args` rebinds `cmd_args` further down the batch loop, and the guard cannot move below it -- the connection-level intercepts it exists to hold back (AUTH, CLIENT, CONFIG, INFO, SELECT, ...) run in between. The discrepancy is not small. A workspace key is `{<32-hex>}:`, and that prefix is a hash TAG, so every key in a workspace routes to ONE shard however the raw names scatter. A mask read off raw names therefore calls a command disjoint from the very shard its own batch's writes are pending on -- moon#507 reopened for exactly the connections that opted into isolation. Measured on the pre-fix build of this branch: 5 of 12 workspace connections had `SET a; SET b; MGET a b` answer `$-1 $-1` for keys the same batch had already acked `+OK`. The same test is green on the commit this branch forked from, so this was introduced here, not uncovered here. Fix: treat every shard as pending when the connection has a workspace, which makes `must_wait_for_pending_remote` answer exactly as it did before the mask existed. Workspace connections lose the batch-cut saving; they keep their data. Tests: `pco14` drives 12 workspace connections and asserts the MGET observes its own batch's writes. It asserts CORRECTNESS rather than a deferral count because the count is only wrong when the workspace's shard is foreign to the connection and SO_REUSEPORT decides that -- the correctness claim holds for every connection, so twelve make placement moot. Run 8x consecutively, green. Verified by mutation: dropping the workspace arm reproduces 5/12 losses. `pco13` gains a single-shard leg -- the shape the mask is most tempted to wave through -- and its final correctness block is relabelled: those keys span two shards, so calling them "co-located" was wrong. Refs #513, #507, #702 author: Tin Dang --- CHANGELOG.md | 16 ++++-- src/server/conn/handler_monoio/mod.rs | 19 ++++++- src/server/conn/handler_sharded/mod.rs | 19 ++++++- tests/pipeline_cross_shard_ordering.rs | 74 +++++++++++++++++++++++++- 4 files changed, 121 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d4f5608..56dc7708 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,9 +38,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 key position holding a non-string, more shards than the mask has bits — also still waits: wrongly waiting costs a batch boundary, wrongly proceeding corrupts data. - A co-located `{tag}` multi-key command still defers, and must: the coordinator executes it - inline rather than slotting it, so skipping the wait would re-open #507. Routing a - single-owner multi-key command into the slotted batch is tracked separately. + A co-located `{tag}` multi-key command still defers whenever its one shard is the shard with + pending work, and must: the coordinator executes a multi-key command inline rather than + slotting it, so skipping the wait there would re-open #507. Routing a single-owner multi-key + command into the slotted batch is the remaining scope of #513. + + Workspace connections keep the old always-wait behaviour. `workspace_rewrite_args` rebinds + the argv *below* this guard, and the guard cannot move down there — the connection-level + intercepts it exists to hold back run in between — so the keys visible to it are the raw + ones. The workspace prefix is a hash *tag*, so every key in a workspace routes to one shard + however the raw names scatter, and a mask read off raw names called commands disjoint from + the very shard their writes were pending on (5 of 12 connections lost an `MGET`'s own batch + writes before this was caught). Treating every shard as pending makes the predicate answer + exactly as it did before the mask existed. ### Added - **`INFO stats` reports what the pipeline ordering guarantee costs** (`total_pipeline_remote_defer`, diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 92a31810..8dcdcfd5 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -1678,12 +1678,29 @@ pub(crate) async fn handle_connection_sharded_monoio< // phase 2 resolves the pending replies first, and the tail // re-parses at the top of the next batch with `remote_groups` // empty, so this cannot loop. + // moon#513: a workspace connection's keys are rewritten to + // `{<32-hex>}:` BELOW this point (the `cmd_args` rebind), and + // this guard cannot move down there — the connection-level + // intercepts it exists to hold back run in between. So the keys + // visible HERE are the raw ones, and they hash to the wrong + // shards. Not by a little: the prefix is a hash TAG, so every key + // in a workspace routes to ONE shard however the raw names + // scatter, and a mask read off the raw names can call a command + // disjoint from the very shard its writes are pending on + // (measured: 5 of 12 connections lost an MGET's own batch writes). + // Treating every shard as pending makes the predicate answer + // exactly as it did before the mask existed. + let effective_pending = if conn.workspace_id.is_some() { + u64::MAX + } else { + pending_mask + }; if pending_mask != 0 && crate::server::conn::shared::must_wait_for_pending_remote( cmd, cmd_args, ctx.num_shards, - pending_mask, + effective_pending, ) { frames[frame_idx - 1] = frame; diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 398d4842..bf81edde 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -818,12 +818,29 @@ pub(crate) async fn handle_connection_sharded_inner< // pending replies first, and the tail re-parses at the top // of the next batch with `remote_groups` empty, so this // cannot loop. + // moon#513: a workspace connection's keys are rewritten to + // `{<32-hex>}:` BELOW this point (the `cmd_args` rebind), and + // this guard cannot move down there — the connection-level + // intercepts it exists to hold back run in between. So the keys + // visible HERE are the raw ones, and they hash to the wrong + // shards. Not by a little: the prefix is a hash TAG, so every key + // in a workspace routes to ONE shard however the raw names + // scatter, and a mask read off the raw names can call a command + // disjoint from the very shard its writes are pending on + // (measured: 5 of 12 connections lost an MGET's own batch writes). + // Treating every shard as pending makes the predicate answer + // exactly as it did before the mask existed. + let effective_pending = if conn.workspace_id.is_some() { + u64::MAX + } else { + pending_mask + }; if pending_mask != 0 && crate::server::conn::shared::must_wait_for_pending_remote( cmd, cmd_args, ctx.num_shards, - pending_mask, + effective_pending, ) { batch[frame_idx - 1] = frame; diff --git a/tests/pipeline_cross_shard_ordering.rs b/tests/pipeline_cross_shard_ordering.rs index b091fac5..f7256322 100644 --- a/tests/pipeline_cross_shard_ordering.rs +++ b/tests/pipeline_cross_shard_ordering.rs @@ -386,8 +386,9 @@ fn pco13_disjoint_shard_multikey_does_not_cut_the_batch() { {overlap_defers} on this same harness, so the guard is reachable)" ); - // The relaxation must not cost the guarantee the guard exists for: a - // co-located MGET reading the keys its own batch just wrote still sees them. + // The relaxation must not cost the guarantee the guard exists for. These two + // keys are the WRITTEN ones — spanning shards 0 and 1, so the read meets + // whatever is pending — and the MGET must observe both values it just acked. let mut c = Conn::open(m.port); let r = c.pipeline(&[ &["SET", &w[0], "11"], @@ -399,6 +400,75 @@ fn pco13_disjoint_shard_multikey_does_not_cut_the_batch() { "an MGET reading the keys its own batch just wrote must observe them — \ moon#507 is the whole reason the guard exists: {r:?}" ); + // And the same claim where the read is confined to ONE shard, which is the + // shape the mask is most tempted to wave through. + let one = on("pco13x", &[1]); + let r = c.pipeline(&[&["SET", &one[0], "33"], &["MGET", &one[0], &one[0]]]); + assert!( + r.contains("$2\r\n33\r\n"), + "a single-shard MGET must observe the write that acked before it in the \ + same batch: {r:?}" + ); +} + +/// moon#513: inside a workspace, the guard must judge the keys the command will +/// ACTUALLY run against. +/// +/// `workspace_rewrite_args` rebinds `cmd_args` further down the batch loop than +/// the ordering guard sits — and it cannot be moved below it, because the +/// connection-level intercepts the guard exists to hold back (`AUTH`, `CLIENT`, +/// `CONFIG`, `INFO`, `SELECT`, …) run in between. So the keys visible AT the +/// guard are the raw ones, while routing later hashes the prefixed ones. +/// +/// That is not a small discrepancy here. A workspace key is +/// `{<32-hex>}:` — a hash TAG — so every key in a workspace routes to one +/// shard no matter how the raw names scatter. A shard mask read off the raw +/// names can therefore say "touches nothing pending" about a command that in +/// truth reads the very shard the batch's pending writes are going to, which is +/// moon#507 reopened for exactly the connections that opted into isolation. +/// +/// Asserted as CORRECTNESS across many connections rather than as a deferral +/// count, because the count is only wrong when the workspace's shard is foreign +/// to the connection, and `SO_REUSEPORT` decides that. The correctness claim +/// holds for every connection, so twelve of them make the placement question +/// moot. +#[test] +fn pco14_workspace_rewritten_keys_still_wait_for_their_own_batch() { + let m = spawn_moon(SHARDS); + let mut c0 = Conn::open(m.port); + let created = c0.send(&["WS", "CREATE", "pco14ws"]); + let Some(ws_id) = created + .strip_prefix('$') + .and_then(|r| r.split_once("\r\n")) + .and_then(|(_, rest)| rest.split("\r\n").next()) + .filter(|id| !id.is_empty()) + .map(str::to_owned) + else { + panic!("WS CREATE did not answer a workspace id: {created:?}"); + }; + + let mut wrong = Vec::new(); + for i in 0..12 { + let mut c = Conn::open(m.port); + let authed = c.send(&["WS", "AUTH", &ws_id]); + assert_eq!(authed, "+OK\r\n", "conn {i}: WS AUTH: {authed:?}"); + // Raw names chosen to scatter across shards; the workspace prefix + // collapses them onto one. Whatever the guard reads, the MGET must + // observe the two writes that acked earlier in its own batch. + let (a, b) = (format!("pco14:{i}:aaaa"), format!("pco14:{i}:zzzz")); + let r = c.pipeline(&[&["SET", &a, "1"], &["SET", &b, "2"], &["MGET", &a, &b]]); + if !r.contains("$1\r\n1\r\n") || !r.contains("$1\r\n2\r\n") { + wrong.push(format!(" conn {i}: {r:?}")); + } + } + assert!( + wrong.is_empty(), + "{}/12 workspace connections had an MGET miss a write from its own \ + batch — the ordering guard judged the RAW keys while the command ran \ + against the workspace-prefixed ones:\n{}", + wrong.len(), + wrong.join("\n") + ); } #[test]