From 4d19a999371190e25f6138916d93e6c75093e4a6 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Sat, 16 Nov 2019 22:24:28 -0800
Subject: [PATCH] runtime: set spawn context on enter (#1780)

---
 tokio/src/runtime/basic_scheduler.rs     |  9 +++++++++
 tokio/src/runtime/handle.rs              |  8 +++++++-
 tokio/src/runtime/thread_pool/mod.rs     |  2 +-
 tokio/src/runtime/thread_pool/spawner.rs |  8 ++++++++
 tokio/tests/rt_common.rs                 | 10 ++++++++++
 5 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/tokio/src/runtime/basic_scheduler.rs b/tokio/src/runtime/basic_scheduler.rs
index cb99cf183..affe2a5e6 100644
--- a/tokio/src/runtime/basic_scheduler.rs
+++ b/tokio/src/runtime/basic_scheduler.rs
@@ -179,6 +179,15 @@ impl Spawner {
         self.scheduler.schedule(task);
         handle
     }
+
+    /// Enter the executor context
+    pub(crate) fn enter<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce() -> R,
+    {
+        use crate::runtime::global;
+        global::with_basic_scheduler(&*self.scheduler, f)
+    }
 }
 
 impl SchedulerPriv {
diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs
index 936448405..533477781 100644
--- a/tokio/src/runtime/handle.rs
+++ b/tokio/src/runtime/handle.rs
@@ -90,7 +90,13 @@ impl Handle {
         self.blocking_spawner.enter(|| {
             let _io = io::set_default(&self.io_handles[0]);
 
-            time::with_default(&self.time_handles[0], &self.clock, f)
+            time::with_default(&self.time_handles[0], &self.clock, || match &self.kind {
+                Kind::Shell => f(),
+                #[cfg(feature = "rt-core")]
+                Kind::Basic(spawner) => spawner.enter(f),
+                #[cfg(feature = "rt-full")]
+                Kind::ThreadPool(spawner) => spawner.enter(f),
+            })
         })
     }
 }
diff --git a/tokio/src/runtime/thread_pool/mod.rs b/tokio/src/runtime/thread_pool/mod.rs
index 599ce5480..b45d37072 100644
--- a/tokio/src/runtime/thread_pool/mod.rs
+++ b/tokio/src/runtime/thread_pool/mod.rs
@@ -117,7 +117,7 @@ impl ThreadPool {
     where
         F: Future,
     {
-        crate::runtime::global::with_thread_pool(self.spawner(), || {
+        self.spawner.enter(|| {
             let mut enter = crate::runtime::enter();
             enter.block_on(future)
         })
diff --git a/tokio/src/runtime/thread_pool/spawner.rs b/tokio/src/runtime/thread_pool/spawner.rs
index e2975313b..4773ea9ad 100644
--- a/tokio/src/runtime/thread_pool/spawner.rs
+++ b/tokio/src/runtime/thread_pool/spawner.rs
@@ -37,6 +37,14 @@ impl Spawner {
         self.workers.spawn_typed(future)
     }
 
+    /// Enter the executor context
+    pub(crate) fn enter<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce() -> R,
+    {
+        crate::runtime::global::with_thread_pool(self, f)
+    }
+
     /// Reference to the worker set. Used by `ThreadPool` to initiate shutdown.
     pub(super) fn workers(&self) -> &slice::Set<Box<dyn Unpark>> {
         &*self.workers
diff --git a/tokio/tests/rt_common.rs b/tokio/tests/rt_common.rs
index 81fe68010..33f779c70 100644
--- a/tokio/tests/rt_common.rs
+++ b/tokio/tests/rt_common.rs
@@ -418,6 +418,16 @@ rt_test! {
         .await
     }
 
+    #[test]
+    fn enter_and_spawn() {
+        let mut rt = rt();
+        let handle = rt.enter(|| {
+            tokio::spawn(async {})
+        });
+
+        assert_ok!(rt.block_on(handle));
+    }
+
     async fn client_server(tx: mpsc::Sender<()>) {
         let mut server = assert_ok!(TcpListener::bind("127.0.0.1:0").await);