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.

1 files changed, 67 insertions, 0 deletions

diff --git a/tokio/tests/task_blocking.rs b/tokio/tests/task_blocking.rs
index 50c070a3..4ca1596e 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();
+}