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