pollable_map/optional/
timeout.rs

1use crate::common::Timed;
2use crate::error::TimedError;
3use crate::optional::Optional;
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;
10
11/// A reusable future or stream based on `Option` that will time out after a specific duration as elapse.
12#[pin_project::pin_project]
13pub struct TimeoutOptional<T> {
14    duration: Duration,
15    #[pin]
16    task: Optional<Timed<T>>,
17}
18
19impl<T> Deref for TimeoutOptional<T> {
20    type Target = Optional<Timed<T>>;
21    fn deref(&self) -> &Self::Target {
22        &self.task
23    }
24}
25
26impl<T> DerefMut for TimeoutOptional<T> {
27    fn deref_mut(&mut self) -> &mut Self::Target {
28        &mut self.task
29    }
30}
31
32impl<T> TimeoutOptional<T> {
33    /// Construct a new [`TimeoutOptional`].
34    pub fn new(duration: Duration) -> Self {
35        Self {
36            duration,
37            task: Optional::default(),
38        }
39    }
40
41    /// Construct a new [`TimeoutOptional`] with an existing [`Future`] or [`Stream`].
42    pub fn new_with_task(duration: Duration, task: T) -> Self {
43        Self {
44            duration,
45            task: Optional::new(Timed::new(task, duration)),
46        }
47    }
48
49    /// Construct a new [`TimeoutOptional`] with an existing [`Future`].
50    pub fn new_with_future(duration: Duration, task: T) -> Self
51    where
52        T: Future,
53    {
54        Self {
55            duration,
56            task: Optional::with_future(Timed::new(task, duration)),
57        }
58    }
59
60    /// Construct a new [`TimeoutOptional`] with an existing [`Stream`].
61    pub fn new_with_stream(duration: Duration, task: T) -> Self
62    where
63        T: Stream,
64    {
65        Self {
66            duration,
67            task: Optional::with_stream(Timed::new(task, duration)),
68        }
69    }
70
71    /// Replaces the current the future or stream with a new one, returning the previous value if present.
72    pub fn replace(&mut self, task: T) -> Option<T> {
73        let prev = self.task.replace(Timed::new(task, self.duration));
74        prev.map(|item| item.into_inner())
75    }
76
77    /// Replaces the current future or stream in place without moving the previous value.
78    pub fn set(self: Pin<&mut Self>, task: T) {
79        let this = self.project();
80        this.task.set(Timed::new(task, *this.duration));
81    }
82}
83
84impl<T: Future> Future for TimeoutOptional<T> {
85    type Output = Result<T::Output, TimedError>;
86    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
87        let this = self.project();
88        this.task.poll(cx).map_err(|_| TimedError)
89    }
90}
91
92impl<T: Stream> Stream for TimeoutOptional<T> {
93    type Item = Result<T::Item, TimedError>;
94    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
95        let this = self.project();
96        this.task.poll_next(cx).map_err(|_| TimedError)
97    }
98}
99
100#[cfg(test)]
101mod test {
102    use crate::optional::timeout::TimeoutOptional;
103    use core::future::pending;
104    use core::pin::Pin;
105    use core::time::Duration;
106    use futures::future::ready;
107
108    #[test]
109    fn test_timeout_optional_ready() {
110        let mut task = TimeoutOptional::new_with_task(Duration::from_secs(1), ready(()));
111        futures::executor::block_on(async move {
112            let fut = Pin::new(&mut task);
113            match fut.await {
114                Ok(_) => assert!(task.is_none()),
115                Err(e) => panic!("unexpected error: {e}"),
116            }
117        })
118    }
119
120    #[test]
121    fn test_timeout_optional_timeout() {
122        let mut task = TimeoutOptional::new_with_task(Duration::from_millis(10), pending::<()>());
123
124        futures::executor::block_on(async move {
125            let fut = Pin::new(&mut task);
126            match fut.await {
127                Ok(_) => unreachable!("should time out"),
128                Err(_) => {
129                    assert!(task.is_none());
130                }
131            }
132        })
133    }
134
135    #[test]
136    fn reusable_pinned_timeout_optional_future() {
137        async fn value(value: i32) -> i32 {
138            value
139        }
140
141        let task = TimeoutOptional::new_with_future(Duration::from_secs(1), value(0));
142        futures::pin_mut!(task);
143
144        futures::executor::block_on(async {
145            assert_eq!(task.as_mut().await.expect("future should not time out"), 0);
146            assert!(task.is_none());
147
148            task.as_mut().set(value(1));
149            assert!(task.is_some());
150
151            assert_eq!(task.as_mut().await.expect("future should not time out"), 1);
152            assert!(task.is_none());
153        });
154    }
155}
pollable_map/optional/timeout.rs

pollable_map/optional/
timeout.rs
