tokio/task/coop/
consume_budget.rs

1/// Consumes a unit of budget and returns the execution back to the Tokio
2/// runtime *if* the task's coop budget was exhausted.
3///
4/// The task will only yield if its entire coop budget has been exhausted.
5/// This function can be used in order to insert optional yield points into long
6/// computations that do not use Tokio resources like sockets or semaphores,
7/// without redundantly yielding to the runtime each time.
8///
9/// # Examples
10///
11/// Make sure that a function which returns a sum of (potentially lots of)
12/// iterated values is cooperative.
13///
14/// ```
15/// async fn sum_iterator(input: &mut impl std::iter::Iterator<Item=i64>) -> i64 {
16///     let mut sum: i64 = 0;
17///     while let Some(i) = input.next() {
18///         sum += i;
19///         tokio::task::consume_budget().await
20///     }
21///     sum
22/// }
23/// ```
24#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
25pub async fn consume_budget() {
26    let mut status = std::task::Poll::Pending;
27
28    std::future::poll_fn(move |cx| {
29        std::task::ready!(crate::trace::trace_leaf(cx));
30        if status.is_ready() {
31            return status;
32        }
33        status = crate::task::coop::poll_proceed(cx).map(|restore| {
34            restore.made_progress();
35        });
36        status
37    })
38    .await
39}