hopr_protocol_session/utils/
skip_queue.rs

1use std::{
2    cmp::Ordering,
3    collections::BTreeSet,
4    pin::Pin,
5    sync::{Arc, atomic::AtomicBool},
6    task::{Context, Poll, Waker},
7    time::{Duration, Instant},
8};
9
10use futures::FutureExt;
11use tracing::instrument;
12
13/// An internal type used by the [`SkipDelayQueue`].
14#[derive(Debug)]
15struct DelayedEntry<T> {
16    item: T,
17    at: Instant,
18    cancelled: AtomicBool,
19}
20
21// The entries are equal only if the items they carry are equal
22impl<T: PartialEq> PartialEq for DelayedEntry<T> {
23    fn eq(&self, other: &Self) -> bool {
24        self.item == other.item
25    }
26}
27
28impl<T: Eq> Eq for DelayedEntry<T> {}
29
30impl<T: Ord> Ord for DelayedEntry<T> {
31    fn cmp(&self, other: &Self) -> Ordering {
32        if other.item != self.item {
33            // If items are not equal, the order is determined by the deadline
34            match self.at.cmp(&other.at) {
35                // If the deadlines are equal, use the natural order of the items.
36                // This should be presumably consistent with their PartialEq and won't
37                // therefore return Ordering::Equal.
38                Ordering::Equal => self.item.cmp(&other.item),
39                x => x,
40            }
41        } else {
42            // Be consistent with PartialEq
43            Ordering::Equal
44        }
45    }
46}
47
48impl<T: Ord> PartialOrd for DelayedEntry<T> {
49    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
50        Some(self.cmp(other))
51    }
52}
53
54/// Internal type used by the [`skip_delay_channel`].
55struct SkipDelayQueue<T> {
56    entries: BTreeSet<DelayedEntry<T>>,
57    next_wakeup: Option<futures_time::task::SleepUntil>,
58    rx_waker: Option<Waker>,
59    is_closed: bool,
60}
61
62/// An item with a deadline, which can be pushed into the [`SkipDelayQueue`].
63///
64/// For convenience, the type implements From traits from
65/// `(T, Instant)`, `(T, Duration)` and `(T, Skip)`.
66#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
67pub enum DelayedItem<T> {
68    /// Adds new (or replaces an existing) item with a deadline.
69    New(T, Instant),
70    /// Cancel a previously added item.
71    Cancel(T),
72}
73
74/// A marker type for canceling items pushed into the [`SkipDelayQueue`].
75#[derive(Debug, Copy, Clone, PartialEq, Eq)]
76pub struct Skip;
77
78impl<T> From<(T, Duration)> for DelayedItem<T> {
79    fn from(value: (T, Duration)) -> Self {
80        Self::New(value.0, Instant::now() + value.1)
81    }
82}
83
84impl<T> From<(T, Instant)> for DelayedItem<T> {
85    fn from(value: (T, Instant)) -> Self {
86        Self::New(value.0, value.1)
87    }
88}
89
90impl<T> From<(T, Skip)> for DelayedItem<T> {
91    fn from(value: (T, Skip)) -> Self {
92        Self::Cancel(value.0)
93    }
94}
95
96impl<T> SkipDelayQueue<T> {
97    const TOLERANCE: Duration = Duration::from_millis(5);
98
99    /// Creates a new instance.
100    ///
101    /// As a common practice, [`futures::StreamExt::split`] can be called to
102    /// get separate sending and receiving part of the queue.
103    pub fn new() -> Self {
104        Self {
105            entries: BTreeSet::new(),
106            next_wakeup: None,
107            rx_waker: None,
108            is_closed: false,
109        }
110    }
111}
112
113/// Receiver part for the [`skip_delay_channel`].
114pub struct SkipDelayReceiver<T>(Arc<std::sync::Mutex<SkipDelayQueue<T>>>);
115
116impl<T> Drop for SkipDelayReceiver<T> {
117    #[instrument(name = "SkipDelayReceiver::drop", level = "trace", skip(self))]
118    fn drop(&mut self) {
119        // When the receiver is dropped, clear the poison and mark the queue as closed.
120        self.0.clear_poison();
121        let mut queue = self.0.lock().expect("cannot panic because poison is cleared");
122        queue.is_closed = true;
123        if let Some(waker) = queue.rx_waker.take() {
124            waker.wake();
125        }
126    }
127}
128
129impl<T: Ord> futures::Stream for SkipDelayReceiver<T> {
130    type Item = T;
131
132    #[instrument(name = "SkipDelayReceiver::poll_next", level = "trace", skip(self, cx))]
133    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
134        let Ok(mut queue) = self.0.lock() else {
135            tracing::error!("poisoned mutex");
136            return Poll::Ready(None);
137        };
138
139        // Wait until the timer is done, if any
140        if let Some(next_wakeup) = queue.next_wakeup.as_mut() {
141            tracing::trace!("polling timer");
142            let _ = futures::ready!(next_wakeup.poll_unpin(cx));
143            queue.next_wakeup = None;
144        }
145
146        tracing::trace!("timer finished");
147
148        let now = Instant::now();
149        while let Some(e) = queue.entries.first() {
150            if !e.cancelled.load(std::sync::atomic::Ordering::SeqCst) {
151                return if e.at.saturating_duration_since(now) < SkipDelayQueue::<T>::TOLERANCE {
152                    // If the item is already in the past, yield it
153                    tracing::trace!("ready");
154                    Poll::Ready(queue.entries.pop_first().map(|e| e.item))
155                } else {
156                    // The next item is in the future, set up the timer and wake us up to start it
157                    tracing::trace!("pending new timer");
158                    queue.next_wakeup = Some(futures_time::task::sleep_until(e.at.into()));
159                    cx.waker().wake_by_ref();
160                    Poll::Pending
161                };
162            } else {
163                // If the item has been canceled, remove it and continue
164                queue.entries.pop_first();
165                tracing::trace!("item cancelled");
166            }
167        }
168
169        if !queue.is_closed {
170            // Need more data, wake us up when some are added
171            tracing::trace!("pending for data");
172            queue.rx_waker = Some(cx.waker().clone());
173            Poll::Pending
174        } else {
175            // We're done
176            Poll::Ready(None)
177        }
178    }
179}
180
181/// Sender part for the [`skip_delay_channel`].
182pub struct SkipDelaySender<T>(Option<Arc<std::sync::Mutex<SkipDelayQueue<T>>>>);
183
184impl<T> Clone for SkipDelaySender<T> {
185    fn clone(&self) -> Self {
186        Self(self.0.clone())
187    }
188}
189
190impl<T> SkipDelaySender<T> {
191    fn ensure_closure(&mut self) {
192        if let Some(queue) = self.0.take() {
193            let count_holders = Arc::strong_count(&queue);
194            tracing::trace!(count_holders, "ensure_closure");
195
196            // Check if the last holders are this instance and (potentially) the receiver
197            if count_holders == 2 {
198                Self::finalize_closure(queue);
199            }
200        }
201    }
202
203    fn finalize_closure(queue: Arc<std::sync::Mutex<SkipDelayQueue<T>>>) {
204        tracing::trace!("finalize_closure");
205        queue.clear_poison();
206        let mut queue = queue.lock().expect("cannot panic because poison is cleared");
207        queue.is_closed = true;
208        if let Some(waker) = queue.rx_waker.take() {
209            waker.wake();
210        }
211    }
212
213    /// Forces closure of the queue (regardless of any remaining senders).
214    pub fn force_close(&mut self) {
215        if let Some(queue) = self.0.take() {
216            Self::finalize_closure(queue);
217        }
218    }
219}
220
221impl<T: Ord> SkipDelaySender<T> {
222    #[instrument(
223        name = "SkipDelaySender::send_internal",
224        level = "trace",
225        skip(self, items, flush),
226        ret
227    )]
228    fn send_internal<I: Iterator<Item = DelayedItem<T>>>(&self, items: I, flush: bool) -> Result<(), std::io::Error> {
229        if let Some(queue) = self.0.as_ref() {
230            let mut queue = queue.lock().map_err(|_| std::io::ErrorKind::BrokenPipe)?;
231
232            // This can happen only when the receiver was dropped.
233            if queue.is_closed {
234                return Err(std::io::ErrorKind::BrokenPipe.into());
235            }
236
237            for item in items {
238                match item {
239                    DelayedItem::New(item, at) => {
240                        tracing::trace!(at =  ?at.saturating_duration_since(Instant::now()), "inserting");
241                        queue.entries.replace(DelayedEntry {
242                            item,
243                            at,
244                            cancelled: AtomicBool::new(false),
245                        });
246                    }
247                    DelayedItem::Cancel(item) => {
248                        tracing::trace!("cancelling");
249                        queue
250                            .entries
251                            .iter()
252                            .filter(|e| item == e.item)
253                            .for_each(|e| e.cancelled.store(true, std::sync::atomic::Ordering::SeqCst));
254                    }
255                }
256            }
257
258            if flush {
259                tracing::trace!("flushing");
260                if let Some(waker) = queue.rx_waker.take() {
261                    waker.wake();
262                }
263            }
264
265            Ok(())
266        } else {
267            Err(std::io::ErrorKind::NotConnected.into())
268        }
269    }
270
271    /// Sends the given single item and flushes the queue.
272    pub fn send_one<I: Into<DelayedItem<T>>>(&mut self, item: I) -> Result<(), std::io::Error> {
273        self.send_internal(std::iter::once(item.into()), true)
274    }
275
276    /// Sends many items at once and then flushes the queue.
277    pub fn send_many<I: IntoIterator<Item = DelayedItem<T>>>(&mut self, items: I) -> Result<(), std::io::Error> {
278        self.send_internal(items.into_iter(), true)
279    }
280}
281
282impl<T> Drop for SkipDelaySender<T> {
283    fn drop(&mut self) {
284        self.ensure_closure();
285    }
286}
287
288impl<T: Ord> futures::Sink<DelayedItem<T>> for SkipDelaySender<T> {
289    type Error = std::io::Error;
290
291    fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
292        if self.0.is_some() {
293            Poll::Ready(Ok(()))
294        } else {
295            Poll::Ready(Err(std::io::ErrorKind::NotConnected.into()))
296        }
297    }
298
299    fn start_send(self: Pin<&mut Self>, item: DelayedItem<T>) -> Result<(), Self::Error> {
300        self.send_internal(std::iter::once(item), false)
301    }
302
303    #[instrument(name = "SkipDelaySender::poll_flush", level = "trace", skip(self), ret)]
304    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
305        if let Some(queue) = self.0.as_ref() {
306            let Ok(mut queue) = queue.lock() else {
307                return Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into()));
308            };
309
310            tracing::trace!("flushing");
311            if let Some(waker) = queue.rx_waker.take() {
312                waker.wake();
313            }
314
315            Poll::Ready(Ok(()))
316        } else {
317            Poll::Ready(Err(std::io::ErrorKind::NotConnected.into()))
318        }
319    }
320
321    fn poll_close(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
322        if self.0.is_none() {
323            return Poll::Ready(Err(std::io::ErrorKind::NotConnected.into()));
324        }
325
326        self.ensure_closure();
327        Poll::Ready(Ok(()))
328    }
329}
330
331/// A MPSC queue of [items](DelayedItem) with attached [`Instant`] that determines a deadline
332/// at which it should be yielded from the [`Stream`](futures::Stream) side of the queue.
333/// The queue also has the cancellation ability: an item that has been pushed into the
334/// queue earlier can be canceled before it meets its deadline.
335/// A canceled item will then be skipped in the output stream.
336///
337/// The items are internally sorted based on their deadline.
338/// In a case when two items have equal deadlines, they are sorted according
339/// to their values; therefore, items must implement [`Ord`].
340///
341/// If equal items are inserted, the deadline of the earlier one inserted is updated.
342pub fn skip_delay_channel<T: Ord>() -> (SkipDelaySender<T>, SkipDelayReceiver<T>) {
343    let queue = Arc::new(std::sync::Mutex::new(SkipDelayQueue::new()));
344    (SkipDelaySender(Some(queue.clone())), SkipDelayReceiver(queue))
345}
346
347#[cfg(test)]
348mod tests {
349    use futures::{SinkExt, StreamExt, pin_mut};
350
351    use super::*;
352
353    #[test_log::test(tokio::test)]
354    async fn skip_delay_queue_should_yield_items() -> anyhow::Result<()> {
355        let (mut tx, rx) = skip_delay_channel();
356        pin_mut!(rx);
357
358        let now = Instant::now();
359        tx.send((1, now + Duration::from_millis(100)).into()).await?;
360        tx.close().await?;
361
362        assert_eq!(Some(1), rx.next().await);
363        assert!(now.elapsed() >= Duration::from_millis(100));
364        assert_eq!(None, rx.next().await);
365
366        Ok(())
367    }
368
369    #[test_log::test(tokio::test)]
370    async fn skip_delay_queue_should_replace_and_yield_items() -> anyhow::Result<()> {
371        let (mut tx, rx) = skip_delay_channel();
372        pin_mut!(rx);
373
374        let now = Instant::now();
375        tx.send((1, now + Duration::from_millis(100)).into()).await?;
376        tx.send((1, now + Duration::from_millis(200)).into()).await?;
377        tx.close().await?;
378
379        assert_eq!(Some(1), rx.next().await);
380        assert!(now.elapsed() >= Duration::from_millis(200));
381        assert_eq!(None, rx.next().await);
382
383        Ok(())
384    }
385
386    #[test_log::test(tokio::test)]
387    async fn skip_delay_queue_should_yield_items_from_multiple_senders() -> anyhow::Result<()> {
388        let (mut tx, rx) = skip_delay_channel();
389        pin_mut!(rx);
390
391        let mut tx2 = tx.clone();
392
393        let now = Instant::now();
394        tx.send((2, now + Duration::from_millis(100)).into()).await?;
395        tx.close().await?;
396
397        tx2.send((1, now + Duration::from_millis(150)).into()).await?;
398        tx2.close().await?;
399
400        assert_eq!(Some(2), rx.next().await);
401        assert!(now.elapsed() >= Duration::from_millis(100));
402        assert_eq!(Some(1), rx.next().await);
403        assert!(now.elapsed() >= Duration::from_millis(150));
404
405        assert_eq!(None, rx.next().await);
406
407        Ok(())
408    }
409
410    #[test_log::test(tokio::test)]
411    async fn skip_delay_queue_yielded_items_should_be_apart() -> anyhow::Result<()> {
412        let (mut tx, rx) = skip_delay_channel();
413        pin_mut!(rx);
414
415        let now1 = Instant::now();
416        tx.send((1, now1 + Duration::from_millis(100)).into()).await?;
417        let now2 = Instant::now();
418        tx.send((2, now2 + Duration::from_millis(200)).into()).await?;
419        tx.close().await?;
420
421        assert_eq!(Some(1), rx.next().await);
422        assert!(now1.elapsed() >= Duration::from_millis(100));
423        assert_eq!(Some(2), rx.next().await);
424        assert!(now2.elapsed() >= Duration::from_millis(200));
425
426        assert_eq!(None, rx.next().await);
427
428        Ok(())
429    }
430
431    #[test_log::test(tokio::test)]
432    async fn skip_delay_queue_should_not_yield_cancelled_items() -> anyhow::Result<()> {
433        let (mut tx, rx) = skip_delay_channel();
434        pin_mut!(rx);
435
436        let now = Instant::now();
437        tx.send((1, now + Duration::from_millis(100)).into()).await?;
438        tx.send((1, Skip).into()).await?;
439        tx.close().await?;
440
441        assert_eq!(None, rx.next().await);
442
443        Ok(())
444    }
445
446    #[test_log::test(tokio::test)]
447    async fn skip_delay_queue_should_yield_past_items_immediately() -> anyhow::Result<()> {
448        let (mut tx, rx) = skip_delay_channel();
449        pin_mut!(rx);
450
451        let now = Instant::now();
452        tx.send((1, now).into()).await?;
453        tx.send((2, now).into()).await?;
454        tx.close().await?;
455
456        let now = Instant::now();
457        assert_eq!(Some(1), rx.next().await);
458        assert_eq!(Some(2), rx.next().await);
459        assert_eq!(None, rx.next().await);
460
461        assert!(now.elapsed() < Duration::from_millis(25));
462
463        Ok(())
464    }
465
466    #[test_log::test(tokio::test)]
467    async fn skip_delay_queue_should_not_yield_future_cancelled_items() -> anyhow::Result<()> {
468        let (mut tx, rx) = skip_delay_channel();
469        pin_mut!(rx);
470
471        let now = Instant::now();
472        tx.send((1, now).into()).await?;
473        tx.send((2, now + Duration::from_millis(100)).into()).await?;
474        tx.send((2, Skip).into()).await?;
475        tx.close().await?;
476
477        assert_eq!(Some(1), rx.next().await);
478        assert_eq!(None, rx.next().await);
479        assert!(now.elapsed() < Duration::from_millis(50));
480
481        Ok(())
482    }
483
484    #[test_log::test(tokio::test)]
485    async fn skip_delay_queue_should_discard_duplicate_entries() -> anyhow::Result<()> {
486        let (mut tx, rx) = skip_delay_channel();
487        pin_mut!(rx);
488
489        let now = Instant::now();
490        tx.send((1, now).into()).await?;
491        tx.send((1, now).into()).await?;
492        tx.close().await?;
493
494        assert_eq!(Some(1), rx.next().await);
495        assert_eq!(None, rx.next().await);
496
497        Ok(())
498    }
499
500    #[test_log::test(tokio::test)]
501    async fn skip_delay_queue_should_yield_items_in_order() -> anyhow::Result<()> {
502        let (mut tx, rx) = skip_delay_channel();
503        pin_mut!(rx);
504
505        let now = Instant::now();
506        tx.send((2, now).into()).await?;
507        tx.send((1, now).into()).await?;
508        tx.close().await?;
509
510        assert_eq!(Some(1), rx.next().await);
511        assert_eq!(Some(2), rx.next().await);
512        assert_eq!(None, rx.next().await);
513
514        Ok(())
515    }
516
517    #[test_log::test(tokio::test)]
518    async fn skip_delay_queue_should_yield_fed_items_in_order() -> anyhow::Result<()> {
519        let (mut tx, rx) = skip_delay_channel();
520        pin_mut!(rx);
521
522        let now = Instant::now();
523        tx.feed((2, now).into()).await?;
524        tx.feed((1, now).into()).await?;
525        tx.flush().await?;
526        tx.close().await?;
527
528        assert_eq!(Some(1), rx.next().await);
529        assert_eq!(Some(2), rx.next().await);
530        assert_eq!(None, rx.next().await);
531
532        Ok(())
533    }
534
535    #[test_log::test(tokio::test)]
536    async fn skip_delay_queue_should_not_send_items_when_closed() -> anyhow::Result<()> {
537        let (mut tx, rx) = skip_delay_channel();
538        pin_mut!(rx);
539        tx.close().await?;
540
541        let now = Instant::now();
542        tx.send((1, now).into()).await.unwrap_err();
543        tx.close().await.unwrap_err();
544
545        assert_eq!(None, rx.next().await);
546
547        Ok(())
548    }
549
550    #[test_log::test(tokio::test)]
551    async fn skip_delay_queue_should_continuously_yield_items() -> anyhow::Result<()> {
552        let (mut tx, rx) = skip_delay_channel();
553
554        let items = [5, 2, 1, 4, 3];
555
556        let now = Instant::now();
557        let timed_items = (0..5)
558            .map(|i| (items[i], now + Duration::from_millis(100) * (i as u32)))
559            .collect::<Vec<_>>();
560
561        let timed_items_clone = timed_items.clone();
562        let jh = hopr_utils::runtime::prelude::spawn(async move {
563            for (n, time) in timed_items_clone {
564                tx.send((n, time).into()).await?;
565                hopr_utils::runtime::prelude::sleep(Duration::from_millis(50)).await;
566            }
567            tx.close().await?;
568            Ok::<_, std::io::Error>(())
569        });
570
571        let collected = rx.map(|item| (item, Instant::now())).collect::<Vec<_>>().await;
572
573        assert_eq!(timed_items.len(), collected.len());
574
575        for (i, (item, received_at)) in collected.into_iter().enumerate() {
576            assert_eq!(timed_items[i].0, item);
577            if received_at < timed_items[i].1 {
578                assert!(timed_items[i].1.saturating_duration_since(received_at) < Duration::from_millis(20));
579            } else {
580                assert!(received_at.saturating_duration_since(timed_items[i].1) < Duration::from_millis(20));
581            }
582        }
583
584        jh.await??;
585        Ok(())
586    }
587}
hopr_protocol_session/utils/skip_queue.rs

hopr_protocol_session/utils/
skip_queue.rs
