Repository: tokio-rs/tokio
File: tokio/src/task/coop/mod.rs (Rust)

Code before:

    #[allow(unused_variables)]
    let maybe_guard = context::budget(|cell| {
        let prev = cell.get();
        cell.set(budget);

        ResetGuard { prev }
    });

    // The function is called regardless even if the budget is not successfully
    // set due to the thread-local being destroyed.
    f()
}

/// Returns `true` if there is still budget left on the task.
///
/// Futures created by the tokio library functions are budget-aware and yield when there is no more
/// budget left, but not all futures will be budget-aware. Consider future `A` that polls two inner
/// futures `B` and `C`, and returns `Poll::Ready` when one of them is ready. If both inner futures
/// were budget-aware, at some point the budget would be depleted which would cause both futures to
/// return `Poll::Pending` resulting in `A` returning `Poll::Pending` as well. Yielding all the way
/// back to the runtime, the budget would be reset, and `B` and `C` could make progress again.
/// Now let's consider `B` is budget-aware, but `C` is not and is always ready. The budget will be
/// depleted as before, but now since `C` always returns `Poll::Ready`, `A` will always return
/// `Poll::Ready` as well. This way, `A` will not yield back to the runtime which will keep the budget
/// depleted and `B` not making progress.
///
/// In these scenarios you could use [`has_budget_remaining`] to check whether the budget has been depleted
/// or not and act accordingly:
/// ```
/// # use std::future::{poll_fn, Future};
/// # use std::task::{Context, Poll};
/// # use std::pin::{pin, Pin};
/// # use std::sync::atomic::{AtomicBool, Ordering};
/// # use tokio::task::coop::{consume_budget, has_budget_remaining};
/// #
/// # #[tokio::main]
/// # async fn main() {
/// struct Greedy;
/// struct Aware;
/// struct Combined {
///     greedy: Greedy,
///     aware: Aware,
/// }
///
/// impl Future for Greedy {
///     type Output = ();
///
///     fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
///         Poll::Ready(())
///     }

Review comment (suggestion):
This explanation is confusing. And more generally, this seems like a niche use-case. Are there not more common use-cases for this that we can explain instead?

E.g. maybe just have this example:
```rs
if !has_budget_remaining() {
    tokio::task::yield_now().await;
}
```
and say that this is similar to `consume_budget`, but doesn't decrement the budget when non-zero.

Code after:

    #[allow(unused_variables)]
    let maybe_guard = context::budget(|cell| {
        let prev = cell.get();
        cell.set(budget);

        ResetGuard { prev }
    });

    // The function is called regardless even if the budget is not successfully
    // set due to the thread-local being destroyed.
    f()
}

/// 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.
///
/// ```
/// # use std::future::Future;
/// # use std::pin::{pin, Pin};
/// # use std::task::{ready, Context, Poll};
/// # use tokio::task::coop;
/// # use tokio::time::Sleep;
/// 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));