kevy_embedded/
ops_atomic.rs

1//! Single-shard read-modify-write closure: `Store::atomic`.
2//!
3//! `atomic(|tx| { ... })` holds the shard's write lock for the
4//! closure body. Reads inside the closure see prior writes inside
5//! the same closure, so read-modify-write loops work as expected.
6//! AOF writes are deferred and batched into a single fsync at
7//! commit time.
8//!
9//! Every key touched inside the closure must hash to the same
10//! shard. For closures that span shards use
11//! [`Store::atomic_all_shards`](crate::Store::atomic_all_shards).
12
13use crate::{KevyError, KevyResult};
14use std::sync::RwLockWriteGuard;
15
16use crate::store::ensure_writable;
17use crate::store::{Inner, Store, commit_write, store_err};
18
19/// One key's pre-transaction state: the key, and what was there before
20/// the transaction first touched it (`None` = the key did not exist).
21type UndoEntry = (Vec<u8>, Option<(kevy_store::Value, Option<u64>)>);
22
23/// Handle passed to the `atomic` closure body. Methods mirror the
24/// equivalent `Store` ops but operate on the already-held write
25/// lock, so reads inside the block see the closure's own writes.
26pub struct AtomicCtx<'a> {
27    inner: &'a mut Inner,
28    log: Vec<Vec<Vec<u8>>>,
29    /// Prior state of every key this transaction has touched, captured
30    /// on FIRST touch: `None` means the key did not exist. Replayed in
31    /// reverse by [`Store::atomic`] when the closure returns `Err`, so a
32    /// rejected transaction leaves neither memory nor the AOF changed.
33    undo: Vec<UndoEntry>,
34    /// Keys already in `undo` — a key is snapshotted once, before its
35    /// first mutation, never after.
36    touched: std::collections::HashSet<Vec<u8>>,
37}
38
39impl AtomicCtx<'_> {
40    // ---- string ops ------------------------------------------------
41
42    /// `SET key value`. Returns `true` (SET always succeeds without
43    /// `NX`/`XX` veto).
44    pub fn set(&mut self, key: &[u8], value: &[u8]) -> bool {
45        self.snap(key);
46        let ok = self
47            .inner
48            .store
49            .set(key, value.to_vec(), None, false, false);
50        self.log_arg(&[b"SET", key, value]);
51        ok
52    }
53
54    /// `GET key`.
55    pub fn get(&mut self, key: &[u8]) -> KevyResult<Option<Vec<u8>>> {
56        self.inner
57            .store
58            .get(key)
59            .map(|opt| opt.as_deref().map(<[u8]>::to_vec))
60            .map_err(store_err)
61    }
62
63    /// `INCR key` — by 1.
64    pub fn incr(&mut self, key: &[u8]) -> KevyResult<i64> {
65        self.snap(key);
66        let n = self.inner.store.incr_by(key, 1).map_err(store_err)?;
67        self.log_arg(&[b"INCR", key]);
68        Ok(n)
69    }
70
71    /// `INCRBY key delta`.
72    pub fn incr_by(&mut self, key: &[u8], delta: i64) -> KevyResult<i64> {
73        self.snap(key);
74        let n = self.inner.store.incr_by(key, delta).map_err(store_err)?;
75        let s = format!("{delta}");
76        self.log_arg(&[b"INCRBY", key, s.as_bytes()]);
77        Ok(n)
78    }
79
80    // ---- hash ops ---------------------------------------------------
81
82    /// `HSET key field value`.
83    pub fn hset(&mut self, key: &[u8], pairs: &[(&[u8], &[u8])]) -> KevyResult<usize> {
84        self.snap(key);
85        let n = self.inner.store.hset(key, pairs).map_err(store_err)?;
86        let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
87        parts.push(b"HSET");
88        parts.push(key);
89        for (f, v) in pairs {
90            parts.push(f);
91            parts.push(v);
92        }
93        self.log_arg(&parts);
94        Ok(n)
95    }
96
97    /// `HGET key field`.
98    pub fn hget(&mut self, key: &[u8], field: &[u8]) -> KevyResult<Option<Vec<u8>>> {
99        Ok(self
100            .inner
101            .store
102            .hget(key, field)
103            .map_err(store_err)?
104            .map(<[u8]>::to_vec))
105    }
106
107    /// `HINCRBY key field delta`.
108    pub fn hincrby(&mut self, key: &[u8], field: &[u8], delta: i64) -> KevyResult<i64> {
109        self.snap(key);
110        let n = self.inner.store.hincrby(key, field, delta).map_err(store_err)?;
111        let s = format!("{delta}");
112        self.log_arg(&[b"HINCRBY", key, field, s.as_bytes()]);
113        Ok(n)
114    }
115
116    // ---- zset ops ---------------------------------------------------
117
118    /// `ZADD key score member`.
119    pub fn zadd(&mut self, key: &[u8], pairs: &[(f64, &[u8])]) -> KevyResult<usize> {
120        self.snap(key);
121        let n = self.inner.store.zadd(key, pairs).map_err(store_err)?;
122        let score_strs: Vec<Vec<u8>> =
123            pairs.iter().map(|(s, _)| format!("{s}").into_bytes()).collect();
124        let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
125        parts.push(b"ZADD");
126        parts.push(key);
127        for (i, (_, m)) in pairs.iter().enumerate() {
128            parts.push(&score_strs[i]);
129            parts.push(m);
130        }
131        self.log_arg(&parts);
132        Ok(n)
133    }
134
135    /// `ZINCRBY key delta member`.
136    pub fn zincrby(&mut self, key: &[u8], delta: f64, member: &[u8]) -> KevyResult<f64> {
137        self.snap(key);
138        let n = self.inner.store.zincrby(key, delta, member).map_err(store_err)?;
139        let s = format!("{delta}");
140        self.log_arg(&[b"ZINCRBY", key, s.as_bytes(), member]);
141        Ok(n)
142    }
143
144    /// `ZSCORE key member`.
145    pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> KevyResult<Option<f64>> {
146        self.inner.store.zscore(key, member).map_err(store_err)
147    }
148
149    // ---- helpers ----------------------------------------------------
150
151    // ---- keyspace ops (Pipeline write parity) ----------------------
152
153    /// `DEL key [key ...]` — every key must hash to this shard.
154    pub fn del(&mut self, keys: &[&[u8]]) -> usize {
155        for k in keys {
156            self.snap(k);
157        }
158        let n = self.inner.store.del(keys);
159        if n > 0 {
160            let mut argv: Vec<&[u8]> = Vec::with_capacity(1 + keys.len());
161            argv.push(b"DEL");
162            argv.extend_from_slice(keys);
163            self.log_arg(&argv);
164        }
165        n
166    }
167
168    /// `EXISTS key [key ...]` — count of the given keys that exist.
169    pub fn exists(&mut self, keys: &[&[u8]]) -> usize {
170        keys.iter()
171            .filter(|k| self.inner.store.key_exists(k))
172            .count()
173    }
174
175    // ---- hash ops --------------------------------------------------
176
177    /// `HDEL key field [field ...]`.
178    pub fn hdel(&mut self, key: &[u8], fields: &[&[u8]]) -> KevyResult<usize> {
179        self.snap(key);
180        let removed = self.inner.store.hdel(key, fields).map_err(store_err)?;
181        if removed > 0 {
182            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + fields.len());
183            argv.push(b"HDEL");
184            argv.push(key);
185            argv.extend_from_slice(fields);
186            self.log_arg(&argv);
187        }
188        Ok(removed)
189    }
190
191    /// `HGETALL key` — `(field, value)` pairs; reads see the
192    /// closure's own prior writes.
193    pub fn hgetall(&mut self, key: &[u8]) -> KevyResult<Vec<(Vec<u8>, Vec<u8>)>> {
194        let flat = self.inner.store.hgetall(key).map_err(store_err)?;
195        let mut out = Vec::with_capacity(flat.len() / 2);
196        let mut it = flat.into_iter();
197        while let (Some(f), Some(v)) = (it.next(), it.next()) {
198            out.push((f, v));
199        }
200        Ok(out)
201    }
202
203    /// `HMGET key field [field ...]` — `None` per absent field.
204    pub fn hmget(&mut self, key: &[u8], fields: &[&[u8]]) -> KevyResult<Vec<Option<Vec<u8>>>> {
205        self.inner.store.hmget(key, fields).map_err(store_err)
206    }
207
208    /// `HEXISTS key field`.
209    pub fn hexists(&mut self, key: &[u8], field: &[u8]) -> KevyResult<bool> {
210        self.inner.store.hexists(key, field).map_err(store_err)
211    }
212
213    // ---- set ops ---------------------------------------------------
214
215    /// `SADD key member [member ...]`.
216    pub fn sadd(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
217        self.snap(key);
218        let added = self.inner.store.sadd(key, members).map_err(store_err)?;
219        if added > 0 {
220            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
221            argv.push(b"SADD");
222            argv.push(key);
223            argv.extend_from_slice(members);
224            self.log_arg(&argv);
225        }
226        Ok(added)
227    }
228
229    /// `SREM key member [member ...]`.
230    pub fn srem(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
231        self.snap(key);
232        let removed = self.inner.store.srem(key, members).map_err(store_err)?;
233        if removed > 0 {
234            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
235            argv.push(b"SREM");
236            argv.push(key);
237            argv.extend_from_slice(members);
238            self.log_arg(&argv);
239        }
240        Ok(removed)
241    }
242
243    // ---- list ops --------------------------------------------------
244
245    /// `LPUSH key value [value ...]` — returns the new list length.
246    pub fn lpush(&mut self, key: &[u8], values: &[&[u8]]) -> KevyResult<usize> {
247        self.snap(key);
248        let len = self.inner.store.lpush(key, values).map_err(store_err)?;
249        let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + values.len());
250        argv.push(b"LPUSH");
251        argv.push(key);
252        argv.extend_from_slice(values);
253        self.log_arg(&argv);
254        Ok(len)
255    }
256
257    /// `RPUSH key value [value ...]` — returns the new list length.
258    pub fn rpush(&mut self, key: &[u8], values: &[&[u8]]) -> KevyResult<usize> {
259        self.snap(key);
260        let len = self.inner.store.rpush(key, values).map_err(store_err)?;
261        let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + values.len());
262        argv.push(b"RPUSH");
263        argv.push(key);
264        argv.extend_from_slice(values);
265        self.log_arg(&argv);
266        Ok(len)
267    }
268
269    // ---- zset ops --------------------------------------------------
270
271    /// `ZREM key member [member ...]`.
272    pub fn zrem(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
273        self.snap(key);
274        let removed = self.inner.store.zrem(key, members).map_err(store_err)?;
275        if removed > 0 {
276            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
277            argv.push(b"ZREM");
278            argv.push(key);
279            argv.extend_from_slice(members);
280            self.log_arg(&argv);
281        }
282        Ok(removed)
283    }
284
285    /// `ZCARD key` — member count; 0 when absent.
286    pub fn zcard(&mut self, key: &[u8]) -> KevyResult<usize> {
287        self.inner.store.zcard(key).map_err(store_err)
288    }
289
290    /// Flags-aware `ZADD`. AOF logs the applied pairs as plain
291    /// `ZADD` — the effect, never the condition (deterministic replay).
292    pub fn zadd_flags(
293        &mut self,
294        key: &[u8],
295        pairs: &[(f64, &[u8])],
296        flags: kevy_store::ZaddFlags,
297    ) -> KevyResult<kevy_store::ZaddReport> {
298        if !flags.valid() {
299            return Err(KevyError::InvalidInput("invalid ZADD flag combo".into()));
300        }
301        let rep = self
302            .inner
303            .store
304            .zadd_flags(key, pairs, flags)
305            .map_err(store_err)?;
306        if !rep.applied.is_empty() {
307            let score_strs: Vec<Vec<u8>> = rep
308                .applied
309                .iter()
310                .map(|(s, _)| format!("{s}").into_bytes())
311                .collect();
312            let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + rep.applied.len() * 2);
313            parts.push(b"ZADD");
314            parts.push(key);
315            for (i, (_, m)) in rep.applied.iter().enumerate() {
316                parts.push(&score_strs[i]);
317                parts.push(m);
318            }
319            self.log_arg(&parts);
320        }
321        Ok(rep)
322    }
323
324
325    // ---- collection reads --------------------------------------------
326    // Requested by a consumer: a set could be written inside a transaction but never read back
327    // inside one, so any child collection a cascade delete must
328    // enumerate had to be modelled as a hash — they reshaped a whole
329    // keyspace around the omission. These hold the shard write lock
330    // already, so there was never a consistency reason to withhold them.
331
332    /// `SMEMBERS key`.
333    pub fn smembers(&mut self, key: &[u8]) -> KevyResult<Vec<Vec<u8>>> {
334        self.inner.store.smembers(key).map_err(store_err)
335    }
336
337    /// `SISMEMBER key member`.
338    pub fn sismember(&mut self, key: &[u8], member: &[u8]) -> KevyResult<bool> {
339        self.inner.store.sismember(key, member).map_err(store_err)
340    }
341
342    /// `LRANGE key start stop` (inclusive, negatives count from the end).
343    pub fn lrange(&mut self, key: &[u8], start: i64, stop: i64) -> KevyResult<Vec<Vec<u8>>> {
344        self.inner.store.lrange(key, start, stop).map_err(store_err)
345    }
346
347    /// `LLEN key`.
348    pub fn llen(&mut self, key: &[u8]) -> KevyResult<usize> {
349        self.inner.store.llen(key).map_err(store_err)
350    }
351
352    /// `SCARD key`.
353    pub fn scard(&mut self, key: &[u8]) -> KevyResult<usize> {
354        self.inner.store.scard(key).map_err(store_err)
355    }
356
357    /// `ZRANGEBYSCORE key min max` — `(member, score)` in score order.
358    pub fn zrangebyscore(
359        &mut self,
360        key: &[u8],
361        min: kevy_store::ScoreBound,
362        max: kevy_store::ScoreBound,
363    ) -> KevyResult<Vec<(Vec<u8>, f64)>> {
364        self.inner.store.zrange_by_score(key, min, max).map_err(store_err)
365    }
366
367    /// Record `key`'s prior state, once, before its first mutation.
368    fn snap(&mut self, key: &[u8]) {
369        if self.touched.contains(key) {
370            return;
371        }
372        let prior = self.inner.store.clone_with_ttl(key);
373        self.touched.insert(key.to_vec());
374        self.undo.push((key.to_vec(), prior));
375    }
376
377    fn log_arg(&mut self, parts: &[&[u8]]) {
378        self.log.push(parts.iter().map(|p| p.to_vec()).collect());
379    }
380}
381
382impl Store {
383    /// Run `body` as a single-shard atomic transaction: it applies
384    /// entirely, or not at all.
385    ///
386    /// Inside the closure every read sees the closure's own previous
387    /// writes. On `Ok`, the queued AOF frames are committed as one
388    /// group — under `Fsync::Always` that is a single fsync for the
389    /// whole block, not one per mutation.
390    ///
391    /// On `Err`, **every write the closure made is rolled back** and
392    /// nothing is appended to the AOF. This is what lets the closure
393    /// act as the enforcement point for an invariant: read, decide,
394    /// write, and return `Err` to reject — the rejection leaves no
395    /// trace. (Before 4.0 the writes stayed live in memory while their
396    /// AOF frames were discarded, so a restarted process disagreed
397    /// with the running one.)
398    ///
399    /// Rollback restores each touched key to the value and TTL it had
400    /// before the transaction — including deleting keys the closure
401    /// created. It is a snapshot of the keys the closure touches, so
402    /// the cost scales with the transaction, not the keyspace.
403    ///
404    /// Constraint: every key touched inside the closure must hash to
405    /// the same shard. The default embedded config uses 1 shard, so
406    /// any key works.
407    pub fn atomic<R>(
408        &self,
409        body: impl FnOnce(&mut AtomicCtx<'_>) -> KevyResult<R>,
410    ) -> KevyResult<R> {
411        ensure_writable(self)?;
412        let mut g: RwLockWriteGuard<'_, Inner> = self.lock();
413        let mut ctx = AtomicCtx {
414            inner: &mut g,
415            log: Vec::new(),
416            undo: Vec::new(),
417            touched: std::collections::HashSet::new(),
418        };
419        let outcome = body(&mut ctx);
420        let log = std::mem::take(&mut ctx.log);
421        let undo = std::mem::take(&mut ctx.undo);
422        let r = match outcome {
423            Ok(r) => r,
424            Err(e) => {
425                rollback(&mut g, undo);
426                return Err(e);
427            }
428        };
429        commit_group(&mut g, log)?;
430        Ok(r)
431    }
432}
433
434/// Parity manifest: command names `AtomicCtx` implements.
435#[cfg_attr(not(test), allow(dead_code))]
436pub(crate) const ATOMIC_OPS: &[&str] = &[
437    "SET", "GET", "INCR", "INCRBY", "HSET", "HGET", "HINCRBY", "ZADD",
438    "ZINCRBY", "ZSCORE", "DEL", "EXISTS", "HDEL", "HGETALL", "HMGET",
439    "HEXISTS", "SADD", "SREM", "LPUSH", "RPUSH", "ZREM", "ZCARD",
440    "SMEMBERS", "SISMEMBER", "LRANGE", "LLEN", "SCARD", "ZRANGEBYSCORE",
441];
442
443/// Undo a rejected transaction.
444///
445/// The closure's writes hit the store as they were made — reads inside
446/// the block have to see them — so a rejected transaction must be undone
447/// here, or the rejected write stays live while its AOF frames are
448/// discarded and a restart disagrees with the running process. Reverse
449/// order so a key touched more than once lands on its earliest recorded
450/// state.
451fn rollback(g: &mut Inner, undo: Vec<UndoEntry>) {
452    for (key, prior) in undo.into_iter().rev() {
453        match prior {
454            Some((value, ttl_ms)) => g.store.put_with_ttl(key, value, ttl_ms),
455            None => {
456                let k: &[u8] = &key;
457                g.store.del(&[k]);
458            }
459        }
460    }
461}
462
463/// Commit the queued AOF frames as ONE bracketed group.
464///
465/// The brackets are what make replay all-or-nothing at any size, and the
466/// group is what makes `Fsync::Always` cost one sync instead of N. See
467/// `kevy_persist::Aof::begin_group`.
468fn commit_group(g: &mut Inner, log: Vec<Vec<Vec<u8>>>) -> KevyResult<()> {
469    #[cfg(feature = "persist")]
470    if let Some(aof) = g.aof.as_mut() {
471        aof.begin_group();
472    }
473    let mut commit = Ok(());
474    for entry in log {
475        let parts: Vec<&[u8]> = entry.iter().map(|v| v.as_slice()).collect();
476        commit = commit_write(g, &parts);
477        if commit.is_err() {
478            break;
479        }
480    }
481    #[cfg(feature = "persist")]
482    if let Some(aof) = g.aof.as_mut() {
483        let synced = aof.end_group().map_err(KevyError::from);
484        commit = commit.and(synced);
485    }
486    commit
487}
kevy_embedded/ops_atomic.rs

kevy_embedded/
ops_atomic.rs
