add a Handle::current() for getting access to TLS Runtime

tokio/src/runtime/blocking/mod.rs

@@ -56,6 +56,12 @@ cfg_not_blocking_impl! {
             self
         }
 
+        #[cfg(any(
+            feature = "blocking",
+            feature = "dns",
+            feature = "fs",
+            feature = "io-std",
+        ))]
         pub(crate) fn enter<F, R>(&self, f: F) -> R
         where
             F: FnOnce() -> R,

tokio/src/runtime/blocking/pool.rs

@@ -8,7 +8,6 @@ use crate::runtime::blocking::task::BlockingTask;
 use crate::runtime::{self, context::ThreadContext, io, time, Builder, Callback};
 use crate::task::{self, JoinHandle};
 
-use std::cell::Cell;
 use std::collections::VecDeque;
 use std::fmt;
 use std::time::Duration;
@@ -68,28 +67,21 @@ struct Shared {
 
 type Task = task::Task<NoopSchedule>;
 
-thread_local! {
-    /// Thread-local tracking the current executor
-    static BLOCKING: Cell<Option<*const Spawner>> = Cell::new(None)
-}
-
 const KEEP_ALIVE: Duration = Duration::from_secs(10);
 
 /// Run the provided function on an executor dedicated to blocking operations.
 pub(crate) fn spawn_blocking<F, R>(func: F) -> JoinHandle<R>
 where
     F: FnOnce() -> R + Send + 'static,
 {
-    BLOCKING.with(|cell| {
-        let schedule = match cell.get() {
-            Some(ptr) => unsafe { &*ptr },
-            None => panic!("not currently running on the Tokio runtime."),
-        };
+    use crate::runtime::context::ThreadContext;
+
+    let schedule =
+        ThreadContext::blocking_spawner().expect("not currently running on the Tokio runtime.");
 
-        let (task, handle) = task::joinable(BlockingTask::new(func));
-        schedule.schedule(task);
-        handle
-    })
+    let (task, handle) = task::joinable(BlockingTask::new(func));
+    schedule.schedule(task);
+    handle
 }
 
 // ===== impl BlockingPool =====
@@ -168,30 +160,10 @@ impl Spawner {
     where
         F: FnOnce() -> R,
     {
-        // While scary, this is safe. The function takes a `&BlockingPool`,
-        // which guarantees that the reference lives for the duration of
-        // `with_pool`.
-        //
-        // Because we are always clearing the TLS value at the end of the
-        // function, we can cast the reference to 'static which thread-local
-        // cells require.
-        BLOCKING.with(|cell| {
-            let was = cell.replace(None);
-
-            // Ensure that the pool is removed from the thread-local context
-            // when leaving the scope. This handles cases that involve panicking.
-            struct Reset<'a>(&'a Cell<Option<*const Spawner>>, Option<*const Spawner>);
-
-            impl Drop for Reset<'_> {
-                fn drop(&mut self) {
-                    self.0.set(self.1);
-                }
-            }
+        let ctx = crate::runtime::context::ThreadContext::clone_current();
+        let _e = ctx.with_blocking_spawner(self.clone()).enter();
 
-            let _reset = Reset(cell, was);
-            cell.set(Some(self as *const Spawner));
-            f()
-        })
+        f()
     }
 
     fn schedule(&self, task: Task) {
@@ -248,6 +220,7 @@ impl Spawner {
             self.inner.io_handle.clone(),
             self.inner.time_handle.clone(),
             Some(self.inner.clock.clone()),
+            Some(self.clone()),
         );
         let spawner = self.clone();
         builder

tokio/src/runtime/context.rs

@@ -20,6 +20,9 @@ pub(crate) struct ThreadContext {
 
     /// Source of `Instant::now()`
     clock: Option<crate::runtime::time::Clock>,
+
+    /// Blocking pool spawner
+    blocking_spawner: Option<crate::runtime::blocking::Spawner>,
 }
 
 impl Default for ThreadContext {
@@ -35,6 +38,7 @@ impl Default for ThreadContext {
             #[cfg(any(not(feature = "time"), loom))]
             time_handle: (),
             clock: None,
+            blocking_spawner: None,
         }
     }
 }
