From 70364b707975b84daf223915d60a1610de2539a2 Mon Sep 17 00:00:00 2001
From: Alice Ryhl <aliceryhl@google.com>
Date: Mon, 15 May 2023 14:40:04 +0200
Subject: [PATCH] runtime: fix possible starvation when using lifo slot (#5686)

---
 tokio/src/runtime/coop.rs                     | 11 +++++++
 .../runtime/scheduler/multi_thread/worker.rs  |  6 ++++
 tokio/tests/rt_threaded.rs                    | 29 ++++++++++++++++++-
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/tokio/src/runtime/coop.rs b/tokio/src/runtime/coop.rs
index f3ed17cff..b1177c2a4 100644
--- a/tokio/src/runtime/coop.rs
+++ b/tokio/src/runtime/coop.rs
@@ -119,6 +119,17 @@ cfg_rt_multi_thread! {
     pub(crate) fn set(budget: Budget) {
         let _ = context::budget(|cell| cell.set(budget));
     }
+
+    /// Consume one unit of progress from the current task's budget.
+    pub(crate) fn consume_one() {
+        let _ = context::budget(|cell| {
+            let mut budget = cell.get();
+            if let Some(ref mut counter) = budget.0 {
+                *counter = counter.saturating_sub(1);
+            }
+            cell.set(budget);
+        });
+    }
 }
 
 cfg_rt! {
diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs
index 148255a85..f83c6bc1e 100644
--- a/tokio/src/runtime/scheduler/multi_thread/worker.rs
+++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs
@@ -479,6 +479,12 @@ impl Context {
                     None => return Ok(core),
                 };
 
+                // Polling a task doesn't necessarily consume any budget, if it
+                // doesn't use any Tokio leaf futures. To prevent such tasks
+                // from using the lifo slot in an infinite loop, we consume an
+                // extra unit of budget between each iteration of the loop.
+                coop::consume_one();
+
                 if coop::has_budget_remaining() {
                     // Run the LIFO task, then loop
                     core.metrics.incr_poll_count();
diff --git a/tokio/tests/rt_threaded.rs b/tokio/tests/rt_threaded.rs
index c5984182c..feb2f9f8c 100644
--- a/tokio/tests/rt_threaded.rs
+++ b/tokio/tests/rt_threaded.rs
@@ -30,7 +30,8 @@ fn single_thread() {
     let _ = runtime::Builder::new_multi_thread()
         .enable_all()
         .worker_threads(1)
-        .build();
+        .build()
+        .unwrap();
 }
 
 #[test]
@@ -160,6 +161,32 @@ fn many_multishot_futures() {
     }
 }
 
+#[test]
+fn lifo_slot_budget() {
+    async fn my_fn() {
+        spawn_another();
+    }
+
+    fn spawn_another() {
+        tokio::spawn(my_fn());
+    }
+
+    let rt = runtime::Builder::new_multi_thread()
+        .enable_all()
+        .worker_threads(1)
+        .build()
+        .unwrap();
+
+    let (send, recv) = oneshot::channel();
+
+    rt.spawn(async move {
+        tokio::spawn(my_fn());
+        let _ = send.send(());
+    });
+
+    let _ = rt.block_on(recv);
+}
+
 #[test]
 fn spawn_shutdown() {
     let rt = rt();