From f3fce7d515e8d833389778038b7eefa3c571026c Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Sat, 19 Sep 2026 01:53:08 +0700 Subject: [PATCH] fix(blocking): multi-key blocking pops serve exactly once; refuse spanning BLMPOP/BZMPOP (moon#989) At --shards >= 2 a multi-key blocking pop (BLMPOP, BZMPOP, BLPOP, BRPOP, BZPOPMIN, BZPOPMAX) over keys co-located under one {hash} tag replied correctly while a SECOND non-empty key lost its head element, delivered to nobody: BLMPOP 0.3 3 {t}a {t}b {t}c LEFT ({t}a empty, {t}b=[B1 B2], {t}c=[C1 C2]) redis 8.6.1 -> {t}b [B1] {t}c=[C1 C2] moon --shards 4 -> {t}b [B1] {t}c=[C2] <-- C1 destroyed Root cause (the triage hypothesis, confirmed from code and reproduction): when the client's connection lives on a different shard than the keys, immediate_scan can see none of them, so the multi-key coordinator sent one BlockRegister per key to their owner. The owner handled each message on its own -- register, see data, serve -- and so served the SAME waiter once per non-empty key. The client kept the first reply and dropped the rest. The same fan-out popped a key named twice twice, and skipped Redis's -WRONGTYPE for an earlier wrong-typed key. It only happens when the connection is off the owner shard, which is why it varied between server instances and never showed at --shards 1. Fix: - The coordinator now sends ONE BlockRegisterGroup per owner shard carrying every key that owner holds, in argument order (blocking_multikey.rs, shared by both runtimes). The owner registers every member under one wait_id and then serves in argument order in one synchronous stretch (blocking::group::register_group); serving runs remove_wait, which unregisters the siblings before the next key is looked at, so the waiter is served at most once. When the group is the whole command it also answers the -WRONGTYPE Redis owes; a partial group skips a wrong-typed key instead of tearing down its sibling registrations. - Spanning BLMPOP/BZMPOP are refused with CROSSSLOT in immediate_scan, before anything is popped or registered, via the same family list moon#962 uses for LMPOP/ZMPOP. That placement also popped the WRONG key (a later local key over an earlier remote one) and could be served by two owners. - BlockRegister now carries single-key waiters only, so its sole_key flag is gone and the owner's type check is unconditional. Not changed: spanning BLPOP/BRPOP/BZPOPMIN/BZPOPMAX can still be served by two owner shards (measured unchanged within noise, 35 vs 40 extra pops of 120 over 10 instances each). Refusing them is a client-visible decision and is moon#1019. Tests: - tests/blocking_multikey_cross_shard.rs: 7 tests. Co-located rows run for tags owned by EVERY shard, so they cannot pass by the connection happening to land on the owner. RED on 43ad387b (bmk1 58/80, bmk3, bmk6 20/32), GREEN with the fix on monoio and tokio; mutating the grouping or the refusal turns them red again. - unit tests for register_group and stage_multikey_wait. - test-consistency.sh: blmpop/bzmpop span+colo rows, colo-only rows for the other four, and answered rows now compare the keyspace with redis, not just the reply. test-commands.sh: co-located BLMPOP/BZMPOP rows with a read of the key they must not touch. Refs: moon#989, moon#962, moon#1019 author: Tin Dang --- CHANGELOG.md | 30 ++ scripts/test-commands.sh | 13 + scripts/test-consistency.sh | 42 +- src/blocking/group.rs | 444 ++++++++++++++++ src/blocking/mod.rs | 11 + src/server/conn/blocking.rs | 168 +++--- src/server/conn/blocking_multikey.rs | 210 ++++++++ src/server/conn/mod.rs | 3 + src/server/conn/shared.rs | 6 + src/shard/dispatch.rs | 56 +- src/shard/spsc_handler.rs | 66 +-- tests/blocking_multikey_cross_shard.rs | 710 +++++++++++++++++++++++++ 12 files changed, 1606 insertions(+), 153 deletions(-) create mode 100644 src/blocking/group.rs create mode 100644 src/server/conn/blocking_multikey.rs create mode 100644 tests/blocking_multikey_cross_shard.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 4099cd6d4..05d8dd33a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,9 +31,39 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `WEIGHTS nan` used to be accepted and poison every aggregated score; it is now `ERR weight value is not a float`. Infinite weights remain legal. A client relying on either form silently doing nothing will now see an error. +- **BEHAVIOUR CHANGE — `BLMPOP`/`BZMPOP` whose keys span shards are refused with + `CROSSSLOT`** at `--shards > 1` (moon#989), the rule moon#962 already applies + to `LMPOP`/`ZMPOP`. They used to answer, and measured at `--shards 4` over 16 + three-shard placements, 20 of 32 probes popped a key the reply did not name: + either the WRONG key (a later local key served over an earlier remote one) or + a second key whose element no client ever received. "Pop from the first + non-empty key in argument order, exactly once" is a property of the whole key + vector that no single shard can see. The refusal is decided from the key + names before anything is touched, so the keyspace is unchanged. Keys under + one `{hash}` tag, and every placement at `--shards 1`, are unaffected. ### Fixed +- **Multi-key blocking pops no longer destroy an element from a key they did + not answer with** (moon#989). `BLMPOP`, `BZMPOP`, `BLPOP`, `BRPOP`, + `BZPOPMIN` and `BZPOPMAX` over keys co-located under one `{hash}` tag replied + correctly while a SECOND non-empty key silently lost its head element: + `BLMPOP 0.3 3 {t}a {t}b {t}c LEFT` answered `{t}b B1` and left `{t}c` at + `[C2]`, with `C1` delivered to nobody. It happened whenever the client's + connection lived on a different shard than the keys: the client's shard + could not see them, so it sent one registration per key to their owner, and + the owner served the same waiter once per non-empty key — the client kept + the first reply and dropped the rest. The same fan-out popped a key named + twice (`BLMPOP 0 2 k k LEFT`) twice, and skipped Redis's `-WRONGTYPE` for an + earlier key of the wrong type. Measured at `--shards 4` (`BLMPOP`, `BZMPOP`, + `BLPOP`, `BZPOPMIN` × 16 tags × 3 server instances, against redis 8.6.1): 139 + of 192 probes destroyed an element before, 0 after; `--shards 1` was and is 0. The client now sends ONE registration per owner shard carrying + every key it owns, and the owner registers, type-checks and serves them in + one synchronous stretch, so a waiter is served at most once there. A + spanning `BLPOP`/`BRPOP`/`BZPOPMIN`/`BZPOPMAX` can still be served by two + owner shards at once; that placement is unchanged by this fix and is + tracked as moon#1019. + - **Commands routed to another shard are counted and timed** (moon#982). At `--shards > 1` a command whose key lives on a shard other than the connection's went through no telemetry probe at all — neither the diff --git a/scripts/test-commands.sh b/scripts/test-commands.sh index 1bb6c38ae..a498db9d6 100755 --- a/scripts/test-commands.sh +++ b/scripts/test-commands.sh @@ -1808,6 +1808,19 @@ if should_run "blocking"; then # moon#570: `{blk}` co-locates the pair -- see the LMOVE row above. rcli RPUSH {blk}:src x y z >/dev/null 2>&1; mcli RPUSH {blk}:src x y z >/dev/null 2>&1 assert_match "BLMOVE (ready)" BLMOVE {blk}:src {blk}:dst LEFT RIGHT 1 + + # moon#989: a multi-key blocking pop serves from the FIRST non-empty key, + # exactly once. The reply alone cannot show a SECOND key losing an element + # (moon answered correctly while destroying it), so each row is followed by + # a read of the key it must not touch. `{blk}` co-locates the three keys. + rcli RPUSH {blk}:mp2 B1 B2 >/dev/null 2>&1; mcli RPUSH {blk}:mp2 B1 B2 >/dev/null 2>&1 + rcli RPUSH {blk}:mp3 C1 C2 >/dev/null 2>&1; mcli RPUSH {blk}:mp3 C1 C2 >/dev/null 2>&1 + assert_match "BLMPOP (ready, 3 co-located keys)" BLMPOP 1 3 {blk}:mp1 {blk}:mp2 {blk}:mp3 LEFT + assert_match "BLMPOP left the later key untouched" LRANGE {blk}:mp3 0 -1 + rcli ZADD {blk}:zp2 1 B1 2 B2 >/dev/null 2>&1; mcli ZADD {blk}:zp2 1 B1 2 B2 >/dev/null 2>&1 + rcli ZADD {blk}:zp3 1 C1 2 C2 >/dev/null 2>&1; mcli ZADD {blk}:zp3 1 C1 2 C2 >/dev/null 2>&1 + assert_match "BZMPOP (ready, 3 co-located keys)" BZMPOP 1 3 {blk}:zp1 {blk}:zp2 {blk}:zp3 MIN + assert_match "BZMPOP left the later key untouched" ZRANGE {blk}:zp3 0 -1 WITHSCORES fi # =========================================================================== diff --git a/scripts/test-consistency.sh b/scripts/test-consistency.sh index 7f753c3f2..06ea0035b 100755 --- a/scripts/test-consistency.sh +++ b/scripts/test-consistency.sh @@ -2013,7 +2013,7 @@ mk_norm() { # wrong-key pop cannot happen. route_probe_multi() { local mode="$1" label="$2" n="$3" seeds="$4" probe="$5" check="${6:-}" - local i j port s p r m before="" after="" + local i j port s p r m before="" after="" after_redis="" local wrong=0 refused=0 local -a keys seedv sv pv for i in $(seq 1 "$MK_TRIALS"); do @@ -2050,6 +2050,7 @@ route_probe_multi() { m="$(mk_norm "$(redis-cli -p "$PORT_RUST" "${pv[@]}" 2>&1)")" if [[ -n "$check" ]]; then after="$(mk_state "$PORT_RUST" "$check" "${keys[@]}")" + after_redis="$(mk_state "$PORT_REDIS" "$check" "${keys[@]}")" fi # A SUBSTRING test, not an anchored `case` pattern. The reply may carry # a leading blank line (see `mk_norm`), and an anchored pattern that @@ -2071,6 +2072,12 @@ route_probe_multi() { elif [[ "$r" != "$m" ]]; then echo " FAIL detail: ${label}[$i] ($mode) answered '$m'; redis says '$r'" wrong=$((wrong + 1)) + elif [[ -n "$check" && "$after" != "$after_redis" ]]; then + # moon#989: the RIGHT reply is not enough. BLMPOP answered exactly + # like redis while popping a second key it never named, and only + # the keyspace after the probe could show it. + echo " FAIL detail: ${label}[$i] ($mode) answered like redis but the keyspace differs: moon '$after' vs redis '$after_redis'" + wrong=$((wrong + 1)) fi done assert_eq "moon#962 ${label} ${mode} (shards=${SHARDS}, ${MK_TRIALS} placements)" \ @@ -2099,9 +2106,33 @@ MK_ROWS=( "touch|3|SET %K v|SET %K v|SET %K v|TOUCH %K1 %K2 %K3|GET %K" "lmpop|3||RPUSH %K B1 B2|RPUSH %K C1 C2|LMPOP 3 %K1 %K2 %K3 LEFT|LRANGE %K 0 -1" "zmpop|3||ZADD %K 1 B1 2 B2|ZADD %K 1 C1 2 C2|ZMPOP 3 %K1 %K2 %K3 MIN|ZRANGE %K 0 -1" + # moon#989: the blocking twins. Data is seeded, so neither blocks -- the + # 0.1s timeout only bounds a regression that would. `colo` is the row that + # caught the defect: the reply matched redis while a second co-located key + # lost its head element, visible only through the per-key check. + "blmpop|3||RPUSH %K B1 B2|RPUSH %K C1 C2|BLMPOP 0.1 3 %K1 %K2 %K3 LEFT|LRANGE %K 0 -1" + "bzmpop|3||ZADD %K 1 B1 2 B2|ZADD %K 1 C1 2 C2|BZMPOP 0.1 3 %K1 %K2 %K3 MIN|ZRANGE %K 0 -1" +) + +# moon#989: the rest of the multi-key blocking-pop family, CO-LOCATED only. +# They shared BLMPOP's double-pop and are fixed with it, so `colo` must agree +# with redis byte for byte. Their SPANNING placement is still a known defect +# (two owner shards can each serve the same waiter) and is deliberately not +# refused yet -- that is a behaviour decision tracked as moon#1019, so a +# `span` row here would only assert the bug. +MK_COLO_ONLY_ROWS=( + "blpop|3||RPUSH %K B1 B2|RPUSH %K C1 C2|BLPOP %K1 %K2 %K3 0.1|LRANGE %K 0 -1" + "brpop|3||RPUSH %K B1 B2|RPUSH %K C1 C2|BRPOP %K1 %K2 %K3 0.1|LRANGE %K 0 -1" + "bzpopmin|3||ZADD %K 1 B1 2 B2|ZADD %K 1 C1 2 C2|BZPOPMIN %K1 %K2 %K3 0.1|ZRANGE %K 0 -1" + "bzpopmax|3||ZADD %K 1 B1 2 B2|ZADD %K 1 C1 2 C2|BZPOPMAX %K1 %K2 %K3 0.1|ZRANGE %K 0 -1" ) -for mk_row in "${MK_ROWS[@]}"; do +for mk_row in "${MK_ROWS[@]}" "${MK_COLO_ONLY_ROWS[@]/#/colo-only:}"; do + mk_modes="span colo" + if [[ "$mk_row" == colo-only:* ]]; then + mk_modes="colo" + mk_row="${mk_row#colo-only:}" + fi IFS='|' read -r -a mk_f <<<"$mk_row" mk_label="${mk_f[0]}"; mk_n="${mk_f[1]}" # fields 2..(2+n-1) are the per-key seeds, then the probe, then the check @@ -2112,8 +2143,9 @@ for mk_row in "${MK_ROWS[@]}"; do mk_seeds="${mk_seeds%|}" mk_probe="${mk_f[$((2 + mk_n))]}" mk_check="${mk_f[$((3 + mk_n))]:-}" - route_probe_multi span "$mk_label" "$mk_n" "$mk_seeds" "$mk_probe" "$mk_check" - route_probe_multi colo "$mk_label" "$mk_n" "$mk_seeds" "$mk_probe" "$mk_check" + for mk_mode in $mk_modes; do + route_probe_multi "$mk_mode" "$mk_label" "$mk_n" "$mk_seeds" "$mk_probe" "$mk_check" + done done # Non-vacuity. At --shards>1 the span sweep MUST have reached the cross-shard @@ -2133,7 +2165,7 @@ assert_eq "moon#962 TOUCH is never refused (shards=$SHARDS)" "0" "$MK_TOUCH_REFU # Tidy up by exact name -- `--scan | xargs -r` is GNU-only and this script runs # on macOS too. -for mk_row in "${MK_ROWS[@]}"; do +for mk_row in "${MK_ROWS[@]}" "${MK_COLO_ONLY_ROWS[@]}"; do IFS='|' read -r -a mk_f <<<"$mk_row" for mk_i in $(seq 1 "$MK_TRIALS"); do for mk_j in $(seq 1 "${mk_f[1]}"); do diff --git a/src/blocking/group.rs b/src/blocking/group.rs new file mode 100644 index 000000000..c1832b065 --- /dev/null +++ b/src/blocking/group.rs @@ -0,0 +1,444 @@ +//! Owner-side registration of a multi-key blocking waiter (moon#989). +//! +//! A multi-key `BLPOP`/`BRPOP`/`BZPOPMIN`/`BZPOPMAX`/`BLMPOP`/`BZMPOP` whose +//! keys are owned by another shard used to reach that owner as one +//! `BlockRegister` per key. Each was handled on its own — register the key, +//! see data, serve the waiter — so the SAME waiter was served once per +//! non-empty key. The client kept the first reply and dropped the rest, and +//! every element in the dropped replies had already left the keyspace: +//! +//! ```text +//! BLMPOP 0.3 3 {t}a {t}b {t}c LEFT ({t}a empty, {t}b=[B1 B2], {t}c=[C1 C2]) +//! redis 8.6.1 -> {t}b [B1] {t}c=[C1 C2] +//! moon --shards 4 -> {t}b [B1] {t}c=[C2] <-- C1 served to nobody +//! ``` +//! +//! [`register_group`] receives every key this shard owns of one waiter in a +//! single message and handles them in one synchronous stretch of the owner's +//! event loop. Nothing can interleave with it, so the waiter is served at most +//! once by construction: every key is registered under one `wait_id` before +//! any wake runs, and the wake that serves the waiter runs `remove_wait`, +//! which unregisters all of its siblings before the next key is looked at. + +use bytes::Bytes; + +use crate::blocking::{BlockedCommand, BlockingRegistry, WaitEntry, WaitFamily}; +use crate::protocol::Frame; +use crate::shard::dispatch::{BlockRegisterGroupPayload, BlockRegisterMember}; +use crate::storage::Database; + +/// Register — and if the data is already there, serve — every key of one +/// multi-key waiter that this shard owns. +/// +/// Keys are visited in the command's argument order, which is the order Redis +/// serves them in: the first non-empty key answers, exactly once. +/// +/// When the group is the WHOLE command (every key co-located here) the owner +/// is the only place that can see all of it, so it also answers the +/// `-WRONGTYPE` Redis owes for the first existing key of the wrong type, and +/// never registers in that case. A partial group cannot: the waiter's other +/// keys live on other shards and may be serving right now. +/// +/// Remote registrations carry no deadline; the client's own timer times the +/// wait out and sends `BlockCancel`, which removes every member at once. +pub fn register_group( + registry: &mut BlockingRegistry, + db: &mut Database, + payload: BlockRegisterGroupPayload, +) { + let BlockRegisterGroupPayload { + db_index, + wait_id, + members, + whole_command, + } = payload; + + // Every member carries the same command, so they share one family. + let family = members.first().map(|m| m.cmd.family()); + + if whole_command && let Some(err) = first_type_error(db, &members) { + // Nothing registered, so nothing to unwind; the client's `BlockCancel` + // on the way out is a no-op. One reply is enough — the client skips + // the members whose senders drop here. + if let Some(first) = members.into_iter().next() { + let _ = first.reply_tx.send(Some(err)); + } + return; + } + + let mut keys: smallvec::SmallVec<[Bytes; 4]> = smallvec::SmallVec::with_capacity(members.len()); + for BlockRegisterMember { + key, + mut cmd, + reply_tx, + } in members + { + // moon#595: bind `$` here, with no suspension point before `register`. + // A no-op for everything that is not an `XREAD ... $`. + cmd.bind_stream_since(db, &key); + registry.register( + db_index, + key.clone(), + WaitEntry { + wait_id, + cmd, + reply_tx, + deadline: None, + }, + ); + keys.push(key); + } + + // Data may already be there (it arrived before the registration, or the + // client's shard simply could not see it). Serve in argument order and + // stop as soon as this waiter is gone — served, or reaped as dead — so a + // later key is never consulted on its behalf. + // + // Only a key holding this waiter's type is offered to the wakers. A waker + // handed a waiter it cannot serve (a `BLPOP` waiter on a string) pops it, + // fails, and runs the served-waiter cleanup — `remove_wait` plus a `None` + // reply — which would silently unregister every SIBLING key on this shard + // too, leaving the client parked on keys nobody is watching. A wrong-typed + // key is skipped instead; for a whole-command group it cannot come before + // the key that serves, because `first_type_error` already answered it. + for key in &keys { + if !registry.is_waiting(wait_id) { + break; + } + if db.exists(key) && family.is_some_and(|f| family_type_error(db, key, f).is_none()) { + crate::blocking::wakeup::try_wake_list_waiter(registry, db, db_index, key); + crate::blocking::wakeup::try_wake_zset_waiter(registry, db, db_index, key); + crate::blocking::wakeup::try_wake_stream_waiter(registry, db, db_index, key); + } + } +} + +/// Redis's pre-block ladder, minus the pop: walk the keys in argument order, +/// answer `-WRONGTYPE` for the first existing key of the wrong type, and stop +/// at the first key that exists with the right one — that key serves, so +/// nothing after it is consulted (Redis does not type-check past it either). +/// +/// Read-only: the list probe uses the shared-borrow accessor so that a type +/// check never flattens a compact encoding (moon#832). +fn first_type_error(db: &mut Database, members: &[BlockRegisterMember]) -> Option { + for m in members { + if let Some(err) = type_error(db, &m.key, &m.cmd) { + return Some(err); + } + if db.exists(&m.key) { + return None; + } + } + None +} + +/// The error `cmd` owes its client for `key` before it may block: a wrong +/// type, or `XREADGROUP`'s missing key or group. +fn type_error(db: &mut Database, key: &Bytes, cmd: &BlockedCommand) -> Option { + match cmd.family() { + WaitFamily::Stream => crate::blocking::wakeup::stream_register_error(db, key, cmd), + family => family_type_error(db, key, family), + } +} + +/// `-WRONGTYPE` when `key` exists holding something `family` cannot pop from. +fn family_type_error(db: &mut Database, key: &Bytes, family: WaitFamily) -> Option { + match family { + WaitFamily::List => { + let now_ms = db.now_ms(); + db.get_list_ref_if_alive(key, now_ms).err() + } + WaitFamily::ZSet => db.get_sorted_set(key).err(), + WaitFamily::Stream => db.get_stream(key).err(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::blocking::Direction; + use crate::runtime::channel::{self, OneshotReceiver}; + + fn b(s: &str) -> Bytes { + Bytes::copy_from_slice(s.as_bytes()) + } + + /// A group for `keys`, all waiting with `cmd()`; returns the receivers in + /// member order. + fn group( + reg: &mut BlockingRegistry, + keys: &[&str], + cmd: impl Fn() -> BlockedCommand, + whole_command: bool, + ) -> ( + BlockRegisterGroupPayload, + Vec>>, + ) { + let wait_id = reg.next_wait_id(); + let mut members = Vec::new(); + let mut rxs = Vec::new(); + for k in keys { + let (tx, rx) = channel::oneshot(); + members.push(BlockRegisterMember { + key: b(k), + cmd: cmd(), + reply_tx: tx, + }); + rxs.push(rx); + } + ( + BlockRegisterGroupPayload { + db_index: 0, + wait_id, + members, + whole_command, + }, + rxs, + ) + } + + /// Every reply the waiter received, across all of its members. + fn replies(rxs: &[OneshotReceiver>]) -> Vec { + rxs.iter() + .filter_map(|rx| rx.try_recv().ok().flatten()) + .collect() + } + + fn list(db: &mut Database, key: &str) -> Vec { + let now_ms = db.now_ms(); + match db.get_list_ref_if_alive(&b(key), now_ms) { + Ok(Some(l)) => l.iter_bytes(), + _ => Vec::new(), + } + } + + fn blmpop() -> BlockedCommand { + BlockedCommand::BLMPop { + dir: Direction::Left, + count: 1, + } + } + + /// The issue's exact shape: two non-empty keys, one waiter, ONE pop. + #[test] + fn two_non_empty_keys_serve_the_waiter_once_from_the_first() { + for whole in [true, false] { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + db.list_push_back(&b("b"), b("B1")); + db.list_push_back(&b("b"), b("B2")); + db.list_push_back(&b("c"), b("C1")); + db.list_push_back(&b("c"), b("C2")); + let (payload, rxs) = group(&mut reg, &["a", "b", "c"], blmpop, whole); + let wait_id = payload.wait_id; + register_group(&mut reg, &mut db, payload); + + let got = replies(&rxs); + assert_eq!( + got.len(), + 1, + "whole={whole}: exactly one reply, got {got:?}" + ); + assert_eq!( + got[0], + Frame::Array(crate::framevec![ + Frame::BulkString(b("b")), + Frame::Array(crate::framevec![Frame::BulkString(b("B1"))]), + ]) + ); + assert_eq!(list(&mut db, "b"), vec![b("B2")]); + assert_eq!( + list(&mut db, "c"), + vec![b("C1"), b("C2")], + "c must be untouched" + ); + assert!( + !reg.is_waiting(wait_id), + "a served waiter is fully unregistered" + ); + } + } + + /// The same key named twice is one pop, not two. + #[test] + fn a_key_named_twice_is_popped_once() { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + db.list_push_back(&b("b"), b("B1")); + db.list_push_back(&b("b"), b("B2")); + let (payload, rxs) = group(&mut reg, &["b", "b"], blmpop, true); + register_group(&mut reg, &mut db, payload); + assert_eq!(replies(&rxs).len(), 1); + assert_eq!(list(&mut db, "b"), vec![b("B2")]); + } + + /// Nothing to serve: every key is registered and a later push wakes the + /// waiter exactly once, from the key that received the data. + #[test] + fn empty_keys_register_and_the_first_push_serves_once() { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + let (payload, rxs) = group(&mut reg, &["a", "b", "c"], blmpop, true); + let wait_id = payload.wait_id; + register_group(&mut reg, &mut db, payload); + assert!(replies(&rxs).is_empty()); + assert!(reg.is_waiting(wait_id)); + for k in ["a", "b", "c"] { + assert!(reg.has_waiters(0, &b(k)), "{k} must be registered"); + } + + db.list_push_back(&b("c"), b("C1")); + assert!(crate::blocking::wakeup::try_wake_list_waiter( + &mut reg, + &mut db, + 0, + &b("c") + )); + db.list_push_back(&b("b"), b("B1")); + assert!(!crate::blocking::wakeup::try_wake_list_waiter( + &mut reg, + &mut db, + 0, + &b("b") + )); + assert_eq!(replies(&rxs).len(), 1); + assert_eq!( + list(&mut db, "b"), + vec![b("B1")], + "b's push had nobody to serve" + ); + assert!(!reg.is_waiting(wait_id)); + } + + /// A whole-command group answers Redis's `-WRONGTYPE` for the first + /// existing key of the wrong type, pops nothing and registers nothing. + #[test] + fn whole_command_wrong_type_before_data_is_an_error_and_pops_nothing() { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + db.set_string(b"s", b("x")); + db.list_push_back(&b("b"), b("B1")); + let (payload, rxs) = group(&mut reg, &["a", "s", "b"], blmpop, true); + let wait_id = payload.wait_id; + register_group(&mut reg, &mut db, payload); + let got = replies(&rxs); + assert_eq!(got.len(), 1); + assert!( + matches!(&got[0], Frame::Error(e) if e.starts_with(b"WRONGTYPE")), + "{got:?}" + ); + assert_eq!(list(&mut db, "b"), vec![b("B1")]); + assert!(!reg.is_waiting(wait_id)); + } + + /// Redis stops type-checking at the key that serves: a wrong-typed key + /// AFTER it is never consulted. + #[test] + fn whole_command_wrong_type_after_data_does_not_block_the_pop() { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + db.list_push_back(&b("b"), b("B1")); + db.set_string(b"s", b("x")); + let (payload, rxs) = group(&mut reg, &["b", "s"], blmpop, true); + register_group(&mut reg, &mut db, payload); + let got = replies(&rxs); + assert_eq!(got.len(), 1); + assert!(matches!(&got[0], Frame::Array(_)), "{got:?}"); + assert!(list(&mut db, "b").is_empty()); + } + + /// A PARTIAL group must not decide the command with an error — its + /// siblings on other shards may be serving right now — and a wrong-typed + /// key must not tear down the waiter's OTHER keys on this shard either: + /// a later push to one of them still serves it. + #[test] + fn partial_group_skips_a_wrong_type_and_keeps_its_siblings() { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + db.set_string(b"s", b("x")); + let (payload, rxs) = group(&mut reg, &["s", "a"], blmpop, false); + let wait_id = payload.wait_id; + register_group(&mut reg, &mut db, payload); + assert!(replies(&rxs).is_empty(), "no error and no nil"); + assert!(reg.is_waiting(wait_id), "still parked on its keys"); + assert!(reg.has_waiters(0, &b("a")), "the sibling survives"); + + db.list_push_back(&b("a"), b("A1")); + assert!(crate::blocking::wakeup::try_wake_list_waiter( + &mut reg, + &mut db, + 0, + &b("a") + )); + assert_eq!(replies(&rxs).len(), 1); + } + + /// A whole-command group whose first servable key is taken by a waiter + /// queued AHEAD of this one: the wrong-typed key after it is skipped, not + /// used to tear this waiter down with a spurious nil. + #[test] + fn whole_group_never_offers_a_wrong_typed_key_to_the_wakers() { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + let (ahead_tx, ahead_rx) = channel::oneshot(); + let ahead_id = reg.next_wait_id(); + reg.register( + 0, + b("b"), + WaitEntry { + wait_id: ahead_id, + cmd: blmpop(), + reply_tx: ahead_tx, + deadline: None, + }, + ); + db.list_push_back(&b("b"), b("B1")); + db.set_string(b"s", b("x")); + let (payload, rxs) = group(&mut reg, &["b", "s"], blmpop, true); + let wait_id = payload.wait_id; + register_group(&mut reg, &mut db, payload); + assert!( + ahead_rx.try_recv().ok().flatten().is_some(), + "FIFO: the waiter ahead is served first" + ); + assert!(replies(&rxs).is_empty(), "ours got neither B1 nor a nil"); + assert!(reg.is_waiting(wait_id), "and is still parked"); + } + + /// A2 carried over: a waiter whose client is already gone must not + /// consume anything. + #[test] + fn a_dead_waiter_consumes_nothing() { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + db.list_push_back(&b("b"), b("B1")); + let (payload, rxs) = group(&mut reg, &["a", "b"], blmpop, true); + let wait_id = payload.wait_id; + drop(rxs); + register_group(&mut reg, &mut db, payload); + assert_eq!(list(&mut db, "b"), vec![b("B1")]); + assert!(!reg.is_waiting(wait_id), "the dead waiter is reaped"); + } + + /// `BZMPOP` over co-located sorted sets: one pop, from the first. + #[test] + fn zset_group_serves_once() { + let mut reg = BlockingRegistry::new(0); + let mut db = Database::new(); + db.zset_restore(&b("b"), b("B1"), 1.0); + db.zset_restore(&b("c"), b("C1"), 1.0); + let (payload, rxs) = group( + &mut reg, + &["a", "b", "c"], + || BlockedCommand::BZMPop { + min: true, + count: 1, + }, + true, + ); + register_group(&mut reg, &mut db, payload); + assert_eq!(replies(&rxs).len(), 1); + assert!(!db.exists(b"b"), "b served its only member"); + assert!(db.exists(b"c"), "c must be untouched"); + } +} diff --git a/src/blocking/mod.rs b/src/blocking/mod.rs index 723cfdf4b..df859031a 100644 --- a/src/blocking/mod.rs +++ b/src/blocking/mod.rs @@ -1,3 +1,4 @@ +pub mod group; pub mod wakeup; use std::collections::{HashMap, VecDeque}; @@ -419,6 +420,16 @@ impl BlockingRegistry { } } + /// Is `wait_id` still registered on at least one key of this shard? + /// + /// `false` once it has been served, cancelled, timed out or reaped as + /// dead — every one of those runs `remove_wait`. moon#989's group + /// registration uses it to stop consulting a waiter's later keys the + /// moment an earlier one has served it. + pub fn is_waiting(&self, wait_id: u64) -> bool { + self.wait_keys.contains_key(&wait_id) + } + /// Check if any waiters exist for this (db_index, key). pub fn has_waiters(&self, db_index: usize, key: &Bytes) -> bool { self.waiters diff --git a/src/server/conn/blocking.rs b/src/server/conn/blocking.rs index 6bda3f298..1bbe054cd 100644 --- a/src/server/conn/blocking.rs +++ b/src/server/conn/blocking.rs @@ -440,7 +440,6 @@ pub(crate) async fn handle_blocking_command( where S: tokio::io::AsyncRead + Unpin, { - use futures::stream::FuturesUnordered; use tokio::io::AsyncReadExt; // Parse timeout (last argument for all blocking commands) @@ -485,9 +484,6 @@ where wait_id, cmd: blocked_cmd_factory(), reply_tx, - // The single-key fast path: this key IS the command, so - // the owner may answer a type error for it (moon#556). - sole_key: true, }, )); if !push_block_msg(shutdown, dispatch_tx, shard_id, target, msg, None).await { @@ -602,56 +598,28 @@ where } // --- Multi-key coordinator: register on ALL keys across local + remote shards --- - // Uses FuturesUnordered for first-wakeup-wins semantics. - let wait_id; - let mut receivers: FuturesUnordered>> = - FuturesUnordered::new(); - let mut registered_remote_shards: Vec = Vec::new(); - - // A4/A5: build every registration under one borrow, but STAGE the remote - // ones and push them only after the borrow is released — the old code - // pushed while holding both borrows, which forced the bare `try_push` - // (no await possible, so no backpressure retry and no shutdown arm). - let mut pending_remote: Vec<(usize, ShardMessage)> = Vec::new(); - { - let mut reg = blocking_registry.borrow_mut(); - wait_id = reg.next_wait_id(); - - for key in &keys { - let target = key_to_shard(key, num_shards); - let (tx, rx) = channel::oneshot::>(); - receivers.push(rx); - - if target == shard_id { - // Local registration - let entry = crate::blocking::WaitEntry { - wait_id, - cmd: local_blocked_cmd(&blocked_cmd_factory, selected_db, key), - reply_tx: tx, - deadline, - }; - reg.register(selected_db, key.clone(), entry); - } else { - // Remote registration via SPSC - let msg = ShardMessage::BlockRegister(Box::new( - crate::shard::dispatch::BlockRegisterPayload { - db_index: selected_db, - key: key.clone(), - wait_id, - cmd: blocked_cmd_factory(), - reply_tx: tx, - // One of several keys — see `sole_key`'s docs for why - // the owner must NOT decide the whole command here. - sole_key: false, - }, - )); - pending_remote.push((target, msg)); - if !registered_remote_shards.contains(&target) { - registered_remote_shards.push(target); - } - } - } - } // borrows dropped -- CRITICAL before await + // Uses FuturesUnordered for first-wakeup-wins semantics. moon#989: remote + // keys travel as ONE group per owner shard — the owner serves this waiter + // at most once only because it sees all of its keys in one message. See + // `blocking_multikey` for the protocol and what it does not cover. + // + // A4/A5: every registration is built under one borrow and the remote ones + // are STAGED, pushed only after the borrow is released. + let super::blocking_multikey::StagedWait { + wait_id, + mut receivers, + pending_remote, + remote_shards: registered_remote_shards, + } = super::blocking_multikey::stage_multikey_wait( + &mut blocking_registry.borrow_mut(), + &keys, + &|key| local_blocked_cmd(&blocked_cmd_factory, selected_db, key), + &*blocked_cmd_factory, + selected_db, + shard_id, + num_shards, + deadline, + ); // borrow dropped -- CRITICAL before await // A5: a silently dropped registration leaves this waiter blocked on a key // nobody is watching — it would sleep through data that IS available. @@ -793,8 +761,6 @@ pub(crate) async fn handle_blocking_command_monoio( where S: super::handler_monoio::idle_park::IdleParkRead, { - use futures::stream::FuturesUnordered; - // Parse timeout (last argument for all blocking commands) let timeout_secs = match parse_blocking_timeout(cmd, args) { Ok(t) => t, @@ -841,9 +807,6 @@ where wait_id, cmd: blocked_cmd_factory(), reply_tx, - // The single-key fast path: this key IS the command, so - // the owner may answer a type error for it (moon#556). - sole_key: true, }, )); if !push_block_msg( @@ -943,54 +906,28 @@ where } // --- Multi-key coordinator: register on ALL keys across local + remote shards --- - // Uses FuturesUnordered for first-wakeup-wins semantics. - let wait_id; - let mut receivers: FuturesUnordered>> = - FuturesUnordered::new(); - let mut registered_remote_shards: Vec = Vec::new(); - - // A4/A5: stage remote registrations, push them after the borrow is - // released. See the tokio twin for the full rationale. - let mut pending_remote: Vec<(usize, ShardMessage)> = Vec::new(); - { - let mut reg = blocking_registry.borrow_mut(); - wait_id = reg.next_wait_id(); - - for key in &keys { - let target = key_to_shard(key, num_shards); - let (tx, rx) = channel::oneshot::>(); - receivers.push(rx); - - if target == shard_id { - // Local registration - let entry = crate::blocking::WaitEntry { - wait_id, - cmd: local_blocked_cmd(&blocked_cmd_factory, selected_db, key), - reply_tx: tx, - deadline, - }; - reg.register(selected_db, key.clone(), entry); - } else { - // Remote registration via SPSC - let msg = ShardMessage::BlockRegister(Box::new( - crate::shard::dispatch::BlockRegisterPayload { - db_index: selected_db, - key: key.clone(), - wait_id, - cmd: blocked_cmd_factory(), - reply_tx: tx, - // One of several keys — see `sole_key`'s docs for why - // the owner must NOT decide the whole command here. - sole_key: false, - }, - )); - pending_remote.push((target, msg)); - if !registered_remote_shards.contains(&target) { - registered_remote_shards.push(target); - } - } - } - } // borrows dropped -- CRITICAL before await + // Uses FuturesUnordered for first-wakeup-wins semantics. moon#989: remote + // keys travel as ONE group per owner shard — the owner serves this waiter + // at most once only because it sees all of its keys in one message. See + // `blocking_multikey` for the protocol and what it does not cover. + // + // A4/A5: every registration is built under one borrow and the remote ones + // are STAGED, pushed only after the borrow is released. + let super::blocking_multikey::StagedWait { + wait_id, + mut receivers, + pending_remote, + remote_shards: registered_remote_shards, + } = super::blocking_multikey::stage_multikey_wait( + &mut blocking_registry.borrow_mut(), + &keys, + &|key| local_blocked_cmd(&blocked_cmd_factory, selected_db, key), + &*blocked_cmd_factory, + selected_db, + shard_id, + num_shards, + deadline, + ); // borrow dropped -- CRITICAL before await let mut registration_failed = false; for (target, msg) in pending_remote { @@ -2055,6 +1992,23 @@ pub(crate) fn immediate_scan( }) { return Some(err); } + // moon#989: a `BLMPOP`/`BZMPOP` whose keys span shards is refused here, + // before anything is popped or registered — the rule moon#962 applies to + // their non-blocking twins. "Pop from the FIRST non-empty key in argument + // order, exactly once" is a property of the whole key vector, and no shard + // can see the whole vector: this scan skips the keys it does not own (and + // so served a LATER local key over an earlier remote one), and two owners + // could each serve the same waiter. Refusing is decided from the key names + // alone, so it cannot lose an element. Co-located keys (`{hash}` tags) are + // unaffected and answer exactly as at `--shards 1`. + // + // The family list lives in `cross_shard_multikey_rejection`, the same + // guard the non-blocking dispatch path and scripts consult; of the + // blocking commands only these two are in it (a blocking stream read + // names one stream, and one key cannot span shards). + if let Some(err) = super::shared::cross_shard_multikey_rejection(cmd, args, num_shards) { + return Some(err); + } // moon#595: a stream read is answered by running the reader itself, not by // the per-key ladder below — its reply names the streams that had data // rather than the one key that served, and `XREADGROUP` needs the group diff --git a/src/server/conn/blocking_multikey.rs b/src/server/conn/blocking_multikey.rs new file mode 100644 index 000000000..920465de9 --- /dev/null +++ b/src/server/conn/blocking_multikey.rs @@ -0,0 +1,210 @@ +//! How a multi-key blocking waiter registers on its keys (moon#989). +//! +//! Shared by both runtimes' `handle_blocking_command*` so the two cannot +//! drift: the registration protocol IS the exactly-once argument, and a +//! second copy of it is a second place for that argument to break. +//! +//! * keys THIS shard owns are registered directly, under one `wait_id`; +//! * keys other shards own travel as ONE `BlockRegisterGroup` per owner, +//! carrying every key that owner holds, in argument order. +//! +//! The per-owner grouping is the fix. The coordinator used to send one +//! `BlockRegister` per key, and the owner served each on arrival — so a waiter +//! whose co-located keys both held data was served once per key, and every +//! reply after the first was dropped with its element already gone. A group is +//! handled in one synchronous stretch of the owner's loop +//! (`blocking::group::register_group`), where serving the waiter unregisters +//! its siblings before the next key is looked at. +//! +//! What grouping does NOT fix: keys owned by two DIFFERENT remote shards can +//! still each serve the same waiter, because nothing orders two shards' +//! wakes. `BLMPOP`/`BZMPOP` refuse that placement up front +//! (`immediate_scan`, CROSSSLOT); the rest of the family is moon#1019. + +use bytes::Bytes; +use futures::stream::FuturesUnordered; + +use crate::blocking::{BlockedCommand, BlockingRegistry, WaitEntry}; +use crate::protocol::Frame; +use crate::runtime::channel; +use crate::shard::dispatch::{ + BlockRegisterGroupPayload, BlockRegisterMember, ShardMessage, key_to_shard, +}; + +/// Everything a multi-key waiter holds once its registrations are staged. +pub(super) struct StagedWait { + pub wait_id: u64, + /// One receiver per key, local and remote. The first `Some(frame)` wins. + pub receivers: FuturesUnordered>>, + /// One `BlockRegisterGroup` per remote owner, not yet pushed — the caller + /// pushes them after releasing its registry borrow (A4/A5). + pub pending_remote: Vec<(usize, ShardMessage)>, + /// The owners in `pending_remote`, for the `BlockCancel` fan-out. + pub remote_shards: Vec, +} + +/// Register `keys` for one waiter: local keys now, remote keys staged as one +/// group per owner shard. +/// +/// `local_cmd` builds the `BlockedCommand` for a LOCAL key (it binds a stream +/// `$` against this shard's view); `remote_cmd` builds one for a remote key, +/// whose owner binds it on arrival. +pub(super) fn stage_multikey_wait( + registry: &mut BlockingRegistry, + keys: &[Bytes], + local_cmd: &dyn Fn(&Bytes) -> BlockedCommand, + remote_cmd: &dyn Fn() -> BlockedCommand, + selected_db: usize, + shard_id: usize, + num_shards: usize, + deadline: Option, +) -> StagedWait { + let wait_id = registry.next_wait_id(); + let receivers = FuturesUnordered::new(); + // (owner, members) in the order each owner first appears in argv. + let mut groups: Vec<(usize, Vec)> = Vec::new(); + let mut any_local = false; + + for key in keys { + let target = key_to_shard(key, num_shards); + let (tx, rx) = channel::oneshot::>(); + receivers.push(rx); + if target == shard_id { + any_local = true; + registry.register( + selected_db, + key.clone(), + WaitEntry { + wait_id, + cmd: local_cmd(key), + reply_tx: tx, + deadline, + }, + ); + continue; + } + let member = BlockRegisterMember { + key: key.clone(), + cmd: remote_cmd(), + reply_tx: tx, + }; + match groups.iter_mut().find(|(owner, _)| *owner == target) { + Some((_, members)) => members.push(member), + None => groups.push((target, vec![member])), + } + } + + // The owner may decide the whole command — including a `-WRONGTYPE` — only + // when it holds every key: nothing is registered here and no other shard + // was sent anything. + let whole_command = !any_local && groups.len() == 1; + let remote_shards: Vec = groups.iter().map(|(owner, _)| *owner).collect(); + let pending_remote = groups + .into_iter() + .map(|(owner, members)| { + ( + owner, + ShardMessage::BlockRegisterGroup(Box::new(BlockRegisterGroupPayload { + db_index: selected_db, + wait_id, + members, + whole_command, + })), + ) + }) + .collect(); + + StagedWait { + wait_id, + receivers, + pending_remote, + remote_shards, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cmd() -> BlockedCommand { + BlockedCommand::BLPop + } + + fn group_of(msg: &ShardMessage) -> (&[BlockRegisterMember], bool) { + match msg { + ShardMessage::BlockRegisterGroup(p) => (&p.members, p.whole_command), + _ => panic!("multi-key registrations travel as BlockRegisterGroup"), + } + } + + /// Co-located keys owned by another shard: ONE message, every key in + /// argument order, flagged as the whole command. + #[test] + fn colocated_remote_keys_travel_as_one_whole_group() { + const N: usize = 4; + let keys: Vec = ["{t}a", "{t}b", "{t}c", "{t}b"] + .iter() + .map(|k| Bytes::from_static(k.as_bytes())) + .collect(); + let owner = key_to_shard(&keys[0], N); + let me = (owner + 1) % N; + let mut reg = BlockingRegistry::new(me); + let staged = stage_multikey_wait(&mut reg, &keys, &|_| cmd(), &cmd, 0, me, N, None); + assert_eq!(staged.remote_shards, vec![owner]); + assert_eq!(staged.pending_remote.len(), 1); + let (members, whole) = group_of(&staged.pending_remote[0].1); + assert!(whole); + let names: Vec<&[u8]> = members.iter().map(|m| m.key.as_ref()).collect(); + assert_eq!(names, [&b"{t}a"[..], b"{t}b", b"{t}c", b"{t}b"]); + assert_eq!(staged.receivers.len(), 4); + assert!( + !reg.is_waiting(staged.wait_id), + "nothing registered locally" + ); + } + + /// Co-located keys owned by THIS shard never leave it. + #[test] + fn local_keys_register_locally_and_send_nothing() { + const N: usize = 4; + let keys: Vec = ["{t}a", "{t}b"] + .iter() + .map(|k| Bytes::from_static(k.as_bytes())) + .collect(); + let me = key_to_shard(&keys[0], N); + let mut reg = BlockingRegistry::new(me); + let staged = stage_multikey_wait(&mut reg, &keys, &|_| cmd(), &cmd, 0, me, N, None); + assert!(staged.pending_remote.is_empty()); + assert!(staged.remote_shards.is_empty()); + assert!(reg.is_waiting(staged.wait_id)); + } + + /// Keys on several shards: one PARTIAL group per remote owner, so no owner + /// can mistake its share for the whole command. + #[test] + fn spanning_keys_make_one_partial_group_per_owner() { + const N: usize = 4; + let me = 0usize; + let pick = |owner: usize, tag: &str| -> Bytes { + (0..10_000) + .map(|i| Bytes::from(format!("{tag}{i}"))) + .find(|k| key_to_shard(k, N) == owner) + .expect("a key for every shard") + }; + let keys = vec![pick(1, "x"), pick(2, "y"), pick(1, "z"), pick(me, "w")]; + let mut reg = BlockingRegistry::new(me); + let staged = stage_multikey_wait(&mut reg, &keys, &|_| cmd(), &cmd, 0, me, N, None); + assert_eq!(staged.remote_shards, vec![1, 2]); + let (m1, whole1) = group_of(&staged.pending_remote[0].1); + let (m2, whole2) = group_of(&staged.pending_remote[1].1); + assert!(!whole1 && !whole2); + assert_eq!(m1.len(), 2, "both of shard 1's keys in ONE message"); + assert_eq!(m1[0].key, keys[0]); + assert_eq!(m1[1].key, keys[2]); + assert_eq!(m2.len(), 1); + assert!( + reg.is_waiting(staged.wait_id), + "the local key is registered" + ); + } +} diff --git a/src/server/conn/mod.rs b/src/server/conn/mod.rs index 47918490c..b477a82f9 100644 --- a/src/server/conn/mod.rs +++ b/src/server/conn/mod.rs @@ -1,6 +1,9 @@ pub mod affinity; pub mod blocking; pub mod blocking_effect; +/// moon#989: the multi-key waiter's registration protocol, shared by both +/// runtimes' blocking handlers. +mod blocking_multikey; /// moon#556/#557: runtime-agnostic tests for the blocking pre-registration /// scan. Separate from `tests` below, which only compiles under monoio. #[cfg(test)] diff --git a/src/server/conn/shared.rs b/src/server/conn/shared.rs index a26482c04..e2a72f1ff 100644 --- a/src/server/conn/shared.rs +++ b/src/server/conn/shared.rs @@ -2191,6 +2191,12 @@ fn touches_a_key_it_did_not_route_on(cmd: &[u8]) -> bool { // path — strictly worse than either endpoint. See the family doc // above for the measured pop. (5, b'l') => cmd.eq_ignore_ascii_case(b"LMPOP"), + // moon#989: the blocking twins. They never reach the pre-routing + // guard — the connection handlers intercept every blocking command + // first — so `blocking::immediate_scan` consults this function itself, + // before it pops or registers anything. Listed HERE so there is one + // family list, not two that can drift. + (6, b'b') => cmd.eq_ignore_ascii_case(b"BLMPOP") || cmd.eq_ignore_ascii_case(b"BZMPOP"), (3, b'l') => cmd.eq_ignore_ascii_case(b"LCS"), (5, b'z') => cmd.eq_ignore_ascii_case(b"ZMPOP") || cmd.eq_ignore_ascii_case(b"ZDIFF"), (6, b's') => cmd.eq_ignore_ascii_case(b"SINTER") || cmd.eq_ignore_ascii_case(b"SUNION"), diff --git a/src/shard/dispatch.rs b/src/shard/dispatch.rs index 1695b55c1..cc29688aa 100644 --- a/src/shard/dispatch.rs +++ b/src/shard/dispatch.rs @@ -348,22 +348,55 @@ pub struct VectorSearchPayload { /// /// `BlockedCommand::XReadGroup` carries Vec + two Bytes + count options, pushing /// the inline variant past 160 B. Boxing collapses it to a pointer. +/// +/// Carries a SINGLE-key waiter only: `key` is the whole command, so the owner +/// answers `-WRONGTYPE` for it at registration time (moon#556). Multi-key +/// waiters register through [`BlockRegisterGroupPayload`] (moon#989). pub struct BlockRegisterPayload { pub db_index: usize, pub key: Bytes, pub wait_id: u64, pub cmd: crate::blocking::BlockedCommand, pub reply_tx: channel::OneshotSender>, - /// moon#556: is `key` the ONLY key this waiter is blocked on? - /// - /// The owning shard answers `-WRONGTYPE` at registration time for a key it - /// finds holding the wrong type — but only when it is the whole command. - /// For a multi-key waiter the other keys are registered on other shards and - /// may be serving concurrently, so an error raised here would race a real - /// wake-up whose element has already left the keyspace. Multi-key remote - /// registrations therefore keep their pre-#556 behaviour (the key is - /// skipped, the client stays blocked on its remaining keys). - pub sole_key: bool, +} + +/// One key of a [`BlockRegisterGroupPayload`]: the key, the command its wake +/// runs, and the channel that wake answers on. +pub struct BlockRegisterMember { + pub key: Bytes, + pub cmd: crate::blocking::BlockedCommand, + pub reply_tx: channel::OneshotSender>, +} + +/// moon#989: every key ONE shard owns of ONE multi-key blocking waiter, +/// delivered as a single message. +/// +/// The multi-key coordinator used to send one [`BlockRegisterPayload`] per +/// key. The owner handled each in isolation — register, see data, serve — so +/// a waiter whose co-located keys `b` and `c` both held data was served TWICE: +/// once when `b`'s registration landed and again when `c`'s did. The client +/// kept the first reply and dropped the second, and the element in it had +/// already left the keyspace. +/// +/// Grouping the keys makes "serve this waiter at most once" a property of one +/// synchronous stretch of the owner's event loop: every member is registered +/// under the same `wait_id` before any wake runs, and the wake that serves it +/// runs `remove_wait`, which unregisters its siblings in the same stretch. +pub struct BlockRegisterGroupPayload { + pub db_index: usize, + pub wait_id: u64, + /// The keys this shard owns, in the command's argument order — the order + /// Redis serves them in. A key named twice appears twice. + pub members: Vec, + /// `true` when `members` is EVERY key of the command, i.e. the keys are + /// co-located on this shard. The owner then decides the whole command, the + /// way `--shards 1` does, including the `-WRONGTYPE` Redis owes for the + /// first existing key of the wrong type. `false` for a waiter whose other + /// keys live on other shards: those may be serving concurrently, so an + /// error raised here would race a real wake-up whose element has already + /// left the keyspace — a wrong-typed key is skipped instead, and the + /// client stays blocked on its remaining keys. + pub whole_command: bool, } /// Portable raw socket file descriptor type. @@ -657,6 +690,9 @@ pub enum ShardMessage { /// Boxed (Phase 177) — `BlockedCommand::XReadGroup` pushes the inline variant /// past 160 B. BlockRegister(Box), + /// Register every key one shard owns of a multi-key blocked client, in + /// one message, so the owner can serve it at most once (moon#989). + BlockRegisterGroup(Box), /// Cancel a blocked client registration (woken by another shard or timed out). BlockCancel { wait_id: u64 }, /// Register a connected replica's per-shard sender channel with this shard. diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index cdd0908ff..05c5a5258 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -2319,7 +2319,6 @@ pub(crate) fn handle_shard_message_shared( wait_id, cmd, reply_tx, - sole_key, } = *payload; // moon#556: THIS shard owns the key, so it is the only one that // can answer the type question for it. A blocking pop on an @@ -2329,38 +2328,33 @@ pub(crate) fn handle_shard_message_shared( // remote case keeps the old behaviour (the waker finds nothing to // pop and answers a null the client reads as "empty"). // - // Gated on `sole_key`: for a multi-key waiter the sibling keys are - // registered on other shards and may be serving right now, and an - // error raised here would race a real wake-up whose element has - // already left the keyspace. Those keep the pre-#556 behaviour. + // Unconditional since moon#989: this message now carries only a + // single-key waiter (the key IS the command). Multi-key waiters + // register through `BlockRegisterGroup`, which makes the same + // decision only when it holds every key of the command. let mut cmd = cmd; - let type_error = if sole_key { - crate::shard::slice::with_shard_db(db_index, |guard| { - match cmd.family() { - // moon#832: a type probe before parking a waiter must - // not rewrite the value it is probing — `get_list` - // (via `get_promoted`) flattened the list's compact - // encoding on every remote blocking registration. - crate::blocking::WaitFamily::List => { - let now_ms = guard.now_ms(); - guard.get_list_ref_if_alive(&key, now_ms).err() - } - crate::blocking::WaitFamily::ZSet => guard.get_sorted_set(&key).err(), - // moon#595: `-WRONGTYPE` for a stream read on the - // wrong type, plus XREADGROUP's missing-key and - // missing-group errors. The client's own scan cannot - // see a key it does not own, so — exactly as moon#556 - // argued for the pops — the check has to happen here - // too or the remote case parks on a key that can never - // serve it. - crate::blocking::WaitFamily::Stream => { - crate::blocking::wakeup::stream_register_error(guard, &key, &cmd) - } + let type_error = crate::shard::slice::with_shard_db(db_index, |guard| { + match cmd.family() { + // moon#832: a type probe before parking a waiter must not + // rewrite the value it is probing — `get_list` (via + // `get_promoted`) flattened the list's compact encoding + // on every remote blocking registration. + crate::blocking::WaitFamily::List => { + let now_ms = guard.now_ms(); + guard.get_list_ref_if_alive(&key, now_ms).err() } - }) - } else { - None - }; + crate::blocking::WaitFamily::ZSet => guard.get_sorted_set(&key).err(), + // moon#595: `-WRONGTYPE` for a stream read on the wrong + // type, plus XREADGROUP's missing-key and missing-group + // errors. The client's own scan cannot see a key it does + // not own, so — exactly as moon#556 argued for the pops — + // the check has to happen here too or the remote case + // parks on a key that can never serve it. + crate::blocking::WaitFamily::Stream => { + crate::blocking::wakeup::stream_register_error(guard, &key, &cmd) + } + } + }); if let Some(err) = type_error { // Never registered, so there is nothing to unwind: the // client's `BlockCancel` on the way out is a no-op. @@ -2396,6 +2390,16 @@ pub(crate) fn handle_shard_message_shared( } }); } + ShardMessage::BlockRegisterGroup(payload) => { + // moon#989: every key this shard owns of one multi-key waiter, in + // ONE message — so registering, type-checking and serving them is + // one synchronous stretch and the waiter is served at most once. + let db_index = payload.db_index; + let mut reg = blocking_registry.borrow_mut(); + crate::shard::slice::with_shard_db(db_index, |guard| { + crate::blocking::group::register_group(&mut reg, guard, *payload); + }); + } ShardMessage::BlockCancel { wait_id } => { blocking_registry.borrow_mut().remove_wait(wait_id); } diff --git a/tests/blocking_multikey_cross_shard.rs b/tests/blocking_multikey_cross_shard.rs new file mode 100644 index 000000000..cf94704f3 --- /dev/null +++ b/tests/blocking_multikey_cross_shard.rs @@ -0,0 +1,710 @@ +//! A multi-key blocking pop serves exactly ONE element, from the first +//! non-empty key in argument order — at every shard count (moon#989). +//! +//! ## The defect +//! +//! `BLMPOP`/`BZMPOP` (and the rest of the multi-key blocking-pop family — +//! `BLPOP`, `BRPOP`, `BZPOPMIN`, `BZPOPMAX`) popped an element from a key they +//! did not answer with. The reply was right; a second key silently lost its +//! head element, which no client ever received: +//! +//! ```text +//! BLMPOP 0.3 3 {t}a {t}b {t}c LEFT ({t}a empty, {t}b=[B1 B2], {t}c=[C1 C2]) +//! redis 8.6.1 -> {t}b [B1] {t}b=[B2] {t}c=[C1 C2] +//! moon --shards 4 -> {t}b [B1] {t}b=[B2] {t}c=[C2] <-- C1 destroyed +//! ``` +//! +//! The keys are co-located under one `{hash}` tag, so this is not routing. It +//! happens whenever the CLIENT's connection lives on a different shard than +//! the keys. The client's own pre-block scan can only see keys its shard owns, +//! so it found nothing and fell through to the multi-key coordinator, which +//! sent one `BlockRegister` per key to the owner. The owner handled each +//! message in isolation: register the key, see data, serve the waiter. `{t}b` +//! served `B1`; then `{t}c`'s registration arrived, found data, and served the +//! SAME waiter again with `C1`. The client took the first reply and dropped +//! the second, with `C1` already gone from the keyspace. +//! +//! With keys on several shards the same fan-out also popped the WRONG key: the +//! local scan skipped a remote non-empty key and served a later local one. +//! +//! ## Why these placements cannot pass vacuously +//! +//! Whether a probe exercises the bug depends on which shard its CONNECTION +//! landed on, which a test cannot choose — on macOS every connection lands on +//! one shard, on Linux the kernel's `SO_REUSEPORT` hash decides. So instead of +//! sampling, every co-located row is run for tags constructed (with the +//! server's own `key_to_shard`) to be owned by EVERY shard in turn. Wherever +//! the connections land, at least `SHARDS - 1` of the owners are remote to +//! them. The issue's own sweep hit a first-placement-clean run exactly because +//! it did not do this. +//! +//! ## The contract asserted +//! +//! * co-located keys (`{hash}` tag) at `--shards 4`: the reply AND the keyspace +//! are byte-identical to redis 8.6.1 / `--shards 1`; +//! * spanning `BLMPOP`/`BZMPOP`: either the redis answer with only the answered +//! key touched, or an error with NOTHING touched — never a pop from a key the +//! reply does not name (the same answer contract as moon#962's `LMPOP`); +//! * `--shards 1`: every row answers exactly as redis does (control). + +mod common; + +use common::Conn; + +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +use moon::shard::dispatch::key_to_shard; + +const SHARDS: usize = 4; +/// Placements per owner shard. Two, not one: the issue observed that the first +/// placement after start can be clean on its own. +const PER_OWNER: usize = 2; + +// --------------------------------------------------------------------------- +// Harness +// --------------------------------------------------------------------------- + +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: usize) -> Moon { + // `CARGO_BIN_EXE_moon` is the binary cargo built for THIS test run; the + // `target/release/moon` fallback has unknown provenance. + let bin = std::path::PathBuf::from(env!("CARGO_BIN_EXE_moon")); + let (child, port) = common::spawn_listening(|port| { + let tmp_dir = std::env::temp_dir().join(format!("moon-bmk-{port}")); + let _ = std::fs::create_dir_all(&tmp_dir); + Command::new(&bin) + .args([ + "--port", + &port.to_string(), + "--shards", + &shards.to_string(), + "--admin-port", + "0", + "--appendonly", + "no", + "--disk-free-min-pct", + "0", + "--dir", + tmp_dir.to_str().unwrap_or("/tmp"), + ]) + .stdout(Stdio::null()) + .stderr( + std::fs::File::create(tmp_dir.join("moon.stderr")).expect("create moon stderr log"), + ) + .spawn() + .expect("spawn moon") + }); + let tmp_dir = std::env::temp_dir().join(format!("moon-bmk-{port}")); + let moon = Moon { + child, + port, + tmp_dir, + }; + let deadline = Instant::now() + Duration::from_secs(30); + 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) + && n > 0 + && buf.starts_with(b"+PONG") + { + return moon; + } + } + } + std::thread::sleep(Duration::from_millis(100)); + } + let log = std::fs::read_to_string(moon.tmp_dir.join("moon.stderr")).unwrap_or_default(); + panic!("moon never became ready on port {port}\n--- stderr ---\n{log}"); +} + +/// Render a RESP reply as stable text: `+OK` -> `OK`, `:3` -> `3`, a bulk -> +/// its bytes, an array -> `[a,b]`, a null -> `nil`, any error -> `!`. +fn canon(s: &str) -> String { + let b = s.as_bytes(); + let mut i = 0usize; + canon_one(b, &mut i) +} + +fn take_line(b: &[u8], i: &mut usize) -> String { + let start = *i; + while *i + 1 < b.len() && !(b[*i] == b'\r' && b[*i + 1] == b'\n') { + *i += 1; + } + let s = String::from_utf8_lossy(&b[start..*i]).into_owned(); + *i = (*i + 2).min(b.len()); + s +} + +fn canon_one(b: &[u8], i: &mut usize) -> String { + if *i >= b.len() { + return "".to_string(); + } + let tag = b[*i]; + *i += 1; + let head = take_line(b, i); + match tag { + b'+' | b':' | b',' | b'#' => head, + b'-' => format!("!{head}"), + b'_' => "nil".to_string(), + b'$' => { + if head.starts_with('-') { + return "nil".to_string(); + } + let n: usize = head.parse().unwrap_or(0); + let end = (*i + n).min(b.len()); + let s = String::from_utf8_lossy(&b[*i..end]).into_owned(); + *i = (end + 2).min(b.len()); + s + } + b'*' => { + if head.starts_with('-') { + return "nil".to_string(); + } + let n: usize = head.parse().unwrap_or(0); + let parts: Vec = (0..n).map(|_| canon_one(b, i)).collect(); + format!("[{}]", parts.join(",")) + } + other => format!("", other as char, head), + } +} + +/// `{hash}`-tagged key names `{bmk::}:1..=3` whose tag is owned by +/// shard `owner`. Found by search with the server's own routing function, so +/// the set of owners a test covers is a fact, not a hope. +fn colocated_owned_by(tag: &str, owner: usize, nth: usize) -> [String; 3] { + let mut found = 0usize; + for n in 0..10_000 { + let hash = format!("bmk:{tag}:{n}"); + if key_to_shard(hash.as_bytes(), SHARDS) != owner { + continue; + } + if found == nth { + let k = |j: u8| format!("{{{hash}}}:{j}"); + let keys = [k(1), k(2), k(3)]; + for key in &keys { + assert_eq!( + key_to_shard(key.as_bytes(), SHARDS), + owner, + "a hash-tagged key must route by its tag" + ); + } + return keys; + } + found += 1; + } + panic!("no tag owned by shard {owner} among 10000 candidates"); +} + +/// Three untagged key names on three DIFFERENT shards. +fn spanning_three(tag: &str, i: usize) -> [String; 3] { + let k1 = format!("bmk:{tag}:{i}:a"); + let o1 = key_to_shard(k1.as_bytes(), SHARDS); + let k2 = (0..1000) + .map(|j| format!("bmk:{tag}:{i}:b{j}")) + .find(|k| key_to_shard(k.as_bytes(), SHARDS) != o1) + .expect("a second shard among 1000 candidates"); + let o2 = key_to_shard(k2.as_bytes(), SHARDS); + let k3 = (0..1000) + .map(|j| format!("bmk:{tag}:{i}:c{j}")) + .find(|k| { + let o = key_to_shard(k.as_bytes(), SHARDS); + o != o1 && o != o2 + }) + .expect("a third shard among 1000 candidates"); + [k1, k2, k3] +} + +// --------------------------------------------------------------------------- +// The family, as rows +// --------------------------------------------------------------------------- + +#[derive(Clone, Copy, PartialEq)] +enum Kind { + List, + Zset, +} + +/// One blocking pop, parameterised over its three key names. `{1}` `{2}` `{3}` +/// are substituted. +struct Row { + label: &'static str, + kind: Kind, + argv: &'static [&'static str], + /// The redis 8.6.1 reply when `{1}` is empty and `{2}`, `{3}` hold + /// `[B1 B2]` / `[C1 C2]` (zsets: `B1`=1 `B2`=2, `C1`=1 `C2`=2). + expect: &'static str, + /// `{2}` after that reply. + k2_after: &'static str, +} + +const ROWS: &[Row] = &[ + Row { + label: "BLMPOP LEFT", + kind: Kind::List, + argv: &["BLMPOP", "0.3", "3", "{1}", "{2}", "{3}", "LEFT"], + expect: "[{2},[B1]]", + k2_after: "[B2]", + }, + Row { + label: "BLMPOP RIGHT COUNT 5", + kind: Kind::List, + argv: &[ + "BLMPOP", "0.3", "3", "{1}", "{2}", "{3}", "RIGHT", "COUNT", "5", + ], + expect: "[{2},[B2,B1]]", + k2_after: "[]", + }, + Row { + label: "BZMPOP MIN", + kind: Kind::Zset, + argv: &["BZMPOP", "0.3", "3", "{1}", "{2}", "{3}", "MIN"], + expect: "[{2},[[B1,1]]]", + k2_after: "[B2]", + }, + Row { + label: "BZMPOP MAX COUNT 5", + kind: Kind::Zset, + argv: &[ + "BZMPOP", "0.3", "3", "{1}", "{2}", "{3}", "MAX", "COUNT", "5", + ], + expect: "[{2},[[B2,2],[B1,1]]]", + k2_after: "[]", + }, + Row { + label: "BLPOP", + kind: Kind::List, + argv: &["BLPOP", "{1}", "{2}", "{3}", "0.3"], + expect: "[{2},B1]", + k2_after: "[B2]", + }, + Row { + label: "BRPOP", + kind: Kind::List, + argv: &["BRPOP", "{1}", "{2}", "{3}", "0.3"], + expect: "[{2},B2]", + k2_after: "[B1]", + }, + Row { + label: "BZPOPMIN", + kind: Kind::Zset, + argv: &["BZPOPMIN", "{1}", "{2}", "{3}", "0.3"], + expect: "[{2},B1,1]", + k2_after: "[B2]", + }, + Row { + label: "BZPOPMAX", + kind: Kind::Zset, + argv: &["BZPOPMAX", "{1}", "{2}", "{3}", "0.3"], + expect: "[{2},B2,2]", + k2_after: "[B1]", + }, + // The same key named twice: one pop, not two. + Row { + label: "BLMPOP same key twice", + kind: Kind::List, + argv: &["BLMPOP", "0.3", "3", "{1}", "{2}", "{2}", "LEFT"], + expect: "[{2},[B1]]", + k2_after: "[B2]", + }, + Row { + label: "BLPOP same key twice", + kind: Kind::List, + argv: &["BLPOP", "{2}", "{2}", "0.3"], + expect: "[{2},B1]", + k2_after: "[B2]", + }, +]; + +fn subst(argv: &[&str], keys: &[String; 3]) -> Vec { + argv.iter() + .map(|a| { + a.replace("{1}", &keys[0]) + .replace("{2}", &keys[1]) + .replace("{3}", &keys[2]) + }) + .collect() +} + +fn send(c: &mut Conn, argv: &[String]) -> String { + let refs: Vec<&str> = argv.iter().map(String::as_str).collect(); + canon(&c.send(&refs)) +} + +fn seed(c: &mut Conn, kind: Kind, keys: &[String; 3]) { + for k in keys { + let _ = c.send(&["DEL", k]); + } + let cmds: [Vec<&str>; 2] = match kind { + Kind::List => [ + vec!["RPUSH", &keys[1], "B1", "B2"], + vec!["RPUSH", &keys[2], "C1", "C2"], + ], + Kind::Zset => [ + vec!["ZADD", &keys[1], "1", "B1", "2", "B2"], + vec!["ZADD", &keys[2], "1", "C1", "2", "C2"], + ], + }; + for cmd in &cmds { + let r = canon(&c.send(cmd)); + assert!(!r.starts_with('!'), "seeding {cmd:?} failed: {r}"); + } +} + +fn contents(c: &mut Conn, kind: Kind, key: &str) -> String { + match kind { + Kind::List => canon(&c.send(&["LRANGE", key, "0", "-1"])), + Kind::Zset => canon(&c.send(&["ZRANGE", key, "0", "-1"])), + } +} + +/// Run every row against `keys`, the blocking probe on a FRESH connection. +/// Returns one line per row whose reply or keyspace differs from redis. +fn run_rows_exact(port: u16, keys: &[String; 3], wrong: &mut Vec) { + let mut admin = Conn::open(port); + for row in ROWS { + seed(&mut admin, row.kind, keys); + let argv = subst(row.argv, keys); + let reply = { + let mut probe = Conn::open(port); + send(&mut probe, &argv) + }; + let expect = row.expect.replace("{2}", &keys[1]); + let k1 = contents(&mut admin, row.kind, &keys[0]); + let k2 = contents(&mut admin, row.kind, &keys[1]); + let k3 = contents(&mut admin, row.kind, &keys[2]); + if reply != expect || k1 != "[]" || k2 != row.k2_after || k3 != "[C1,C2]" { + wrong.push(format!( + " {:<22} owner={} reply={reply} (want {expect}) {}={k2} (want {}) {}={k3} \ + (want [C1,C2]){}", + row.label, + key_to_shard(keys[1].as_bytes(), SHARDS), + keys[1], + row.k2_after, + keys[2], + if k3 != "[C1,C2]" && reply.contains(&keys[1]) { + " <-- ELEMENT DESTROYED: popped from a key the reply does not name" + } else { + "" + } + )); + } + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +/// The issue's shape, for every member of the family and every owner shard. +#[test] +fn bmk1_colocated_multikey_blocking_pop_serves_exactly_once() { + let m = spawn_moon(SHARDS); + let mut wrong = Vec::new(); + for owner in 0..SHARDS { + for nth in 0..PER_OWNER { + let keys = colocated_owned_by("imm", owner, nth); + run_rows_exact(m.port, &keys, &mut wrong); + } + } + assert!( + wrong.is_empty(), + "{} of {} co-located multi-key blocking pops at --shards {SHARDS} differ from \ + redis 8.6.1 (moon#989):\n{}", + wrong.len(), + SHARDS * PER_OWNER * ROWS.len(), + wrong.join("\n") + ); +} + +/// Control: the same rows at `--shards 1`, where the defect never existed. +/// A fix that broke the family outright fails here. +#[test] +fn bmk2_single_shard_control_every_row_answers_like_redis() { + let m = spawn_moon(1); + let mut wrong = Vec::new(); + for owner in 0..SHARDS { + let keys = colocated_owned_by("s1", owner, 0); + run_rows_exact(m.port, &keys, &mut wrong); + let keys = spanning_three("s1", owner); + run_rows_exact(m.port, &keys, &mut wrong); + } + assert!( + wrong.is_empty(), + "--shards 1 no longer answers like redis:\n{}", + wrong.join("\n") + ); +} + +/// Redis type-checks each key in argument order and answers `-WRONGTYPE` for +/// the first existing key of the wrong type, even when a LATER key could +/// serve. Co-located keys owned by another shard used to skip the check (the +/// owner could not tell one key of several from the whole command) and pop +/// the later key instead. +#[test] +fn bmk3_colocated_wrong_type_is_answered_like_redis() { + let m = spawn_moon(SHARDS); + let mut wrong = Vec::new(); + let cases: &[(&[&str], Kind)] = &[ + (&["BLMPOP", "0.3", "2", "{1}", "{2}", "LEFT"], Kind::List), + (&["BZMPOP", "0.3", "2", "{1}", "{2}", "MIN"], Kind::Zset), + (&["BLPOP", "{3}", "{1}", "{2}", "0.3"], Kind::List), + (&["BZPOPMIN", "{3}", "{1}", "{2}", "0.3"], Kind::Zset), + ]; + for owner in 0..SHARDS { + for nth in 0..PER_OWNER { + let keys = colocated_owned_by("wt", owner, nth); + for (argv, kind) in cases { + let mut admin = Conn::open(m.port); + for k in &keys { + let _ = admin.send(&["DEL", k]); + } + let _ = admin.send(&["SET", &keys[0], "x"]); + let _ = match kind { + Kind::List => admin.send(&["RPUSH", &keys[1], "B1"]), + Kind::Zset => admin.send(&["ZADD", &keys[1], "1", "B1"]), + }; + let reply = { + let mut probe = Conn::open(m.port); + send(&mut probe, &subst(argv, &keys)) + }; + let k2 = contents(&mut admin, *kind, &keys[1]); + if !reply.starts_with("!WRONGTYPE") || k2 != "[B1]" { + wrong.push(format!( + " {} owner={owner}: reply={reply} (want !WRONGTYPE ...) {}={k2} (want [B1])", + argv[0], keys[1] + )); + } + } + } + } + assert!( + wrong.is_empty(), + "co-located blocking pops skipped redis's type check:\n{}", + wrong.join("\n") + ); +} + +/// Block first, then push: the waiter is woken by the first key to receive +/// data and served exactly once. A control for the registration protocol — +/// the wake path must still find a waiter registered by the owner shard. +#[test] +fn bmk4_colocated_block_then_wake_serves_once() { + let m = spawn_moon(SHARDS); + let mut wrong = Vec::new(); + for owner in 0..SHARDS { + let keys = colocated_owned_by("wake", owner, 0); + for (label, argv, kind, expect) in [ + ( + "BLMPOP", + vec!["BLMPOP", "5", "3", "{1}", "{2}", "{3}", "LEFT"], + Kind::List, + "[{2},[B1]]", + ), + ( + "BZMPOP", + vec!["BZMPOP", "5", "3", "{1}", "{2}", "{3}", "MIN"], + Kind::Zset, + "[{2},[[B1,1]]]", + ), + ( + "BLPOP", + vec!["BLPOP", "{1}", "{2}", "{3}", "5"], + Kind::List, + "[{2},B1]", + ), + ] { + let mut admin = Conn::open(m.port); + for k in &keys { + let _ = admin.send(&["DEL", k]); + } + let argv = subst(&argv, &keys); + let port = m.port; + let waiter = std::thread::spawn(move || { + let mut probe = Conn::open(port); + send(&mut probe, &argv) + }); + // Registered = counted in `blocked_clients`. Polling beats a sleep: + // a slow box cannot turn this into a push-before-block race. + let deadline = Instant::now() + Duration::from_secs(4); + loop { + let info = admin.send(&["INFO", "clients"]); + if info.contains("blocked_clients:1") { + break; + } + assert!( + Instant::now() < deadline, + "{label}: waiter never registered: {info}" + ); + std::thread::sleep(Duration::from_millis(10)); + } + let (push2, push3): (Vec<&str>, Vec<&str>) = match kind { + Kind::List => ( + vec!["RPUSH", &keys[1], "B1", "B2"], + vec!["RPUSH", &keys[2], "C1", "C2"], + ), + Kind::Zset => ( + vec!["ZADD", &keys[1], "1", "B1", "2", "B2"], + vec!["ZADD", &keys[2], "1", "C1", "2", "C2"], + ), + }; + let _ = admin.send(&push2); + let _ = admin.send(&push3); + let reply = waiter.join().expect("waiter thread"); + let expect = expect.replace("{2}", &keys[1]); + let k2 = contents(&mut admin, kind, &keys[1]); + let k3 = contents(&mut admin, kind, &keys[2]); + if reply != expect || k2 != "[B2]" || k3 != "[C1,C2]" { + wrong.push(format!( + " {label} owner={owner}: reply={reply} (want {expect}) k2={k2} (want [B2]) \ + k3={k3} (want [C1,C2])" + )); + } + } + } + assert!(wrong.is_empty(), "block-then-wake:\n{}", wrong.join("\n")); +} + +/// A co-located waiter that TIMES OUT must leave nothing registered on the +/// owner: a later push stays in the key instead of feeding a ghost. +#[test] +fn bmk5_colocated_timeout_leaves_no_ghost_waiter() { + let m = spawn_moon(SHARDS); + let mut wrong = Vec::new(); + for owner in 0..SHARDS { + let keys = colocated_owned_by("ghost", owner, 0); + let mut admin = Conn::open(m.port); + for k in &keys { + let _ = admin.send(&["DEL", k]); + } + let reply = { + let mut probe = Conn::open(m.port); + send( + &mut probe, + &subst(&["BLMPOP", "0.1", "3", "{1}", "{2}", "{3}", "LEFT"], &keys), + ) + }; + let _ = admin.send(&["RPUSH", &keys[1], "B1"]); + let _ = admin.send(&["RPUSH", &keys[2], "C1"]); + let k2 = contents(&mut admin, Kind::List, &keys[1]); + let k3 = contents(&mut admin, Kind::List, &keys[2]); + if reply != "nil" || k2 != "[B1]" || k3 != "[C1]" { + wrong.push(format!( + " owner={owner}: reply={reply} (want nil) k2={k2} k3={k3} (want [B1] [C1])" + )); + } + } + assert!(wrong.is_empty(), "ghost waiters:\n{}", wrong.join("\n")); +} + +/// Keys on three different shards: `BLMPOP`/`BZMPOP` must either answer +/// exactly as redis does, or refuse with the keyspace untouched — never pop a +/// key the reply does not name, and never pop two. +#[test] +fn bmk6_spanning_blmpop_bzmpop_never_pop_an_unanswered_key() { + let m = spawn_moon(SHARDS); + let mut wrong = Vec::new(); + for i in 0..(SHARDS * PER_OWNER * 2) { + let keys = spanning_three("span", i); + for (label, argv, kind, expect) in [ + ( + "BLMPOP", + &["BLMPOP", "0.3", "3", "{1}", "{2}", "{3}", "LEFT"][..], + Kind::List, + "[{2},[B1]]", + ), + ( + "BZMPOP", + &["BZMPOP", "0.3", "3", "{1}", "{2}", "{3}", "MIN"][..], + Kind::Zset, + "[{2},[[B1,1]]]", + ), + ] { + let mut admin = Conn::open(m.port); + seed(&mut admin, kind, &keys); + let reply = { + let mut probe = Conn::open(m.port); + send(&mut probe, &subst(argv, &keys)) + }; + let k2 = contents(&mut admin, kind, &keys[1]); + let k3 = contents(&mut admin, kind, &keys[2]); + let ok = if reply.starts_with('!') { + k2 == "[B1,B2]" && k3 == "[C1,C2]" + } else { + reply == expect.replace("{2}", &keys[1]) && k2 == "[B2]" && k3 == "[C1,C2]" + }; + if !ok { + wrong.push(format!( + " {label} [{} | {} | {}]: reply={reply} k2={k2} k3={k3}", + keys[0], keys[1], keys[2] + )); + } + } + } + assert!( + wrong.is_empty(), + "{} spanning BLMPOP/BZMPOP placements popped a key they did not answer with:\n{}", + wrong.len(), + wrong.join("\n") + ); +} + +/// Inside MULTI the blocking pop is queued as its non-blocking twin and runs +/// at EXEC. Co-located keys answer exactly like redis; spanning keys follow +/// the same answer contract. +#[test] +fn bmk7_multi_exec_blocking_pop_is_exactly_once() { + let m = spawn_moon(SHARDS); + let mut wrong = Vec::new(); + for owner in 0..SHARDS { + let keys = colocated_owned_by("txn", owner, 0); + let spanning = spanning_three("txn", owner); + for (keys, colocated) in [(keys, true), (spanning, false)] { + let mut c = Conn::open(m.port); + seed(&mut c, Kind::List, &keys); + assert_eq!(c.send(&["MULTI"]), "+OK\r\n"); + let q = send( + &mut c, + &subst(&["BLMPOP", "0", "3", "{1}", "{2}", "{3}", "LEFT"], &keys), + ); + let reply = if q.starts_with('!') { + let _ = c.send(&["DISCARD"]); + q + } else { + canon(&c.send(&["EXEC"])) + }; + let k2 = contents(&mut c, Kind::List, &keys[1]); + let k3 = contents(&mut c, Kind::List, &keys[2]); + let refused = reply.contains('!'); + let ok = if refused { + !colocated && k2 == "[B1,B2]" && k3 == "[C1,C2]" + } else { + reply == format!("[[{},[B1]]]", keys[1]) && k2 == "[B2]" && k3 == "[C1,C2]" + }; + if !ok { + wrong.push(format!( + " colocated={colocated} owner={owner}: reply={reply} k2={k2} k3={k3}" + )); + } + } + } + assert!(wrong.is_empty(), "MULTI/EXEC:\n{}", wrong.join("\n")); +}