tokio::task::coop

Function has_budget_remaining

Source
pub fn has_budget_remaining() -> bool
Expand description

Returns true if there is still budget left on the task.

§Examples

This example defines a Timeout future that requires a given future to complete before the specified duration elapses. If it does, its result is returned; otherwise, an error is returned and the future is canceled.

Note that the future could exhaust the budget before we evaluate the timeout. Using has_budget_remaining, we can detect this scenario and ensure the timeout is always checked.

pub struct Timeout<T> {
    future: T,
    delay: Pin<Box<Sleep>>,
}

impl<T> Future for Timeout<T>
where
    T: Future + Unpin,
{
    type Output = Result<T::Output, ()>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = Pin::into_inner(self);
        let future = Pin::new(&mut this.future);
        let delay = Pin::new(&mut this.delay);

        // check if the future is ready
        let had_budget_before = coop::has_budget_remaining();
        if let Poll::Ready(v) = future.poll(cx) {
            return Poll::Ready(Ok(v));
        }
        let has_budget_now = coop::has_budget_remaining();

        // evaluate the timeout
        if let (true, false) = (had_budget_before, has_budget_now) {
            // it is the underlying future that exhausted the budget
            ready!(pin!(coop::unconstrained(delay)).poll(cx));
        } else {
            ready!(delay.poll(cx));
        }
        return Poll::Ready(Err(()));
    }
}