Skip to main content

kevy_client/
cluster_coll.rs

1//! [`ClusterClient`] collection commands (hash / list / set / sorted set).
2//! Every collection command is single-key, so each routes to its key's owner
3//! shard — split out of `cluster.rs` for the 500-LOC house rule. The multi-key
4//! set-combine ops (SINTER/SUNION/SDIFF) route by their first key: like Redis
5//! Cluster, they require all keys in one slot (use a `{hashtag}`), else the
6//! server answers `-MOVED`.
7
8use std::io;
9
10use kevy_resp::Reply;
11
12use crate::cluster::ClusterClient;
13use crate::collections::{list_pop, list_push, remote_set_combine, set_multi};
14use crate::{array_to_bulks, string, unexpected, vec2, vec3};
15
16/// `match reply` → non-negative count, shared by the many `*LEN`/`*ADD`/`*REM`
17/// ops that reply `:N`.
18fn count(reply: Reply) -> io::Result<usize> {
19    match reply {
20        Reply::Int(n) if n >= 0 => Ok(n as usize),
21        Reply::Error(e) => Err(io::Error::other(string(e))),
22        other => Err(unexpected(other)),
23    }
24}
25
26/// `match reply` → bulk array (HGETALL/SMEMBERS/LRANGE/…).
27fn bulks(reply: Reply) -> io::Result<Vec<Vec<u8>>> {
28    match reply {
29        Reply::Array(items) => array_to_bulks(items),
30        Reply::Error(e) => Err(io::Error::other(string(e))),
31        other => Err(unexpected(other)),
32    }
33}
34
35impl ClusterClient {
36    // ===== Hash =====
37
38    /// `HSET key field value [field value ...]` — count of newly added fields.
39    pub fn hset(&mut self, key: &[u8], pairs: &[(&[u8], &[u8])]) -> io::Result<usize> {
40        let mut args = Vec::with_capacity(2 + pairs.len() * 2);
41        args.push(b"HSET".to_vec());
42        args.push(key.to_vec());
43        for (f, v) in pairs {
44            args.push(f.to_vec());
45            args.push(v.to_vec());
46        }
47        count(self.request_keyed(key, &args)?)
48    }
49
50    /// `HGET key field`. `None` when key or field absent.
51    pub fn hget(&mut self, key: &[u8], field: &[u8]) -> io::Result<Option<Vec<u8>>> {
52        match self.request_keyed(key, &vec3(b"HGET", key, field))? {
53            Reply::Bulk(v) => Ok(Some(v)),
54            Reply::Nil => Ok(None),
55            Reply::Error(e) => Err(io::Error::other(string(e))),
56            other => Err(unexpected(other)),
57        }
58    }
59
60    /// `HDEL key field [field ...]` — count of fields removed.
61    pub fn hdel(&mut self, key: &[u8], fields: &[&[u8]]) -> io::Result<usize> {
62        let mut args = Vec::with_capacity(fields.len() + 2);
63        args.push(b"HDEL".to_vec());
64        args.push(key.to_vec());
65        args.extend(fields.iter().map(|f| f.to_vec()));
66        count(self.request_keyed(key, &args)?)
67    }
68
69    /// `HLEN key`.
70    pub fn hlen(&mut self, key: &[u8]) -> io::Result<usize> {
71        count(self.request_keyed(key, &vec2(b"HLEN", key))?)
72    }
73
74    /// `HGETALL key` — flat `[f0, v0, f1, v1, …]`.
75    pub fn hgetall(&mut self, key: &[u8]) -> io::Result<Vec<Vec<u8>>> {
76        bulks(self.request_keyed(key, &vec2(b"HGETALL", key))?)
77    }
78
79    /// `HKEYS key`.
80    pub fn hkeys(&mut self, key: &[u8]) -> io::Result<Vec<Vec<u8>>> {
81        bulks(self.request_keyed(key, &vec2(b"HKEYS", key))?)
82    }
83
84    /// `HVALS key`.
85    pub fn hvals(&mut self, key: &[u8]) -> io::Result<Vec<Vec<u8>>> {
86        bulks(self.request_keyed(key, &vec2(b"HVALS", key))?)
87    }
88
89    // ===== List =====
90
91    /// `LPUSH key value [value ...]` — new list length.
92    pub fn lpush(&mut self, key: &[u8], values: &[&[u8]]) -> io::Result<usize> {
93        list_push(self.route_mut(key), b"LPUSH", key, values)
94    }
95
96    /// `RPUSH key value [value ...]` — new list length.
97    pub fn rpush(&mut self, key: &[u8], values: &[&[u8]]) -> io::Result<usize> {
98        list_push(self.route_mut(key), b"RPUSH", key, values)
99    }
100
101    /// `LPOP key count`.
102    pub fn lpop(&mut self, key: &[u8], n: usize) -> io::Result<Vec<Vec<u8>>> {
103        list_pop(self.route_mut(key), b"LPOP", key, n)
104    }
105
106    /// `RPOP key count`.
107    pub fn rpop(&mut self, key: &[u8], n: usize) -> io::Result<Vec<Vec<u8>>> {
108        list_pop(self.route_mut(key), b"RPOP", key, n)
109    }
110
111    /// `LLEN key`.
112    pub fn llen(&mut self, key: &[u8]) -> io::Result<usize> {
113        count(self.request_keyed(key, &vec2(b"LLEN", key))?)
114    }
115
116    /// `LRANGE key start stop` (Redis-style negative indices).
117    pub fn lrange(&mut self, key: &[u8], start: i64, stop: i64) -> io::Result<Vec<Vec<u8>>> {
118        let args = vec![
119            b"LRANGE".to_vec(),
120            key.to_vec(),
121            start.to_string().into_bytes(),
122            stop.to_string().into_bytes(),
123        ];
124        bulks(self.request_keyed(key, &args)?)
125    }
126
127    // ===== Set =====
128
129    /// `SADD key member [member ...]` — count newly added.
130    pub fn sadd(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
131        set_multi(self.route_mut(key), b"SADD", key, members)
132    }
133
134    /// `SREM key member [member ...]` — count removed.
135    pub fn srem(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
136        set_multi(self.route_mut(key), b"SREM", key, members)
137    }
138
139    /// `SMEMBERS key`.
140    pub fn smembers(&mut self, key: &[u8]) -> io::Result<Vec<Vec<u8>>> {
141        bulks(self.request_keyed(key, &vec2(b"SMEMBERS", key))?)
142    }
143
144    /// `SCARD key`.
145    pub fn scard(&mut self, key: &[u8]) -> io::Result<usize> {
146        count(self.request_keyed(key, &vec2(b"SCARD", key))?)
147    }
148
149    /// `SISMEMBER key member`.
150    pub fn sismember(&mut self, key: &[u8], member: &[u8]) -> io::Result<bool> {
151        match self.request_keyed(key, &vec3(b"SISMEMBER", key, member))? {
152            Reply::Int(1) => Ok(true),
153            Reply::Int(0) => Ok(false),
154            Reply::Error(e) => Err(io::Error::other(string(e))),
155            other => Err(unexpected(other)),
156        }
157    }
158
159    /// `SINTER key [key ...]` — all keys must share a slot (route by the first).
160    pub fn sinter(&mut self, keys: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
161        self.set_combine(b"SINTER", keys)
162    }
163
164    /// `SUNION key [key ...]` — same-slot.
165    pub fn sunion(&mut self, keys: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
166        self.set_combine(b"SUNION", keys)
167    }
168
169    /// `SDIFF key [key ...]` — same-slot.
170    pub fn sdiff(&mut self, keys: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
171        self.set_combine(b"SDIFF", keys)
172    }
173
174    fn set_combine(&mut self, verb: &[u8], keys: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
175        let Some(first) = keys.first() else {
176            return Ok(Vec::new());
177        };
178        remote_set_combine(self.route_mut(first), verb, keys)
179    }
180
181    // ===== Sorted set =====
182
183    /// `ZADD key score member [score member ...]` — count newly added.
184    pub fn zadd(&mut self, key: &[u8], pairs: &[(f64, &[u8])]) -> io::Result<usize> {
185        let mut args = Vec::with_capacity(2 + pairs.len() * 2);
186        args.push(b"ZADD".to_vec());
187        args.push(key.to_vec());
188        for (score, m) in pairs {
189            args.push(score.to_string().into_bytes());
190            args.push(m.to_vec());
191        }
192        count(self.request_keyed(key, &args)?)
193    }
194
195    /// `ZREM key member [member ...]` — count removed.
196    pub fn zrem(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
197        set_multi(self.route_mut(key), b"ZREM", key, members)
198    }
199
200    /// `ZSCORE key member`. `None` if absent.
201    pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> io::Result<Option<f64>> {
202        match self.request_keyed(key, &vec3(b"ZSCORE", key, member))? {
203            Reply::Bulk(v) => {
204                let s = std::str::from_utf8(&v)
205                    .map_err(|_| io::Error::other("non-utf8 score reply"))?;
206                let n: f64 = s
207                    .parse()
208                    .map_err(|_| io::Error::other(format!("bad score float: {s}")))?;
209                Ok(Some(n))
210            }
211            Reply::Nil => Ok(None),
212            Reply::Error(e) => Err(io::Error::other(string(e))),
213            other => Err(unexpected(other)),
214        }
215    }
216
217    /// `ZCARD key`.
218    pub fn zcard(&mut self, key: &[u8]) -> io::Result<usize> {
219        count(self.request_keyed(key, &vec2(b"ZCARD", key))?)
220    }
221
222    /// `ZRANGE key start stop` (ascending score; negative indices from tail).
223    pub fn zrange(&mut self, key: &[u8], start: i64, stop: i64) -> io::Result<Vec<Vec<u8>>> {
224        let args = vec![
225            b"ZRANGE".to_vec(),
226            key.to_vec(),
227            start.to_string().into_bytes(),
228            stop.to_string().into_bytes(),
229        ];
230        bulks(self.request_keyed(key, &args)?)
231    }
232}