Skip to content

Commit

Permalink
yield_now: add maybe_yield_now for cooperative scheduling
Browse files Browse the repository at this point in the history
For cpu-only computations that do not use any Tokio resources,
budgeting does not really kick in in order to yield and prevent
other tasks from starvation. The new mechanism - maybe_yield_now,
performs a budget check, consumes a unit of it, and yields only
if the task exceeded the budget. That allows cpu-intenstive
computations to define yield points that yield only conditionally,
since unconditional yielding to runtime is a potentially heavy operation.
  • Loading branch information
psarna committed Feb 14, 2022
1 parent 8fb15da commit 5c0bbc6
Showing 1 changed file with 38 additions and 0 deletions.
38 changes: 38 additions & 0 deletions tokio/src/task/yield_now.rs
Expand Up @@ -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
}

0 comments on commit 5c0bbc6

Please sign in to comment.