1use crate::{KevyError, KevyResult};
14
15use kevy_embedded::FeedError;
16use kevy_resp::Reply;
17use kevy_resp_client::RespClient;
18
19use crate::{Connection, string, unexpected};
20
21#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct FeedFrame {
24 pub offset: u64,
26 pub argv: Vec<Vec<u8>>,
29}
30
31#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct FeedBatch {
34 pub generation: u64,
36 pub next_offset: u64,
38 pub frames: Vec<FeedFrame>,
40}
41
42impl Connection {
43 pub fn feed_shards(&mut self) -> KevyResult<usize> {
45 match self {
46 Self::Embedded(s) => Ok(s.feed_shards()),
47 Self::Remote(c) => match c.request_borrowed(&[b"FEED.SHARDS"])? {
48 Reply::Int(n) if n >= 0 => Ok(n as usize),
49 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
50 other => Err(unexpected(other)),
51 },
52 }
53 }
54
55 pub fn feed_tail(&mut self, shard: usize) -> KevyResult<(u64, u64)> {
59 match self {
60 Self::Embedded(s) => {
61 check_embedded_shard(shard)?;
62 s.changes_tail().map_err(feed_err)
63 }
64 Self::Remote(c) => {
65 let sh = shard.to_string();
66 match c.request_borrowed(&[b"FEED.TAIL", sh.as_bytes()])? {
67 Reply::Array(items) if items.len() == 2 => {
68 let mut it = items.into_iter();
69 match (it.next().unwrap(), it.next().unwrap()) {
70 (Reply::Int(g), Reply::Int(o)) => Ok((g as u64, o as u64)),
71 (a, _) => Err(unexpected(a)),
72 }
73 }
74 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
75 other => Err(unexpected(other)),
76 }
77 }
78 }
79 }
80
81 pub fn feed_read(
87 &mut self,
88 shard: usize,
89 generation: u64,
90 offset: u64,
91 count: Option<usize>,
92 prefixes: &[&[u8]],
93 ) -> KevyResult<FeedBatch> {
94 match self {
95 Self::Embedded(s) => {
96 check_embedded_shard(shard)?;
97 let batch = s
98 .changes_since(generation, offset, count.unwrap_or(256), prefixes)
99 .map_err(feed_err)?;
100 Ok(FeedBatch {
101 generation: batch.next.0,
102 next_offset: batch.next.1,
103 frames: batch
104 .changes
105 .into_iter()
106 .map(|ch| FeedFrame { offset: ch.offset, argv: ch.argv })
107 .collect(),
108 })
109 }
110 Self::Remote(c) => {
111 parse_batch(feed_read_request(c, shard, generation, offset, count, prefixes)?)
112 }
113 }
114 }
115}
116
117fn check_embedded_shard(shard: usize) -> KevyResult<()> {
119 if shard != 0 {
120 return Err(KevyError::InvalidInput(
121 "embedded feed is single-shard: shard must be 0".into(),
122 ));
123 }
124 Ok(())
125}
126
127fn feed_err(e: FeedError) -> KevyError {
130 match e {
131 FeedError::Resync { generation, tail } => {
132 KevyError::Protocol(format!("FEEDRESYNC {generation} {tail}"))
133 }
134 FeedError::Future => KevyError::Protocol("ERR feed cursor ahead of stream".into()),
135 FeedError::Disabled => KevyError::Unsupported(
136 "feed disabled: open the embedded store with Config::with_feed".into(),
137 ),
138 }
139}
140
141fn feed_read_request(
142 c: &mut RespClient,
143 shard: usize,
144 generation: u64,
145 offset: u64,
146 count: Option<usize>,
147 prefixes: &[&[u8]],
148) -> KevyResult<Reply> {
149 let mut args: Vec<Vec<u8>> = vec![
150 b"FEED.READ".to_vec(),
151 shard.to_string().into_bytes(),
152 generation.to_string().into_bytes(),
153 offset.to_string().into_bytes(),
154 ];
155 if let Some(n) = count {
156 args.push(b"COUNT".to_vec());
157 args.push(n.to_string().into_bytes());
158 }
159 for p in prefixes {
160 args.push(b"PREFIX".to_vec());
161 args.push(p.to_vec());
162 }
163 Ok(c.request(&args)?)
164}
165
166fn parse_batch(reply: Reply) -> KevyResult<FeedBatch> {
169 let Reply::Array(items) = reply else {
170 return match reply {
171 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
172 other => Err(unexpected(other)),
173 };
174 };
175 if items.len() != 3 {
176 return Err(KevyError::Protocol("FEED.READ: expected [gen, next, frames]".into()));
177 }
178 let mut it = items.into_iter();
179 let (Reply::Int(g), Reply::Int(next)) = (it.next().unwrap(), it.next().unwrap()) else {
180 return Err(KevyError::Protocol("FEED.READ: non-integer cursor".into()));
181 };
182 let Reply::Array(raw_frames) = it.next().unwrap() else {
183 return Err(KevyError::Protocol("FEED.READ: frames not an array".into()));
184 };
185 let frames = raw_frames.into_iter().map(parse_frame).collect::<KevyResult<_>>()?;
186 Ok(FeedBatch { generation: g as u64, next_offset: next as u64, frames })
187}
188
189fn parse_frame(frame: Reply) -> KevyResult<FeedFrame> {
190 let Reply::Array(cells) = frame else {
191 return Err(KevyError::Protocol("FEED.READ: frame not an array".into()));
192 };
193 let mut it = cells.into_iter();
194 let (Some(Reply::Int(off)), Some(Reply::Array(argv_raw))) = (it.next(), it.next()) else {
195 return Err(KevyError::Protocol("FEED.READ: frame shape != [offset, argv]".into()));
196 };
197 let argv = argv_raw
198 .into_iter()
199 .map(|a| match a {
200 Reply::Bulk(b) | Reply::Simple(b) => Ok(b),
201 other => Err(unexpected(other)),
202 })
203 .collect::<KevyResult<_>>()?;
204 Ok(FeedFrame { offset: off as u64, argv })
205}
206
207#[cfg(test)]
208mod tests {
209 use super::*;
210
211 #[test]
212 fn embedded_without_feed_config_is_unsupported() {
213 let mut c = Connection::connect("mem://").unwrap();
215 assert_eq!(c.feed_shards().unwrap(), 1);
216 let err = c.feed_tail(0).unwrap_err();
217 assert!(matches!(err, KevyError::Unsupported(_)));
218 let err = c.feed_read(0, 1, 0, None, &[]).unwrap_err();
219 assert!(matches!(err, KevyError::Unsupported(_)));
220 }
221
222 #[test]
223 fn embedded_nonzero_shard_rejected() {
224 let mut c = Connection::connect("mem://").unwrap();
225 let err = c.feed_tail(1).unwrap_err();
226 assert!(matches!(err, KevyError::InvalidInput(_)));
227 }
228
229 #[test]
230 fn batch_parser_maps_frames() {
231 let reply = Reply::Array(vec![
232 Reply::Int(1),
233 Reply::Int(42),
234 Reply::Array(vec![Reply::Array(vec![
235 Reply::Int(41),
236 Reply::Array(vec![
237 Reply::Bulk(b"SET".to_vec()),
238 Reply::Bulk(b"k".to_vec()),
239 Reply::Bulk(b"v".to_vec()),
240 ]),
241 ])]),
242 ]);
243 let batch = parse_batch(reply).unwrap();
244 assert_eq!(batch.generation, 1);
245 assert_eq!(batch.next_offset, 42);
246 assert_eq!(batch.frames.len(), 1);
247 assert_eq!(batch.frames[0].offset, 41);
248 assert_eq!(batch.frames[0].argv[0], b"SET");
249 }
250}