diff --git a/tokio/src/task/yield_now.rs b/tokio/src/task/yield_now.rs index 148e3dc0c87..a370a1a6748 100644 --- a/tokio/src/task/yield_now.rs +++ b/tokio/src/task/yield_now.rs @@ -56,3 +56,41 @@ pub async fn yield_now() { YieldNow { yielded: false }.await } + +/// Yields 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. +/// +/// See also the usage example in the [task module](index.html#yield_now). +#[must_use = "maybe_yield_now does nothing unless polled/`await`-ed"] +pub async fn maybe_yield_now() { + struct MaybeYieldNow { + status: Poll<()>, + } + + impl Future for MaybeYieldNow { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + if self.status.is_ready() { + return self.status; + } + self.status = match crate::coop::poll_proceed(cx) { + Poll::Ready(restore) => { + restore.made_progress(); + Poll::Ready(()) + }, + Poll::Pending => { + cx.waker().wake_by_ref(); + Poll::Pending + }, + }; + self.status + } + } + + MaybeYieldNow { status: Poll::Pending }.await +} \ No newline at end of file