future_wrap/
lib.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use pin_project::pin_project;
6
7#[pin_project]
8pub struct WrappedFuture<Fut, F, O>
9    where
10        Fut: Future,
11        F: FnMut(Pin<&mut Fut>, &mut Context) -> Poll<O>,
12{
13    #[pin]
14    inner: Fut,
15    f: F,
16}
17
18impl<Fut, F, O> Future for WrappedFuture<Fut, F, O>
19    where
20        Fut: Future,
21        F: FnMut(Pin<&mut Fut>, &mut Context) -> Poll<O>,
22{
23    type Output = O;
24
25    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
26        let this = self.project();
27
28        (this.f)(this.inner, cx)
29    }
30}
31
32/// Adds a `wrap` function to all types that implement `Future`
33/// Lets you track each poll, interrupt the execution of the Future and change the return type
34/// ```
35/// use future_wrap::WrapFuture;
36/// use std::future::Future;
37/// use std::task::Poll;
38/// use std::time::{Duration, Instant};
39///
40/// let fut = some_async_fn();
41///
42/// let mut remaining_time = Duration::from_millis(10);
43/// fut.wrap(|fut, cx| {
44///     let poll_start = Instant::now();
45///
46///     println!("Poll start");
47///     let res = fut.poll(cx);
48///     println!("Poll end");
49///
50///     remaining_time = remaining_time.saturating_sub(poll_start.elapsed());
51///     if remaining_time.is_zero() {
52///         println!("Too much time spent on polls :(");
53///         Poll::Ready(None)
54///     } else {
55///         res.map(|v| Some(v))
56///     }
57/// }).await;
58/// ```
59pub trait WrapFuture<O>: Sized + Future {
60    fn wrap<F>(self, f: F) -> WrappedFuture<Self, F, O>
61        where
62            F: FnMut(Pin<&mut Self>, &mut Context) -> Poll<O>,
63    {
64        WrappedFuture {
65            inner: self,
66            f,
67        }
68    }
69}
70
71impl<F: Future, O> WrapFuture<O> for F {}
future_wrap/lib.rs

future_wrap/
lib.rs
