pollable_map/futures/
timeout_set.rs1use 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 pub fn new(duration: Duration) -> Self {
36 Self {
37 duration,
38 set: FutureSet::new(),
39 }
40 }
41
42 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}