From 0193df3a593cb69d23414109118784de2948024c Mon Sep 17 00:00:00 2001 From: Benjamin Fry Date: Mon, 6 Jan 2020 11:32:21 -0800 Subject: [PATCH] rt: add a Handle::current() (#2040) Adds `Handle::current()` for accessing a handle to the runtime associated with the current thread. This handle can then be passed to other threads in order to spawn or perform other runtime related tasks. --- tokio/src/runtime/blocking/mod.rs | 6 ++++ tokio/src/runtime/blocking/pool.rs | 49 +++++++----------------------- tokio/src/runtime/context.rs | 39 ++++++++++++++++++++---- tokio/src/runtime/handle.rs | 43 +++++++++++++++++++++++++- 4 files changed, 92 insertions(+), 45 deletions(-) diff --git a/tokio/src/runtime/blocking/mod.rs b/tokio/src/runtime/blocking/mod.rs index fa4c77a12..70832d46d 100644 --- a/tokio/src/runtime/blocking/mod.rs +++ b/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(&self, f: F) -> R where F: FnOnce() -> R, diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 2aa4c8bf9..320bd7f2a 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/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,11 +67,6 @@ struct Shared { type Task = task::Task; -thread_local! { - /// Thread-local tracking the current executor - static BLOCKING: Cell> = Cell::new(None) -} - const KEEP_ALIVE: Duration = Duration::from_secs(10); /// Run the provided function on an executor dedicated to blocking operations. @@ -80,16 +74,14 @@ pub(crate) fn spawn_blocking(func: F) -> JoinHandle 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 (task, handle) = task::joinable(BlockingTask::new(func)); - schedule.schedule(task); - handle - }) + 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 } // ===== 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); + let ctx = crate::runtime::context::ThreadContext::clone_current(); + let _e = ctx.with_blocking_spawner(self.clone()).enter(); - // 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>); - - impl Drop for Reset<'_> { - fn drop(&mut self) { - self.0.set(self.1); - } - } - - 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 diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index d3bd69822..07c311edd 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -20,6 +20,9 @@ pub(crate) struct ThreadContext { /// Source of `Instant::now()` clock: Option, + + /// Blocking pool spawner + blocking_spawner: Option, } 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, + blocking_spawner: Option, ) -> Self { ThreadContext { spawner, @@ -60,6 +65,7 @@ impl ThreadContext { #[cfg(any(not(feature = "time"), loom))] time_handle, clock, + blocking_spawner, } } @@ -81,23 +87,20 @@ 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 { CONTEXT.with(|ctx| match *ctx.borrow() { Some(ref ctx) => Some(ctx.spawner.clone()), @@ -105,7 +108,6 @@ impl ThreadContext { }) } - #[cfg(all(feature = "test-util", feature = "time"))] pub(crate) fn clock() -> Option { 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 { + 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. diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index ed3348df8..723650256 100644 --- a/tokio/src/runtime/handle.rs +++ b/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."), + } } }