pollable_map/futures/
timeout_map.rs

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

pollable_map/futures/
timeout_map.rs
