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 std::io;
14use std::sync::RwLockWriteGuard;
15
16#[cfg(not(target_arch = "wasm32"))]
17use crate::replica_glue::ensure_writable;
18use crate::store::{Inner, Store, commit_write, store_err};
19
20#[cfg(target_arch = "wasm32")]
21fn ensure_writable(_s: &Store) -> io::Result<()> { Ok(()) }
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}
30
31impl AtomicCtx<'_> {
32    // ---- string ops ------------------------------------------------
33
34    /// `SET key value`. Returns `true` (SET always succeeds without
35    /// `NX`/`XX` veto).
36    pub fn set(&mut self, key: &[u8], value: &[u8]) -> bool {
37        let ok = self
38            .inner
39            .store
40            .set(key, value.to_vec(), None, false, false);
41        self.log_arg(&[b"SET", key, value]);
42        ok
43    }
44
45    /// `GET key`.
46    pub fn get(&mut self, key: &[u8]) -> io::Result<Option<Vec<u8>>> {
47        self.inner
48            .store
49            .get(key)
50            .map(|opt| opt.as_deref().map(<[u8]>::to_vec))
51            .map_err(store_err)
52    }
53
54    /// `INCR key` — by 1.
55    pub fn incr(&mut self, key: &[u8]) -> io::Result<i64> {
56        let n = self.inner.store.incr_by(key, 1).map_err(store_err)?;
57        self.log_arg(&[b"INCR", key]);
58        Ok(n)
59    }
60
61    /// `INCRBY key delta`.
62    pub fn incr_by(&mut self, key: &[u8], delta: i64) -> io::Result<i64> {
63        let n = self.inner.store.incr_by(key, delta).map_err(store_err)?;
64        let s = format!("{delta}");
65        self.log_arg(&[b"INCRBY", key, s.as_bytes()]);
66        Ok(n)
67    }
68
69    // ---- hash ops ---------------------------------------------------
70
71    /// `HSET key field value`.
72    pub fn hset(&mut self, key: &[u8], pairs: &[(&[u8], &[u8])]) -> io::Result<usize> {
73        let owned: Vec<(Vec<u8>, Vec<u8>)> = pairs
74            .iter()
75            .map(|(f, v)| (f.to_vec(), v.to_vec()))
76            .collect();
77        let n = self.inner.store.hset(key, &owned).map_err(store_err)?;
78        let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
79        parts.push(b"HSET");
80        parts.push(key);
81        for (f, v) in pairs {
82            parts.push(f);
83            parts.push(v);
84        }
85        self.log_arg(&parts);
86        Ok(n)
87    }
88
89    /// `HGET key field`.
90    pub fn hget(&mut self, key: &[u8], field: &[u8]) -> io::Result<Option<Vec<u8>>> {
91        Ok(self
92            .inner
93            .store
94            .hget(key, field)
95            .map_err(store_err)?
96            .map(<[u8]>::to_vec))
97    }
98
99    /// `HINCRBY key field delta`.
100    pub fn hincrby(&mut self, key: &[u8], field: &[u8], delta: i64) -> io::Result<i64> {
101        let n = self.inner.store.hincrby(key, field, delta).map_err(store_err)?;
102        let s = format!("{delta}");
103        self.log_arg(&[b"HINCRBY", key, field, s.as_bytes()]);
104        Ok(n)
105    }
106
107    // ---- zset ops ---------------------------------------------------
108
109    /// `ZADD key score member`.
110    pub fn zadd(&mut self, key: &[u8], pairs: &[(f64, &[u8])]) -> io::Result<usize> {
111        let owned: Vec<(f64, Vec<u8>)> =
112            pairs.iter().map(|(s, m)| (*s, m.to_vec())).collect();
113        let n = self.inner.store.zadd(key, &owned).map_err(store_err)?;
114        let score_strs: Vec<Vec<u8>> =
115            pairs.iter().map(|(s, _)| format!("{s}").into_bytes()).collect();
116        let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
117        parts.push(b"ZADD");
118        parts.push(key);
119        for (i, (_, m)) in pairs.iter().enumerate() {
120            parts.push(&score_strs[i]);
121            parts.push(m);
122        }
123        self.log_arg(&parts);
124        Ok(n)
125    }
126
127    /// `ZINCRBY key delta member`.
128    pub fn zincrby(&mut self, key: &[u8], delta: f64, member: &[u8]) -> io::Result<f64> {
129        let n = self.inner.store.zincrby(key, delta, member).map_err(store_err)?;
130        let s = format!("{delta}");
131        self.log_arg(&[b"ZINCRBY", key, s.as_bytes(), member]);
132        Ok(n)
133    }
134
135    /// `ZSCORE key member`.
136    pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> io::Result<Option<f64>> {
137        self.inner.store.zscore(key, member).map_err(store_err)
138    }
139
140    // ---- helpers ----------------------------------------------------
141
142    // ---- keyspace ops (v2.1 — Pipeline write parity) ---------------
143
144    /// `DEL key [key ...]` — every key must hash to this shard.
145    pub fn del(&mut self, keys: &[&[u8]]) -> usize {
146        let n = self.inner.store.del_borrowed(keys);
147        if n > 0 {
148            let mut argv: Vec<&[u8]> = Vec::with_capacity(1 + keys.len());
149            argv.push(b"DEL");
150            argv.extend_from_slice(keys);
151            self.log_arg(&argv);
152        }
153        n
154    }
155
156    /// `EXISTS key [key ...]` — count of the given keys that exist.
157    pub fn exists(&mut self, keys: &[&[u8]]) -> usize {
158        keys.iter()
159            .filter(|k| self.inner.store.key_exists(k))
160            .count()
161    }
162
163    // ---- hash ops --------------------------------------------------
164
165    /// `HDEL key field [field ...]`.
166    pub fn hdel(&mut self, key: &[u8], fields: &[&[u8]]) -> io::Result<usize> {
167        let owned: Vec<Vec<u8>> = fields.iter().map(|f| f.to_vec()).collect();
168        let removed = self.inner.store.hdel(key, &owned).map_err(store_err)?;
169        if removed > 0 {
170            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + fields.len());
171            argv.push(b"HDEL");
172            argv.push(key);
173            argv.extend_from_slice(fields);
174            self.log_arg(&argv);
175        }
176        Ok(removed)
177    }
178
179    /// `HGETALL key` — `(field, value)` pairs; reads see the
180    /// closure's own prior writes.
181    pub fn hgetall(&mut self, key: &[u8]) -> io::Result<Vec<(Vec<u8>, Vec<u8>)>> {
182        let flat = self.inner.store.hgetall(key).map_err(store_err)?;
183        let mut out = Vec::with_capacity(flat.len() / 2);
184        let mut it = flat.into_iter();
185        while let (Some(f), Some(v)) = (it.next(), it.next()) {
186            out.push((f, v));
187        }
188        Ok(out)
189    }
190
191    /// `HMGET key field [field ...]` — `None` per absent field.
192    pub fn hmget(&mut self, key: &[u8], fields: &[&[u8]]) -> io::Result<Vec<Option<Vec<u8>>>> {
193        self.inner.store.hmget_borrowed(key, fields).map_err(store_err)
194    }
195
196    /// `HEXISTS key field`.
197    pub fn hexists(&mut self, key: &[u8], field: &[u8]) -> io::Result<bool> {
198        self.inner.store.hexists(key, field).map_err(store_err)
199    }
200
201    // ---- set ops ---------------------------------------------------
202
203    /// `SADD key member [member ...]`.
204    pub fn sadd(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
205        let owned: Vec<Vec<u8>> = members.iter().map(|m| m.to_vec()).collect();
206        let added = self.inner.store.sadd(key, &owned).map_err(store_err)?;
207        if added > 0 {
208            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
209            argv.push(b"SADD");
210            argv.push(key);
211            argv.extend_from_slice(members);
212            self.log_arg(&argv);
213        }
214        Ok(added)
215    }
216
217    /// `SREM key member [member ...]`.
218    pub fn srem(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
219        let owned: Vec<Vec<u8>> = members.iter().map(|m| m.to_vec()).collect();
220        let removed = self.inner.store.srem(key, &owned).map_err(store_err)?;
221        if removed > 0 {
222            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
223            argv.push(b"SREM");
224            argv.push(key);
225            argv.extend_from_slice(members);
226            self.log_arg(&argv);
227        }
228        Ok(removed)
229    }
230
231    // ---- list ops --------------------------------------------------
232
233    /// `LPUSH key value [value ...]` — returns the new list length.
234    pub fn lpush(&mut self, key: &[u8], values: &[&[u8]]) -> io::Result<usize> {
235        let owned: Vec<Vec<u8>> = values.iter().map(|v| v.to_vec()).collect();
236        let len = self.inner.store.lpush(key, &owned).map_err(store_err)?;
237        let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + values.len());
238        argv.push(b"LPUSH");
239        argv.push(key);
240        argv.extend_from_slice(values);
241        self.log_arg(&argv);
242        Ok(len)
243    }
244
245    /// `RPUSH key value [value ...]` — returns the new list length.
246    pub fn rpush(&mut self, key: &[u8], values: &[&[u8]]) -> io::Result<usize> {
247        let owned: Vec<Vec<u8>> = values.iter().map(|v| v.to_vec()).collect();
248        let len = self.inner.store.rpush(key, &owned).map_err(store_err)?;
249        let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + values.len());
250        argv.push(b"RPUSH");
251        argv.push(key);
252        argv.extend_from_slice(values);
253        self.log_arg(&argv);
254        Ok(len)
255    }
256
257    // ---- zset ops --------------------------------------------------
258
259    /// `ZREM key member [member ...]`.
260    pub fn zrem(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
261        let owned: Vec<Vec<u8>> = members.iter().map(|m| m.to_vec()).collect();
262        let removed = self.inner.store.zrem(key, &owned).map_err(store_err)?;
263        if removed > 0 {
264            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
265            argv.push(b"ZREM");
266            argv.push(key);
267            argv.extend_from_slice(members);
268            self.log_arg(&argv);
269        }
270        Ok(removed)
271    }
272
273    /// `ZCARD key` — member count; 0 when absent.
274    pub fn zcard(&mut self, key: &[u8]) -> io::Result<usize> {
275        self.inner.store.zcard(key).map_err(store_err)
276    }
277
278    /// Flags-aware `ZADD` (v2.1). AOF logs the applied pairs as plain
279    /// `ZADD` — the effect, never the condition (deterministic replay).
280    pub fn zadd_flags(
281        &mut self,
282        key: &[u8],
283        pairs: &[(f64, &[u8])],
284        flags: kevy_store::ZaddFlags,
285    ) -> io::Result<kevy_store::ZaddReport> {
286        if !flags.valid() {
287            return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid ZADD flag combo"));
288        }
289        let rep = self
290            .inner
291            .store
292            .zadd_flags_borrowed(key, pairs, flags)
293            .map_err(store_err)?;
294        if !rep.applied.is_empty() {
295            let score_strs: Vec<Vec<u8>> = rep
296                .applied
297                .iter()
298                .map(|(s, _)| format!("{s}").into_bytes())
299                .collect();
300            let mut parts: Vec<&[u8]> = Vec::with_capacity(2 + rep.applied.len() * 2);
301            parts.push(b"ZADD");
302            parts.push(key);
303            for (i, (_, m)) in rep.applied.iter().enumerate() {
304                parts.push(&score_strs[i]);
305                parts.push(m);
306            }
307            self.log_arg(&parts);
308        }
309        Ok(rep)
310    }
311
312    fn log_arg(&mut self, parts: &[&[u8]]) {
313        self.log.push(parts.iter().map(|p| p.to_vec()).collect());
314    }
315}
316
317impl Store {
318    /// Run `body` as a single-shard atomic transaction. Inside the
319    /// closure every read sees previous writes; on closure return
320    /// the queued AOF writes are committed under one fsync.
321    ///
322    /// Constraint: every key touched inside the closure must hash to
323    /// the same shard. The default embedded config uses 1 shard, so
324    /// any key works.
325    pub fn atomic<R>(
326        &self,
327        body: impl FnOnce(&mut AtomicCtx<'_>) -> io::Result<R>,
328    ) -> io::Result<R> {
329        ensure_writable(self)?;
330        let mut g: RwLockWriteGuard<'_, Inner> = self.lock();
331        let mut ctx = AtomicCtx { inner: &mut g, log: Vec::new() };
332        let r = body(&mut ctx)?;
333        // Commit queued AOF writes — one append per op, one fsync at
334        // the end via `commit_write`'s standard path.
335        let log = std::mem::take(&mut ctx.log);
336        for entry in log {
337            let parts: Vec<&[u8]> = entry.iter().map(|v| v.as_slice()).collect();
338            commit_write(&mut g, &parts)?;
339        }
340        Ok(r)
341    }
342}
343
344/// Parity manifest (v2.1): command names `AtomicCtx` implements.
345#[cfg_attr(not(test), allow(dead_code))]
346pub(crate) const ATOMIC_OPS: &[&str] = &[
347    "SET", "GET", "INCR", "INCRBY", "HSET", "HGET", "HINCRBY", "ZADD",
348    "ZINCRBY", "ZSCORE", "DEL", "EXISTS", "HDEL", "HGETALL", "HMGET",
349    "HEXISTS", "SADD", "SREM", "LPUSH", "RPUSH", "ZREM", "ZCARD",
350];
kevy_embedded/ops_atomic.rs

kevy_embedded/
ops_atomic.rs
