From 70402f86a5d6ef61f27b3afb1868c6255360deac Mon Sep 17 00:00:00 2001 From: noah Date: Fri, 29 May 2026 15:55:36 -0500 Subject: [PATCH] address feedback from alice --- tokio/src/runtime/blocking/mod.rs | 2 +- tokio/src/runtime/blocking/pool.rs | 6 +- tokio/src/runtime/mod.rs | 2 +- .../runtime/scheduler/current_thread/mod.rs | 79 +++++------- .../runtime/scheduler/multi_thread/handle.rs | 42 +++---- .../runtime/scheduler/multi_thread/worker.rs | 4 +- tokio/src/runtime/task/core.rs | 80 +++++++------ tokio/src/runtime/task/harness.rs | 112 +++++++----------- tokio/src/runtime/task/list.rs | 57 +++++++-- tokio/src/runtime/task/mod.rs | 14 ++- tokio/src/runtime/task/raw.rs | 4 +- tokio/src/runtime/task_hooks.rs | 12 ++ tokio/tests/task_hooks.rs | 28 +++++ 13 files changed, 248 insertions(+), 194 deletions(-) diff --git a/tokio/src/runtime/blocking/mod.rs b/tokio/src/runtime/blocking/mod.rs index 77ff53029..a87edf724 100644 --- a/tokio/src/runtime/blocking/mod.rs +++ b/tokio/src/runtime/blocking/mod.rs @@ -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! { diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 99bfb3150..5c5e1ed8b 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -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(func: F) -> JoinHandle +pub(crate) fn spawn_blocking_skip_hooks(func: F) -> JoinHandle 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(&self, rt: &Handle, func: F) -> JoinHandle + pub(crate) fn spawn_blocking_skip_hooks(&self, rt: &Handle, func: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index 22c8fb73d..6cba1cd45 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -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; diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index 6ec7d0e4e..e8018ab49 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/tokio/src/runtime/scheduler/current_thread/mod.rs @@ -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>| { - // 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 { 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); } diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index d53d0b801..7fefc577c 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -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 { 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); diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 9e5cacf48..469f837b5 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -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)); } } } diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index da5e59d94..d4c53638d 100644 --- a/tokio/src/runtime/task/core.rs +++ b/tokio/src/runtime/task/core.rs @@ -350,29 +350,67 @@ impl Drop for TaskIdGuard { } } -#[cfg(tokio_unstable)] struct TaskContextGuard { parent_task_id: Option, + #[cfg(tokio_unstable)] 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)); + #[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>, +} + +#[cfg(tokio_unstable)] +impl CurrentTaskGuard { + fn enter(header: NonNull
) -> 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 Core { // 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 Core { } } -#[cfg(tokio_unstable)] -pub(crate) struct CurrentTaskGuard { - parent_task: Option>, -} - -#[cfg(tokio_unstable)] -impl CurrentTaskGuard { - fn enter(header: NonNull
) -> 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>) { self.queue_next.with_mut(|ptr| *ptr = next); diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 0dee88f6b..4816784a0 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/tokio/src/runtime/task/harness.rs @@ -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::(&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( +unsafe fn task_meta<'meta, T: Future, S: Schedule>( core: &Core, header: NonNull
, - hook_panic: Box, -) { - 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(core: &Core, header: NonNull
) { + 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(core: &Core, header: NonNull
) { + 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( header: NonNull
, 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( 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, diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs index 281b397f1..884507675 100644 --- a/tokio/src/runtime/task/list.rs +++ b/tokio/src/runtime/task/list.rs @@ -112,23 +112,29 @@ impl OwnedTasks { (join, notified) } - #[cfg(tokio_unstable)] pub(crate) fn bind_with_spawn_hook( &self, task: T, scheduler: S, id: super::Id, spawned_at: SpawnLocation, - user_data: Option, - before_bind: impl FnOnce(&Task), + #[cfg(tokio_unstable)] user_data: Option, + task_hooks: &crate::runtime::TaskHooks, ) -> (JoinHandle, Option>) 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 OwnedTasks { /// # Safety /// /// Only use this in `LocalRuntime` where the task cannot move. - #[cfg(tokio_unstable)] pub(crate) unsafe fn bind_local_with_spawn_hook( &self, task: T, scheduler: S, id: super::Id, spawned_at: SpawnLocation, - user_data: Option, - before_bind: impl FnOnce(&Task), + #[cfg(tokio_unstable)] user_data: Option, + task_hooks: &crate::runtime::TaskHooks, ) -> (JoinHandle, Option>) 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 OwnedTasks { } } +#[cfg(tokio_unstable)] +fn run_spawn_hook( + task: &Task, + _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( + _task: &Task, + 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 OwnedTasks { /// Locks the tasks, and calls `f` on an iterator over them. diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index 451893b64..0b7af459f 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -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<'_>) {} } diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs index 766a7b418..f6a640f34 100644 --- a/tokio/src/runtime/task/raw.rs +++ b/tokio/src/runtime/task/raw.rs @@ -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. diff --git a/tokio/src/runtime/task_hooks.rs b/tokio/src/runtime/task_hooks.rs index 27dc457b7..3b9f7b77a 100644 --- a/tokio/src/runtime/task_hooks.rs +++ b/tokio/src/runtime/task_hooks.rs @@ -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<'_>) { diff --git a/tokio/tests/task_hooks.rs b/tokio/tests/task_hooks.rs index 1d649803b..8295d14f9 100644 --- a/tokio/tests/task_hooks.rs +++ b/tokio/tests/task_hooks.rs @@ -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,