tang_rs/util/
timeout.rs

1use core::future::Future;
2use core::pin::Pin;
3use core::task::{Context, Poll};
4
5#[derive(Debug)]
6pub struct ManagerTimeout<T, D> {
7    future: T,
8    delay: D,
9}
10
11impl<T, D> ManagerTimeout<T, D> {
12    pub fn new(future: T, delay: D) -> Self {
13        Self { future, delay }
14    }
15}
16
17impl<T, D, E> Future for ManagerTimeout<T, D>
18where
19    T: Future,
20    D: Future<Output = E>,
21{
22    type Output = Result<T::Output, E>;
23
24    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
25        unsafe {
26            let p = self.as_mut().map_unchecked_mut(|this| &mut this.future);
27            if let Poll::Ready(fut) = p.poll(cx) {
28                return Poll::Ready(Ok(fut));
29            }
30        }
31
32        unsafe {
33            match self.map_unchecked_mut(|this| &mut this.delay).poll(cx) {
34                Poll::Ready(e) => Poll::Ready(Err(e)),
35                Poll::Pending => Poll::Pending,
36            }
37        }
38    }
39}