1use std::io;
14
15use crate::store::Store;
16
17pub struct Pipeline<'a> {
21 store: &'a Store,
22 ops: Vec<PendingOp>,
23}
24
25enum PendingOp {
26 Set { key: Vec<u8>, value: Vec<u8> },
27 Del { keys: Vec<Vec<u8>> },
28 Incr { key: Vec<u8> },
29 IncrBy { key: Vec<u8>, delta: i64 },
30 HSet { key: Vec<u8>, pairs: Vec<(Vec<u8>, Vec<u8>)> },
31 HDel { key: Vec<u8>, fields: Vec<Vec<u8>> },
32 HIncrBy { key: Vec<u8>, field: Vec<u8>, delta: i64 },
33 ZAdd { key: Vec<u8>, pairs: Vec<(f64, Vec<u8>)> },
34 ZAddFlags { key: Vec<u8>, pairs: Vec<(f64, Vec<u8>)>, flags: kevy_store::ZaddFlags },
35 ZRem { key: Vec<u8>, members: Vec<Vec<u8>> },
36 ZIncrBy { key: Vec<u8>, delta: f64, member: Vec<u8> },
37 SAdd { key: Vec<u8>, members: Vec<Vec<u8>> },
38 SRem { key: Vec<u8>, members: Vec<Vec<u8>> },
39 LPush { key: Vec<u8>, values: Vec<Vec<u8>> },
40 RPush { key: Vec<u8>, values: Vec<Vec<u8>> },
41}
42
43impl<'a> Pipeline<'a> {
44 pub(crate) fn new(store: &'a Store) -> Self {
45 Self { store, ops: Vec::new() }
46 }
47
48 pub fn len(&self) -> usize {
50 self.ops.len()
51 }
52
53 pub fn is_empty(&self) -> bool {
55 self.ops.is_empty()
56 }
57
58 pub fn set(mut self, key: &[u8], value: &[u8]) -> Self {
63 self.ops.push(PendingOp::Set {
64 key: key.to_vec(),
65 value: value.to_vec(),
66 });
67 self
68 }
69
70 pub fn del(mut self, keys: &[&[u8]]) -> Self {
72 self.ops.push(PendingOp::Del {
73 keys: keys.iter().map(|k| k.to_vec()).collect(),
74 });
75 self
76 }
77
78 pub fn incr(mut self, key: &[u8]) -> Self {
80 self.ops.push(PendingOp::Incr { key: key.to_vec() });
81 self
82 }
83
84 pub fn incr_by(mut self, key: &[u8], delta: i64) -> Self {
86 self.ops.push(PendingOp::IncrBy { key: key.to_vec(), delta });
87 self
88 }
89
90 pub fn hset(mut self, key: &[u8], pairs: &[(&[u8], &[u8])]) -> Self {
92 self.ops.push(PendingOp::HSet {
93 key: key.to_vec(),
94 pairs: pairs.iter().map(|(f, v)| (f.to_vec(), v.to_vec())).collect(),
95 });
96 self
97 }
98
99 pub fn hdel(mut self, key: &[u8], fields: &[&[u8]]) -> Self {
101 self.ops.push(PendingOp::HDel {
102 key: key.to_vec(),
103 fields: fields.iter().map(|f| f.to_vec()).collect(),
104 });
105 self
106 }
107
108 pub fn hincrby(mut self, key: &[u8], field: &[u8], delta: i64) -> Self {
110 self.ops.push(PendingOp::HIncrBy {
111 key: key.to_vec(),
112 field: field.to_vec(),
113 delta,
114 });
115 self
116 }
117
118 pub fn zadd(mut self, key: &[u8], pairs: &[(f64, &[u8])]) -> Self {
120 self.ops.push(PendingOp::ZAdd {
121 key: key.to_vec(),
122 pairs: pairs.iter().map(|(s, m)| (*s, m.to_vec())).collect(),
123 });
124 self
125 }
126
127 pub fn zadd_flags(
129 mut self,
130 key: &[u8],
131 pairs: &[(f64, &[u8])],
132 flags: kevy_store::ZaddFlags,
133 ) -> Self {
134 self.ops.push(PendingOp::ZAddFlags {
135 key: key.to_vec(),
136 pairs: pairs.iter().map(|(s, m)| (*s, m.to_vec())).collect(),
137 flags,
138 });
139 self
140 }
141
142 pub fn zrem(mut self, key: &[u8], members: &[&[u8]]) -> Self {
144 self.ops.push(PendingOp::ZRem {
145 key: key.to_vec(),
146 members: members.iter().map(|m| m.to_vec()).collect(),
147 });
148 self
149 }
150
151 pub fn zincrby(mut self, key: &[u8], delta: f64, member: &[u8]) -> Self {
153 self.ops.push(PendingOp::ZIncrBy {
154 key: key.to_vec(),
155 delta,
156 member: member.to_vec(),
157 });
158 self
159 }
160
161 pub fn sadd(mut self, key: &[u8], members: &[&[u8]]) -> Self {
163 self.ops.push(PendingOp::SAdd {
164 key: key.to_vec(),
165 members: members.iter().map(|m| m.to_vec()).collect(),
166 });
167 self
168 }
169
170 pub fn srem(mut self, key: &[u8], members: &[&[u8]]) -> Self {
172 self.ops.push(PendingOp::SRem {
173 key: key.to_vec(),
174 members: members.iter().map(|m| m.to_vec()).collect(),
175 });
176 self
177 }
178
179 pub fn lpush(mut self, key: &[u8], values: &[&[u8]]) -> Self {
181 self.ops.push(PendingOp::LPush {
182 key: key.to_vec(),
183 values: values.iter().map(|v| v.to_vec()).collect(),
184 });
185 self
186 }
187
188 pub fn rpush(mut self, key: &[u8], values: &[&[u8]]) -> Self {
190 self.ops.push(PendingOp::RPush {
191 key: key.to_vec(),
192 values: values.iter().map(|v| v.to_vec()).collect(),
193 });
194 self
195 }
196
197 pub fn commit(mut self) -> io::Result<()> {
202 let ops = std::mem::take(&mut self.ops);
203 for op in ops {
204 self.apply_one(op)?;
205 }
206 Ok(())
207 }
208
209 fn apply_one(&self, op: PendingOp) -> io::Result<()> {
213 match op {
214 PendingOp::Set { key, value } => {
215 self.store.set(&key, &value)?;
216 }
217 PendingOp::Del { keys } => {
218 let refs: Vec<&[u8]> = keys.iter().map(|k| k.as_slice()).collect();
219 self.store.del(&refs)?;
220 }
221 PendingOp::Incr { key } => {
222 self.store.incr(&key)?;
223 }
224 PendingOp::IncrBy { key, delta } => {
225 self.store.incr_by(&key, delta)?;
226 }
227 PendingOp::HSet { key, pairs } => {
228 let refs: Vec<(&[u8], &[u8])> =
229 pairs.iter().map(|(f, v)| (f.as_slice(), v.as_slice())).collect();
230 self.store.hset(&key, &refs)?;
231 }
232 PendingOp::HDel { key, fields } => {
233 let refs: Vec<&[u8]> = fields.iter().map(|f| f.as_slice()).collect();
234 self.store.hdel(&key, &refs)?;
235 }
236 PendingOp::HIncrBy { key, field, delta } => {
237 self.store.hincrby(&key, &field, delta)?;
238 }
239 PendingOp::ZAdd { key, pairs } => {
240 let refs: Vec<(f64, &[u8])> =
241 pairs.iter().map(|(s, m)| (*s, m.as_slice())).collect();
242 self.store.zadd(&key, &refs)?;
243 }
244 PendingOp::ZAddFlags { key, pairs, flags } => {
245 let refs: Vec<(f64, &[u8])> =
246 pairs.iter().map(|(s, m)| (*s, m.as_slice())).collect();
247 self.store.zadd_flags(&key, &refs, flags)?;
248 }
249 PendingOp::ZRem { key, members } => {
250 let refs: Vec<&[u8]> = members.iter().map(|m| m.as_slice()).collect();
251 self.store.zrem(&key, &refs)?;
252 }
253 PendingOp::ZIncrBy { key, delta, member } => {
254 self.store.zincrby(&key, delta, &member)?;
255 }
256 PendingOp::SAdd { key, members } => {
257 let refs: Vec<&[u8]> = members.iter().map(|m| m.as_slice()).collect();
258 self.store.sadd(&key, &refs)?;
259 }
260 PendingOp::SRem { key, members } => {
261 let refs: Vec<&[u8]> = members.iter().map(|m| m.as_slice()).collect();
262 self.store.srem(&key, &refs)?;
263 }
264 PendingOp::LPush { key, values } => {
265 let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect();
266 self.store.lpush(&key, &refs)?;
267 }
268 PendingOp::RPush { key, values } => {
269 let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect();
270 self.store.rpush(&key, &refs)?;
271 }
272 }
273 Ok(())
274 }
275}
276
277impl Store {
278 pub fn pipeline(&self) -> Pipeline<'_> {
281 Pipeline::new(self)
282 }
283}
284
285#[cfg_attr(not(test), allow(dead_code))]
289pub(crate) const PIPELINE_OPS: &[&str] = &[
290 "SET", "DEL", "INCR", "INCRBY", "HSET", "HDEL", "HINCRBY", "ZADD",
291 "ZREM", "ZINCRBY", "SADD", "SREM", "LPUSH", "RPUSH",
292];