kevy_embedded/
ops_atomic_all.rs

1//! Cross-shard read-modify-write closure:
2//! `Store::atomic_all_shards`.
3//!
4//! `atomic_all_shards(|tx| { ... })` holds a write lock on every
5//! shard for the closure body. Operations inside the closure are
6//! routed to their owning shards, and AOF writes are batched
7//! per-shard with one fsync per shard at commit time.
8//!
9//! Heavier than [`Store::atomic`](crate::Store::atomic): every
10//! reader and writer on the affected shards blocks until the
11//! closure returns. Use it only when the closure genuinely needs
12//! more than one shard and atomicity across them is required.
13
14use crate::{KevyError, KevyResult};
15use std::sync::RwLockWriteGuard;
16
17use crate::shard::shard_idx;
18use crate::store::{Inner, Store, commit_write, store_err};
19
20use crate::store::ensure_writable;
21
22/// One key's pre-transaction state, plus the shard it lives on.
23/// `None` prior = the key did not exist.
24type ShardUndoEntry = (usize, Vec<u8>, Option<(kevy_store::Value, Option<u64>)>);
25
26/// Context handed to the `atomic_all_shards` closure body. Methods
27/// route to the right shard by hashing the key.
28pub struct AtomicAllShards<'a> {
29    pub(crate) guards: Vec<RwLockWriteGuard<'a, Inner>>,
30    /// (shard_idx, serialised RESP-frame parts) queued for AOF commit.
31    log: Vec<(usize, Vec<Vec<u8>>)>,
32    /// `(shard_idx, key, prior)` captured on first touch; `None` prior
33    /// means the key did not exist. See [`Store::atomic`] — same
34    /// rollback contract, and worse to get wrong here because a
35    /// rejected transaction would otherwise diverge several shards at
36    /// once.
37    undo: Vec<ShardUndoEntry>,
38    touched: std::collections::HashSet<Vec<u8>>,
39    /// The index catalog, for the transaction-scoped index reads in
40    /// `ops_atomic_all_index.rs`. Held as a handle rather than reached
41    /// through `Store` because those reads must use the guards above,
42    /// not take the shard locks again.
43    #[cfg(feature = "index")]
44    pub(crate) indexes: std::sync::Arc<crate::ops_index::IndexReg>,
45}
46
47impl<'a> AtomicAllShards<'a> {
48    pub(crate) fn idx(&self, key: &[u8]) -> usize {
49        shard_idx(key, self.guards.len())
50    }
51
52    /// Record `key`'s prior state, once, before its first mutation.
53    fn snap(&mut self, key: &[u8]) {
54        if self.touched.contains(key) {
55            return;
56        }
57        let i = self.idx(key);
58        let prior = self.guards[i].store.clone_with_ttl(key);
59        self.touched.insert(key.to_vec());
60        self.undo.push((i, key.to_vec(), prior));
61    }
62
63    fn log_arg(&mut self, idx: usize, parts: &[&[u8]]) {
64        self.log.push((idx, parts.iter().map(|p| p.to_vec()).collect()));
65    }
66
67    // ---- string ops -----------------------------------------------
68
69    /// `SET key value` — always succeeds.
70    pub fn set(&mut self, key: &[u8], value: &[u8]) -> bool {
71        self.snap(key);
72        let i = self.idx(key);
73        let ok = self.guards[i].store.set(key, value.to_vec(), None, false, false);
74        self.log_arg(i, &[b"SET", key, value]);
75        ok
76    }
77
78    /// `GET key`.
79    pub fn get(&mut self, key: &[u8]) -> KevyResult<Option<Vec<u8>>> {
80        let i = self.idx(key);
81        self.guards[i]
82            .store
83            .get(key)
84            .map(|opt| opt.as_deref().map(<[u8]>::to_vec))
85            .map_err(store_err)
86    }
87
88    /// `INCR key`.
89    pub fn incr(&mut self, key: &[u8]) -> KevyResult<i64> {
90        self.snap(key);
91        let i = self.idx(key);
92        let n = self.guards[i].store.incr_by(key, 1).map_err(store_err)?;
93        self.log_arg(i, &[b"INCR", key]);
94        Ok(n)
95    }
96
97    /// `INCRBY key delta`.
98    pub fn incr_by(&mut self, key: &[u8], delta: i64) -> KevyResult<i64> {
99        self.snap(key);
100        let i = self.idx(key);
101        let n = self.guards[i].store.incr_by(key, delta).map_err(store_err)?;
102        let s = format!("{delta}");
103        self.log_arg(i, &[b"INCRBY", key, s.as_bytes()]);
104        Ok(n)
105    }
106
107    // ---- hash ops --------------------------------------------------
108
109    /// `HSET key field value [field value ...]`. Returns count newly
110    /// added (existing fields are overwritten but not counted).
111    pub fn hset(&mut self, key: &[u8], pairs: &[(&[u8], &[u8])]) -> KevyResult<usize> {
112        self.snap(key);
113        let i = self.idx(key);
114        let n = self.guards[i].store.hset(key, pairs).map_err(store_err)?;
115        let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
116        parts.push(b"HSET");
117        parts.push(key);
118        for (f, v) in pairs {
119            parts.push(f);
120            parts.push(v);
121        }
122        self.log_arg(i, &parts);
123        Ok(n)
124    }
125
126    /// `HGET key field` — `None` when the key or field is absent.
127    pub fn hget(&mut self, key: &[u8], field: &[u8]) -> KevyResult<Option<Vec<u8>>> {
128        let i = self.idx(key);
129        Ok(self.guards[i].store.hget(key, field).map_err(store_err)?.map(<[u8]>::to_vec))
130    }
131
132    /// `HINCRBY key field delta` — returns the field's new value.
133    pub fn hincrby(&mut self, key: &[u8], field: &[u8], delta: i64) -> KevyResult<i64> {
134        self.snap(key);
135        let i = self.idx(key);
136        let n = self.guards[i].store.hincrby(key, field, delta).map_err(store_err)?;
137        let s = format!("{delta}");
138        self.log_arg(i, &[b"HINCRBY", key, field, s.as_bytes()]);
139        Ok(n)
140    }
141
142    // ---- zset ops --------------------------------------------------
143
144    /// `ZADD key score member [score member ...]`. Returns count newly
145    /// added (score updates of existing members are not counted).
146    pub fn zadd(&mut self, key: &[u8], pairs: &[(f64, &[u8])]) -> KevyResult<usize> {
147        self.snap(key);
148        let i = self.idx(key);
149        let n = self.guards[i].store.zadd(key, pairs).map_err(store_err)?;
150        let score_strs: Vec<Vec<u8>> =
151            pairs.iter().map(|(s, _)| format!("{s}").into_bytes()).collect();
152        let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
153        parts.push(b"ZADD");
154        parts.push(key);
155        for (j, (_, m)) in pairs.iter().enumerate() {
156            parts.push(&score_strs[j]);
157            parts.push(m);
158        }
159        self.log_arg(i, &parts);
160        Ok(n)
161    }
162
163    /// `ZINCRBY key delta member` — returns the member's new score.
164    pub fn zincrby(&mut self, key: &[u8], delta: f64, member: &[u8]) -> KevyResult<f64> {
165        self.snap(key);
166        let i = self.idx(key);
167        let n = self.guards[i].store.zincrby(key, delta, member).map_err(store_err)?;
168        let s = format!("{delta}");
169        self.log_arg(i, &[b"ZINCRBY", key, s.as_bytes(), member]);
170        Ok(n)
171    }
172
173    /// `ZSCORE key member` (parity with [`super::ops_atomic::AtomicCtx`]).
174    pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> KevyResult<Option<f64>> {
175        let i = self.idx(key);
176        self.guards[i].store.zscore(key, member).map_err(store_err)
177    }
178
179    // ---- keyspace ops (Pipeline write parity) ----------------------
180
181    /// `DEL key [key ...]` — keys may span shards; each key's delete
182    /// is applied and AOF-logged on its own shard.
183    pub fn del(&mut self, keys: &[&[u8]]) -> usize {
184        for k in keys {
185            self.snap(k);
186        }
187        let mut n = 0;
188        for k in keys {
189            let i = self.idx(k);
190            if self.guards[i].store.del(&[k]) > 0 {
191                n += 1;
192                self.log_arg(i, &[b"DEL", k]);
193            }
194        }
195        n
196    }
197
198    /// `EXISTS key [key ...]` — count of the given keys that exist.
199    pub fn exists(&mut self, keys: &[&[u8]]) -> usize {
200        keys.iter()
201            .filter(|k| {
202                let i = self.idx(k);
203                self.guards[i].store.key_exists(k)
204            })
205            .count()
206    }
207
208    // ---- hash ops --------------------------------------------------
209
210    /// `HDEL key field [field ...]`.
211    pub fn hdel(&mut self, key: &[u8], fields: &[&[u8]]) -> KevyResult<usize> {
212        self.snap(key);
213        let i = self.idx(key);
214        let removed = self.guards[i].store.hdel(key, fields).map_err(store_err)?;
215        if removed > 0 {
216            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + fields.len());
217            argv.push(b"HDEL");
218            argv.push(key);
219            argv.extend_from_slice(fields);
220            self.log_arg(i, &argv);
221        }
222        Ok(removed)
223    }
224
225    /// `HGETALL key` — `(field, value)` pairs.
226    pub fn hgetall(&mut self, key: &[u8]) -> KevyResult<Vec<(Vec<u8>, Vec<u8>)>> {
227        let i = self.idx(key);
228        let flat = self.guards[i].store.hgetall(key).map_err(store_err)?;
229        let mut out = Vec::with_capacity(flat.len() / 2);
230        let mut it = flat.into_iter();
231        while let (Some(f), Some(v)) = (it.next(), it.next()) {
232            out.push((f, v));
233        }
234        Ok(out)
235    }
236
237    /// `HMGET key field [field ...]` — `None` per absent field.
238    pub fn hmget(&mut self, key: &[u8], fields: &[&[u8]]) -> KevyResult<Vec<Option<Vec<u8>>>> {
239        let i = self.idx(key);
240        self.guards[i].store.hmget(key, fields).map_err(store_err)
241    }
242
243    /// `HEXISTS key field`.
244    pub fn hexists(&mut self, key: &[u8], field: &[u8]) -> KevyResult<bool> {
245        let i = self.idx(key);
246        self.guards[i].store.hexists(key, field).map_err(store_err)
247    }
248
249    // ---- set ops ---------------------------------------------------
250
251    /// `SADD key member [member ...]`.
252    pub fn sadd(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
253        self.snap(key);
254        let i = self.idx(key);
255        let added = self.guards[i].store.sadd(key, members).map_err(store_err)?;
256        if added > 0 {
257            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
258            argv.push(b"SADD");
259            argv.push(key);
260            argv.extend_from_slice(members);
261            self.log_arg(i, &argv);
262        }
263        Ok(added)
264    }
265
266    /// `SREM key member [member ...]`.
267    pub fn srem(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
268        self.snap(key);
269        let i = self.idx(key);
270        let removed = self.guards[i].store.srem(key, members).map_err(store_err)?;
271        if removed > 0 {
272            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
273            argv.push(b"SREM");
274            argv.push(key);
275            argv.extend_from_slice(members);
276            self.log_arg(i, &argv);
277        }
278        Ok(removed)
279    }
280
281    // ---- list ops --------------------------------------------------
282
283    /// `LPUSH key value [value ...]` — returns the new list length.
284    pub fn lpush(&mut self, key: &[u8], values: &[&[u8]]) -> KevyResult<usize> {
285        self.snap(key);
286        let i = self.idx(key);
287        let len = self.guards[i].store.lpush(key, values).map_err(store_err)?;
288        let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + values.len());
289        argv.push(b"LPUSH");
290        argv.push(key);
291        argv.extend_from_slice(values);
292        self.log_arg(i, &argv);
293        Ok(len)
294    }
295
296    /// `RPUSH key value [value ...]` — returns the new list length.
297    pub fn rpush(&mut self, key: &[u8], values: &[&[u8]]) -> KevyResult<usize> {
298        self.snap(key);
299        let i = self.idx(key);
300        let len = self.guards[i].store.rpush(key, values).map_err(store_err)?;
301        let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + values.len());
302        argv.push(b"RPUSH");
303        argv.push(key);
304        argv.extend_from_slice(values);
305        self.log_arg(i, &argv);
306        Ok(len)
307    }
308
309    // ---- zset ops --------------------------------------------------
310
311    /// `ZREM key member [member ...]`.
312    pub fn zrem(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
313        self.snap(key);
314        let i = self.idx(key);
315        let removed = self.guards[i].store.zrem(key, members).map_err(store_err)?;
316        if removed > 0 {
317            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
318            argv.push(b"ZREM");
319            argv.push(key);
320            argv.extend_from_slice(members);
321            self.log_arg(i, &argv);
322        }
323        Ok(removed)
324    }
325
326    /// `ZCARD key` — member count; 0 when absent.
327    pub fn zcard(&mut self, key: &[u8]) -> KevyResult<usize> {
328        let i = self.idx(key);
329        self.guards[i].store.zcard(key).map_err(store_err)
330    }
331
332    /// Flags-aware `ZADD`. AOF logs the applied pairs as plain
333    /// `ZADD` — the effect, never the condition (deterministic replay).
334    pub fn zadd_flags(
335        &mut self,
336        key: &[u8],
337        pairs: &[(f64, &[u8])],
338        flags: kevy_store::ZaddFlags,
339    ) -> KevyResult<kevy_store::ZaddReport> {
340        if !flags.valid() {
341            return Err(KevyError::InvalidInput("invalid ZADD flag combo".into()));
342        }
343        let i = self.idx(key);
344        let rep = self.guards[i].store.zadd_flags(key, pairs, flags).map_err(store_err)?;
345        if !rep.applied.is_empty() {
346            let score_strs: Vec<Vec<u8>> =
347                rep.applied.iter().map(|(s, _)| format!("{s}").into_bytes()).collect();
348            let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + rep.applied.len() * 2);
349            parts.push(b"ZADD");
350            parts.push(key);
351            for (j, (_, m)) in rep.applied.iter().enumerate() {
352                parts.push(&score_strs[j]);
353                parts.push(m);
354            }
355            self.log_arg(i, &parts);
356        }
357        Ok(rep)
358    }
359}
360
361impl Store {
362    /// Run `body` as a transaction holding write locks on EVERY
363    /// shard for the closure's duration. Reads inside the closure
364    /// see prior writes (full read-modify-write). On closure
365    /// return, AOF writes commit with one fsync per shard.
366    ///
367    /// Cost: blocks every other writer + reader on this Store for
368    /// the closure body. Use when atomic multi-shard semantics are
369    /// required; otherwise prefer the single-shard `atomic`.
370    pub fn atomic_all_shards<R>(
371        &self,
372        body: impl FnOnce(&mut AtomicAllShards<'_>) -> KevyResult<R>,
373    ) -> KevyResult<R> {
374        ensure_writable(self)?;
375        // Take every shard's write lock in shard-index order
376        // (deterministic order avoids deadlock).
377        let guards: Vec<RwLockWriteGuard<'_, Inner>> =
378            self.shards.iter().map(|s| s.write().expect("lock poisoned")).collect();
379        let mut ctx = AtomicAllShards {
380            guards,
381            log: Vec::new(),
382            undo: Vec::new(),
383            touched: std::collections::HashSet::new(),
384            #[cfg(feature = "index")]
385            indexes: std::sync::Arc::clone(&self.indexes),
386        };
387        let outcome = body(&mut ctx);
388        let log = std::mem::take(&mut ctx.log);
389        let undo = std::mem::take(&mut ctx.undo);
390        let r = match outcome {
391            Ok(r) => r,
392            Err(e) => {
393                rollback_all(&mut ctx.guards, undo);
394                return Err(e);
395            }
396        };
397        commit_group_all(&mut ctx.guards, log)?;
398        Ok(r)
399    }
400}
401
402/// Parity manifest: command names `AtomicAllShards` implements.
403/// MUST stay identical to `ops_atomic::ATOMIC_OPS` (the two ctxs
404/// drifted before — zscore was missing here).
405#[cfg_attr(not(test), allow(dead_code))]
406pub(crate) const ATOMIC_ALL_OPS: &[&str] = &[
407    "SET",
408    "GET",
409    "INCR",
410    "INCRBY",
411    "HSET",
412    "HGET",
413    "HINCRBY",
414    "ZADD",
415    "ZINCRBY",
416    "ZSCORE",
417    "DEL",
418    "EXISTS",
419    "HDEL",
420    "HGETALL",
421    "HMGET",
422    "HEXISTS",
423    "SADD",
424    "SREM",
425    "LPUSH",
426    "RPUSH",
427    "ZREM",
428    "ZCARD",
429    "SMEMBERS",
430    "SISMEMBER",
431    "LRANGE",
432    "LLEN",
433    "SCARD",
434    "ZRANGEBYSCORE",
435];
436
437/// Undo a rejected cross-shard transaction. See `Store::atomic`; reverse
438/// order so a key touched more than once lands on its earliest state.
439fn rollback_all(guards: &mut [RwLockWriteGuard<'_, Inner>], undo: Vec<ShardUndoEntry>) {
440    for (idx, key, prior) in undo.into_iter().rev() {
441        let g = &mut guards[idx];
442        match prior {
443            Some((value, ttl_ms)) => g.store.put_with_ttl(key, value, ttl_ms),
444            None => {
445                let k: &[u8] = &key;
446                g.store.del(&[k]);
447            }
448        }
449    }
450}
451
452/// Bracket and group-commit each shard's queued frames. The brackets make
453/// replay all-or-nothing at any size; the group makes `Fsync::Always`
454/// cost one sync per shard instead of one per frame.
455fn commit_group_all(
456    guards: &mut [RwLockWriteGuard<'_, Inner>],
457    log: Vec<(usize, Vec<Vec<u8>>)>,
458) -> KevyResult<()> {
459    #[cfg(feature = "persist")]
460    for g in guards.iter_mut() {
461        if let Some(aof) = g.aof.as_mut() {
462            aof.begin_group();
463        }
464    }
465    let mut commit = Ok(());
466    for (idx, parts) in log {
467        let g = &mut guards[idx];
468        let refs: Vec<&[u8]> = parts.iter().map(|v| v.as_slice()).collect();
469        commit = commit_write(g, &refs);
470        if commit.is_err() {
471            break;
472        }
473    }
474    #[cfg(feature = "persist")]
475    for g in guards.iter_mut() {
476        if let Some(aof) = g.aof.as_mut() {
477            let synced = aof.end_group().map_err(KevyError::from);
478            if commit.is_ok() {
479                commit = synced;
480            }
481        }
482    }
483    commit
484}
kevy_embedded/ops_atomic_all.rs

kevy_embedded/
ops_atomic_all.rs
