Skip to main content

pollable_map/futures/
timeout_set.rs

1use crate::common::Timed;
2use crate::error::TimedError;
3use crate::futures::set::FutureSet;
4use core::future::Future;
5use core::ops::{Deref, DerefMut};
6use core::pin::Pin;
7use core::task::{Context, Poll};
8use core::time::Duration;
9use futures::stream::FusedStream;
10use futures::{Stream, StreamExt};
11
12pub struct TimeoutFutureSet<S> {
13    duration: Duration,
14    set: FutureSet<Timed<S>>,
15}
16
17impl<S> Deref for TimeoutFutureSet<S> {
18    type Target = FutureSet<Timed<S>>;
19    fn deref(&self) -> &Self::Target {
20        &self.set
21    }
22}
23
24impl<S> DerefMut for TimeoutFutureSet<S> {
25    fn deref_mut(&mut self) -> &mut Self::Target {
26        &mut self.set
27    }
28}
29
30impl<F> TimeoutFutureSet<F>
31where
32    F: Future,
33{
34    /// Create an empty [`TimeoutFutureSet`]
35    pub fn new(duration: Duration) -> Self {
36        Self {
37            duration,
38            set: FutureSet::new(),
39        }
40    }
41
42    /// Insert a future into the set of futures.
43    pub fn insert(&mut self, future: F) -> bool {
44        self.set.insert(Timed::new(future, self.duration))
45    }
46}
47
48impl<F> Stream for TimeoutFutureSet<F>
49where
50    F: Future,
51{
52    type Item = Result<F::Output, TimedError>;
53    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
54        self.set.poll_next_unpin(cx)
55    }
56
57    fn size_hint(&self) -> (usize, Option<usize>) {
58        self.set.size_hint()
59    }
60}
61
62impl<F> FusedStream for TimeoutFutureSet<F>
63where
64    F: Future,
65{
66    fn is_terminated(&self) -> bool {
67        self.set.is_terminated()
68    }
69}
70
71#[cfg(test)]
72mod test {
73    use crate::{error::TimedError, futures::timeout_set::TimeoutFutureSet};
74    use futures::StreamExt;
75    use std::time::Duration;
76
77    #[test]
78    fn timeout_set() {
79        let mut list = TimeoutFutureSet::new(Duration::from_millis(100));
80        assert!(list.insert(futures::future::pending::<()>()));
81
82        futures::executor::block_on(async move {
83            let result = list.next().await;
84            let Some(Err(e)) = result else {
85                unreachable!("result is err");
86            };
87
88            assert_eq!(e, TimedError);
89        });
90    }
91
92    #[test]
93    fn valid_stream() {
94        let mut list = TimeoutFutureSet::new(Duration::from_secs(10));
95        assert!(list.insert(futures::future::ready(0)));
96
97        futures::executor::block_on(async move {
98            let result = list.next().await;
99            let Some(Ok(val)) = result else {
100                unreachable!("result is err");
101            };
102
103            assert_eq!(val, 0);
104        });
105    }
106}