madsim_real_tokio/task/
consume_budget.rs

1use std::task::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/// **Note**: This is an [unstable API][unstable]. The public API of this type
12/// may break in 1.x releases. See [the documentation on unstable
13/// features][unstable] for details.
14///
15/// # Examples
16///
17/// Make sure that a function which returns a sum of (potentially lots of)
18/// iterated values is cooperative.
19///
20/// ```
21/// async fn sum_iterator(input: &mut impl std::iter::Iterator<Item=i64>) -> i64 {
22///     let mut sum: i64 = 0;
23///     while let Some(i) = input.next() {
24///         sum += i;
25///         tokio::task::consume_budget().await
26///     }
27///     sum
28/// }
29/// ```
30/// [unstable]: crate#unstable-features
31#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]
32pub async fn consume_budget() {
33    let mut status = Poll::Pending;
34
35    crate::future::poll_fn(move |cx| {
36        ready!(crate::trace::trace_leaf(cx));
37        if status.is_ready() {
38            return status;
39        }
40        status = crate::runtime::coop::poll_proceed(cx).map(|restore| {
41            restore.made_progress();
42        });
43        status
44    })
45    .await
46}
madsim_real_tokio/task/consume_budget.rs

madsim_real_tokio/task/
consume_budget.rs
