hopr_protocol_session/processing/
sequencer.rs

1//! This module defines the [`Sequencer`] stream adaptor.
2
3use std::{
4    collections::BinaryHeap,
5    future::Future,
6    pin::Pin,
7    task::{Context, Poll},
8    time::{Duration, Instant},
9};
10
11use futures_time::future::Timer;
12use tracing::instrument;
13
14use crate::{errors::SessionError, protocol::FrameId};
15
16/// Sequencer is an adaptor for streams, that yield elements that have a natural ordering and
17/// can be compared with [`FrameId`] and puts them in the correct sequence starting with
18/// `FrameId` equal to 1.
19///
20/// Sequencer internally maintains a `FrameId` to be yielded next, polls the underlying stream
21/// and yields elements only when they match the next `FrameId` to be yielded, incrementing the
22/// value on each yield.
23///
24/// The Sequencer takes to arguments: `max_wait` and `capacity`:
25///
26/// The `max_wait` indicates the maximum amount of time to wait for a certain `FrameId` to
27/// be yielded from the underlying stream.
28/// If this does not happen, the Segmenter yields an error,
29/// indicating that the given frame was discarded.
30///
31/// The `capacity` parameter sets the maximum number of buffered elements inside the Sequencer.
32/// If this value is reached, the Sequencer will stop polling the underlying stream, waiting for the
33/// next element to expire.
34///
35/// By definition, Sequencer is a fallible stream, yielding either `Ok(Item)`, `Err(`[`SessionError::FrameDiscarded`]`)`
36/// or `Ok(None)` when the underlying stream is closed and no more elements can be yielded.
37///
38/// Use [`SequencerExt`] methods to turn a stream into a sequenced stream.
39#[must_use = "streams do nothing unless polled"]
40#[pin_project::pin_project]
41pub struct Sequencer<S: futures::Stream> {
42    #[pin]
43    inner: S,
44    #[pin]
45    timer: futures_time::task::Sleep,
46    buffer: BinaryHeap<std::cmp::Reverse<S::Item>>,
47    next_id: FrameId,
48    last_emitted: Instant,
49    max_wait: Duration,
50    state: State,
51}
52
53impl<S> Sequencer<S>
54where
55    S: futures::Stream,
56    S::Item: Ord + PartialOrd<FrameId>,
57{
58    /// Creates a new instance, wrapping the given `inner` Segment sink.
59    ///
60    /// The `frame_size` value will be clamped into the `[C, (C - SessionMessage::SEGMENT_OVERHEAD) * SeqIndicator::MAX
61    /// + 1]` interval.
62    fn new(inner: S, max_wait: Duration, capacity: usize) -> Self {
63        assert!(capacity > 0, "capacity should be positive");
64        Self {
65            inner,
66            buffer: BinaryHeap::with_capacity(capacity),
67            timer: futures_time::task::sleep(max_wait.max(Duration::from_millis(1)).into()),
68            next_id: 1,
69            last_emitted: Instant::now(),
70            max_wait,
71            state: State::Polling,
72        }
73    }
74}
75
76#[derive(Clone, Copy, Debug, PartialEq, Eq)]
77enum State {
78    Polling,
79    BufferUpdated,
80    Done,
81}
82
83impl<S> futures::Stream for Sequencer<S>
84where
85    S: futures::Stream,
86    S::Item: Ord + PartialOrd<FrameId>,
87{
88    type Item = Result<S::Item, SessionError>;
89
90    #[instrument(name = "Sequencer::poll_next", level = "trace", skip(self, cx), fields(next_frame_id = self.next_id, state = ?self.state))]
91    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
92        let mut this = self.project();
93        if *this.next_id == 0 {
94            tracing::debug!("end of frame sequence reached");
95            return Poll::Ready(None);
96        }
97
98        loop {
99            match *this.state {
100                State::Polling => {
101                    if this.buffer.len() < this.buffer.capacity() {
102                        // We still have capacity available, poll the underlying stream
103                        let stream_poll = this.inner.as_mut().poll_next(cx);
104
105                        // Only poll timer if there's something in the buffer
106                        let timer_poll = if !this.buffer.is_empty() {
107                            let poll = this.timer.as_mut().poll(cx);
108                            if poll.is_ready() {
109                                this.timer.as_mut().reset_timer();
110                            }
111                            poll
112                        } else {
113                            Poll::Pending
114                        };
115
116                        match (stream_poll, timer_poll) {
117                            (Poll::Pending, Poll::Pending) => {
118                                tracing::trace!("pending");
119                                *this.state = State::Polling;
120                                return Poll::Pending;
121                            }
122                            (Poll::Ready(Some(item)), _) => {
123                                // We have to reset the last emitted timestamp if
124                                // the buffer was empty until now
125                                if this.buffer.is_empty() {
126                                    *this.last_emitted = Instant::now();
127                                }
128
129                                if item.lt(this.next_id) {
130                                    // Do not accept older frame ids
131                                    tracing::error!("old item");
132                                    *this.state = State::Polling;
133                                } else {
134                                    // Push new item to the buffer
135                                    tracing::trace!("new item");
136                                    this.buffer.push(std::cmp::Reverse(item));
137                                    *this.state = State::BufferUpdated;
138                                }
139                            }
140                            (Poll::Ready(None), _) => {
141                                tracing::trace!(len = this.buffer.len(), "stream is done");
142                                *this.state = State::Done
143                            }
144                            (_, Poll::Ready(_)) => {
145                                // Simulate buffer update when the timer elapses
146                                tracing::trace!("timer elapsed");
147                                *this.state = State::BufferUpdated;
148                            }
149                        }
150                    } else {
151                        // Simulate buffer update when at capacity
152                        tracing::warn!("sequencer buffer is full");
153                        *this.state = State::BufferUpdated;
154                    }
155                }
156                State::BufferUpdated => {
157                    // The buffer has been updated, check if we can yield something
158                    if let Some(next) = this.buffer.peek().map(|item| &item.0) {
159                        if next.eq(this.next_id) {
160                            *this.next_id = this.next_id.wrapping_add(1);
161                            *this.last_emitted = Instant::now();
162                            *this.state = State::BufferUpdated;
163
164                            tracing::trace!("emit next frame");
165
166                            return Poll::Ready(this.buffer.pop().map(|item| Ok(item.0)));
167                        } else if this.last_emitted.elapsed() >= *this.max_wait
168                            || this.buffer.len() == this.buffer.capacity()
169                        {
170                            let discarded = *this.next_id;
171                            *this.next_id = this.next_id.wrapping_add(1);
172                            // `last_emitted` is intentionally NOT reset here: it is only reset
173                            // when an actual frame is emitted. Resetting it per discarded id would
174                            // drain a contiguous gap of K missing frames at 1 frame per `max_wait`
175                            // (a K x max_wait delivery stall of frames already sitting in the
176                            // buffer), instead of flushing the whole gap once `max_wait` elapses.
177                            *this.state = State::BufferUpdated;
178
179                            tracing::trace!(discarded, "discard frame");
180
181                            return Poll::Ready(Some(Err(SessionError::FrameDiscarded(discarded))));
182                        }
183                    } else {
184                        tracing::trace!("buffer is empty");
185                    }
186
187                    // Nothing to yield, keep on polling
188                    *this.state = State::Polling;
189                }
190                State::Done => {
191                    // The underlying stream is done, drain what we have in the internal buffer
192                    return if let Some(next) = this.buffer.peek().map(|item| &item.0) {
193                        if next.lt(this.next_id) {
194                            tracing::error!("old item");
195                            this.buffer.pop();
196                            continue;
197                        } else if next.eq(this.next_id) {
198                            *this.next_id = this.next_id.wrapping_add(1);
199                            tracing::trace!("emit next frame when done");
200
201                            Poll::Ready(this.buffer.pop().map(|item| Ok(item.0)))
202                        } else {
203                            let discarded = *this.next_id;
204                            *this.next_id = this.next_id.wrapping_add(1);
205                            tracing::trace!(discarded, "discard frame when done");
206
207                            Poll::Ready(Some(Err(SessionError::FrameDiscarded(discarded))))
208                        }
209                    } else {
210                        tracing::trace!("buffer is empty and done");
211                        Poll::Ready(None)
212                    };
213                }
214            }
215        }
216    }
217}
218
219/// Stream extensions methods for item sequencing.
220pub trait SequencerExt: futures::Stream {
221    /// Attaches a [`Sequencer`] to the underlying stream, given the item `timeout` and `capacity`
222    /// of items.
223    fn sequencer(self, timeout: Duration, capacity: usize) -> Sequencer<Self>
224    where
225        Self::Item: Ord + PartialOrd<FrameId>,
226        Self: Sized,
227    {
228        Sequencer::new(self, timeout, capacity)
229    }
230}
231
232impl<T: ?Sized> SequencerExt for T where T: futures::Stream {}
233
234#[cfg(test)]
235mod tests {
236    use futures::{SinkExt, StreamExt, TryStreamExt, pin_mut};
237    use futures_time::future::FutureExt;
238
239    use super::*;
240
241    #[test_log::test(tokio::test)]
242    async fn sequencer_should_return_entries_in_order() -> anyhow::Result<()> {
243        let mut expected = vec![4u32, 1, 5, 7, 8, 6, 2, 3];
244
245        let actual: Vec<u32> = futures::stream::iter(expected.clone())
246            .sequencer(Duration::from_secs(5), 4096)
247            .try_collect()
248            .timeout(futures_time::time::Duration::from_secs(5))
249            .await??;
250
251        expected.sort();
252        assert_eq!(expected, actual);
253
254        Ok(())
255    }
256
257    #[test_log::test(tokio::test)]
258    async fn sequencer_should_not_allow_emitted_entries() -> anyhow::Result<()> {
259        let (seq_sink, seq_stream) = futures::channel::mpsc::unbounded();
260
261        let seq_stream = seq_stream.sequencer(Duration::from_secs(1), 4096);
262
263        pin_mut!(seq_sink);
264        pin_mut!(seq_stream);
265
266        seq_sink.send(1u32).await?;
267        assert_eq!(Some(1), seq_stream.try_next().await?);
268
269        seq_sink.send(2u32).await?;
270        assert_eq!(Some(2), seq_stream.try_next().await?);
271
272        seq_sink.send(2u32).await?;
273        seq_sink.send(1u32).await?;
274
275        seq_sink.send(3u32).await?;
276        assert_eq!(Some(3), seq_stream.try_next().await?);
277
278        Ok(())
279    }
280
281    #[test_log::test(tokio::test)]
282    async fn sequencer_should_discard_entry_on_timeout() -> anyhow::Result<()> {
283        let timeout = Duration::from_millis(25);
284        let (mut seq_sink, seq_stream) = futures::channel::mpsc::unbounded();
285
286        let input = vec![2u32, 1, 4, 5, 8, 7, 9, 11, 10];
287
288        let input_clone = input.clone();
289        let jh = hopr_utils::runtime::prelude::spawn(async move {
290            for v in input_clone {
291                seq_sink
292                    .feed(v)
293                    .delay(futures_time::time::Duration::from_millis(5))
294                    .await?;
295            }
296            seq_sink.flush().await?;
297            seq_sink.close().await
298        });
299
300        let seq_stream = seq_stream.sequencer(timeout, 4096);
301
302        pin_mut!(seq_stream);
303
304        assert_eq!(Some(1), seq_stream.try_next().await?);
305        assert_eq!(Some(2), seq_stream.try_next().await?);
306
307        let now = Instant::now();
308        assert!(matches!(
309            seq_stream.try_next().await,
310            Err(SessionError::FrameDiscarded(3))
311        ));
312        assert!(now.elapsed() >= timeout);
313
314        assert_eq!(Some(4), seq_stream.try_next().await?);
315        assert_eq!(Some(5), seq_stream.try_next().await?);
316
317        assert!(matches!(
318            seq_stream.try_next().await,
319            Err(SessionError::FrameDiscarded(6))
320        ));
321
322        assert_eq!(Some(7), seq_stream.try_next().await?);
323        assert_eq!(Some(8), seq_stream.try_next().await?);
324        assert_eq!(Some(9), seq_stream.try_next().await?);
325        assert_eq!(Some(10), seq_stream.try_next().await?);
326        assert_eq!(Some(11), seq_stream.try_next().await?);
327
328        assert_eq!(None, seq_stream.try_next().await?);
329
330        let _ = jh.await?;
331        Ok(())
332    }
333
334    #[test_log::test(tokio::test)]
335    async fn sequencer_should_discard_entry_close() -> anyhow::Result<()> {
336        let (seq_sink, seq_stream) = futures::channel::mpsc::unbounded();
337
338        let input = vec![2u32, 1, 3, 5, 4, 8, 11];
339
340        hopr_utils::runtime::prelude::spawn(futures::stream::iter(input.clone()).map(Ok).forward(seq_sink)).await??;
341
342        let seq_stream = seq_stream.sequencer(Duration::from_millis(25), 4096);
343
344        pin_mut!(seq_stream);
345
346        assert_eq!(Some(1), seq_stream.try_next().await?);
347        assert_eq!(Some(2), seq_stream.try_next().await?);
348        assert_eq!(Some(3), seq_stream.try_next().await?);
349        assert_eq!(Some(4), seq_stream.try_next().await?);
350        assert_eq!(Some(5), seq_stream.try_next().await?);
351        assert!(matches!(
352            seq_stream.try_next().await,
353            Err(SessionError::FrameDiscarded(6))
354        ));
355        assert!(matches!(
356            seq_stream.try_next().await,
357            Err(SessionError::FrameDiscarded(7))
358        ));
359        assert_eq!(Some(8), seq_stream.try_next().await?);
360        assert!(matches!(
361            seq_stream.try_next().await,
362            Err(SessionError::FrameDiscarded(9))
363        ));
364        assert!(matches!(
365            seq_stream.try_next().await,
366            Err(SessionError::FrameDiscarded(10))
367        ));
368        assert_eq!(Some(11), seq_stream.try_next().await?);
369        assert_eq!(None, seq_stream.try_next().await?);
370
371        Ok(())
372    }
373
374    #[test_log::test(tokio::test)]
375    async fn sequencer_should_discard_entry_when_inner_stream_pending() -> anyhow::Result<()> {
376        let sent = vec![4u32, 1, 7, 8, 6, 2, 3];
377        let (tx, rx) = futures::channel::mpsc::unbounded();
378
379        pin_mut!(tx);
380        tx.send_all(&mut futures::stream::iter(sent.clone()).map(Ok)).await?;
381
382        let rx = rx.sequencer(Duration::from_millis(10), 4096);
383        pin_mut!(rx);
384
385        assert!(matches!(rx.next().await, Some(Ok(1))));
386        assert!(matches!(rx.next().await, Some(Ok(2))));
387        assert!(matches!(rx.next().await, Some(Ok(3))));
388        assert!(matches!(rx.next().await, Some(Ok(4))));
389        assert!(matches!(rx.next().await, Some(Err(SessionError::FrameDiscarded(5)))));
390        assert!(matches!(rx.next().await, Some(Ok(6))));
391        assert!(matches!(rx.next().await, Some(Ok(7))));
392        assert!(matches!(rx.next().await, Some(Ok(8))));
393
394        Ok(())
395    }
396
397    #[test_log::test(tokio::test)]
398    async fn sequencer_should_discard_entry_when_capacity_is_reached() -> anyhow::Result<()> {
399        let sent = vec![4u32, 5, 7, 8, 2, 6, 3];
400        let (tx, rx) = futures::channel::mpsc::unbounded();
401
402        pin_mut!(tx);
403        tx.send_all(&mut futures::stream::iter(sent.clone()).map(Ok)).await?;
404
405        let rx = rx.sequencer(Duration::from_millis(10), 4);
406        pin_mut!(rx);
407
408        assert!(matches!(rx.next().await, Some(Err(SessionError::FrameDiscarded(1)))));
409        assert!(matches!(rx.next().await, Some(Err(SessionError::FrameDiscarded(2)))));
410        assert!(matches!(rx.next().await, Some(Err(SessionError::FrameDiscarded(3)))));
411        assert!(matches!(rx.next().await, Some(Ok(4))));
412        assert!(matches!(rx.next().await, Some(Ok(5))));
413        assert!(matches!(rx.next().await, Some(Ok(6))));
414        assert!(matches!(rx.next().await, Some(Ok(7))));
415        assert!(matches!(rx.next().await, Some(Ok(8))));
416
417        Ok(())
418    }
419
420    #[test_log::test(tokio::test)]
421    async fn sequencer_should_drain_contiguous_gap_within_single_timeout_window() -> anyhow::Result<()> {
422        let timeout = Duration::from_millis(50);
423        let (tx, rx) = futures::channel::mpsc::unbounded();
424
425        pin_mut!(tx);
426        tx.send_all(&mut futures::stream::iter([1u32, 2, 10, 11, 12]).map(Ok))
427            .await?;
428
429        let rx = rx.sequencer(timeout, 4096);
430        pin_mut!(rx);
431
432        assert_eq!(Some(1), rx.try_next().await?);
433        assert_eq!(Some(2), rx.try_next().await?);
434
435        let now = Instant::now();
436        for expected in 3u32..=9 {
437            assert!(matches!(
438                rx.next().await,
439                Some(Err(SessionError::FrameDiscarded(id))) if id == expected
440            ));
441        }
442        assert_eq!(Some(10), rx.try_next().await?);
443        assert_eq!(Some(11), rx.try_next().await?);
444        assert_eq!(Some(12), rx.try_next().await?);
445
446        // The 7-frame gap must be flushed after one timeout window,
447        // not at a rate of one frame per window.
448        assert!(
449            now.elapsed() < 3 * timeout,
450            "gap drain took {:?}, expected well under {:?}",
451            now.elapsed(),
452            7 * timeout
453        );
454
455        Ok(())
456    }
457
458    #[test_log::test(tokio::test)]
459    async fn sequencer_must_terminate_on_last_frame_id() -> anyhow::Result<()> {
460        let (tx, rx) = futures::channel::mpsc::unbounded();
461
462        pin_mut!(tx);
463        tx.send_all(&mut futures::stream::iter([FrameId::MAX - 1, FrameId::MAX, 1, 2]).map(Ok))
464            .await?;
465
466        let mut rx = rx.sequencer(Duration::from_millis(10), 1024);
467        rx.next_id = FrameId::MAX - 1;
468        pin_mut!(rx);
469
470        const LAST_ID: FrameId = FrameId::MAX - 1;
471        assert!(matches!(rx.next().await, Some(Ok(LAST_ID))));
472        assert!(matches!(rx.next().await, Some(Ok(FrameId::MAX))));
473        assert!(rx.next().await.is_none());
474
475        Ok(())
476    }
477
478    #[test_log::test(tokio::test(flavor = "multi_thread"))]
479    async fn sequencer_must_not_discard_frames_when_buffer_was_empty_after_timeout() -> anyhow::Result<()> {
480        let (tx, rx) = futures::channel::mpsc::unbounded();
481
482        let jh = tokio::task::spawn(async move {
483            tokio::time::sleep(Duration::from_millis(2)).await;
484            pin_mut!(tx);
485            tx.send_all(&mut futures::stream::iter([3, 1, 2, 4]).map(Ok)).await?;
486
487            tokio::time::sleep(Duration::from_millis(150)).await;
488
489            tx.send_all(&mut futures::stream::iter([6, 5, 7]).map(Ok)).await?;
490
491            anyhow::Ok(())
492        });
493
494        let chunks = rx
495            .sequencer(Duration::from_millis(50), 1024)
496            .try_ready_chunks(10)
497            .try_collect::<Vec<Vec<_>>>()
498            .await?;
499
500        assert_eq!(chunks, vec![vec![1, 2, 3, 4], vec![5, 6, 7]]);
501        jh.await??;
502
503        Ok(())
504    }
505}
hopr_protocol_session/processing/sequencer.rs

hopr_protocol_session/processing/
sequencer.rs
