kevy_client/
cluster.rs

1//! Cluster-aware client: one connection per shard, routing each key to its
2//! owner shard by CRC16 slot — so no command pays the server-side cross-shard
3//! forwarding hop (the low-load tail / high-load throughput cost measured on
4//! the kevy-server role). The topology is discovered once at connect via
5//! `CLUSTER SLOTS`, which advertises each shard's slot range + address, so the
6//! client routes against the server's *actual* partition (no need to replicate
7//! its `slot → shard` arithmetic).
8//!
9//! Requires the server to run in cluster mode (`--cluster`): each shard then
10//! binds a deterministic listener and answers a wrong-shard key with `-MOVED`
11//! rather than forwarding. Correct routing here means `-MOVED` never fires.
12
13use std::io;
14use std::time::Duration;
15
16use kevy_hash::key_hash_slot;
17use kevy_resp::Reply;
18use kevy_resp_client::RespClient;
19
20use crate::{string, unexpected, vec2, vec3};
21
22/// Redis-cluster keyspace size: every key hashes to one of these slots.
23const NUM_SLOTS: usize = 16384;
24
25/// One open connection per distinct shard node, with a slot → shard index so a
26/// single-key command goes straight to its owner.
27pub struct ClusterClient {
28    /// Per distinct shard node, in first-advertised order.
29    shards: Vec<RespClient>,
30    /// `slot_to_shard[slot]` = index into [`Self::shards`]. Length [`NUM_SLOTS`].
31    slot_to_shard: Vec<u16>,
32}
33
34/// `(distinct shard nodes in advertised order, slot → shard-index)`.
35type Topology = (Vec<(String, u16)>, Vec<u16>);
36
37/// One `[start, end, host, port]` entry parsed out of `CLUSTER SLOTS`.
38struct SlotRange {
39    start: u16,
40    end: u16,
41    host: String,
42    port: u16,
43}
44
45impl ClusterClient {
46    /// Connect via a seed node, discover the topology (`CLUSTER SLOTS`), and
47    /// open one connection per shard.
48    pub fn connect(host: &str, port: u16) -> io::Result<Self> {
49        let mut seed = RespClient::connect(host, port)?;
50        let reply = seed.request(&[b"CLUSTER".to_vec(), b"SLOTS".to_vec()])?;
51        let ranges = parse_cluster_slots(reply)?;
52        let (nodes, slot_to_shard) = build_topology(&ranges)?;
53        let shards = nodes
54            .iter()
55            .map(|(h, p)| RespClient::connect(h, *p))
56            .collect::<io::Result<Vec<_>>>()?;
57        Ok(Self { shards, slot_to_shard })
58    }
59
60    /// Owner-shard index of `key`.
61    #[inline]
62    fn shard_for(&self, key: &[u8]) -> usize {
63        self.slot_to_shard[key_hash_slot(key) as usize] as usize
64    }
65
66    /// The connection to `key`'s owner shard — for callers (collection ops in
67    /// `cluster_coll`) that build a request via a shared helper.
68    #[inline]
69    pub(crate) fn route_mut(&mut self, key: &[u8]) -> &mut RespClient {
70        let i = self.shard_for(key);
71        &mut self.shards[i]
72    }
73
74    /// Route a single-key command (`args`) to the shard owning `key`.
75    pub fn request_keyed(&mut self, key: &[u8], args: &[Vec<u8>]) -> io::Result<Reply> {
76        let i = self.shard_for(key);
77        self.shards[i].request(args)
78    }
79
80    /// A keyless command (PING / etc.) — answered identically by any shard, so
81    /// send it to the first.
82    pub fn request_unkeyed(&mut self, args: &[Vec<u8>]) -> io::Result<Reply> {
83        self.shards[0].request(args)
84    }
85
86    /// Number of distinct shard nodes.
87    pub fn shard_count(&self) -> usize {
88        self.shards.len()
89    }
90
91}
92
93// ───────────── command surface (routed) ─────────────
94//
95// One connection per shard, so each single-key command goes straight to its
96// owner — no `-MOVED`, no server forwarding hop. Multi-key DEL/EXISTS route
97// per key and sum; keyspace-wide DBSIZE/FLUSHALL fan out to every shard.
98
99impl ClusterClient {
100    /// `PING` — answered by any shard.
101    pub fn ping(&mut self) -> io::Result<()> {
102        match self.request_unkeyed(&[b"PING".to_vec()])? {
103            Reply::Simple(s) if s == b"PONG" || s == b"OK" => Ok(()),
104            Reply::Error(e) => Err(io::Error::other(string(e))),
105            other => Err(unexpected(other)),
106        }
107    }
108
109    /// `PUBLISH channel message` — returns the number of subscribers reached.
110    /// kevy's pub/sub is process-global (delivered to subscribers on every
111    /// core), so a cluster PUBLISH goes to any one shard.
112    pub fn publish(&mut self, channel: &[u8], message: &[u8]) -> io::Result<usize> {
113        match self.request_unkeyed(&vec3(b"PUBLISH", channel, message))? {
114            Reply::Int(n) if n >= 0 => Ok(n as usize),
115            Reply::Error(e) => Err(io::Error::other(string(e))),
116            other => Err(unexpected(other)),
117        }
118    }
119
120    /// `SET key value`.
121    pub fn set(&mut self, key: &[u8], value: &[u8]) -> io::Result<()> {
122        match self.request_keyed(key, &vec3(b"SET", key, value))? {
123            Reply::Simple(s) if s == b"OK" => Ok(()),
124            Reply::Error(e) => Err(io::Error::other(string(e))),
125            other => Err(unexpected(other)),
126        }
127    }
128
129    /// `SET key value PX ttl_ms` — value with an expiry.
130    pub fn set_with_ttl(&mut self, key: &[u8], value: &[u8], ttl: Duration) -> io::Result<()> {
131        let ms = ttl.as_millis().min(i64::MAX as u128) as i64;
132        let args = vec![
133            b"SET".to_vec(),
134            key.to_vec(),
135            value.to_vec(),
136            b"PX".to_vec(),
137            ms.to_string().into_bytes(),
138        ];
139        match self.request_keyed(key, &args)? {
140            Reply::Simple(s) if s == b"OK" => Ok(()),
141            Reply::Error(e) => Err(io::Error::other(string(e))),
142            other => Err(unexpected(other)),
143        }
144    }
145
146    /// `GET key`. `None` if absent or expired.
147    pub fn get(&mut self, key: &[u8]) -> io::Result<Option<Vec<u8>>> {
148        match self.request_keyed(key, &vec2(b"GET", key))? {
149            Reply::Bulk(v) => Ok(Some(v)),
150            Reply::Nil => Ok(None),
151            Reply::Error(e) => Err(io::Error::other(string(e))),
152            other => Err(unexpected(other)),
153        }
154    }
155
156    /// `INCR key`. Returns the post-increment value.
157    pub fn incr(&mut self, key: &[u8]) -> io::Result<i64> {
158        match self.request_keyed(key, &vec2(b"INCR", key))? {
159            Reply::Int(n) => Ok(n),
160            Reply::Error(e) => Err(io::Error::other(string(e))),
161            other => Err(unexpected(other)),
162        }
163    }
164
165    /// `INCRBY key delta`.
166    pub fn incr_by(&mut self, key: &[u8], delta: i64) -> io::Result<i64> {
167        let args = vec![b"INCRBY".to_vec(), key.to_vec(), delta.to_string().into_bytes()];
168        match self.request_keyed(key, &args)? {
169            Reply::Int(n) => Ok(n),
170            Reply::Error(e) => Err(io::Error::other(string(e))),
171            other => Err(unexpected(other)),
172        }
173    }
174
175    /// `PEXPIRE key ttl_ms`. Whether the key existed and got a TTL.
176    pub fn expire(&mut self, key: &[u8], ttl: Duration) -> io::Result<bool> {
177        let ms = ttl.as_millis().min(i64::MAX as u128) as i64;
178        let args = vec![b"PEXPIRE".to_vec(), key.to_vec(), ms.to_string().into_bytes()];
179        match self.request_keyed(key, &args)? {
180            Reply::Int(1) => Ok(true),
181            Reply::Int(0) => Ok(false),
182            Reply::Error(e) => Err(io::Error::other(string(e))),
183            other => Err(unexpected(other)),
184        }
185    }
186
187    /// `PERSIST key`. Whether a TTL was removed.
188    pub fn persist(&mut self, key: &[u8]) -> io::Result<bool> {
189        match self.request_keyed(key, &vec2(b"PERSIST", key))? {
190            Reply::Int(1) => Ok(true),
191            Reply::Int(0) => Ok(false),
192            Reply::Error(e) => Err(io::Error::other(string(e))),
193            other => Err(unexpected(other)),
194        }
195    }
196
197    /// `PTTL key`. ms remaining, -2 no key, -1 no TTL.
198    pub fn ttl_ms(&mut self, key: &[u8]) -> io::Result<i64> {
199        match self.request_keyed(key, &vec2(b"PTTL", key))? {
200            Reply::Int(n) => Ok(n),
201            Reply::Error(e) => Err(io::Error::other(string(e))),
202            other => Err(unexpected(other)),
203        }
204    }
205
206    /// `DEL key [key ...]` — routed per key (each to its owner) and summed, so
207    /// keys spanning shards work without a same-slot constraint.
208    pub fn del(&mut self, keys: &[&[u8]]) -> io::Result<usize> {
209        let mut removed = 0;
210        for k in keys {
211            match self.request_keyed(k, &vec2(b"DEL", k))? {
212                Reply::Int(n) if n >= 0 => removed += n as usize,
213                Reply::Error(e) => return Err(io::Error::other(string(e))),
214                other => return Err(unexpected(other)),
215            }
216        }
217        Ok(removed)
218    }
219
220    /// `EXISTS key [key ...]` — routed per key and summed (a repeated key
221    /// counts each time, matching Redis).
222    pub fn exists(&mut self, keys: &[&[u8]]) -> io::Result<usize> {
223        let mut count = 0;
224        for k in keys {
225            match self.request_keyed(k, &vec2(b"EXISTS", k))? {
226                Reply::Int(n) if n >= 0 => count += n as usize,
227                Reply::Error(e) => return Err(io::Error::other(string(e))),
228                other => return Err(unexpected(other)),
229            }
230        }
231        Ok(count)
232    }
233
234    /// `DBSIZE` — the cluster-wide total. kevy answers DBSIZE by fanning out
235    /// across shards internally (`Route::Dbsize`), so a single shard already
236    /// reports the whole-cluster count; no client-side summing.
237    pub fn dbsize(&mut self) -> io::Result<usize> {
238        match self.request_unkeyed(&[b"DBSIZE".to_vec()])? {
239            Reply::Int(n) if n >= 0 => Ok(n as usize),
240            Reply::Error(e) => Err(io::Error::other(string(e))),
241            other => Err(unexpected(other)),
242        }
243    }
244
245    /// `FLUSHALL` — clears every shard. kevy fans FLUSHALL out internally
246    /// (`Route::Flush`), so one call wipes the whole cluster.
247    pub fn flushall(&mut self) -> io::Result<()> {
248        match self.request_unkeyed(&[b"FLUSHALL".to_vec()])? {
249            Reply::Simple(s) if s == b"OK" => Ok(()),
250            Reply::Error(e) => Err(io::Error::other(string(e))),
251            other => Err(unexpected(other)),
252        }
253    }
254}
255
256/// Build the `(distinct nodes, slot → shard-index)` topology from the parsed
257/// ranges. Distinct nodes are kept in first-advertised order; a slot left
258/// uncovered (a gap a healthy cluster never has) defaults to shard 0. Pure —
259/// no I/O — so the routing is unit-testable without a live server.
260fn build_topology(ranges: &[SlotRange]) -> io::Result<Topology> {
261    if ranges.is_empty() {
262        return Err(io::Error::new(
263            io::ErrorKind::InvalidData,
264            "CLUSTER SLOTS returned no ranges",
265        ));
266    }
267    let mut nodes: Vec<(String, u16)> = Vec::new();
268    let mut slot_to_shard = vec![0u16; NUM_SLOTS];
269    for r in ranges {
270        let idx = if let Some(i) = nodes.iter().position(|(h, p)| h == &r.host && *p == r.port) { i } else {
271            nodes.push((r.host.clone(), r.port));
272            nodes.len() - 1
273        } as u16;
274        for slot in r.start..=r.end {
275            slot_to_shard[slot as usize] = idx;
276        }
277    }
278    Ok((nodes, slot_to_shard))
279}
280
281/// Parse a `CLUSTER SLOTS` reply: `[[start, end, [host, port, id, []], …], …]`.
282fn parse_cluster_slots(reply: Reply) -> io::Result<Vec<SlotRange>> {
283    fn bad() -> io::Error {
284        io::Error::new(io::ErrorKind::InvalidData, "malformed CLUSTER SLOTS reply")
285    }
286    let Reply::Array(rows) = reply else {
287        return Err(bad());
288    };
289    let mut out = Vec::with_capacity(rows.len());
290    for row in rows {
291        let Reply::Array(cols) = row else { return Err(bad()) };
292        if cols.len() < 3 {
293            return Err(bad());
294        }
295        let start = as_int(&cols[0]).ok_or_else(bad)?;
296        let end = as_int(&cols[1]).ok_or_else(bad)?;
297        let Reply::Array(node) = &cols[2] else {
298            return Err(bad());
299        };
300        if node.len() < 2 {
301            return Err(bad());
302        }
303        let host = as_str(&node[0]).ok_or_else(bad)?;
304        let port = as_int(&node[1]).ok_or_else(bad)?;
305        if !(0..=i64::from(u16::MAX)).contains(&start)
306            || !(0..=i64::from(u16::MAX)).contains(&end)
307            || !(0..=i64::from(u16::MAX)).contains(&port)
308        {
309            return Err(bad());
310        }
311        out.push(SlotRange {
312            start: start as u16,
313            end: end as u16,
314            host,
315            port: port as u16,
316        });
317    }
318    Ok(out)
319}
320
321fn as_int(r: &Reply) -> Option<i64> {
322    match r {
323        Reply::Int(n) => Some(*n),
324        _ => None,
325    }
326}
327
328fn as_str(r: &Reply) -> Option<String> {
329    match r {
330        Reply::Bulk(b) | Reply::Simple(b) => Some(String::from_utf8_lossy(b).into_owned()),
331        _ => None,
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338
339    /// A 4-shard even split mirroring the server's contiguous slot ranges,
340    /// encoded as a `CLUSTER SLOTS` reply, then parsed + routed.
341    fn four_shard_reply() -> Reply {
342        // 16384 / 4 = 4096 slots each.
343        let row = |start: i64, end: i64, port: i64| {
344            Reply::Array(vec![
345                Reply::Int(start),
346                Reply::Int(end),
347                Reply::Array(vec![
348                    Reply::Bulk(b"127.0.0.1".to_vec()),
349                    Reply::Int(port),
350                    Reply::Bulk(b"nodeid".to_vec()),
351                    Reply::Array(vec![]),
352                ]),
353            ])
354        };
355        Reply::Array(vec![
356            row(0, 4095, 7001),
357            row(4096, 8191, 7002),
358            row(8192, 12287, 7003),
359            row(12288, 16383, 7004),
360        ])
361    }
362
363    #[test]
364    fn parse_and_build_4_shard_topology() {
365        let ranges = parse_cluster_slots(four_shard_reply()).unwrap();
366        assert_eq!(ranges.len(), 4);
367        let (nodes, slot_to_shard) = build_topology(&ranges).unwrap();
368        assert_eq!(nodes.len(), 4);
369        assert_eq!(nodes[0], ("127.0.0.1".to_string(), 7001));
370        assert_eq!(nodes[3], ("127.0.0.1".to_string(), 7004));
371        // Every slot maps to the advertised shard.
372        assert_eq!(slot_to_shard[0], 0);
373        assert_eq!(slot_to_shard[4095], 0);
374        assert_eq!(slot_to_shard[4096], 1);
375        assert_eq!(slot_to_shard[8192], 2);
376        assert_eq!(slot_to_shard[16383], 3);
377    }
378
379    #[test]
380    fn keys_route_to_their_slot_owner() {
381        let ranges = parse_cluster_slots(four_shard_reply()).unwrap();
382        let (_, slot_to_shard) = build_topology(&ranges).unwrap();
383        // A key's shard = its CRC16 slot's owner — the same mapping the server
384        // enforces, so `-MOVED` never fires.
385        for k in ["k0", "k1", "user:42", "rate:10.0.0.1", "gl:abc"] {
386            let slot = key_hash_slot(k.as_bytes()) as usize;
387            let shard = slot_to_shard[slot] as usize;
388            // Sanity: slot's owner is the contiguous-range shard.
389            assert_eq!(shard, slot / 4096, "key {k} slot {slot}");
390        }
391    }
392
393    #[test]
394    fn rejects_empty_and_malformed() {
395        assert!(parse_cluster_slots(Reply::Int(1)).is_err());
396        assert!(build_topology(&[]).is_err());
397        assert!(parse_cluster_slots(Reply::Array(vec![Reply::Array(vec![Reply::Int(0)])])).is_err());
398    }
399}
kevy_client/cluster.rs

kevy_client/
cluster.rs
