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.
24pub struct LoadedGroup {
25    /// Group name.
26    pub name: Vec<u8>,
27    /// `last_delivered_id` as `(ms, seq)`.
28    pub last_delivered: (u64, u64),
29    /// `(name, last_seen_ms)` per known consumer. `pel_count` is
30    /// recomputed from `pel` on import.
31    pub consumers: Vec<(Vec<u8>, u64)>,
32    /// Every PEL row, including tombstones (entries XDEL'd while
33    /// pending) — snapshot keeps those; AOF rewrite filters them.
34    pub pel: Vec<LoadedPelEntry>,
35}
36
37impl StreamData {
38    /// Does an entry with `id` currently exist? AOF rewrite uses this
39    /// to filter tombstone PEL rows (XCLAIM can't re-create those).
40    pub fn contains_entry(&self, id: StreamId) -> bool {
41        self.entries.contains_key(&id)
42    }
43
44    /// Dump every group into the primitive exchange form.
45    pub fn export_groups(&self) -> Vec<LoadedGroup> {
46        self.groups
47            .iter()
48            .map(|(name, g)| LoadedGroup {
49                name: name.to_vec(),
50                last_delivered: (g.last_delivered_id.ms, g.last_delivered_id.seq),
51                consumers: g
52                    .consumers
53                    .iter()
54                    .map(|(c, cs)| (c.to_vec(), cs.last_seen_ms))
55                    .collect(),
56                pel: g
57                    .pel
58                    .iter()
59                    .map(|(id, p)| {
60                        (id.ms, id.seq, p.consumer.to_vec(), p.delivery_time_ms, p.delivery_count)
61                    })
62                    .collect(),
63            })
64            .collect()
65    }
66
67    /// Rebuild the group map from the exchange form (loader-side twin of
68    /// [`Self::export_groups`]). Per-consumer `pel_count` is recomputed;
69    /// a PEL owner missing from the consumer roster (hand-built or
70    /// corrupt file) gets a roster slot rather than a panic.
71    pub fn import_groups(&mut self, groups: Vec<LoadedGroup>) {
72        for lg in groups {
73            let mut consumers: KevyMap<SmallBytes, Box<ConsumerState>> = KevyMap::default();
74            for (name, last_seen_ms) in lg.consumers {
75                let name = SmallBytes::from_vec(name);
76                consumers.insert(
77                    name.clone(),
78                    Box::new(ConsumerState { name, last_seen_ms, pel_count: 0 }),
79                );
80            }
81            let mut pel: BTreeMap<StreamId, PelEntry> = BTreeMap::new();
82            for (ms, seq, consumer, delivery_time_ms, delivery_count) in lg.pel {
83                let consumer = SmallBytes::from_vec(consumer);
84                if consumers.get(consumer.as_slice()).is_none() {
85                    consumers.insert(
86                        consumer.clone(),
87                        Box::new(ConsumerState {
88                            name: consumer.clone(),
89                            last_seen_ms: 0,
90                            pel_count: 0,
91                        }),
92                    );
93                }
94                if let Some(cs) = consumers.get_mut(consumer.as_slice()) {
95                    cs.pel_count += 1;
96                }
97                pel.insert(
98                    StreamId { ms, seq },
99                    PelEntry { consumer, delivery_time_ms, delivery_count },
100                );
101            }
102            self.groups.insert(
103                SmallBytes::from_vec(lg.name),
104                Box::new(ConsumerGroup {
105                    last_delivered_id: StreamId {
106                        ms: lg.last_delivered.0,
107                        seq: lg.last_delivered.1,
108                    },
109                    pel,
110                    consumers,
111                }),
112            );
113        }
114    }
115
116    /// `XSETID key last-id [ENTRIESADDED n] [MAXDELETEDID id]` — overwrite
117    /// the stream's scalar state. Rejects a `last_id` below the current
118    /// top entry (Redis: "smaller than the target stream top item").
119    pub fn xsetid(
120        &mut self,
121        last_id: StreamId,
122        entries_added: Option<u64>,
123        max_deleted_id: Option<StreamId>,
124    ) -> Result<(), StoreError> {
125        if let Some((top, _)) = self.entries.iter().next_back()
126            && last_id < *top
127        {
128            return Err(StoreError::OutOfRange);
129        }
130        self.last_id = last_id;
131        if let Some(n) = entries_added {
132            self.entries_added = n;
133        }
134        if let Some(id) = max_deleted_id {
135            self.max_deleted_id = id;
136        }
137        Ok(())
138    }
139}