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

@@ -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

@@ -475,20 +475,20 @@ impl Handle {
         F::Output: Send + 'static,
     {
         #[cfg(tokio_unstable)]
-        let parent = task::current_task_meta();
+        let (handle, notified) = task::with_current_task_meta(|parent| {
-        #[cfg(tokio_unstable)]
+            me.shared.owned.bind_with_spawn_hook(
-        let (handle, notified) = me.shared.owned.bind_with_spawn_hook(
+                future,
-            future,
+                me.clone(),
-            me.clone(),
+                id,
-            id,
+                spawned_at,
-            spawned_at,
+                user_data,
-            user_data,
+                |task| {
-            |task| {
+                    // Safety: the task is freshly allocated and not published yet.
-                // Safety: the task is freshly allocated and not published yet.
+                    let mut meta = unsafe { task.task_meta() };
-                let mut meta = unsafe { task.task_meta() };
+                    me.task_hooks.spawn(&mut meta, parent);
-                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();
+        let (handle, notified) = task::with_current_task_meta(|parent| {
-        #[cfg(tokio_unstable)]
+            let before_bind = |task: &Task<Arc<Handle>>| {
-        let before_bind = |task: &Task<Arc<Handle>>| {
+                // Safety: the task is freshly allocated and not published yet.
-            // Safety: the task is freshly allocated and not published yet.
+                let mut meta = unsafe { task.task_meta() };
-            let mut meta = unsafe { task.task_meta() };
+                me.task_hooks.spawn(&mut meta, parent);
-            me.task_hooks.spawn(&mut meta, parent);
+            };
-        };
+            // Safety: the caller guarantees that this is only called on a `LocalRuntime`.
-        #[cfg(tokio_unstable)]
+            unsafe {
-        let (handle, notified) = unsafe {
+                me.shared.owned.bind_local_with_spawn_hook(
-            me.shared.owned.bind_local_with_spawn_hook(
+                    future,
-                future,
+                    me.clone(),
-                me.clone(),
+                    id,
-                id,
+                    spawned_at,
-                spawned_at,
+                    user_data,
-                user_data,
+                    before_bind,
-                before_bind,
+                )
-            )
+            }
-        };
+        });
         #[cfg(not(tokio_unstable))]
         let (handle, notified) = unsafe {
             me.shared

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();
+        let (handle, notified) = task::with_current_task_meta(|parent| {
-        #[cfg(tokio_unstable)]
+            me.shared.owned.bind_with_spawn_hook(
-        let (handle, notified) = me.shared.owned.bind_with_spawn_hook(
+                future,
-            future,
+                me.clone(),
-            me.clone(),
+                id,
-            id,
+                spawned_at,
-            spawned_at,
+                user_data,
-            user_data,
+                |task| {
-            |task| {
+                    // Safety: the task is freshly allocated and not published yet.
-                // Safety: the task is freshly allocated and not published yet.
+                    let mut meta = unsafe { task.task_meta() };
-                let mut meta = unsafe { task.task_meta() };
+                    me.task_hooks.spawn(&mut meta, parent);
-                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

@@ -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

@@ -222,8 +222,9 @@ where
 
                 #[cfg(tokio_unstable)]
                 {
-                    // Safety: the task is in the RUNNING state, which excludes
+                    // Safety: the task is in the RUNNING state, so shutdown
-                    // concurrent shutdown and termination metadata access.
+                    // 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(
+                // Safety: `transition_to_running` succeeded, so this thread has
-                    self.core(),
+                // exclusive access to the future/output storage. The header pointer
-                    #[cfg(tokio_unstable)]
+                // comes from this harness and remains live while the task is running.
-                    header_ptr,
+                let res = unsafe {
-                    cx,
+                    poll_future(
-                );
+                        self.core(),
+                        #[cfg(tokio_unstable)]
+                        header_ptr,
+                        cx,
+                    )
+                };
 
                 #[cfg(tokio_unstable)]
                 {
-                    // Safety: the task is still in the RUNNING state, which
+                    // Safety: the task is still in the RUNNING state, so
-                    // excludes concurrent shutdown and termination metadata
+                    // shutdown cannot take ownership of the task contents and
-                    // access.
+                    // 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(
+        // Safety: the caller guarantees the mutual-exclusion requirements of
-            #[cfg(tokio_unstable)]
+        // `Core::poll` and that `header` identifies this live task allocation.
-            header,
+        let res = unsafe {
-            cx,
+            guard.core.poll(
-        );
+                #[cfg(tokio_unstable)]
+                header,
+                cx,
+            )
+        };
         mem::forget(guard);
         res
     }));

tokio/src/runtime/task/mod.rs

@@ -669,13 +669,18 @@ impl SpawnLocation {
 }
 
 #[cfg(tokio_unstable)]
-pub(crate) fn current_task_meta<'meta>() -> Option<crate::runtime::TaskMetaRef<'meta>> {
+pub(crate) fn with_current_task_meta<R>(
-    let ptr = crate::runtime::context::current_task()?;
+    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
+    // Safety: parent metadata is exposed read-only during this synchronous call
-    // hook invocation while no mutable parent hook metadata is live.
+    // while no mutable parent hook metadata is live. The closure-bound lifetime
-    Some(unsafe { raw.task_meta_ref() })
+    // prevents references exposed through the metadata from escaping.
+    f(Some(unsafe { raw.task_meta_ref() }))
 }

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
+        // Safety: the caller guarantees that the task allocation is live and that
-        // metadata only for the current hook invocation.
+        // 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
+        // Safety: the caller guarantees that the task allocation is live and that
-        // shared access to task data.
+        // 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

@@ -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
+    /// If `user_data` is present, it must point to live task storage for the duration
-    /// duration of any references exposed through this metadata value.
+    /// 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

@@ -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 (terminated_tx, terminated_rx) = std::sync::mpsc::channel();
-    let terminated2 = Arc::clone(&terminated);
 
     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(