diff --git a/tokio/src/task/consume_budget.rs b/tokio/src/task/consume_budget.rs new file mode 100644 index 00000000000..fbab96e82d9 --- /dev/null +++ b/tokio/src/task/consume_budget.rs @@ -0,0 +1,33 @@ +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// Consumes a unit of budget and returns the execution back to the Tokio +/// runtime *if* the task went out ouf budget. +/// +/// The task will only yield if it ran out of its coop budget. +/// It 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 runtime each time. +#[cfg_attr(docsrs, doc(cfg(feature = "rt")))] +pub async fn consume_budget() { + struct ConsumeBudget { + status: Poll<()>, + } + + impl Future for ConsumeBudget { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + if self.status.is_ready() { + return self.status; + } + self.status = crate::coop::poll_proceed(cx).map(|restore| { + restore.made_progress(); + }); + self.status + } + } + + ConsumeBudget { status: Poll::Pending }.await +} diff --git a/tokio/src/task/mod.rs b/tokio/src/task/mod.rs index 7d254190148..c5431d3fb4b 100644 --- a/tokio/src/task/mod.rs +++ b/tokio/src/task/mod.rs @@ -291,6 +291,9 @@ cfg_rt! { mod yield_now; pub use yield_now::yield_now; + mod consume_budget; + pub use consume_budget::consume_budget; + mod local; pub use local::{spawn_local, LocalSet};