rt: allow block_on inside block_in_place inside block_on (#2645)

A fast path in block_on_place was failing to call exit() in the case where we
were in a block_on call.

Fixes: #2639

Co-authored-by: Bryan Donlan <[email protected]>

tokio/src/runtime/thread_pool/worker.rs

@@ -200,16 +200,19 @@ cfg_blocking! {
         }
 
         let mut had_core = false;
+        let mut had_entered = false;
 
         CURRENT.with(|maybe_cx| {
             match (crate::runtime::enter::context(),  maybe_cx.is_some()) {
                 (EnterContext::Entered { .. }, true) => {
                     // We are on a thread pool runtime thread, so we just need to set up blocking.
+                    had_entered = true;
                 }
                 (EnterContext::Entered { allow_blocking }, false) => {
                     // We are on an executor, but _not_ on the thread pool.
                     // That is _only_ okay if we are in a thread pool runtime's block_on method:
                     if allow_blocking {
+                        had_entered = true;
                         return;
                     } else {
                         // This probably means we are on the basic_scheduler or in a LocalSet,
@@ -256,12 +259,13 @@ cfg_blocking! {
             runtime::spawn_blocking(move || run(worker));
         });
 
-
         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 if had_entered {
             crate::runtime::enter::exit(f)
         } else {
             f()

tokio/tests/task_blocking.rs

@@ -96,3 +96,22 @@ fn no_block_in_basic_block_on() {
         task::block_in_place(|| {});
     });
 }
+
+#[test]
+fn can_enter_basic_rt_from_within_block_in_place() {
+    let mut outer = tokio::runtime::Builder::new()
+        .threaded_scheduler()
+        .build()
+        .unwrap();
+
+    outer.block_on(async {
+        tokio::task::block_in_place(|| {
+            let mut inner = tokio::runtime::Builder::new()
+                .basic_scheduler()
+                .build()
+                .unwrap();
+
+            inner.block_on(async {})
+        })
+    });
+}