From 0366a3e6d1aa4e7bf4a1c717680dd0947589264b Mon Sep 17 00:00:00 2001
From: Jon Gjengset <jon@thesquareplanet.com>
Date: Tue, 28 Jul 2020 19:58:33 -0400
Subject: [PATCH] Reset coop budget when blocking in block_on (#2711)

Previously, we would fail to reset the coop budget in this case, making
it so that `coop::poll_proceed` would perpetually yield `Poll::Pending`
in nested executers even when run in `block_in_place`.

This is also a further improvement on #2645.
---
 tokio/src/runtime/thread_pool/worker.rs |  6 +--
 tokio/tests/task_blocking.rs            | 67 +++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs
index abe20da59..c53c9384b 100644
--- a/tokio/src/runtime/thread_pool/worker.rs
+++ b/tokio/src/runtime/thread_pool/worker.rs
@@ -199,7 +199,6 @@ cfg_blocking! {
             }
         }
 
-        let mut had_core = false;
         let mut had_entered = false;
 
         CURRENT.with(|maybe_cx| {
@@ -248,7 +247,6 @@ cfg_blocking! {
             //
             // First, move the core back into the worker's shared core slot.
             cx.worker.core.set(core);
-            had_core = true;
 
             // Next, clone the worker handle and send it to a new thread for
             // processing.
@@ -259,13 +257,11 @@ cfg_blocking! {
             runtime::spawn_blocking(move || run(worker));
         });
 
-        if had_core {
+        if had_entered {
             // 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 if had_entered {
             crate::runtime::enter::exit(f)
         } else {
             f()
diff --git a/tokio/tests/task_blocking.rs b/tokio/tests/task_blocking.rs
index 50c070a35..4ca1596e0 100644
--- a/tokio/tests/task_blocking.rs
+++ b/tokio/tests/task_blocking.rs
@@ -176,3 +176,70 @@ fn can_shutdown_now_in_runtime() {
         rt.shutdown_background();
     });
 }
+
+#[test]
+fn coop_disabled_in_block_in_place() {
+    let mut outer = tokio::runtime::Builder::new()
+        .threaded_scheduler()
+        .enable_time()
+        .build()
+        .unwrap();
+
+    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    for i in 0..200 {
+        tx.send(i).unwrap();
+    }
+    drop(tx);
+
+    outer.block_on(async move {
+        let jh = tokio::spawn(async move {
+            tokio::task::block_in_place(move || {
+                futures::executor::block_on(async move {
+                    use tokio::stream::StreamExt;
+                    assert_eq!(rx.fold(0, |n, _| n + 1).await, 200);
+                })
+            })
+        });
+
+        tokio::time::timeout(Duration::from_secs(1), jh)
+            .await
+            .expect("timed out (probably hanging)")
+            .unwrap()
+    });
+}
+
+#[test]
+fn coop_disabled_in_block_in_place_in_block_on() {
+    let (done_tx, done_rx) = std::sync::mpsc::channel();
+    let done = done_tx.clone();
+    thread::spawn(move || {
+        let mut outer = tokio::runtime::Builder::new()
+            .threaded_scheduler()
+            .build()
+            .unwrap();
+
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        for i in 0..200 {
+            tx.send(i).unwrap();
+        }
+        drop(tx);
+
+        outer.block_on(async move {
+            tokio::task::block_in_place(move || {
+                futures::executor::block_on(async move {
+                    use tokio::stream::StreamExt;
+                    assert_eq!(rx.fold(0, |n, _| n + 1).await, 200);
+                })
+            })
+        });
+
+        let _ = done.send(Ok(()));
+    });
+
+    thread::spawn(move || {
+        thread::sleep(Duration::from_secs(1));
+        let _ = done_tx.send(Err("timed out (probably hanging)"));
+    });
+
+    done_rx.recv().unwrap().unwrap();
+}