1 use 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")))]
consume_budget()27 pub 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 }
42