Skip to main content

kevy_store/stream/
load.rs

1//! Consumer-group exchange types + the `XSETID` scalar setter — the
2//! pieces persistence (snapshot v4, AOF rewrite, reshard's `load_value`
3//! redistribution) needs to carry group/PEL state across a dump/load
4//! boundary. Split from `stream/mod.rs` to stay under the 500-LOC cap.
5
6#[cfg(not(feature = "std"))]
7use crate::nostd_prelude::*;
8use alloc::collections::BTreeMap;
9
10use kevy_map::KevyMap;
11
12use super::group::{ConsumerGroup, ConsumerState, PelEntry};
13use super::{StreamData, StreamId};
14use crate::StoreError;
15use crate::value::SmallBytes;
16
17/// One PEL row in primitive form: `(ms, seq, consumer, delivery_time_ms,
18/// delivery_count)`. The persist crate serializes these verbatim.
19pub type LoadedPelEntry = (u64, u64, Vec<u8>, u64, u32);
20
21/// One consumer group decoded into primitive tuples — the dump/load wire
22/// form shared by snapshot v4, AOF-rewrite filtering, and reshard's
23/// in-memory redistribution.
24#[derive(Debug)]
25pub struct LoadedGroup {
26    /// Group name.
27    pub name: Vec<u8>,
28    /// `last_delivered_id` as `(ms, seq)`.
29    pub last_delivered: (u64, u64),
30    /// `(name, last_seen_ms)` per known consumer. `pel_count` is
31    /// recomputed from `pel` on import.
32    pub consumers: Vec<(Vec<u8>, u64)>,
33    /// Every PEL row, including tombstones (entries XDEL'd while
34    /// pending) — snapshot keeps those; AOF rewrite filters them.
35    pub pel: Vec<LoadedPelEntry>,
36}
37
38impl StreamData {
39    /// Does an entry with `id` currently exist? AOF rewrite uses this
40    /// to filter tombstone PEL rows (XCLAIM can't re-create those).
41    pub fn contains_entry(&self, id: StreamId) -> bool {
42        self.entries.contains_key(&id)
43    }
44
45    /// Dump every group into the primitive exchange form.
46    pub fn export_groups(&self) -> Vec<LoadedGroup> {
47        self.groups
48            .iter()
49            .map(|(name, g)| LoadedGroup {
50                name: name.to_vec(),
51                last_delivered: (g.last_delivered_id.ms, g.last_delivered_id.seq),
52                consumers: g
53                    .consumers
54                    .iter()
55                    .map(|(c, cs)| (c.to_vec(), cs.last_seen_ms))
56                    .collect(),
57                pel: g
58                    .pel
59                    .iter()
60                    .map(|(id, p)| {
61                        (id.ms, id.seq, p.consumer.to_vec(), p.delivery_time_ms, p.delivery_count)
62                    })
63                    .collect(),
64            })
65            .collect()
66    }
67
68    /// Rebuild the group map from the exchange form (loader-side twin of
69    /// [`Self::export_groups`]). Per-consumer `pel_count` is recomputed;
70    /// a PEL owner missing from the consumer roster (hand-built or
71    /// corrupt file) gets a roster slot rather than a panic.
72    pub fn import_groups(&mut self, groups: Vec<LoadedGroup>) {
73        for lg in groups {
74            let mut consumers: KevyMap<SmallBytes, Box<ConsumerState>> = KevyMap::default();
75            for (name, last_seen_ms) in lg.consumers {
76                let name = SmallBytes::from_vec(name);
77                consumers.insert(
78                    name.clone(),
79                    Box::new(ConsumerState { name, last_seen_ms, pel_count: 0 }),
80                );
81            }
82            let mut pel: BTreeMap<StreamId, PelEntry> = BTreeMap::new();
83            for (ms, seq, consumer, delivery_time_ms, delivery_count) in lg.pel {
84                let consumer = SmallBytes::from_vec(consumer);
85                if consumers.get(consumer.as_slice()).is_none() {
86                    consumers.insert(
87                        consumer.clone(),
88                        Box::new(ConsumerState {
89                            name: consumer.clone(),
90                            last_seen_ms: 0,
91                            pel_count: 0,
92                        }),
93                    );
94                }
95                if let Some(cs) = consumers.get_mut(consumer.as_slice()) {
96                    cs.pel_count += 1;
97                }
98                pel.insert(
99                    StreamId { ms, seq },
100                    PelEntry { consumer, delivery_time_ms, delivery_count },
101                );
102            }
103            self.groups.insert(
104                SmallBytes::from_vec(lg.name),
105                Box::new(ConsumerGroup {
106                    last_delivered_id: StreamId {
107                        ms: lg.last_delivered.0,
108                        seq: lg.last_delivered.1,
109                    },
110                    pel,
111                    consumers,
112                }),
113            );
114        }
115    }
116
117    /// `XSETID key last-id [ENTRIESADDED n] [MAXDELETEDID id]` — overwrite
118    /// the stream's scalar state. Rejects a `last_id` below the current
119    /// top entry (Redis: "smaller than the target stream top item").
120    pub fn xsetid(
121        &mut self,
122        last_id: StreamId,
123        entries_added: Option<u64>,
124        max_deleted_id: Option<StreamId>,
125    ) -> Result<(), StoreError> {
126        if let Some((top, _)) = self.entries.iter().next_back()
127            && last_id < *top
128        {
129            return Err(StoreError::OutOfRange);
130        }
131        self.last_id = last_id;
132        if let Some(n) = entries_added {
133            self.entries_added = n;
134        }
135        if let Some(id) = max_deleted_id {
136            self.max_deleted_id = id;
137        }
138        Ok(())
139    }
140}