@@ -48,6 +52,7 @@ impl ThreadContext {
         io_handle: crate::runtime::io::Handle,
         time_handle: crate::runtime::time::Handle,
         clock: Option<crate::runtime::time::Clock>,
+        blocking_spawner: Option<crate::runtime::blocking::Spawner>,
     ) -> Self {
         ThreadContext {
             spawner,
@@ -60,6 +65,7 @@ impl ThreadContext {
             #[cfg(any(not(feature = "time"), loom))]
             time_handle,
             clock,
+            blocking_spawner,
         }
     }
 
@@ -81,31 +87,27 @@ impl ThreadContext {
         })
     }
 
-    #[cfg(all(feature = "io-driver", not(loom)))]
     pub(crate) fn io_handle() -> crate::runtime::io::Handle {
         CONTEXT.with(|ctx| match *ctx.borrow() {
             Some(ref ctx) => ctx.io_handle.clone(),
-            None => None,
+            None => Default::default(),
         })
     }
 
-    #[cfg(all(feature = "time", not(loom)))]
     pub(crate) fn time_handle() -> crate::runtime::time::Handle {
         CONTEXT.with(|ctx| match *ctx.borrow() {
             Some(ref ctx) => ctx.time_handle.clone(),
-            None => None,
+            None => Default::default(),
         })
     }
 
-    #[cfg(feature = "rt-core")]
     pub(crate) fn spawn_handle() -> Option<Spawner> {
         CONTEXT.with(|ctx| match *ctx.borrow() {
             Some(ref ctx) => Some(ctx.spawner.clone()),
             None => None,
         })
     }
 
-    #[cfg(all(feature = "test-util", feature = "time"))]
     pub(crate) fn clock() -> Option<crate::runtime::time::Clock> {
         CONTEXT.with(
             |ctx| match ctx.borrow().as_ref().map(|ctx| ctx.clock.clone()) {
@@ -114,6 +116,31 @@ impl ThreadContext {
             },
         )
     }
+
+    pub(crate) fn blocking_spawner() -> Option<crate::runtime::blocking::Spawner> {
+        CONTEXT.with(|ctx| {
+            match ctx
+                .borrow()
+                .as_ref()
+                .map(|ctx| ctx.blocking_spawner.clone())
+            {
+                Some(Some(blocking_spawner)) => Some(blocking_spawner),
+                _ => None,
+            }
+        })
+    }
+}
+
+cfg_blocking_impl! {
+    impl ThreadContext {
+        pub(crate) fn with_blocking_spawner(
+            mut self,
+            blocking_spawner: crate::runtime::blocking::Spawner,
+        ) -> Self {
+            self.blocking_spawner.replace(blocking_spawner);
+            self
+        }
+    }
 }
 
 /// [`ThreadContextDropGuard`] will replace the `previous` thread context on drop.

tokio/src/runtime/handle.rs

@@ -35,9 +35,50 @@ impl Handle {
             self.io_handle.clone(),
             self.time_handle.clone(),
             Some(self.clock.clone()),
+            Some(self.blocking_spawner.clone()),
         )
         .enter();
-        self.blocking_spawner.enter(|| f())
+
+        f()
+    }
+
+    /// Returns a Handle view over the currently running Runtime
+    ///
+    /// # Panic
+    ///
+    /// A Runtime must have been started or this will panic
+    ///
+    /// # Examples
+    ///
+    /// This allows for the current handle to be gotten when running in a `#`
+    ///
+    /// ```
+    /// # use tokio::runtime::Runtime;
+    ///
+    /// # fn dox() {
+    /// # let rt = Runtime::new().unwrap();
+    /// # rt.spawn(async {
+    /// use tokio::runtime::Handle;
+    ///
+    /// let handle = Handle::current();
+    /// handle.spawn(async {
+    ///     println!("now running in the existing Runtime");
+    /// })
+    /// # });
+    /// # }
+    /// ```
+    pub fn current() -> Self {
+        use crate::runtime::context::ThreadContext;
+
+        Handle {
+            spawner: ThreadContext::spawn_handle()
+                .expect("not currently running on the Tokio runtime."),
+            io_handle: ThreadContext::io_handle(),
+            time_handle: ThreadContext::time_handle(),
+            clock: ThreadContext::clock().expect("not currently running on the Tokio runtime."),
+            blocking_spawner: ThreadContext::blocking_spawner()
+                .expect("not currently running on the Tokio runtime."),
+        }
     }
 }