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