From bf3b15c4716b7d7bb0e8e7871e4680ca82ec5221 Mon Sep 17 00:00:00 2001 From: noah Date: Sun, 10 May 2026 10:41:41 -0500 Subject: [PATCH] cleanup: review and cleanup unsafe code I did some things the other day that I don't like the next morning. Cleaning those up in this commit. Also fixes a flaky test I added. --- tokio/src/runtime/handle.rs | 4 ++ .../runtime/scheduler/current_thread/mod.rs | 65 +++++++++--------- .../runtime/scheduler/multi_thread/handle.rs | 28 ++++---- tokio/src/runtime/task/core.rs | 6 +- tokio/src/runtime/task/harness.rs | 68 +++++++++++++------ tokio/src/runtime/task/mod.rs | 17 +++-- tokio/src/runtime/task/raw.rs | 8 +-- tokio/src/runtime/task_hooks.rs | 10 ++- tokio/tests/task_hooks.rs | 10 +-- 9 files changed, 131 insertions(+), 85 deletions(-) diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index 2d78dee50..2b25a83c2 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -475,6 +475,10 @@ impl Handle { unsafe { self.spawn_local_named_inner(future, meta, user_data) } } + /// # Safety + /// + /// This must only be called in `LocalRuntime` if the runtime has been verified to be owned + /// by the current thread. #[track_caller] unsafe fn spawn_local_named_inner( &self, diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index d12c2e464..6ec7d0e4e 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/tokio/src/runtime/scheduler/current_thread/mod.rs @@ -475,20 +475,20 @@ impl Handle { F::Output: Send + 'static, { #[cfg(tokio_unstable)] - let parent = task::current_task_meta(); - #[cfg(tokio_unstable)] - let (handle, notified) = 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); - }, - ); + 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); @@ -524,26 +524,25 @@ impl Handle { F: crate::future::Future + 'static, F::Output: 'static, { - // Safety: the caller guarantees that this is only called on a `LocalRuntime`. #[cfg(tokio_unstable)] - let parent = task::current_task_meta(); - #[cfg(tokio_unstable)] - 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); - }; - #[cfg(tokio_unstable)] - let (handle, notified) = unsafe { - me.shared.owned.bind_local_with_spawn_hook( - future, - me.clone(), - id, - spawned_at, - user_data, - before_bind, - ) - }; + 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))] let (handle, notified) = unsafe { me.shared diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index 4f809fed4..d53d0b801 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -98,20 +98,20 @@ impl Handle { T::Output: Send + 'static, { #[cfg(tokio_unstable)] - let parent = task::current_task_meta(); - #[cfg(tokio_unstable)] - let (handle, notified) = 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); - }, - ); + 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); diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index 58a7592ce..79c7fb318 100644 --- a/tokio/src/runtime/task/core.rs +++ b/tokio/src/runtime/task/core.rs @@ -364,7 +364,11 @@ impl Core { /// /// `self` must also be pinned. This is handled by storing the task on the /// heap. - pub(super) fn 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. + pub(super) unsafe fn poll( &self, #[cfg(tokio_unstable)] header: NonNull
, mut cx: Context<'_>, diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 46298995e..1f08aa233 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/tokio/src/runtime/task/harness.rs @@ -222,8 +222,9 @@ where #[cfg(tokio_unstable)] { - // Safety: the task is in the RUNNING state, which excludes - // concurrent shutdown and termination metadata access. + // 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() @@ -232,7 +233,9 @@ where })); if let Err(panic) = res { - poll_hook_panic(self.core(), panic); + // 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(), panic) }; return PollFuture::Complete; } @@ -245,18 +248,23 @@ where let header_ptr = self.header_ptr(); let waker_ref = waker_ref::(&header_ptr); let cx = Context::from_waker(&waker_ref); - let res = poll_future( - self.core(), - #[cfg(tokio_unstable)] - header_ptr, - cx, - ); + // 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, + ) + }; #[cfg(tokio_unstable)] { - // Safety: the task is still in the RUNNING state, which - // excludes concurrent shutdown and termination metadata - // access. + // 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() @@ -265,7 +273,9 @@ where })); if let Err(panic) = hook_res { - poll_hook_panic(self.core(), panic); + // 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(), panic) }; return PollFuture::Complete; } } @@ -578,8 +588,14 @@ 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)] -fn poll_hook_panic( +unsafe fn poll_hook_panic( core: &Core, hook_panic: Box, ) { @@ -602,7 +618,15 @@ fn poll_hook_panic( /// Polls the future. If the future completes, the output is written to the /// stage field. -fn poll_future( +/// +/// # Safety +/// +/// 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. +unsafe fn poll_future( core: &Core, #[cfg(tokio_unstable)] header: NonNull
, cx: Context<'_>, @@ -620,11 +644,15 @@ fn poll_future( } } let guard = Guard { core }; - let res = guard.core.poll( - #[cfg(tokio_unstable)] - header, - cx, - ); + // 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, + ) + }; mem::forget(guard); res })); diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index d83b1ac98..eaaaae767 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -669,13 +669,18 @@ impl SpawnLocation { } #[cfg(tokio_unstable)] -pub(crate) fn current_task_meta<'meta>() -> Option> { - let ptr = crate::runtime::context::current_task()?; +pub(crate) fn with_current_task_meta( + f: impl for<'meta> FnOnce(Option>) -> R, +) -> R { + let Some(ptr) = crate::runtime::context::current_task() else { + return f(None); + }; // Safety: the context stores this pointer only while the referenced task is - // being polled, so the allocation is alive for the duration of this call. + // being polled, so the allocation is alive for this synchronous call. let raw = unsafe { RawTask::from_raw(ptr.cast()) }; - // Safety: parent metadata is exposed read-only during synchronous spawn - // hook invocation while no mutable parent hook metadata is live. - Some(unsafe { raw.task_meta_ref() }) + // 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. + f(Some(unsafe { raw.task_meta_ref() })) } diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs index ca409f72e..766a7b418 100644 --- a/tokio/src/runtime/task/raw.rs +++ b/tokio/src/runtime/task/raw.rs @@ -278,8 +278,8 @@ impl RawTask { /// 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. pub(crate) unsafe fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { - // Safety: `self` holds a live task reference, and callers use the - // metadata only for the current hook invocation. + // Safety: the caller guarantees that the task allocation is live and that + // the returned metadata has exclusive access to hook data. unsafe { crate::runtime::TaskMeta::new( Header::get_id(self.ptr), @@ -295,8 +295,8 @@ impl RawTask { /// The task allocation must be live, and hook data must not be mutated while /// references exposed through the returned metadata are live. pub(crate) unsafe fn task_meta_ref<'meta>(&self) -> crate::runtime::TaskMetaRef<'meta> { - // Safety: `self` holds a live task reference, and this only exposes - // shared access to task data. + // Safety: the caller guarantees that the task allocation is live and that + // hook data is not mutated while exposed references are live. unsafe { crate::runtime::TaskMetaRef::new( Header::get_id(self.ptr), diff --git a/tokio/src/runtime/task_hooks.rs b/tokio/src/runtime/task_hooks.rs index e1cfbcac5..27dc457b7 100644 --- a/tokio/src/runtime/task_hooks.rs +++ b/tokio/src/runtime/task_hooks.rs @@ -215,8 +215,11 @@ pub struct TaskMetaRef<'a> { impl<'a> TaskMetaRef<'a> { /// # Safety /// - /// If `user_data` is present, it must point to live task storage for the - /// duration of any references exposed through this metadata value. + /// If `user_data` is present, it must point to live task storage for the duration + /// of any references exposed through this metadata value. + /// + /// While any references exposed through this metadata value are live, the task + /// data must not be mutably accessed, replaced, cleared, or dropped. #[cfg(tokio_unstable)] pub(crate) unsafe fn new( id: super::task::Id, @@ -251,7 +254,8 @@ impl<'a> TaskMetaRef<'a> { let user_data = self.user_data?; // Safety: `TaskMetaRef` is only constructed while the task allocation is - // known to be alive, and it does not expose mutation. + // known to be alive. Its constructor requires that the data is not mutated + // while exposed references are live. unsafe { user_data.as_ref().as_ref()?.downcast_ref::() } } } diff --git a/tokio/tests/task_hooks.rs b/tokio/tests/task_hooks.rs index 75502d44b..e1252b264 100644 --- a/tokio/tests/task_hooks.rs +++ b/tokio/tests/task_hooks.rs @@ -549,13 +549,12 @@ fn task_builder_data_is_visible_to_hooks() { #[cfg(feature = "tracing")] #[test] fn task_builder_data_is_not_dropped_for_spawn_blocking() { - let terminated = Arc::new(Mutex::new(Vec::new())); - let terminated2 = Arc::clone(&terminated); + let (terminated_tx, terminated_rx) = std::sync::mpsc::channel(); let runtime = Builder::new_current_thread() .on_task_terminate(move |meta| { if let Some(value) = meta.take_data::() { - terminated2.lock().unwrap().push(*value); + terminated_tx.send(*value).unwrap(); } }) .build() @@ -570,7 +569,10 @@ fn task_builder_data_is_not_dropped_for_spawn_blocking() { .unwrap(); }); - assert_eq!(*terminated.lock().unwrap(), vec![7]); + let terminated = terminated_rx + .recv_timeout(Duration::from_secs(5)) + .expect("spawn_blocking task terminate hook did not receive task data"); + assert_eq!(terminated, 7); } fn mk_spawn_location_hook(