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

cleanup: review and cleanup unsafe code

Browse Source 
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.
This commit is contained in:
noah
2026-05-29 16:25:29 -05:00
parent 4f2f942e18
commit bf3b15c471
9 changed files with 131 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/runtime/handle.rs
+4
View File
@@ -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<F>(
&self,
tokio/src/runtime/scheduler/current_thread/mod.rs
+32 -33
View File
@@ -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<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);
};
#[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<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))]
let (handle, notified) = unsafe {
me.shared
tokio/src/runtime/scheduler/multi_thread/handle.rs
+14 -14
View File
@@ -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);
tokio/src/runtime/task/core.rs
+5 -1
View File
@@ -364,7 +364,11 @@ impl<T: Future, S: Schedule> Core<T, S> {
///
/// `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<Header>,
mut cx: Context<'_>,
tokio/src/runtime/task/harness.rs
+48 -20
View File
@@ -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::<S>(&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<T: Future, S: Schedule>(
unsafe fn poll_hook_panic<T: Future, S: Schedule>(
core: &Core<T, S>,
hook_panic: Box<dyn Any + Send + 'static>,
) {
@@ -602,7 +618,15 @@ fn poll_hook_panic<T: Future, S: Schedule>(
/// Polls the future. If the future completes, the output is written to the
/// stage field.
fn poll_future<T: Future, S: Schedule>(
///
/// # 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<T: Future, S: Schedule>(
core: &Core<T, S>,
#[cfg(tokio_unstable)] header: NonNull<Header>,
cx: Context<'_>,
@@ -620,11 +644,15 @@ fn poll_future<T: Future, S: Schedule>(
}
}
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
}));
tokio/src/runtime/task/mod.rs
+11 -6
View File
@@ -669,13 +669,18 @@ impl SpawnLocation {
}
#[cfg(tokio_unstable)]
pub(crate) fn current_task_meta<'meta>() -> Option<crate::runtime::TaskMetaRef<'meta>> {
let ptr = crate::runtime::context::current_task()?;
pub(crate) fn with_current_task_meta<R>(
f: impl for<'meta> FnOnce(Option<crate::runtime::TaskMetaRef<'meta>>) -> 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() }))
}
tokio/src/runtime/task/raw.rs
+4 -4
View File
@@ -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),
tokio/src/runtime/task_hooks.rs
+7 -3
View File
@@ -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::<T>() }
}
}
tokio/tests/task_hooks.rs
+6 -4
View File
@@ -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::<usize>() {
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(