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

kevy_embedded/
ops_atomic_all.rs
