ruststream_fred/
subscriber.rs

1//! Redis Streams subscriber driving `XREADGROUP` (fresh tail) or `XAUTOCLAIM` (reclaim).
2
3use std::collections::{HashMap, VecDeque};
4use std::fmt::{Debug, Formatter};
5use std::sync::Arc;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::time::Duration;
8
9use fred::clients::Pool;
10use fred::interfaces::StreamsInterface;
11use fred::types::streams::XReadValue;
12use futures::Stream;
13use futures::stream::unfold;
14use ruststream::{BatchSubscriber, Seekable, Subscriber};
15
16use crate::convert::{HEADER_PREFIX, parts_from_fields};
17use crate::deadletter::{
18    self, DELIVERY_COUNT_HEADER, IDLE_MS_HEADER, PoisonPolicy, REASON_MAX_DELIVERIES,
19};
20use crate::delay::{self, DelayConfig};
21use crate::seek::{EntryId, RedisGroupSeeker};
22use crate::{error::RedisError, message::RedisMessage, stream::ReadMode};
23
24/// One decoded stream entry: its ID and field map.
25type Entry = (String, HashMap<String, Vec<u8>>);
26
27/// `XREADGROUP` reply shape parsed as nested arrays rather than maps: the RESP2 reply is an array of
28/// `[key, [[id, [field, value, ...]], ...]]`, which does not convert to fred's map-based
29/// `XReadResponse` (the outer array is not a flat key/value list). Pairing into tuples does work, so
30/// we collect the entry fields into a map ourselves.
31type RawStreams = Vec<(String, Vec<(String, Vec<(String, Vec<u8>)>)>)>;
32
33/// Cursor a fresh reclaim scan starts from (the whole pending list).
34const RECLAIM_START: &str = "0-0";
35
36fn duration_to_millis(d: Duration) -> u64 {
37    u64::try_from(d.as_millis()).unwrap_or(u64::MAX)
38}
39
40/// A Redis Streams subscription bound to a consumer group.
41///
42/// Constructed by [`crate::ConnectedRedisBroker::subscribe`] from a [`crate::RedisStream`]
43/// descriptor. The read mode (fresh tail vs reclaim) is fixed at construction.
44pub struct RedisSubscriber {
45    pool: Pool,
46    key: String,
47    group: String,
48    consumer: String,
49    count: u64,
50    block: Duration,
51    mode: ReadMode,
52    policy: PoisonPolicy,
53    /// Set when the subscription opted into a durable ZSET delay queue; drives native `nack_after`
54    /// on each delivery and the due-entry sweep on each fetch.
55    delay: Option<DelayConfig>,
56    /// Reclaim cursor; advances across `XAUTOCLAIM` calls, unused in fresh mode.
57    cursor: String,
58    /// Entries fetched but not yet yielded.
59    buffer: VecDeque<Entry>,
60    /// Bumped by every [`RedisGroupSeeker`] minted off this subscription. Shared, because a seek
61    /// can land while this subscriber is parked in a blocking read.
62    generation: Arc<AtomicU64>,
63    /// The generation the buffered entries were selected under. A mismatch means a seek moved the
64    /// cursor after they were chosen, so they belong to the position the group left.
65    buffer_generation: u64,
66}
67
68impl Debug for RedisSubscriber {
69    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
70        f.debug_struct("RedisSubscriber")
71            .field("key", &self.key)
72            .field("group", &self.group)
73            .field("consumer", &self.consumer)
74            .field("mode", &self.mode)
75            .finish_non_exhaustive()
76    }
77}
78
79impl RedisSubscriber {
80    #[allow(
81        clippy::too_many_arguments,
82        reason = "internal constructor mirroring the descriptor"
83    )]
84    pub(crate) fn new(
85        pool: Pool,
86        key: String,
87        group: String,
88        consumer: String,
89        count: u64,
90        block: Duration,
91        mode: ReadMode,
92        policy: PoisonPolicy,
93        delay: Option<DelayConfig>,
94    ) -> Self {
95        Self {
96            pool,
97            key,
98            group,
99            consumer,
100            count,
101            block,
102            mode,
103            policy,
104            delay,
105            cursor: RECLAIM_START.to_owned(),
106            buffer: VecDeque::new(),
107            generation: Arc::new(AtomicU64::new(0)),
108            buffer_generation: 0,
109        }
110    }
111
112    /// Builds the delivery for one fetched entry.
113    ///
114    /// # Errors
115    ///
116    /// Returns [`RedisError::Stream`] when the server sent an id that is not a well-formed
117    /// `<milliseconds>-<sequence>` pair, which would leave the delivery unable to report its
118    /// position.
119    fn message(
120        &self,
121        id: String,
122        fields: HashMap<String, Vec<u8>>,
123    ) -> Result<RedisMessage, RedisError> {
124        let entry: EntryId = id.parse()?;
125        let (payload, headers) = parts_from_fields(fields);
126        Ok(RedisMessage::new(
127            self.pool.clone(),
128            self.key.clone(),
129            self.group.clone(),
130            id,
131            entry,
132            payload,
133            headers,
134            self.policy.clone(),
135            self.delay.clone(),
136        ))
137    }
138
139    /// Drops entries selected before a seek moved the group cursor: they belong to the position
140    /// the group left, so delivering them would contradict the reposition.
141    fn discard_stale(&mut self) {
142        let current = self.generation.load(Ordering::Acquire);
143        if self.buffer_generation != current {
144            self.buffer.clear();
145            self.buffer_generation = current;
146        }
147    }
148
149    /// Fetches the next batch of entries into the buffer. A read that timed out with nothing
150    /// pending leaves the buffer empty (the caller loops and reads again).
151    async fn fetch(&mut self) -> Result<(), RedisError> {
152        // Replay any due delayed-retry entries before reading, so they re-enter the stream and get
153        // delivered through the normal read path. Granularity is the read block interval.
154        if let Some(cfg) = &self.delay {
155            delay::sweep_due(&self.pool, cfg, &self.key).await?;
156        }
157        // Captured before the read, not after: a blocking read selects its entries against the
158        // cursor as it stood when the read started, so a seek that lands mid-read invalidates
159        // whatever comes back.
160        let selected_at = self.generation.load(Ordering::Acquire);
161        let entries = match self.mode.clone() {
162            ReadMode::Fresh => self.fetch_fresh().await?,
163            ReadMode::Reclaim { min_idle } => self.fetch_reclaim(min_idle).await?,
164        };
165        if selected_at != self.generation.load(Ordering::Acquire) {
166            // A seek overtook this read; drop its entries and let the caller read again.
167            return Ok(());
168        }
169        self.buffer_generation = selected_at;
170        self.buffer.extend(entries);
171        Ok(())
172    }
173
174    async fn fetch_fresh(&self) -> Result<Vec<Entry>, RedisError> {
175        let resp: RawStreams = self
176            .pool
177            .xreadgroup(
178                self.group.as_str(),
179                self.consumer.as_str(),
180                Some(self.count),
181                Some(duration_to_millis(self.block)),
182                false,
183                self.key.as_str(),
184                ">",
185            )
186            .await
187            .map_err(RedisError::stream)?;
188        let entries = resp
189            .into_iter()
190            .find(|(key, _)| key == &self.key)
191            .map(|(_, entries)| entries)
192            .unwrap_or_default();
193        Ok(entries
194            .into_iter()
195            .map(|(id, fields)| (id, fields.into_iter().collect()))
196            .collect())
197    }
198
199    async fn fetch_reclaim(&mut self, min_idle: Duration) -> Result<Vec<Entry>, RedisError> {
200        let (cursor, entries): (String, Vec<XReadValue<String, String, Vec<u8>>>) = self
201            .pool
202            .xautoclaim_values(
203                self.key.as_str(),
204                self.group.as_str(),
205                self.consumer.as_str(),
206                duration_to_millis(min_idle),
207                self.cursor.as_str(),
208                Some(self.count),
209                false,
210            )
211            .await
212            .map_err(RedisError::stream)?;
213        self.cursor = cursor;
214        // Nothing left to reclaim this pass: avoid a hot loop until more entries go stale.
215        if entries.is_empty() {
216            tokio::time::sleep(self.block).await;
217            return Ok(entries);
218        }
219        // Plain reclaim with no poison policy: skip the extra XPENDING and deliver as-is.
220        if !self.policy.is_active() {
221            return Ok(entries);
222        }
223        self.enrich_reclaimed(entries).await
224    }
225
226    /// Annotates reclaimed entries with their native delivery count and idle time, and dead-letters
227    /// (or drops) any that have exceeded `max_deliveries` instead of redelivering them.
228    async fn enrich_reclaimed(&self, entries: Vec<Entry>) -> Result<Vec<Entry>, RedisError> {
229        let meta = self.pending_meta().await?;
230        let mut out = Vec::with_capacity(entries.len());
231        for (id, mut fields) in entries {
232            let (idle, count) = meta.get(&id).copied().unwrap_or((0, 0));
233            if self.policy.is_poison(count) {
234                self.dead_letter_reclaimed(&id, &fields).await?;
235                continue;
236            }
237            insert_meta_header(&mut fields, DELIVERY_COUNT_HEADER, count);
238            insert_meta_header(&mut fields, IDLE_MS_HEADER, idle);
239            out.push((id, fields));
240        }
241        Ok(out)
242    }
243
244    /// Maps each of this consumer's pending entry IDs to its `(idle_ms, delivery_count)` via
245    /// extended `XPENDING`, which - unlike `XAUTOCLAIM` - reports the native delivery count.
246    async fn pending_meta(&self) -> Result<HashMap<String, (u64, u64)>, RedisError> {
247        let rows: Vec<(String, String, u64, u64)> = self
248            .pool
249            .xpending(
250                self.key.as_str(),
251                self.group.as_str(),
252                (0_u64, "-", "+", self.count, self.consumer.as_str()),
253            )
254            .await
255            .map_err(RedisError::stream)?;
256        Ok(rows
257            .into_iter()
258            .map(|(id, _consumer, idle, count)| (id, (idle, count)))
259            .collect())
260    }
261
262    /// Routes a poison reclaimed entry to its dead-letter stream (or discards it when none is set),
263    /// then `XACK`s it so it leaves the pending list.
264    async fn dead_letter_reclaimed(
265        &self,
266        id: &str,
267        fields: &HashMap<String, Vec<u8>>,
268    ) -> Result<(), RedisError> {
269        let (payload, headers) = parts_from_fields(fields.clone());
270        deadletter::settle_poison_stream(
271            &self.pool,
272            &self.policy,
273            &payload,
274            &headers,
275            REASON_MAX_DELIVERIES,
276        )
277        .await
278        .map_err(RedisError::stream)?;
279        let _: i64 = self
280            .pool
281            .xack(self.key.as_str(), self.group.as_str(), id)
282            .await
283            .map_err(RedisError::stream)?;
284        Ok(())
285    }
286}
287
288/// Injects a `u64`-valued well-known header into an entry's raw field map (under the `h:` prefix),
289/// so it surfaces as a [`Headers`](ruststream::Headers) entry on the delivered message.
290fn insert_meta_header(fields: &mut HashMap<String, Vec<u8>>, name: &str, value: u64) {
291    fields.insert(
292        format!("{HEADER_PREFIX}{name}"),
293        value.to_string().into_bytes(),
294    );
295}
296
297impl Subscriber for RedisSubscriber {
298    type Message = RedisMessage;
299    type Error = RedisError;
300
301    /// Yields one message per entry, refilling from Redis when the local buffer drains.
302    ///
303    /// # Cancel safety
304    ///
305    /// Dropping the returned stream between items is safe. Dropping it while a read is in flight
306    /// drops the read future; entries already delivered to this consumer but not yet acked stay in
307    /// the group's pending list and are redelivered (fresh mode) or reclaimable (reclaim mode).
308    fn stream(&mut self) -> impl Stream<Item = Result<Self::Message, Self::Error>> + Send + '_ {
309        unfold(self, |s| async move {
310            loop {
311                s.discard_stale();
312                if let Some((id, fields)) = s.buffer.pop_front() {
313                    return Some((s.message(id, fields), s));
314                }
315                // An empty fetch (a blocking read that timed out) just loops and reads again.
316                if let Err(err) = s.fetch().await {
317                    return Some((Err(err), s));
318                }
319            }
320        })
321    }
322}
323
324/// Repositioning the subscription moves its consumer group's cursor, which is shared by every
325/// consumer of that group; see [`RedisGroupSeeker`] for the full contract.
326impl Seekable for RedisSubscriber {
327    type Seeker = RedisGroupSeeker;
328
329    fn seeker(&self) -> RedisGroupSeeker {
330        RedisGroupSeeker::new(
331            self.pool.clone(),
332            self.key.clone(),
333            self.group.clone(),
334            Arc::clone(&self.generation),
335        )
336    }
337}
338
339impl BatchSubscriber for RedisSubscriber {
340    type Batch = Vec<RedisMessage>;
341
342    /// Yields one batch per non-empty read (`XREADGROUP COUNT` / `XAUTOCLAIM`), up to
343    /// [`RedisStream::count`](crate::RedisStream::count) entries. Never yields an empty batch.
344    ///
345    /// # Cancel safety
346    ///
347    /// Same as [`Subscriber::stream`]: dropping the stream mid-read leaves fetched-but-unacked
348    /// entries in the pending list.
349    fn batches(&mut self) -> impl Stream<Item = Result<Self::Batch, Self::Error>> + Send + '_ {
350        unfold(self, |s| async move {
351            loop {
352                s.discard_stale();
353                if !s.buffer.is_empty() {
354                    // Move the buffer out first so `s.message` can borrow `s` without overlapping
355                    // a live mutable borrow of `s.buffer`.
356                    let entries = std::mem::take(&mut s.buffer);
357                    let batch = entries
358                        .into_iter()
359                        .map(|(id, fields)| s.message(id, fields))
360                        .collect::<Result<Vec<_>, _>>();
361                    return Some((batch, s));
362                }
363                if let Err(err) = s.fetch().await {
364                    return Some((Err(err), s));
365                }
366            }
367        })
368    }
369}
ruststream_fred/subscriber.rs

ruststream_fred/
subscriber.rs
