rt(unstable): add spawn Location to TaskMeta (#7417)

As described in issue #7411, task spawning APIs are currently annotated
with `#[track_caller]`, allowing us to capture the location in the user
source code where the task was spawned. This is used for `tracing`
events used by `tokio-console` and friends. However, this information is
*not* exposed to the runtime `on_task_spawn`, `on_before_task_poll`,
`on_after_task_poll`, and `on_task_terminate` hooks, which is a shame,
as it would be useful there as well.

This branch adds the task's spawn location to the `TaskMeta` struct
provided to the runtime's task hooks. This is implemented by storing a
`&'static Location<'static>` in the task's `Core` alongside the
`task::Id`. In [this comment][1], @ADD-SP suggested storing the
`Location` in the task's `Trailer`.

I opted to store it in the `Core` instead, as the `Trailer` is intended
to store "cold" data that is only accessed when the task _completes_,
and not on every poll. Since the task meta is passed to the
`on_before_task_poll` and `on_after_task_poll` hooks, we would be
accessing the `Trailer` on polls if we stored the `Location` there.
Therefore, I put it in the `Core`, instead, which contains data that we
access every time the task is polled.

Closes #7411

[1]: #7411 (comment)

Author: hawkw

tokio/src/runtime/blocking/pool.rs

@@ -379,7 +379,12 @@ impl Spawner {
         let fut =
             blocking_task::<F, BlockingTask<F>>(BlockingTask::new(func), spawn_meta, id.as_u64());
 
-        let (task, handle) = task::unowned(fut, BlockingSchedule::new(rt), id);
+        let (task, handle) = task::unowned(
+            fut,
+            BlockingSchedule::new(rt),
+            id,
+            task::SpawnLocation::capture(),
+        );
 
         let spawned = self.spawn_task(Task::new(task, is_mandatory), rt);
         (handle, spawned)

tokio/src/runtime/scheduler/current_thread/mod.rs

@@ -15,6 +15,7 @@ use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef};
 use std::cell::RefCell;
 use std::collections::VecDeque;
 use std::future::{poll_fn, Future};
+use std::panic::Location;
 use std::sync::atomic::Ordering::{AcqRel, Release};
 use std::task::Poll::{Pending, Ready};
 use std::task::Waker;
@@ -445,6 +446,7 @@ impl Context {
 
 impl Handle {
     /// Spawns a future onto the `CurrentThread` scheduler
+    #[track_caller]
     pub(crate) fn spawn<F>(
         me: &Arc<Self>,
         future: F,
@@ -454,10 +456,15 @@ impl Handle {
         F: crate::future::Future + Send + 'static,
         F::Output: Send + 'static,
     {
-        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id);
+        let spawned_at = Location::caller();
+        let (handle, notified) = me
+            .shared
+            .owned
+            .bind(future, me.clone(), id, spawned_at.into());
 
         me.task_hooks.spawn(&TaskMeta {
             id,
+            spawned_at,
             _phantom: Default::default(),
         });
 
@@ -474,6 +481,7 @@ impl Handle {
     /// This should only be used when this is a `LocalRuntime` or in another case where the runtime
     /// provably cannot be driven from or moved to different threads from the one on which the task
     /// is spawned.
+    #[track_caller]
     pub(crate) unsafe fn spawn_local<F>(
         me: &Arc<Self>,
         future: F,
@@ -483,10 +491,15 @@ impl Handle {
         F: crate::future::Future + 'static,
         F::Output: 'static,
     {
-        let (handle, notified) = me.shared.owned.bind_local(future, me.clone(), id);
+        let spawned_at = Location::caller();
+        let (handle, notified) =
+            me.shared
+                .owned
+                .bind_local(future, me.clone(), id, spawned_at.into());
 
         me.task_hooks.spawn(&TaskMeta {
             id,
+            spawned_at,
             _phantom: Default::default(),
         });
 
@@ -771,16 +784,16 @@ impl CoreGuard<'_> {
                     let task = context.handle.shared.owned.assert_owner(task);
 
                     #[cfg(tokio_unstable)]
-                    let task_id = task.task_id();
+                    let task_meta = task.task_meta();
 
                     let (c, ()) = context.run_task(core, || {
                         #[cfg(tokio_unstable)]
-                        context.handle.task_hooks.poll_start_callback(task_id);
+                        context.handle.task_hooks.poll_start_callback(&task_meta);
 
                         task.run();
 
                         #[cfg(tokio_unstable)]
-                        context.handle.task_hooks.poll_stop_callback(task_id);
+                        context.handle.task_hooks.poll_stop_callback(&task_meta);
                     });
 
                     core = c;

tokio/src/runtime/scheduler/mod.rs

@@ -117,6 +117,7 @@ cfg_rt! {
             }
         }
 
+        #[track_caller]
         pub(crate) fn spawn<F>(&self, future: F, id: Id) -> JoinHandle<F::Output>
         where
             F: Future + Send + 'static,
@@ -136,6 +137,7 @@ cfg_rt! {
         /// This should only be called in `LocalRuntime` if the runtime has been verified to be owned
         /// by the current thread.
         #[allow(irrefutable_let_patterns)]
+        #[track_caller]
         pub(crate) unsafe fn spawn_local<F>(&self, future: F, id: Id) -> JoinHandle<F::Output>
         where
             F: Future + 'static,

tokio/src/runtime/scheduler/multi_thread/handle.rs

@@ -10,6 +10,7 @@ use crate::runtime::{
 use crate::util::RngSeedGenerator;
 
 use std::fmt;
+use std::panic::Location;
 
 mod metrics;
 
@@ -37,6 +38,7 @@ pub(crate) struct Handle {
 
 impl Handle {
     /// Spawns a future onto the thread pool
+    #[track_caller]
     pub(crate) fn spawn<F>(me: &Arc<Self>, future: F, id: task::Id) -> JoinHandle<F::Output>
     where
         F: crate::future::Future + Send + 'static,
@@ -49,15 +51,21 @@ impl Handle {
         self.close();
     }
 
+    #[track_caller]
     pub(super) fn bind_new_task<T>(me: &Arc<Self>, future: T, id: task::Id) -> JoinHandle<T::Output>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id);
+        let spawned_at = Location::caller();
+        let (handle, notified) = me
+            .shared
+            .owned
+            .bind(future, me.clone(), id, spawned_at.into());
 
         me.task_hooks.spawn(&TaskMeta {
             id,
+            spawned_at,
             _phantom: Default::default(),
         });
 

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -568,7 +568,7 @@ impl Context {
 
     fn run_task(&self, task: Notified, mut core: Box<Core>) -> RunResult {
         #[cfg(tokio_unstable)]
-        let task_id = task.task_id();
+        let task_meta = task.task_meta();
 
         let task = self.worker.handle.shared.owned.assert_owner(task);
 
@@ -592,12 +592,15 @@ impl Context {
             // Unlike the poll time above, poll start callback is attached to the task id,
             // so it is tightly associated with the actual poll invocation.
             #[cfg(tokio_unstable)]
-            self.worker.handle.task_hooks.poll_start_callback(task_id);
+            self.worker
+                .handle
+                .task_hooks
+                .poll_start_callback(&task_meta);
 
             task.run();
 
             #[cfg(tokio_unstable)]
-            self.worker.handle.task_hooks.poll_stop_callback(task_id);
+            self.worker.handle.task_hooks.poll_stop_callback(&task_meta);
 
             let mut lifo_polls = 0;
 
@@ -663,15 +666,18 @@ impl Context {
                 let task = self.worker.handle.shared.owned.assert_owner(task);
 
                 #[cfg(tokio_unstable)]
-                let task_id = task.task_id();
+                let task_meta = task.task_meta();
 
                 #[cfg(tokio_unstable)]
-                self.worker.handle.task_hooks.poll_start_callback(task_id);
+                self.worker
+                    .handle
+                    .task_hooks
+                    .poll_start_callback(&task_meta);
 
                 task.run();
 
                 #[cfg(tokio_unstable)]
-                self.worker.handle.task_hooks.poll_stop_callback(task_id);
+                self.worker.handle.task_hooks.poll_stop_callback(&task_meta);
             }
         })
     }

tokio/src/runtime/task/core.rs

@@ -18,6 +18,8 @@ use crate::runtime::task::{Id, Schedule, TaskHarnessScheduleHooks};
 use crate::util::linked_list;
 
 use std::num::NonZeroU64;
+#[cfg(tokio_unstable)]
+use std::panic::Location;
 use std::pin::Pin;
 use std::ptr::NonNull;
 use std::task::{Context, Poll, Waker};
@@ -141,6 +143,13 @@ pub(super) struct Core<T: Future, S> {
     /// The task's ID, used for populating `JoinError`s.
     pub(super) task_id: Id,
 
+    /// The source code location where the task was spawned.
+    ///
+    /// This is used for populating the `TaskMeta` passed to the task runtime
+    /// hooks.
+    #[cfg(tokio_unstable)]
+    pub(super) spawned_at: &'static Location<'static>,
+
     /// Either the future or the output.
     pub(super) stage: CoreStage<T>,
 }
@@ -208,7 +217,13 @@ pub(super) enum Stage<T: Future> {
 impl<T: Future, S: Schedule> Cell<T, S> {
     /// Allocates a new task cell, containing the header, trailer, and core
     /// structures.
-    pub(super) fn new(future: T, scheduler: S, state: State, task_id: Id) -> Box<Cell<T, S>> {
+    pub(super) fn new(
+        future: T,
+        scheduler: S,
+        state: State,
+        task_id: Id,
+        #[cfg(tokio_unstable)] spawned_at: &'static Location<'static>,
+    ) -> Box<Cell<T, S>> {
         // Separated into a non-generic function to reduce LLVM codegen
         fn new_header(
             state: State,
@@ -242,13 +257,21 @@ impl<T: Future, S: Schedule> Cell<T, S> {
                     stage: UnsafeCell::new(Stage::Running(future)),
                 },
                 task_id,
+                #[cfg(tokio_unstable)]
+                spawned_at,
             },
         });
 
         #[cfg(debug_assertions)]
         {
             // Using a separate function for this code avoids instantiating it separately for every `T`.
-            unsafe fn check<S>(header: &Header, trailer: &Trailer, scheduler: &S, task_id: &Id) {
+            unsafe fn check<S>(
+                header: &Header,
+                trailer: &Trailer,
+                scheduler: &S,
+                task_id: &Id,
+                #[cfg(tokio_unstable)] spawn_location: &&'static Location<'static>,
+            ) {
                 let trailer_addr = trailer as *const Trailer as usize;
                 let trailer_ptr = unsafe { Header::get_trailer(NonNull::from(header)) };
                 assert_eq!(trailer_addr, trailer_ptr.as_ptr() as usize);
@@ -260,13 +283,24 @@ impl<T: Future, S: Schedule> Cell<T, S> {
                 let id_addr = task_id as *const Id as usize;
                 let id_ptr = unsafe { Header::get_id_ptr(NonNull::from(header)) };
                 assert_eq!(id_addr, id_ptr.as_ptr() as usize);
+
+                #[cfg(tokio_unstable)]
+                {
+                    let spawn_location_addr =
+                        spawn_location as *const &'static Location<'static> as usize;
+                    let spawn_location_ptr =
+                        unsafe { Header::get_spawn_location_ptr(NonNull::from(header)) };
+                    assert_eq!(spawn_location_addr, spawn_location_ptr.as_ptr() as usize);
+                }
             }
             unsafe {
                 check(
                     &result.header,
                     &result.trailer,
                     &result.core.scheduler,
                     &result.core.task_id,
+                    #[cfg(tokio_unstable)]
+                    &result.core.spawned_at,
                 );
             }
         }
@@ -450,6 +484,37 @@ impl Header {
         *ptr
     }
 
+    /// Gets a pointer to the source code location where the task containing
+    /// this `Header` was spawned.
+    ///
+    /// # Safety
+    ///
+    /// The provided raw pointer must point at the header of a task.
+    #[cfg(tokio_unstable)]
+    pub(super) unsafe fn get_spawn_location_ptr(
+        me: NonNull<Header>,
+    ) -> NonNull<&'static Location<'static>> {
+        let offset = me.as_ref().vtable.spawn_location_offset;
+        let spawned_at = me
+            .as_ptr()
+            .cast::<u8>()
+            .add(offset)
+            .cast::<&'static Location<'static>>();
+        NonNull::new_unchecked(spawned_at)
+    }
+
+    /// Gets the source code location where the task containing
+    /// this `Header` was spawned
+    ///
+    /// # Safety
+    ///
+    /// The provided raw pointer must point at the header of a task.
+    #[cfg(tokio_unstable)]
+    pub(super) unsafe fn get_spawn_location(me: NonNull<Header>) -> &'static Location<'static> {
+        let ptr = Header::get_spawn_location_ptr(me).as_ptr();
+        *ptr
+    }
+
     /// Gets the tracing id of the task containing this `Header`.
     ///
     /// # Safety

tokio/src/runtime/task/harness.rs

@@ -4,6 +4,7 @@ use crate::runtime::task::state::{Snapshot, State};
 use crate::runtime::task::waker::waker_ref;
 use crate::runtime::task::{Id, JoinError, Notified, RawTask, Schedule, Task};
 
+#[cfg(tokio_unstable)]
 use crate::runtime::TaskMeta;
 use std::any::Any;
 use std::mem;
@@ -367,10 +368,12 @@ where
         //
         // We call this in a separate block so that it runs after the task appears to have
         // completed and will still run if the destructor panics.
+        #[cfg(tokio_unstable)]
         if let Some(f) = self.trailer().hooks.task_terminate_callback.as_ref() {
             let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
                 f(&TaskMeta {
                     id: self.core().task_id,
+                    spawned_at: self.core().spawned_at,
                     _phantom: Default::default(),
                 })
             }));

tokio/src/runtime/task/list.rs

@@ -8,7 +8,7 @@
 
 use crate::future::Future;
 use crate::loom::cell::UnsafeCell;
-use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, Task};
+use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, SpawnLocation, Task};
 use crate::util::linked_list::{Link, LinkedList};
 use crate::util::sharded_list;
 
@@ -91,13 +91,14 @@ impl<S: 'static> OwnedTasks<S> {
         task: T,
         scheduler: S,
         id: super::Id,
+        spawned_at: SpawnLocation,
     ) -> (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);
+        let (task, notified, join) = super::new_task(task, scheduler, id, spawned_at);
         let notified = unsafe { self.bind_inner(task, notified) };
         (join, notified)
     }
@@ -111,13 +112,14 @@ impl<S: 'static> OwnedTasks<S> {
         task: T,
         scheduler: S,
         id: super::Id,
+        spawned_at: SpawnLocation,
     ) -> (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);
+        let (task, notified, join) = super::new_task(task, scheduler, id, spawned_at);
         let notified = unsafe { self.bind_inner(task, notified) };
         (join, notified)
     }
@@ -258,13 +260,14 @@ impl<S: 'static> LocalOwnedTasks<S> {
         task: T,
         scheduler: S,
         id: super::Id,
+        spawned_at: SpawnLocation,
     ) -> (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);
+        let (task, notified, join) = super::new_task(task, scheduler, id, spawned_at);
 
         unsafe {
             // safety: We just created the task, so we have exclusive access