From af6f2322f5eba358f55883c815711867eb6b9eb1 Mon Sep 17 00:00:00 2001
From: Bryan Donlan <bdonlan@amazon.com>
Date: Mon, 6 Jul 2020 23:51:49 +0000
Subject: [PATCH] Allow block_on inside block_in_place inside block_on

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
---
 tokio/src/runtime/thread_pool/worker.rs |  6 +++++-
 tokio/tests/task_blocking.rs            | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs
index 0777857e208..abe20da59cf 100644
--- a/tokio/src/runtime/thread_pool/worker.rs
+++ b/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()
diff --git a/tokio/tests/task_blocking.rs b/tokio/tests/task_blocking.rs
index 72fed01e961..b4019868506 100644
--- a/tokio/tests/task_blocking.rs
+++ b/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 {})
+        })
+    });
+}