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

kevy_embedded/
ops_atomic_all.rs
