diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index d201a06df..f03a840e3 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -20,6 +20,8 @@ cfg_rt! { use scoped::Scoped; use crate::runtime::{scheduler, task::Id}; + #[cfg(tokio_unstable)] + use crate::runtime::task::Header; #[cfg(tokio_unstable)] use std::ptr::NonNull; @@ -52,7 +54,7 @@ struct Context { current_task_id: Cell>, #[cfg(all(feature = "rt", tokio_unstable))] - current_task: Cell>>, + current_task: Cell>>, /// Tracks if the current thread is currently driving a runtime. /// Note, that if this is set to "entered", the current scheduler @@ -168,14 +170,28 @@ cfg_rt! { } #[cfg(tokio_unstable)] - pub(crate) fn set_current_task(task: Option>) -> Option> { + pub(crate) fn set_current_task(task: Option>) -> Option> { CONTEXT .try_with(|ctx| ctx.current_task.replace(task)) .unwrap_or(None) } #[cfg(tokio_unstable)] - pub(crate) fn current_task() -> Option> { + pub(crate) fn set_current_task_id_and_task( + id: Option, + task: Option>, + ) -> (Option, Option>) { + CONTEXT + .try_with(|ctx| { + let parent_task_id = ctx.current_task_id.replace(id); + let parent_task = ctx.current_task.replace(task); + (parent_task_id, parent_task) + }) + .unwrap_or((None, None)) + } + + #[cfg(tokio_unstable)] + pub(crate) fn current_task() -> Option> { CONTEXT .try_with(|ctx| ctx.current_task.get()) .unwrap_or(None) diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index b679eadff..da5e59d94 100644 --- a/tokio/src/runtime/task/core.rs +++ b/tokio/src/runtime/task/core.rs @@ -350,6 +350,32 @@ impl Drop for TaskIdGuard { } } +#[cfg(tokio_unstable)] +struct TaskContextGuard { + parent_task_id: Option, + parent_task: Option>, +} + +#[cfg(tokio_unstable)] +impl TaskContextGuard { + fn enter(id: Id, header: NonNull
) -> Self { + let (parent_task_id, parent_task) = + context::set_current_task_id_and_task(Some(id), Some(header)); + + TaskContextGuard { + parent_task_id, + parent_task, + } + } +} + +#[cfg(tokio_unstable)] +impl Drop for TaskContextGuard { + fn drop(&mut self) { + context::set_current_task_id_and_task(self.parent_task_id, self.parent_task); + } +} + impl Core { /// Polls the future. /// @@ -365,12 +391,11 @@ impl Core { /// `self` must also be pinned. This is handled by storing the task on the /// heap. /// - /// When `tokio_unstable` is enabled, `header` must point to the header for - /// this exact task allocation, and the allocation must remain live until - /// this function returns. + /// `header` must point to the header for this exact task allocation, and + /// the allocation must remain live until this function returns. pub(super) unsafe fn poll( &self, - #[cfg(tokio_unstable)] header: NonNull
, + header: NonNull
, mut cx: Context<'_>, ) -> Poll { let res = { @@ -384,18 +409,16 @@ impl Core { // Safety: The caller ensures the future is pinned. let future = unsafe { Pin::new_unchecked(future) }; - let _guard = TaskIdGuard::enter(self.task_id); #[cfg(tokio_unstable)] - let _current_task = CurrentTaskGuard::enter(header); + let _guard = TaskContextGuard::enter(self.task_id, header); + #[cfg(not(tokio_unstable))] + let _guard = TaskIdGuard::enter(self.task_id); future.poll(&mut cx) }) }; if res.is_ready() { - self.drop_future_or_output( - #[cfg(tokio_unstable)] - header, - ); + self.drop_future_or_output(header); } res @@ -406,7 +429,10 @@ impl Core { /// # Safety /// /// The caller must ensure it is safe to mutate the `stage` field. - pub(super) fn drop_future_or_output(&self, #[cfg(tokio_unstable)] header: NonNull
) { + pub(super) fn drop_future_or_output(&self, header: NonNull
) { + #[cfg(not(tokio_unstable))] + let _ = header; + #[cfg(tokio_unstable)] let _current_task = { let dropping_future = self.stage.stage.with(|ptr| { @@ -464,14 +490,14 @@ impl Core { #[cfg(tokio_unstable)] pub(crate) struct CurrentTaskGuard { - parent_task: Option>, + parent_task: Option>, } #[cfg(tokio_unstable)] impl CurrentTaskGuard { fn enter(header: NonNull
) -> Self { CurrentTaskGuard { - parent_task: context::set_current_task(Some(header.cast())), + parent_task: context::set_current_task(Some(header)), } } } diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 6a9f90988..51ccb0cb8 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/tokio/src/runtime/task/harness.rs @@ -242,11 +242,7 @@ where } if self.state().load().is_cancelled() { - cancel_task( - self.core(), - #[cfg(tokio_unstable)] - header_ptr, - ); + cancel_task(self.core(), header_ptr); return PollFuture::Complete; } } @@ -256,14 +252,7 @@ where // Safety: `transition_to_running` succeeded, so this thread has // exclusive access to the future/output storage. The header pointer // comes from this harness and remains live while the task is running. - let res = unsafe { - poll_future( - self.core(), - #[cfg(tokio_unstable)] - header_ptr, - cx, - ) - }; + let res = unsafe { poll_future(self.core(), header_ptr, cx) }; #[cfg(tokio_unstable)] { @@ -294,20 +283,12 @@ where if let TransitionToIdle::Cancelled = transition_res { // The transition to idle failed because the task was // cancelled during the poll. - cancel_task( - self.core(), - #[cfg(tokio_unstable)] - header_ptr, - ); + cancel_task(self.core(), header_ptr); } transition_result_to_poll_future(transition_res) } TransitionToRunning::Cancelled => { - cancel_task( - self.core(), - #[cfg(tokio_unstable)] - self.header_ptr(), - ); + cancel_task(self.core(), self.header_ptr()); PollFuture::Complete } TransitionToRunning::Failed => PollFuture::Done, @@ -330,11 +311,7 @@ where // By transitioning the lifecycle to `Running`, we have permission to // drop the future. - cancel_task( - self.core(), - #[cfg(tokio_unstable)] - self.header_ptr(), - ); + cancel_task(self.core(), self.header_ptr()); self.complete(); } @@ -390,10 +367,7 @@ where // they are dropping the `JoinHandle`, we assume they are not // interested in the panic and swallow it. let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { - self.core().drop_future_or_output( - #[cfg(tokio_unstable)] - self.header_ptr(), - ); + self.core().drop_future_or_output(self.header_ptr()); })); } @@ -434,10 +408,7 @@ where // this task. It is our responsibility to drop the // output. The join waker was already dropped by the // `JoinHandle` before. - self.core().drop_future_or_output( - #[cfg(tokio_unstable)] - self.header_ptr(), - ); + self.core().drop_future_or_output(self.header_ptr()); } else if snapshot.is_join_waker_set() { // Notify the waker. Reading the waker field is safe per rule 4 // in task/mod.rs, since the JOIN_WAKER bit is set and the call @@ -592,16 +563,10 @@ enum PollFuture { } /// Cancels the task and store the appropriate error in the stage field. -fn cancel_task( - core: &Core, - #[cfg(tokio_unstable)] header: NonNull
, -) { +fn cancel_task(core: &Core, header: NonNull
) { // Drop the future from a panic guard. let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { - core.drop_future_or_output( - #[cfg(tokio_unstable)] - header, - ); + core.drop_future_or_output(header); })); core.store_output(Err(panic_result_to_join_error(core.task_id, res))); @@ -653,45 +618,33 @@ unsafe fn poll_hook_panic( /// /// The caller must satisfy the mutual-exclusion requirements of `Core::poll`. /// -/// When `tokio_unstable` is enabled, `header` must point to the header for this -/// exact task allocation, and the allocation must remain live until this -/// function returns. +/// `header` must point to the header for this exact task allocation, and the +/// allocation must remain live until this function returns. unsafe fn poll_future( core: &Core, - #[cfg(tokio_unstable)] header: NonNull
, + header: NonNull
, cx: Context<'_>, ) -> Poll<()> { // Poll the future. let output = panic::catch_unwind(panic::AssertUnwindSafe(|| { struct Guard<'a, T: Future, S: Schedule> { core: &'a Core, - #[cfg(tokio_unstable)] header: NonNull
, } impl<'a, T: Future, S: Schedule> Drop for Guard<'a, T, S> { fn drop(&mut self) { // If the future panics on poll, we drop it inside the panic // guard. - self.core.drop_future_or_output( - #[cfg(tokio_unstable)] - self.header, - ); + self.core.drop_future_or_output(self.header); } } let guard = Guard { core, - #[cfg(tokio_unstable)] header, }; // Safety: the caller guarantees the mutual-exclusion requirements of // `Core::poll` and that `header` identifies this live task allocation. - let res = unsafe { - guard.core.poll( - #[cfg(tokio_unstable)] - header, - cx, - ) - }; + let res = unsafe { guard.core.poll(header, cx) }; mem::forget(guard); res })); diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index d5a1de4eb..451893b64 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -186,7 +186,7 @@ mod core; use self::core::Cell; -use self::core::Header; +pub(crate) use self::core::Header; mod error; pub use self::error::JoinError; @@ -687,7 +687,7 @@ pub(crate) fn with_current_task_meta( // Safety: the context stores this pointer only while the referenced task is // being polled, so the allocation is alive for this synchronous call. - let raw = unsafe { RawTask::from_raw(ptr.cast()) }; + let raw = unsafe { RawTask::from_raw(ptr) }; // Safety: parent metadata is exposed read-only during this synchronous call // while no mutable parent hook metadata is live. The closure-bound lifetime // prevents references exposed through the metadata from escaping.