diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index b8c19f4bb9e..c09e5d286bf 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -31,7 +31,7 @@ cfg_rt! { use crate::runtime::TryCurrentError; #[derive(Debug)] - pub(crate) struct EnterGuard { + pub(crate) struct SetCurrentGuard { old_handle: Option, old_seed: RngSeed, } @@ -47,21 +47,21 @@ cfg_rt! { /// Sets this [`Handle`] as the current active [`Handle`]. /// /// [`Handle`]: crate::runtime::scheduler::Handle - pub(crate) fn try_enter(handle: &scheduler::Handle) -> Option { + pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option { let rng_seed = handle.seed_generator().next_seed(); CONTEXT.try_with(|ctx| { let old_handle = ctx.scheduler.borrow_mut().replace(handle.clone()); let old_seed = ctx.rng.replace_seed(rng_seed); - EnterGuard { + SetCurrentGuard { old_handle, old_seed, } }).ok() } - impl Drop for EnterGuard { + impl Drop for SetCurrentGuard { fn drop(&mut self) { CONTEXT.with(|ctx| { *ctx.scheduler.borrow_mut() = self.old_handle.take(); diff --git a/tokio/src/runtime/enter.rs b/tokio/src/runtime/enter.rs index 221864eb992..be860711b03 100644 --- a/tokio/src/runtime/enter.rs +++ b/tokio/src/runtime/enter.rs @@ -6,7 +6,7 @@ use std::marker::PhantomData; pub(crate) enum EnterContext { #[cfg_attr(not(feature = "rt"), allow(dead_code))] Entered { - allow_blocking: bool, + allow_block_in_place: bool, }, NotEntered, } @@ -30,8 +30,8 @@ cfg_rt! { /// Marks the current thread as being within the dynamic extent of an /// executor. #[track_caller] - pub(crate) fn enter(allow_blocking: bool) -> Enter { - if let Some(enter) = try_enter(allow_blocking) { + pub(crate) fn enter(allow_block_in_place: bool) -> Enter { + if let Some(enter) = try_enter(allow_block_in_place) { return enter; } @@ -45,12 +45,12 @@ cfg_rt! { /// Tries to enter a runtime context, returns `None` if already in a runtime /// context. - pub(crate) fn try_enter(allow_blocking: bool) -> Option { + pub(crate) fn try_enter(allow_block_in_place: bool) -> Option { ENTERED.with(|c| { if c.get().is_entered() { None } else { - c.set(EnterContext::Entered { allow_blocking }); + c.set(EnterContext::Entered { allow_block_in_place }); Some(Enter { _p: PhantomData }) } }) @@ -92,35 +92,35 @@ cfg_rt_multi_thread! { cfg_rt! { /// Disallows blocking in the current runtime context until the guard is dropped. - pub(crate) fn disallow_blocking() -> DisallowBlockingGuard { + pub(crate) fn disallow_block_in_place() -> DisallowBlockInPlaceGuard { let reset = ENTERED.with(|c| { if let EnterContext::Entered { - allow_blocking: true, + allow_block_in_place: true, } = c.get() { c.set(EnterContext::Entered { - allow_blocking: false, + allow_block_in_place: false, }); true } else { false } }); - DisallowBlockingGuard(reset) + DisallowBlockInPlaceGuard(reset) } - pub(crate) struct DisallowBlockingGuard(bool); - impl Drop for DisallowBlockingGuard { + pub(crate) struct DisallowBlockInPlaceGuard(bool); + impl Drop for DisallowBlockInPlaceGuard { fn drop(&mut self) { if self.0 { // XXX: Do we want some kind of assertion here, or is "best effort" okay? ENTERED.with(|c| { if let EnterContext::Entered { - allow_blocking: false, + allow_block_in_place: false, } = c.get() { c.set(EnterContext::Entered { - allow_blocking: true, + allow_block_in_place: true, }); } }) diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index 99f07c4e683..a4f34b87d4c 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -29,7 +29,7 @@ use std::{error, fmt}; #[derive(Debug)] #[must_use = "Creating and dropping a guard does nothing"] pub struct EnterGuard<'a> { - _guard: context::EnterGuard, + _guard: context::SetCurrentGuard, _handle_lifetime: PhantomData<&'a Handle>, } @@ -44,7 +44,10 @@ impl Handle { /// [`tokio::spawn`]: fn@crate::spawn pub fn enter(&self) -> EnterGuard<'_> { EnterGuard { - _guard: self.inner.enter(), + _guard: match context::try_set_current(&self.inner) { + Some(guard) => guard, + None => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + }, _handle_lifetime: PhantomData, } } diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index 0c59cf30b7d..0ec9d683a5e 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -636,7 +636,7 @@ cfg_rt! { Scheduler::CurrentThread(current_thread) => { // This ensures that tasks spawned on the current-thread // runtime are dropped inside the runtime's context. - match context::try_enter(&self.handle.inner) { + match context::try_set_current(&self.handle.inner) { Some(guard) => current_thread.set_context_guard(guard), None => { // The context thread-local has already been destroyed. diff --git a/tokio/src/runtime/scheduler/current_thread.rs b/tokio/src/runtime/scheduler/current_thread.rs index d11d93253ad..d5124d87383 100644 --- a/tokio/src/runtime/scheduler/current_thread.rs +++ b/tokio/src/runtime/scheduler/current_thread.rs @@ -1,7 +1,7 @@ use crate::future::poll_fn; use crate::loom::sync::atomic::AtomicBool; use crate::loom::sync::{Arc, Mutex}; -use crate::runtime::context::EnterGuard; +use crate::runtime::context::SetCurrentGuard; use crate::runtime::driver::{self, Driver}; use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task}; use crate::runtime::{blocking, Config}; @@ -34,7 +34,7 @@ pub(crate) struct CurrentThread { /// scheduler, it is changed to `Some` with the context being the runtime's /// own context. This ensures that any tasks dropped in the `CurrentThread`'s /// destructor run in that runtime's context. - context_guard: Option, + context_guard: Option, } /// Handle to the current thread scheduler @@ -201,7 +201,7 @@ impl CurrentThread { }) } - pub(crate) fn set_context_guard(&mut self, guard: EnterGuard) { + pub(crate) fn set_context_guard(&mut self, guard: SetCurrentGuard) { self.context_guard = Some(guard); } } diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs index 6a685a6594a..18ac474fa75 100644 --- a/tokio/src/runtime/scheduler/mod.rs +++ b/tokio/src/runtime/scheduler/mod.rs @@ -45,7 +45,7 @@ cfg_rt! { use crate::future::Future; use crate::loom::sync::Arc; use crate::runtime::{blocking, task::Id}; - use crate::runtime::context::{self, EnterGuard}; + use crate::runtime::context; use crate::task::JoinHandle; use crate::util::RngSeedGenerator; @@ -58,16 +58,6 @@ cfg_rt! { } } - /// Sets this [`Handle`] as the current active [`Handle`]. - /// - /// [`Handle`]: Handle - pub(crate) fn enter(&self) -> EnterGuard { - match context::try_enter(self) { - Some(guard) => guard, - None => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), - } - } - pub(crate) fn blocking_spawner(&self) -> &blocking::Spawner { match self { Handle::CurrentThread(h) => &h.blocking_spawner, diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 3407186b30e..1e489b7e5c0 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -283,11 +283,16 @@ where // set up blocking. had_entered = true; } - (EnterContext::Entered { allow_blocking }, false) => { + ( + EnterContext::Entered { + allow_block_in_place, + }, + false, + ) => { // We are on an executor, but _not_ on the thread pool. That is // _only_ okay if we are in a thread pool runtime's block_on // method: - if allow_blocking { + if allow_block_in_place { had_entered = true; return Ok(()); } else { diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 185ac9d2679..739d0623b5c 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -890,7 +890,7 @@ impl Future for RunUntil<'_, T> { .waker .register_by_ref(cx.waker()); - let _no_blocking = crate::runtime::enter::disallow_blocking(); + let _no_blocking = crate::runtime::enter::disallow_block_in_place(); let f = me.future; if let Poll::Ready(output) = crate::coop::budget(|| f.poll(cx)) {