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: Poll<()> = std::task::Poll::Pending;
27
28 stdPollFn) -> …>::future::poll_fn(move |cx: &mut Context<'_>| {
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: RestoreOnPending| {
34 restore.made_progress();
35 });
36 status
37 })
38 .await
39}
40