From 147de9a228b8ec824b4910ff451f2e4bc73cb020 Mon Sep 17 00:00:00 2001
From: Piotr Sarna <sarna@scylladb.com>
Date: Tue, 15 Feb 2022 10:58:05 +0100
Subject: [PATCH] task: add consume_budget 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 - consume_budget,
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 points in the program which indicate that
some significant work was performed. It will yield only if the budget
is gone, which is a much better alternative to unconditional yielding,
which is a potentially heavy operation.
---
 tokio/src/task/consume_budget.rs | 33 ++++++++++++++++++++++++++++++++
 tokio/src/task/mod.rs            |  3 +++
 2 files changed, 36 insertions(+)
 create mode 100644 tokio/src/task/consume_budget.rs

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};