From 5c0bbc65fa5f66359d5eea827a33be112972580f Mon Sep 17 00:00:00 2001
From: Piotr Sarna <sarna@scylladb.com>
Date: Mon, 14 Feb 2022 16:24:55 +0100
Subject: [PATCH] yield_now: add maybe_yield_now for cooperative scheduling

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.
---
 tokio/src/task/yield_now.rs | 38 +++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

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