pollable_map/futures/
timeout_map.rs

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

pollable_map/futures/
timeout_map.rs
