runtime: set spawn context on enter (#1780)

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 {

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),
+            })
         })
     }
 }

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)
         })

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

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);