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

kevy_embedded/
ops_atomic_all.rs
