rt: set task budget after block_in_place call (#2502)

In some cases, when a call to `block_in_place` completes, the runtime is
reinstated on the thread. In this case, the task budget must also be set
in order to avoid starving other tasks on the worker.

tokio/src/coop.rs

@@ -92,10 +92,20 @@ cfg_rt_threaded! {
 /// Run the given closure with a cooperative task budget. When the function
 /// returns, the budget is reset to the value prior to calling the function.
 #[inline(always)]
-pub(crate) fn budget<F, R>(f: F) -> R
-where
-    F: FnOnce() -> R,
-{
+pub(crate) fn budget<R>(f: impl FnOnce() -> R) -> R {
+    with_budget(Budget::initial(), f)
+}
+
+cfg_rt_threaded! {
+    /// Set the current task's budget
+    #[cfg(feature = "blocking")]
+    pub(crate) fn set(budget: Budget) {
+        CURRENT.with(|cell| cell.set(budget))
+    }
+}
+
+#[inline(always)]
+fn with_budget<R>(budget: Budget, f: impl FnOnce() -> R) -> R {
     struct ResetGuard<'a> {
         cell: &'a Cell<Budget>,
         prev: Budget,
@@ -110,7 +120,7 @@ where
     CURRENT.with(move |cell| {
         let prev = cell.get();
 
-        cell.set(Budget::initial());
+        cell.set(budget);
 
         let _guard = ResetGuard { cell, prev };
 
@@ -127,10 +137,14 @@ cfg_rt_threaded! {
 
 cfg_blocking_impl! {
     /// Forcibly remove the budgeting constraints early.
-    pub(crate) fn stop() {
+    ///
+    /// Returns the remaining budget
+    pub(crate) fn stop() -> Budget {
         CURRENT.with(|cell| {
+            let prev = cell.get();
             cell.set(Budget::unconstrained());
-        });
+            prev
+        })
     }
 }
 

tokio/src/runtime/thread_pool/worker.rs

@@ -4,6 +4,7 @@
 //! "core" is handed off to a new thread allowing the scheduler to continue to
 //! make progress while the originating thread blocks.
 
+use crate::coop;
 use crate::loom::rand::seed;
 use crate::loom::sync::{Arc, Mutex};
 use crate::park::{Park, Unpark};
@@ -179,31 +180,27 @@ cfg_blocking! {
         F: FnOnce() -> R,
     {
         // Try to steal the worker core back
-        struct Reset(bool);
+        struct Reset(coop::Budget);
 
         impl Drop for Reset {
             fn drop(&mut self) {
                 CURRENT.with(|maybe_cx| {
-                    if !self.0 {
-                        // We were not the ones to give away the core,
-                        // so we do not get to restore it either.
-                        // This is necessary so that with a nested
-                        // block_in_place, the inner block_in_place
-                        // does not restore the core.
-                        return;
-                    }
-
                     if let Some(cx) = maybe_cx {
                         let core = cx.worker.core.take();
                         let mut cx_core = cx.core.borrow_mut();
                         assert!(cx_core.is_none());
                         *cx_core = core;
+
+                        // Reset the task budget as we are re-entering the
+                        // runtime.
+                        coop::set(self.0);
                     }
                 });
             }
         }
 
         let mut had_core = false;
+
         CURRENT.with(|maybe_cx| {
             match (crate::runtime::enter::context(),  maybe_cx.is_some()) {
                 (EnterContext::Entered { .. }, true) => {
@@ -231,16 +228,12 @@ cfg_blocking! {
                     return;
                 }
             }
+
             let cx = maybe_cx.expect("no .is_some() == false cases above should lead here");
 
             // Get the worker core. If none is set, then blocking is fine!
             let core = match cx.core.borrow_mut().take() {
-                Some(core) => {
-                    // We are effectively leaving the executor, so we need to
-                    // forcibly end budgeting.
-                    crate::coop::stop();
-                    core
-                },
+                Some(core) => core,
                 None => return,
             };
 
@@ -263,9 +256,12 @@ cfg_blocking! {
             runtime::spawn_blocking(move || run(worker));
         });
 
-        let _reset = Reset(had_core);
 
         if had_core {
+            // Unset the current task's budget. Blocking sections are not
+            // constrained by task budgets.
+            let _reset = Reset(coop::stop());
+
             crate::runtime::enter::exit(f)
         } else {
             f()
@@ -349,7 +345,7 @@ impl Context {
         *self.core.borrow_mut() = Some(core);
 
         // Run the task
-        crate::coop::budget(|| {
+        coop::budget(|| {
             task.run();
 
             // As long as there is budget remaining and a task exists in the
@@ -368,7 +364,7 @@ impl Context {
                     None => return Ok(core),
                 };
 
-                if crate::coop::has_budget_remaining() {
+                if coop::has_budget_remaining() {
                     // Run the LIFO task, then loop
                     *self.core.borrow_mut() = Some(core);
                     task.run();

tokio/tests/rt_threaded.rs

@@ -7,6 +7,7 @@ use tokio::runtime::{self, Runtime};
 use tokio::sync::oneshot;
 use tokio_test::{assert_err, assert_ok};
 
+use futures::future::poll_fn;
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::atomic::AtomicUsize;
@@ -322,6 +323,64 @@ fn multi_threadpool() {
     done_rx.recv().unwrap();
 }
 
+// When `block_in_place` returns, it attempts to reclaim the yielded runtime
+// worker. In this case, the remainder of the task is on the runtime worker and
+// must take part in the cooperative task budgeting system.
+//
+// The test ensures that, when this happens, attempting to consume from a
+// channel yields occasionally even if there are values ready to receive.
+#[test]
+fn coop_and_block_in_place() {
+    use tokio::sync::mpsc;
+
+    let mut rt = tokio::runtime::Builder::new()
+        .threaded_scheduler()
+        // Setting max threads to 1 prevents another thread from claiming the
+        // runtime worker yielded as part of `block_in_place` and guarantees the
+        // same thread will reclaim the worker at the end of the
+        // `block_in_place` call.
+        .max_threads(1)
+        .build()
+        .unwrap();
+
+    rt.block_on(async move {
+        let (mut tx, mut rx) = mpsc::channel(1024);
+
+        // Fill the channel
+        for _ in 0..1024 {
+            tx.send(()).await.unwrap();
+        }
+
+        drop(tx);
+
+        tokio::spawn(async move {
+            // Block in place without doing anything
+            tokio::task::block_in_place(|| {});
+
+            // Receive all the values, this should trigger a `Pending` as the
+            // coop limit will be reached.
+            poll_fn(|cx| {
+                while let Poll::Ready(v) = {
+                    tokio::pin! {
+                        let fut = rx.recv();
+                    }
+
+                    Pin::new(&mut fut).poll(cx)
+                } {
+                    if v.is_none() {
+                        panic!("did not yield");
+                    }
+                }
+
+                Poll::Ready(())
+            })
+            .await
+        })
+        .await
+        .unwrap();
+    });
+}
+
 // Testing this does not panic
 #[test]
 fn max_threads() {