stop_token/
future.rs

1//! Extension methods and types for the `Future` trait.
2
3use crate::{deadline::TimedOutError, Deadline};
4use core::future::Future;
5use core::pin::Pin;
6
7use pin_project_lite::pin_project;
8use std::task::{Context, Poll};
9
10/// Extend the `Future` trait with the `until` method.
11pub trait FutureExt: Future {
12    /// Run a future until it resolves, or until a deadline is hit.
13    fn timeout_at<T>(self, target: T) -> TimeoutAt<Self>
14    where
15        Self: Sized,
16        T: Into<Deadline>,
17    {
18        TimeoutAt {
19            deadline: target.into(),
20            future: self,
21        }
22    }
23}
24
25impl<F: Future> FutureExt for F {}
26
27pin_project! {
28    /// Run a future until it resolves, or until a deadline is hit.
29    ///
30    /// This method is returned by [`FutureExt::deadline`].
31    #[must_use = "Futures do nothing unless polled or .awaited"]
32    #[derive(Debug)]
33    pub struct TimeoutAt<F> {
34        #[pin]
35        future: F,
36        #[pin]
37        deadline: Deadline,
38    }
39}
40
41impl<F> Future for TimeoutAt<F>
42where
43    F: Future,
44{
45    type Output = Result<F::Output, TimedOutError>;
46
47    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
48        let this = self.project();
49        if let Poll::Ready(()) = this.deadline.poll(cx) {
50            return Poll::Ready(Err(TimedOutError::new()));
51        }
52        match this.future.poll(cx) {
53            Poll::Pending => Poll::Pending,
54            Poll::Ready(it) => Poll::Ready(Ok(it)),
55        }
56    }
57}
stop_token/future.rs

stop_token/
future.rs
