diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index 48e3f6087..ea26e235f 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -182,9 +182,13 @@ cfg_rt! { CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None) } - pub(crate) fn try_current() -> Result { - match CONTEXT.try_with(|ctx| ctx.handle.borrow().clone()) { - Ok(Some(handle)) => Ok(handle), + pub(crate) fn with_current(f: F) -> Result + where + F: FnOnce(&scheduler::Handle) -> R, + { + + match CONTEXT.try_with(|ctx| ctx.handle.borrow().as_ref().map(f)) { + Ok(Some(ret)) => Ok(ret), Ok(None) => Err(TryCurrentError::new_no_context()), Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()), } diff --git a/tokio/src/runtime/coop.rs b/tokio/src/runtime/coop.rs index f3ed17cff..2dba24615 100644 --- a/tokio/src/runtime/coop.rs +++ b/tokio/src/runtime/coop.rs @@ -200,9 +200,9 @@ cfg_coop! { cfg_metrics! { #[inline(always)] fn inc_budget_forced_yield_count() { - if let Ok(handle) = context::try_current() { + let _ = context::with_current(|handle| { handle.scheduler_metrics().inc_budget_forced_yield_count(); - } + }); } } diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index e44711f87..a4dc437db 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -109,7 +109,9 @@ impl Handle { /// /// Contrary to `current`, this never panics pub fn try_current() -> Result { - context::try_current().map(|inner| Handle { inner }) + context::with_current(|inner| Handle { + inner: inner.clone(), + }) } /// Spawns a future onto the Tokio runtime. diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs index 0ea207e1e..fa25dabe2 100644 --- a/tokio/src/runtime/scheduler/mod.rs +++ b/tokio/src/runtime/scheduler/mod.rs @@ -52,7 +52,7 @@ cfg_rt! { impl Handle { #[track_caller] pub(crate) fn current() -> Handle { - match context::try_current() { + match context::with_current(Clone::clone) { Ok(handle) => handle, Err(e) => panic!("{}", e), } diff --git a/tokio/src/task/spawn.rs b/tokio/src/task/spawn.rs index 00334549e..20f74d346 100644 --- a/tokio/src/task/spawn.rs +++ b/tokio/src/task/spawn.rs @@ -1,4 +1,3 @@ -use crate::runtime::Handle; use crate::task::JoinHandle; use std::future::Future; @@ -178,7 +177,8 @@ cfg_rt! { T: Future + Send + 'static, T::Output: Send + 'static, { - use crate::runtime::task; + use crate::runtime::{context, task}; + #[cfg(all( tokio_unstable, tokio_taskdump, @@ -193,7 +193,10 @@ cfg_rt! { let future = task::trace::Trace::root(future); let id = task::Id::next(); let task = crate::util::trace::task(future, "task", name, id.as_u64()); - let handle = Handle::current(); - handle.inner.spawn(task, id) + + match context::with_current(|handle| handle.spawn(task, id)) { + Ok(join_handle) => join_handle, + Err(e) => panic!("{}", e), + } } }