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

address feedback from alice

Browse Source
This commit is contained in:
noah
2026-05-29 16:25:29 -05:00
parent 8b66a4af58
commit 70402f86a5
13 changed files with 248 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/runtime/blocking/mod.rs
+1 -1
View File
@@ -5,7 +5,7 @@
mod pool;
#[cfg(feature = "rt-multi-thread")]
pub(crate) use pool::spawn_blocking_internal;
pub(crate) use pool::spawn_blocking_skip_hooks;
pub(crate) use pool::{spawn_blocking, BlockingPool, Spawner};
cfg_fs! {
tokio/src/runtime/blocking/pool.rs
+3 -3
View File
@@ -190,7 +190,7 @@ where
#[track_caller]
#[cfg(feature = "rt-multi-thread")]
#[cfg_attr(target_os = "wasi", allow(dead_code))]
pub(crate) fn spawn_blocking_internal<F, R>(func: F) -> JoinHandle<R>
pub(crate) fn spawn_blocking_skip_hooks<F, R>(func: F) -> JoinHandle<R>
where
F: FnOnce() -> R + Send + 'static,
R: Send + 'static,
@@ -198,7 +198,7 @@ where
let rt = Handle::current();
rt.inner
.blocking_spawner()
.spawn_blocking_internal(&rt, func)
.spawn_blocking_skip_hooks(&rt, func)
}
cfg_fs! {
@@ -326,7 +326,7 @@ impl Spawner {
#[track_caller]
#[cfg(feature = "rt-multi-thread")]
pub(crate) fn spawn_blocking_internal<F, R>(&self, rt: &Handle, func: F) -> JoinHandle<R>
pub(crate) fn spawn_blocking_skip_hooks<F, R>(&self, rt: &Handle, func: F) -> JoinHandle<R>
where
F: FnOnce() -> R + Send + 'static,
R: Send + 'static,
tokio/src/runtime/mod.rs
+1 -1
View File
@@ -531,7 +531,7 @@ cfg_rt! {
#[cfg_attr(target_os = "wasi", allow(unused_imports))]
pub(crate) use blocking::spawn_blocking;
#[cfg(feature = "rt-multi-thread")]
pub(crate) use blocking::spawn_blocking_internal;
pub(crate) use blocking::spawn_blocking_skip_hooks;
cfg_trace! {
pub(crate) use blocking::Mandatory;
tokio/src/runtime/scheduler/current_thread/mod.rs
+27 -52
View File
@@ -474,29 +474,15 @@ impl Handle {
F: crate::future::Future + Send + 'static,
F::Output: Send + 'static,
{
#[cfg(tokio_unstable)]
let (handle, notified) = task::with_current_task_meta(|parent| {
me.shared.owned.bind_with_spawn_hook(
future,
me.clone(),
id,
spawned_at,
user_data,
|task| {
// Safety: the task is freshly allocated and not published yet.
let mut meta = unsafe { task.task_meta() };
me.task_hooks.spawn(&mut meta, parent);
},
)
});
#[cfg(not(tokio_unstable))]
let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
#[cfg(not(tokio_unstable))]
{
let mut meta = TaskMeta::new(id, spawned_at);
me.task_hooks.spawn(&mut meta, None);
}
let (handle, notified) = me.shared.owned.bind_with_spawn_hook(
future,
me.clone(),
id,
spawned_at,
#[cfg(tokio_unstable)]
user_data,
&me.task_hooks,
);
if let Some(notified) = notified {
me.schedule(notified);
@@ -524,38 +510,19 @@ impl Handle {
F: crate::future::Future + 'static,
F::Output: 'static,
{
#[cfg(tokio_unstable)]
let (handle, notified) = task::with_current_task_meta(|parent| {
let before_bind = |task: &Task<Arc<Handle>>| {
// Safety: the task is freshly allocated and not published yet.
let mut meta = unsafe { task.task_meta() };
me.task_hooks.spawn(&mut meta, parent);
};
// Safety: the caller guarantees that this is only called on a `LocalRuntime`.
unsafe {
me.shared.owned.bind_local_with_spawn_hook(
future,
me.clone(),
id,
spawned_at,
user_data,
before_bind,
)
}
});
#[cfg(not(tokio_unstable))]
// Safety: the caller guarantees that this is only called on a `LocalRuntime`.
let (handle, notified) = unsafe {
me.shared
.owned
.bind_local(future, me.clone(), id, spawned_at)
me.shared.owned.bind_local_with_spawn_hook(
future,
me.clone(),
id,
spawned_at,
#[cfg(tokio_unstable)]
user_data,
&me.task_hooks,
)
};
#[cfg(not(tokio_unstable))]
{
let mut meta = TaskMeta::new(id, spawned_at);
me.task_hooks.spawn(&mut meta, None);
}
if let Some(notified) = notified {
me.schedule(notified);
}
@@ -727,10 +694,18 @@ impl Schedule for Arc<Handle> {
self.task_hooks.task_terminate_callback(meta);
}
fn has_task_poll_start_callback(&self) -> bool {
self.task_hooks.has_poll_start_callback()
}
fn task_poll_start_callback(&self, meta: &mut TaskMeta<'_>) {
self.task_hooks.poll_start_callback(meta);
}
fn has_task_poll_stop_callback(&self) -> bool {
self.task_hooks.has_poll_stop_callback()
}
fn task_poll_stop_callback(&self, meta: &mut TaskMeta<'_>) {
self.task_hooks.poll_stop_callback(meta);
}
tokio/src/runtime/scheduler/multi_thread/handle.rs
+19 -23
View File
@@ -97,29 +97,15 @@ impl Handle {
T: Future + Send + 'static,
T::Output: Send + 'static,
{
#[cfg(tokio_unstable)]
let (handle, notified) = task::with_current_task_meta(|parent| {
me.shared.owned.bind_with_spawn_hook(
future,
me.clone(),
id,
spawned_at,
user_data,
|task| {
// Safety: the task is freshly allocated and not published yet.
let mut meta = unsafe { task.task_meta() };
me.task_hooks.spawn(&mut meta, parent);
},
)
});
#[cfg(not(tokio_unstable))]
let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
#[cfg(not(tokio_unstable))]
{
let mut meta = TaskMeta::new(id, spawned_at);
me.task_hooks.spawn(&mut meta, None);
}
let (handle, notified) = me.shared.owned.bind_with_spawn_hook(
future,
me.clone(),
id,
spawned_at,
#[cfg(tokio_unstable)]
user_data,
&me.task_hooks,
);
me.schedule_option_task_without_yield(notified);
@@ -145,11 +131,21 @@ impl task::Schedule for Arc<Handle> {
self.task_hooks.task_terminate_callback(meta);
}
#[cfg(tokio_unstable)]
fn has_task_poll_start_callback(&self) -> bool {
self.task_hooks.has_poll_start_callback()
}
#[cfg(tokio_unstable)]
fn task_poll_start_callback(&self, meta: &mut TaskMeta<'_>) {
self.task_hooks.poll_start_callback(meta);
}
#[cfg(tokio_unstable)]
fn has_task_poll_stop_callback(&self) -> bool {
self.task_hooks.has_poll_stop_callback()
}
#[cfg(tokio_unstable)]
fn task_poll_stop_callback(&self, meta: &mut TaskMeta<'_>) {
self.task_hooks.poll_stop_callback(meta);
tokio/src/runtime/scheduler/multi_thread/worker.rs
+2 -2
View File
@@ -475,7 +475,7 @@ where
// Once the blocking task is done executing, we will attempt to
// steal the core back.
let worker = cx.worker.clone();
runtime::spawn_blocking_internal(move || run(worker));
runtime::spawn_blocking_skip_hooks(move || run(worker));
Ok(())
});
@@ -500,7 +500,7 @@ where
impl Launch {
pub(crate) fn launch(mut self) {
for worker in self.0.drain(..) {
runtime::spawn_blocking_internal(move || run(worker));
runtime::spawn_blocking_skip_hooks(move || run(worker));
}
}
}
tokio/src/runtime/task/core.rs
+47 -33
View File
@@ -350,29 +350,67 @@ impl Drop for TaskIdGuard {
}
}
#[cfg(tokio_unstable)]
struct TaskContextGuard {
parent_task_id: Option<Id>,
#[cfg(tokio_unstable)]
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));
#[cfg(tokio_unstable)]
{
let (parent_task_id, parent_task) =
context::set_current_task_id_and_task(Some(id), Some(header));
TaskContextGuard {
parent_task_id,
parent_task,
TaskContextGuard {
parent_task_id,
parent_task,
}
}
#[cfg(not(tokio_unstable))]
{
let _ = header;
TaskContextGuard {
parent_task_id: context::set_current_task_id(Some(id)),
}
}
}
}
impl Drop for TaskContextGuard {
fn drop(&mut self) {
#[cfg(tokio_unstable)]
{
context::set_current_task_id_and_task(self.parent_task_id, self.parent_task);
}
#[cfg(not(tokio_unstable))]
{
context::set_current_task_id(self.parent_task_id);
}
}
}
#[cfg(tokio_unstable)]
impl Drop for TaskContextGuard {
struct CurrentTaskGuard {
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)),
}
}
}
#[cfg(tokio_unstable)]
impl Drop for CurrentTaskGuard {
fn drop(&mut self) {
context::set_current_task_id_and_task(self.parent_task_id, self.parent_task);
context::set_current_task(self.parent_task);
}
}
@@ -409,10 +447,7 @@ impl<T: Future, S: Schedule> Core<T, S> {
// Safety: The caller ensures the future is pinned.
let future = unsafe { Pin::new_unchecked(future) };
#[cfg(tokio_unstable)]
let _guard = TaskContextGuard::enter(self.task_id, header);
#[cfg(not(tokio_unstable))]
let _guard = TaskIdGuard::enter(self.task_id);
future.poll(&mut cx)
})
};
@@ -488,27 +523,6 @@ impl<T: Future, S: Schedule> Core<T, S> {
}
}
#[cfg(tokio_unstable)]
pub(crate) struct CurrentTaskGuard {
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)),
}
}
}
#[cfg(tokio_unstable)]
impl Drop for CurrentTaskGuard {
fn drop(&mut self) {
context::set_current_task(self.parent_task);
}
}
impl Header {
pub(super) unsafe fn set_next(&self, next: Option<NonNull<Header>>) {
self.queue_next.with_mut(|ptr| *ptr = next);
tokio/src/runtime/task/harness.rs
+46 -66
View File
@@ -222,31 +222,6 @@ where
let header_ptr = self.header_ptr();
#[cfg(tokio_unstable)]
{
// Safety: the task is in the RUNNING state, so shutdown
// cannot take ownership of the task contents and termination
// cannot access hook data concurrently.
let mut task_meta = unsafe { self.task_meta() };
let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
self.core()
.scheduler
.task_poll_start_callback(&mut task_meta);
}));
if let Err(panic) = res {
// Safety: the task is still in the RUNNING state, so we
// have exclusive access to the future/output storage.
unsafe { poll_hook_panic(self.core(), header_ptr, panic) };
return PollFuture::Complete;
}
if self.state().load().is_cancelled() {
cancel_task(self.core(), header_ptr);
return PollFuture::Complete;
}
}
let waker_ref = waker_ref::<S>(&header_ptr);
let cx = Context::from_waker(&waker_ref);
// Safety: `transition_to_running` succeeded, so this thread has
@@ -254,26 +229,6 @@ where
// comes from this harness and remains live while the task is running.
let res = unsafe { poll_future(self.core(), header_ptr, cx) };
#[cfg(tokio_unstable)]
{
// Safety: the task is still in the RUNNING state, so
// shutdown cannot take ownership of the task contents and
// termination cannot access hook data concurrently.
let mut task_meta = unsafe { self.task_meta() };
let hook_res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
self.core()
.scheduler
.task_poll_stop_callback(&mut task_meta);
}));
if let Err(panic) = hook_res {
// Safety: the task is still in the RUNNING state, so we
// have exclusive access to the future/output storage.
unsafe { poll_hook_panic(self.core(), header_ptr, panic) };
return PollFuture::Complete;
}
}
if res == Poll::Ready(()) {
// The future completed. Move on to complete the task.
return PollFuture::Complete;
@@ -582,32 +537,42 @@ fn panic_result_to_join_error(
}
}
/// Convert a poll hook panic into the task output.
///
/// # Safety
///
/// The caller must have exclusive access to the task's future/output storage,
/// such as by holding the task in the RUNNING state.
#[cfg(tokio_unstable)]
unsafe fn poll_hook_panic<T: Future, S: Schedule>(
unsafe fn task_meta<'meta, T: Future, S: Schedule>(
core: &Core<T, S>,
header: NonNull<Header>,
hook_panic: Box<dyn Any + Send + 'static>,
) {
let drop_res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
core.drop_future_or_output(header);
}));
let join_error = match drop_res {
Ok(()) => panic_to_error(&core.scheduler, core.task_id, hook_panic),
Err(drop_panic) => panic_to_error(&core.scheduler, core.task_id, drop_panic),
};
) -> TaskMeta<'meta> {
// Safety: `header` points to this live task allocation.
let trailer = unsafe { Header::get_trailer(header).as_ref() };
// Safety: the task is in the RUNNING state, so shutdown cannot take
// ownership of the task contents and termination cannot access hook data
// concurrently.
unsafe {
TaskMeta::new(
core.task_id,
core.spawned_at.into(),
Some(trailer.user_data_ptr()),
)
}
}
let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
core.store_output(Err(join_error));
}));
#[cfg(tokio_unstable)]
unsafe fn poll_start_hook<T: Future, S: Schedule>(core: &Core<T, S>, header: NonNull<Header>) {
if core.scheduler.has_task_poll_start_callback() {
let mut task_meta = unsafe { task_meta(core, header) };
let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
core.scheduler.task_poll_start_callback(&mut task_meta);
}));
}
}
if res.is_err() {
core.scheduler.unhandled_panic();
#[cfg(tokio_unstable)]
unsafe fn poll_stop_hook<T: Future, S: Schedule>(core: &Core<T, S>, header: NonNull<Header>) {
if core.scheduler.has_task_poll_stop_callback() {
let mut task_meta = unsafe { task_meta(core, header) };
let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
core.scheduler.task_poll_stop_callback(&mut task_meta);
}));
}
}
@@ -625,6 +590,16 @@ unsafe fn poll_future<T: Future, S: Schedule>(
header: NonNull<Header>,
cx: Context<'_>,
) -> Poll<()> {
#[cfg(tokio_unstable)]
{
unsafe { poll_start_hook(core, header) };
if unsafe { header.as_ref() }.state.load().is_cancelled() {
cancel_task(core, header);
return Poll::Ready(());
}
}
// Poll the future.
let output = panic::catch_unwind(panic::AssertUnwindSafe(|| {
struct Guard<'a, T: Future, S: Schedule> {
@@ -646,6 +621,11 @@ unsafe fn poll_future<T: Future, S: Schedule>(
res
}));
#[cfg(tokio_unstable)]
unsafe {
poll_stop_hook(core, header);
}
// Prepare output for being placed in the core stage.
let output = match output {
Ok(Poll::Pending) => return Poll::Pending,
tokio/src/runtime/task/list.rs
+47 -10
View File
@@ -112,23 +112,29 @@ impl<S: 'static> OwnedTasks<S> {
(join, notified)
}
#[cfg(tokio_unstable)]
pub(crate) fn bind_with_spawn_hook<T>(
&self,
task: T,
scheduler: S,
id: super::Id,
spawned_at: SpawnLocation,
user_data: Option<crate::runtime::TaskData>,
before_bind: impl FnOnce(&Task<S>),
#[cfg(tokio_unstable)] user_data: Option<crate::runtime::TaskData>,
task_hooks: &crate::runtime::TaskHooks,
) -> (JoinHandle<T::Output>, Option<Notified<S>>)
where
S: Schedule,
T: Future + Send + 'static,
T::Output: Send + 'static,
{
let (task, notified, join) = super::new_task(task, scheduler, id, spawned_at, user_data);
before_bind(&task);
let (task, notified, join) = super::new_task(
task,
scheduler,
id,
spawned_at,
#[cfg(tokio_unstable)]
user_data,
);
run_spawn_hook(&task, id, spawned_at, task_hooks);
let notified = unsafe { self.bind_inner(task, notified) };
(join, notified)
}
@@ -167,23 +173,29 @@ impl<S: 'static> OwnedTasks<S> {
/// # Safety
///
/// Only use this in `LocalRuntime` where the task cannot move.
#[cfg(tokio_unstable)]
pub(crate) unsafe fn bind_local_with_spawn_hook<T>(
&self,
task: T,
scheduler: S,
id: super::Id,
spawned_at: SpawnLocation,
user_data: Option<crate::runtime::TaskData>,
before_bind: impl FnOnce(&Task<S>),
#[cfg(tokio_unstable)] user_data: Option<crate::runtime::TaskData>,
task_hooks: &crate::runtime::TaskHooks,
) -> (JoinHandle<T::Output>, Option<Notified<S>>)
where
S: Schedule,
T: Future + 'static,
T::Output: 'static,
{
let (task, notified, join) = super::new_task(task, scheduler, id, spawned_at, user_data);
before_bind(&task);
let (task, notified, join) = super::new_task(
task,
scheduler,
id,
spawned_at,
#[cfg(tokio_unstable)]
user_data,
);
run_spawn_hook(&task, id, spawned_at, task_hooks);
let notified = unsafe { self.bind_inner(task, notified) };
(join, notified)
}
@@ -297,6 +309,31 @@ impl<S: 'static> OwnedTasks<S> {
}
}
#[cfg(tokio_unstable)]
fn run_spawn_hook<S: 'static>(
task: &Task<S>,
_id: super::Id,
_spawned_at: SpawnLocation,
task_hooks: &crate::runtime::TaskHooks,
) {
super::with_current_task_meta(|parent| {
// Safety: the task is freshly allocated and not published yet.
let mut meta = unsafe { task.task_meta() };
task_hooks.spawn(&mut meta, parent);
});
}
#[cfg(not(tokio_unstable))]
fn run_spawn_hook<S: 'static>(
_task: &Task<S>,
id: super::Id,
spawned_at: SpawnLocation,
task_hooks: &crate::runtime::TaskHooks,
) {
let mut meta = crate::runtime::TaskMeta::new(id, spawned_at);
task_hooks.spawn(&mut meta, None);
}
cfg_taskdump! {
impl<S: 'static> OwnedTasks<S> {
/// Locks the tasks, and calls `f` on an iterator over them.
tokio/src/runtime/task/mod.rs
+13 -1
View File
@@ -123,7 +123,9 @@
//! The spawn hook runs before the task is scheduled, poll hooks run while
//! the task holds the RUNNING lock but outside the actual future poll, and
//! the terminate hook runs after completion. Parent task metadata exposed to
//! spawn hooks is read-only.
//! spawn hooks is read-only. If the data is not taken or cleared by a hook,
//! it is dropped when the task allocation is deallocated, after the terminate
//! hook has run and all task references are gone.
//!
//! All other fields are immutable and can be accessed immutably without
//! synchronization by anyone.
@@ -298,9 +300,19 @@ pub(crate) trait Schedule: Sync + Sized + 'static {
#[cfg(tokio_unstable)]
fn task_terminate_callback(&self, _meta: &mut crate::runtime::TaskMeta<'_>) {}
#[cfg(tokio_unstable)]
fn has_task_poll_start_callback(&self) -> bool {
false
}
#[cfg(tokio_unstable)]
fn task_poll_start_callback(&self, _meta: &mut crate::runtime::TaskMeta<'_>) {}
#[cfg(tokio_unstable)]
fn has_task_poll_stop_callback(&self) -> bool {
false
}
#[cfg(tokio_unstable)]
fn task_poll_stop_callback(&self, _meta: &mut crate::runtime::TaskMeta<'_>) {}
}
tokio/src/runtime/task/raw.rs
+2 -2
View File
@@ -272,11 +272,11 @@ impl RawTask {
unsafe { &*self.trailer_ptr().as_ptr() }
}
#[cfg(tokio_unstable)]
/// # Safety
///
/// The task allocation must be live, and the returned metadata must have
/// exclusive access to hook data for as long as it can expose mutable references.
#[cfg(tokio_unstable)]
pub(crate) unsafe fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> {
// Safety: the caller guarantees that the task allocation is live and that
// the returned metadata has exclusive access to hook data.
@@ -289,11 +289,11 @@ impl RawTask {
}
}
#[cfg(tokio_unstable)]
/// # Safety
///
/// The task allocation must be live, and hook data must not be mutated while
/// references exposed through the returned metadata are live.
#[cfg(tokio_unstable)]
pub(crate) unsafe fn task_meta_ref<'meta>(&self) -> crate::runtime::TaskMetaRef<'meta> {
// Safety: the caller guarantees that the task allocation is live and that
// hook data is not mutated while exposed references are live.
tokio/src/runtime/task_hooks.rs
+12
View File
@@ -28,6 +28,12 @@ impl TaskHooks {
}
}
#[cfg(tokio_unstable)]
#[inline]
pub(crate) fn has_poll_start_callback(&self) -> bool {
self.before_poll_callback.is_some()
}
#[cfg(tokio_unstable)]
#[inline]
pub(crate) fn poll_start_callback(&self, meta: &mut TaskMeta<'_>) {
@@ -36,6 +42,12 @@ impl TaskHooks {
}
}
#[cfg(tokio_unstable)]
#[inline]
pub(crate) fn has_poll_stop_callback(&self) -> bool {
self.after_poll_callback.is_some()
}
#[cfg(tokio_unstable)]
#[inline]
pub(crate) fn poll_stop_callback(&self, meta: &mut TaskMeta<'_>) {
tokio/tests/task_hooks.rs
+28
View File
@@ -534,6 +534,34 @@ fn abort_during_before_poll_hook_does_not_poll_future() {
assert_eq!(polls.load(Ordering::SeqCst), 0);
}
#[test]
fn poll_hook_panics_do_not_kill_task() {
let before = Arc::new(AtomicUsize::new(0));
let before2 = Arc::clone(&before);
let after = Arc::new(AtomicUsize::new(0));
let after2 = Arc::clone(&after);
let runtime = Builder::new_current_thread()
.on_before_task_poll(move |_meta| {
if before2.fetch_add(1, Ordering::SeqCst) == 0 {
panic!("before poll hook panic");
}
})
.on_after_task_poll(move |_meta| {
if after2.fetch_add(1, Ordering::SeqCst) == 0 {
panic!("after poll hook panic");
}
})
.build()
.unwrap();
let output = runtime.block_on(async { tokio::spawn(async { 17usize }).await.unwrap() });
assert_eq!(output, 17);
assert_eq!(before.load(Ordering::SeqCst), 1);
assert_eq!(after.load(Ordering::SeqCst), 1);
}
#[derive(Clone, Debug, Eq, PartialEq)]
struct Lineage {
depth: usize,