pollable_map/stream/
timeout_map.rs

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

pollable_map/stream/
timeout_map.rs
