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