pollable_map/stream/
timeout_set.rs

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

pollable_map/stream/
timeout_set.rs
