perf(pipeline): defer only for shards the batch has pending work on (#513)

CHANGELOG.md

@@ -6,6 +6,52 @@ 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 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`,
   and the `moon_pipeline_remote_defer_total` Prometheus counter) — groundwork for #513.

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,30 @@ 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)
+            // moon#513: a workspace connection's keys are rewritten to
+            // `{<32-hex>}:<key>` 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,
+                    effective_pending,
+                )
             {
                 frames[frame_idx - 1] = frame;
                 crate::admin::metrics_setup::record_pipeline_remote_defer();
@@ -3331,6 +3357,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();
             }
         }

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<usize> = Vec::new();
                 let mut should_quit = false;
                 let mut remote_groups: HashMap<usize, Vec<(usize, std::sync::Arc<Frame>, Option<Bytes>, usize, Option<crate::tracking::invalidation::TrackedWriteKeys>, 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,30 @@ 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)
+                    // moon#513: a workspace connection's keys are rewritten to
+                    // `{<32-hex>}:<key>` 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,
+                            effective_pending,
+                        )
                     {
                         batch[frame_idx - 1] = frame;
                         crate::admin::metrics_setup::record_pipeline_remote_defer();
@@ -2657,6 +2683,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);
                     }
                 }

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<u64> {
+    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<Frame> {
+        let found: Vec<Frame> = (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"
+        );
+    }
+}