tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-15 00:00:15 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix alice's findings

Browse Source
This commit is contained in:
noah
2026-05-29 16:25:29 -05:00
parent 422941abea
commit ce372a5cd4
4 changed files with 74 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/runtime/context.rs
+19 -3
View File
@@ -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<Option<Id>>,
#[cfg(all(feature = "rt", tokio_unstable))]
current_task: Cell<Option<NonNull<()>>>,
current_task: Cell<Option<NonNull<Header>>>,
/// 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<NonNull<()>>) -> Option<NonNull<()>> {
pub(crate) fn set_current_task(task: Option<NonNull<Header>>) -> Option<NonNull<Header>> {
CONTEXT
.try_with(|ctx| ctx.current_task.replace(task))
.unwrap_or(None)
}
#[cfg(tokio_unstable)]
pub(crate) fn current_task() -> Option<NonNull<()>> {
pub(crate) fn set_current_task_id_and_task(
id: Option<Id>,
task: Option<NonNull<Header>>,
) -> (Option<Id>, Option<NonNull<Header>>) {
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<NonNull<Header>> {
CONTEXT
.try_with(|ctx| ctx.current_task.get())
.unwrap_or(None)
tokio/src/runtime/task/core.rs
+39 -13
View File
@@ -350,6 +350,32 @@ impl Drop for TaskIdGuard {
}
}
#[cfg(tokio_unstable)]
struct TaskContextGuard {
parent_task_id: Option<Id>,
parent_task: Option<NonNull<Header>>,
}
#[cfg(tokio_unstable)]
impl TaskContextGuard {
fn enter(id: Id, header: NonNull<Header>) -> 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<T: Future, S: Schedule> Core<T, S> {
/// Polls the future.
///
@@ -365,12 +391,11 @@ impl<T: Future, S: Schedule> Core<T, S> {
/// `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>,
header: NonNull<Header>,
mut cx: Context<'_>,
) -> Poll<T::Output> {
let res = {
@@ -384,18 +409,16 @@ impl<T: Future, S: Schedule> Core<T, S> {
// 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<T: Future, S: Schedule> Core<T, S> {
/// # 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<Header>) {
pub(super) fn drop_future_or_output(&self, header: NonNull<Header>) {
#[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<T: Future, S: Schedule> Core<T, S> {
#[cfg(tokio_unstable)]
pub(crate) struct CurrentTaskGuard {
parent_task: Option<NonNull<()>>,
parent_task: Option<NonNull<Header>>,
}
#[cfg(tokio_unstable)]
impl CurrentTaskGuard {
fn enter(header: NonNull<Header>) -> Self {
CurrentTaskGuard {
parent_task: context::set_current_task(Some(header.cast())),
parent_task: context::set_current_task(Some(header)),
}
}
}
tokio/src/runtime/task/harness.rs
+14 -61
View File
@@ -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<T: Future, S: Schedule>(
core: &Core<T, S>,
#[cfg(tokio_unstable)] header: NonNull<Header>,
) {
fn cancel_task<T: Future, S: Schedule>(core: &Core<T, S>, header: NonNull<Header>) {
// 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<T: Future, S: Schedule>(
///
/// 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<T: Future, S: Schedule>(
core: &Core<T, S>,
#[cfg(tokio_unstable)] header: NonNull<Header>,
header: NonNull<Header>,
cx: Context<'_>,
) -> Poll<()> {
// Poll the future.
let output = panic::catch_unwind(panic::AssertUnwindSafe(|| {
struct Guard<'a, T: Future, S: Schedule> {
core: &'a Core<T, S>,
#[cfg(tokio_unstable)]
header: NonNull<Header>,
}
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
}));
tokio/src/runtime/task/mod.rs
+2 -2
View File
@@ -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<R>(
// 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.