tokio/task/
consume_budget.rs

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

tokio/task/
consume_budget.rs
