Skip to main content

kevy_store/stream/
store.rs

1//! `Store::xadd` / `xlen` / `xrange` / `xrevrange` / `xread` / `xdel` /
2//! `xtrim_*` — the per-keyspace surface for sprint A of v2-7 Streams.
3//! Kept separate from `stream/mod.rs` (which owns the `StreamData` /
4//! `StreamId` types + entry-side ops) so each file stays under the
5//! project's ≤500-LOC rule.
6
7use super::group::{AutoclaimResult, ReadGroupId};
8use super::{
9    GroupCreateMode, PendingExtended, PendingSummary, StreamData, StreamId, XAddIdSpec, XClaimOpts,
10};
11#[cfg(not(feature = "std"))]
12use crate::nostd_prelude::*;
13use crate::value::{SmallBytes, Value};
14use crate::{Entry, Store, StoreError};
15use alloc::sync::Arc;
16
17/// Cloned-out view of stream entries, the cross-module wire form. Keeps
18/// the same shape Redis sends and lets the callers stay decoupled from
19/// the `SmallBytes` interning the store uses internally.
20pub type EntryBatch = Vec<(StreamId, Vec<(Vec<u8>, Vec<u8>)>)>;
21
22impl Store {
23    fn stream_mut(
24        &mut self,
25        key: &[u8],
26        create: bool,
27    ) -> Result<Option<&mut StreamData>, StoreError> {
28        if self.live_entry_mut(key).is_none() {
29            if !create {
30                return Ok(None);
31            }
32            self.insert_entry(
33                SmallBytes::from_slice(key),
34                Entry::new(Value::Stream(Arc::default()), None),
35            );
36        }
37        match &mut self.map.get_mut(key).expect("present").value {
38            Value::Stream(s) => Ok(Some(Arc::make_mut(s))),
39            _ => Err(StoreError::WrongType),
40        }
41    }
42
43    fn stream_ref(&mut self, key: &[u8]) -> Result<Option<&StreamData>, StoreError> {
44        match self.live_entry(key) {
45            None => Ok(None),
46            Some(e) => match &e.value {
47                Value::Stream(s) => Ok(Some(s.as_ref())),
48                _ => Err(StoreError::WrongType),
49            },
50        }
51    }
52
53    /// Read-only access to a stream's `StreamData`, used by `XINFO`
54    /// to inspect entries / groups / consumers without going through
55    /// the wrapper layer. Returns `Ok(None)` for a missing key,
56    /// `WrongType` for a non-stream value at `key`.
57    pub fn stream_view(&mut self, key: &[u8]) -> Result<Option<&StreamData>, StoreError> {
58        self.stream_ref(key)
59    }
60
61    /// `XADD key <spec> field value [field value ...]`. Returns the
62    /// assigned ID. `nomkstream` matches Redis's `NOMKSTREAM` flag —
63    /// suppress key creation, returning `Ok(None)`. `now_ms` is the
64    /// wall-clock used for `XAddIdSpec::AutoAll`.
65    pub fn xadd(
66        &mut self,
67        key: &[u8],
68        spec: XAddIdSpec,
69        fields: Vec<(Vec<u8>, Vec<u8>)>,
70        nomkstream: bool,
71        now_ms: u64,
72    ) -> Result<Option<StreamId>, StoreError> {
73        if nomkstream && self.live_entry(key).is_none() {
74            return Ok(None);
75        }
76        let id;
77        let weight_delta;
78        {
79            let s = self.stream_mut(key, true)?.expect("created");
80            id = s.resolve_xadd_id(spec, now_ms)?;
81            let smb_fields: Vec<(SmallBytes, SmallBytes)> = fields
82                .into_iter()
83                .map(|(f, v)| (SmallBytes::from_slice(&f), SmallBytes::from_slice(&v)))
84                .collect();
85            weight_delta = super::stream_entry_weight(&smb_fields);
86            s.insert(id, smb_fields);
87        }
88        self.bump_if_watched(key);
89        self.account_delta(key, weight_delta as i64);
90        Ok(Some(id))
91    }
92
93    /// `XLEN key`. Returns 0 for a missing key.
94    pub fn xlen(&mut self, key: &[u8]) -> Result<u64, StoreError> {
95        Ok(self.stream_ref(key)?.map_or(0, super::StreamData::length))
96    }
97
98    /// `XRANGE key start end [COUNT n]`.
99    pub fn xrange(
100        &mut self,
101        key: &[u8],
102        start: StreamId,
103        end: StreamId,
104        count: Option<usize>,
105    ) -> Result<EntryBatch, StoreError> {
106        Ok(self
107            .stream_ref(key)?
108            .map_or_else(Vec::new, |s| super::clone_entries(s.range(start, end, count))))
109    }
110
111    /// `XREVRANGE key end start [COUNT n]`.
112    pub fn xrevrange(
113        &mut self,
114        key: &[u8],
115        start: StreamId,
116        end: StreamId,
117        count: Option<usize>,
118    ) -> Result<EntryBatch, StoreError> {
119        Ok(self
120            .stream_ref(key)?
121            .map_or_else(Vec::new, |s| super::clone_entries(s.revrange(start, end, count))))
122    }
123
124    /// `XREAD ... STREAMS key last_seen [...]` — per-key part.
125    pub fn xread(
126        &mut self,
127        key: &[u8],
128        last_seen: StreamId,
129        count: Option<usize>,
130    ) -> Result<EntryBatch, StoreError> {
131        Ok(self
132            .stream_ref(key)?
133            .map_or_else(Vec::new, |s| super::clone_entries(s.read_after(last_seen, count))))
134    }
135
136    /// Resolve `$` as XREAD's "last-seen" to the stream's current last
137    /// ID. Returns `MIN` for a missing key.
138    pub fn xread_dollar_last_id(&mut self, key: &[u8]) -> Result<StreamId, StoreError> {
139        Ok(self.stream_ref(key)?.map_or(StreamId::MIN, super::StreamData::last_id))
140    }
141
142    /// `XDEL key id [...]`. Returns count actually removed.
143    pub fn xdel(&mut self, key: &[u8], ids: &[StreamId]) -> Result<u64, StoreError> {
144        let n;
145        {
146            let Some(s) = self.stream_mut(key, false)? else {
147                return Ok(0);
148            };
149            n = s.del_ids(ids);
150        }
151        if n > 0 {
152            self.bump_if_watched(key);
153            self.reweigh_entry(key);
154        }
155        Ok(n as u64)
156    }
157
158    /// `XTRIM key MAXLEN n`. Returns number removed.
159    pub fn xtrim_maxlen(&mut self, key: &[u8], maxlen: u64) -> Result<u64, StoreError> {
160        let n;
161        {
162            let Some(s) = self.stream_mut(key, false)? else {
163                return Ok(0);
164            };
165            n = s.trim_maxlen(maxlen as usize);
166        }
167        if n > 0 {
168            self.bump_if_watched(key);
169            self.reweigh_entry(key);
170        }
171        Ok(n as u64)
172    }
173
174    /// `XTRIM key MINID id`. Returns number removed.
175    pub fn xtrim_minid(&mut self, key: &[u8], minid: StreamId) -> Result<u64, StoreError> {
176        let n;
177        {
178            let Some(s) = self.stream_mut(key, false)? else {
179                return Ok(0);
180            };
181            n = s.trim_minid(minid);
182        }
183        if n > 0 {
184            self.bump_if_watched(key);
185            self.reweigh_entry(key);
186        }
187        Ok(n as u64)
188    }
189
190    /// `XSETID key last-id [ENTRIESADDED n] [MAXDELETEDID id]`. Returns
191    /// `NoSuchKey` for a missing key (dispatch maps it to Redis's
192    /// "requires the key to exist" wording), `OutOfRange` when `last_id`
193    /// is below the stream's top entry.
194    pub fn xsetid(
195        &mut self,
196        key: &[u8],
197        last_id: StreamId,
198        entries_added: Option<u64>,
199        max_deleted_id: Option<StreamId>,
200    ) -> Result<(), StoreError> {
201        {
202            let Some(s) = self.stream_mut(key, false)? else {
203                return Err(StoreError::NoSuchKey);
204            };
205            s.xsetid(last_id, entries_added, max_deleted_id)?;
206        }
207        self.bump_if_watched(key);
208        Ok(())
209    }
210
211    // ─────── consumer-group surface (sprint B) ───────
212
213    /// `XGROUP CREATE key group <id|$> [MKSTREAM]`. Returns `Ok(true)`
214    /// when a fresh group was added; `Ok(false)` if the group already
215    /// existed (caller emits `-BUSYGROUP`). `mkstream` matches Redis:
216    /// auto-create the stream key when missing.
217    pub fn xgroup_create(
218        &mut self,
219        key: &[u8],
220        group: &[u8],
221        mode: GroupCreateMode,
222        mkstream: bool,
223    ) -> Result<bool, StoreError> {
224        let exists = self.live_entry(key).is_some();
225        if !exists && !mkstream {
226            return Err(StoreError::NoSuchKey);
227        }
228        let s = self.stream_mut(key, true)?.expect("created");
229        let created = s.group_create(group, mode)?;
230        self.bump_if_watched(key);
231        self.reweigh_entry(key);
232        Ok(created)
233    }
234
235    /// `XGROUP DESTROY key group`. Returns `true` if a group was dropped.
236    pub fn xgroup_destroy(&mut self, key: &[u8], group: &[u8]) -> Result<bool, StoreError> {
237        let dropped;
238        {
239            let Some(s) = self.stream_mut(key, false)? else {
240                return Ok(false);
241            };
242            dropped = s.group_destroy(group);
243        }
244        if dropped {
245            self.bump_if_watched(key);
246            self.reweigh_entry(key);
247        }
248        Ok(dropped)
249    }
250
251    /// `XGROUP SETID key group <id|$>`.
252    pub fn xgroup_setid(
253        &mut self,
254        key: &[u8],
255        group: &[u8],
256        mode: GroupCreateMode,
257    ) -> Result<bool, StoreError> {
258        let touched;
259        {
260            let Some(s) = self.stream_mut(key, false)? else {
261                return Ok(false);
262            };
263            touched = s.group_setid(group, mode);
264        }
265        if touched {
266            self.bump_if_watched(key);
267        }
268        Ok(touched)
269    }
270
271    /// `XGROUP CREATECONSUMER key group consumer`.
272    pub fn xgroup_create_consumer(
273        &mut self,
274        key: &[u8],
275        group: &[u8],
276        consumer: &[u8],
277        now_ms: u64,
278    ) -> Result<bool, StoreError> {
279        let Some(s) = self.stream_mut(key, false)? else {
280            return Ok(false);
281        };
282        Ok(s.group_create_consumer(group, consumer, now_ms))
283    }
284
285    /// `XGROUP DELCONSUMER key group consumer`. Returns dropped PEL count.
286    pub fn xgroup_del_consumer(
287        &mut self,
288        key: &[u8],
289        group: &[u8],
290        consumer: &[u8],
291    ) -> Result<u64, StoreError> {
292        let Some(s) = self.stream_mut(key, false)? else {
293            return Ok(0);
294        };
295        Ok(s.group_del_consumer(group, consumer))
296    }
297
298    /// `XREADGROUP GROUP g c [COUNT n] [NOACK] STREAMS key id`.
299    #[allow(clippy::too_many_arguments)]
300    pub fn xreadgroup(
301        &mut self,
302        key: &[u8],
303        group: &[u8],
304        consumer: &[u8],
305        last_seen: ReadGroupId,
306        count: Option<usize>,
307        noack: bool,
308        now_ms: u64,
309    ) -> Result<EntryBatch, StoreError> {
310        let result;
311        {
312            let Some(s) = self.stream_mut(key, false)? else {
313                return Err(StoreError::NoSuchKey);
314            };
315            result = s.readgroup(group, consumer, last_seen, count, noack, now_ms)?;
316        }
317        if !result.is_empty() {
318            self.bump_if_watched(key);
319        }
320        Ok(result)
321    }
322
323    /// Non-destructive: would `XREADGROUP … STREAMS key >` yield new
324    /// entries for `group` right now? True iff the stream's last id is
325    /// past the group's last-delivered id. Used by the cross-shard BLOCK
326    /// arbiter's readiness peek — never advances the group cursor. False
327    /// for a missing key / group.
328    pub fn xreadgroup_has_new(&mut self, key: &[u8], group: &[u8]) -> Result<bool, StoreError> {
329        Ok(self
330            .stream_ref(key)?
331            .and_then(|s| s.group(group).map(|g| s.last_id() > g.last_delivered_id()))
332            .unwrap_or(false))
333    }
334
335    /// `XACK key group id [id ...]`. Returns count of PEL removals.
336    pub fn xack(&mut self, key: &[u8], group: &[u8], ids: &[StreamId]) -> Result<u64, StoreError> {
337        let n;
338        {
339            let Some(s) = self.stream_mut(key, false)? else {
340                return Ok(0);
341            };
342            n = s.ack(group, ids);
343        }
344        if n > 0 {
345            self.bump_if_watched(key);
346        }
347        Ok(n)
348    }
349
350    /// `XPENDING key group` — summary form.
351    pub fn xpending_summary(
352        &mut self,
353        key: &[u8],
354        group: &[u8],
355    ) -> Result<Option<PendingSummary>, StoreError> {
356        Ok(self.stream_ref(key)?.and_then(|s| s.pending_summary(group)))
357    }
358
359    /// `XPENDING key group [IDLE ms] start end count [consumer]` —
360    /// extended form.
361    #[allow(clippy::too_many_arguments)]
362    pub fn xpending_extended(
363        &mut self,
364        key: &[u8],
365        group: &[u8],
366        idle_min_ms: Option<u64>,
367        start: StreamId,
368        end: StreamId,
369        count: usize,
370        consumer_filter: Option<&[u8]>,
371        now_ms: u64,
372    ) -> Result<Option<PendingExtended>, StoreError> {
373        Ok(self.stream_ref(key)?.and_then(|s| {
374            s.pending_extended(group, idle_min_ms, start, end, count, consumer_filter, now_ms)
375        }))
376    }
377
378    /// `XCLAIM key group consumer min-idle-ms id [id ...] [...]`.
379    /// Returns the (id, field-value) pairs successfully claimed —
380    /// dispatcher trims to ID-only when `JUSTID` is set.
381    pub fn xclaim(
382        &mut self,
383        key: &[u8],
384        group: &[u8],
385        new_owner: &[u8],
386        ids: &[StreamId],
387        opts: &XClaimOpts,
388        now_ms: u64,
389    ) -> Result<EntryBatch, StoreError> {
390        let claimed;
391        let payloads;
392        {
393            let Some(s) = self.stream_mut(key, false)? else {
394                return Err(StoreError::NoSuchKey);
395            };
396            claimed = s.claim(group, new_owner, ids, opts, now_ms)?;
397            payloads = s.payloads_for(&claimed);
398        }
399        if !claimed.is_empty() {
400            self.bump_if_watched(key);
401        }
402        Ok(payloads)
403    }
404
405    /// `XAUTOCLAIM key group consumer min-idle-ms start [COUNT n]
406    /// [JUSTID]`. Returns the cursor + claimed payloads + deleted IDs.
407    #[allow(clippy::too_many_arguments)]
408    pub fn xautoclaim(
409        &mut self,
410        key: &[u8],
411        group: &[u8],
412        new_owner: &[u8],
413        min_idle_ms: u64,
414        start: StreamId,
415        count: usize,
416        justid: bool,
417        now_ms: u64,
418    ) -> Result<(StreamId, EntryBatch, Vec<StreamId>), StoreError> {
419        let payloads;
420        let next_cursor;
421        let deleted_ids;
422        {
423            let Some(s) = self.stream_mut(key, false)? else {
424                return Err(StoreError::NoSuchKey);
425            };
426            let AutoclaimResult { next_cursor: nc, claimed_ids, deleted_ids: di } =
427                s.autoclaim(group, new_owner, min_idle_ms, start, count, justid, now_ms)?;
428            payloads = s.payloads_for(&claimed_ids);
429            next_cursor = nc;
430            deleted_ids = di;
431        }
432        if !payloads.is_empty() {
433            self.bump_if_watched(key);
434        }
435        Ok((next_cursor, payloads, deleted_ids))
436    }
437}