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]: https://github.com/tokio-rs/tokio/issues/7411#issuecomment-2993377045

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,6 +283,15 @@ 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(
@@ -267,6 +299,8 @@ impl<T: Future, S: Schedule> Cell<T, S> {
                     &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

tokio/src/runtime/task/mod.rs

@@ -216,6 +216,8 @@ use self::state::State;
 
 mod waker;
 
+pub(crate) use self::spawn_location::SpawnLocation;
+
 cfg_taskdump! {
     pub(crate) mod trace;
 }
@@ -226,6 +228,7 @@ use crate::util::sharded_list;
 
 use crate::runtime::TaskCallback;
 use std::marker::PhantomData;
+use std::panic::Location;
 use std::ptr::NonNull;
 use std::{fmt, mem};
 
@@ -243,6 +246,14 @@ unsafe impl<S> Sync for Task<S> {}
 #[repr(transparent)]
 pub(crate) struct Notified<S: 'static>(Task<S>);
 
+impl<S> Notified<S> {
+    #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))]
+    #[inline]
+    pub(crate) fn task_meta<'task, 'meta>(&'task self) -> crate::runtime::TaskMeta<'meta> {
+        self.0.task_meta()
+    }
+}
+
 // safety: This type cannot be used to touch the task without first verifying
 // that the value is on a thread where it is safe to poll the task.
 unsafe impl<S: Schedule> Send for Notified<S> {}
@@ -258,8 +269,9 @@ pub(crate) struct LocalNotified<S: 'static> {
 
 impl<S> LocalNotified<S> {
     #[cfg(tokio_unstable)]
-    pub(crate) fn task_id(&self) -> Id {
-        self.task.id()
+    #[inline]
+    pub(crate) fn task_meta<'task, 'meta>(&'task self) -> crate::runtime::TaskMeta<'meta> {
+        self.task.task_meta()
     }
 }
 
@@ -317,13 +329,19 @@ cfg_rt! {
         task: T,
         scheduler: S,
         id: Id,
+        spawned_at: SpawnLocation,
     ) -> (Task<S>, Notified<S>, JoinHandle<T::Output>)
     where
         S: Schedule,
         T: Future + 'static,
         T::Output: 'static,
     {
-        let raw = RawTask::new::<T, S>(task, scheduler, id);
+        let raw = RawTask::new::<T, S>(
+            task,
+            scheduler,
+            id,
+            spawned_at,
+        );
         let task = Task {
             raw,
             _p: PhantomData,
@@ -341,13 +359,23 @@ cfg_rt! {
     /// only when the task is not going to be stored in an `OwnedTasks` list.
     ///
     /// Currently only blocking tasks use this method.
-    pub(crate) fn unowned<T, S>(task: T, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>)
+    pub(crate) fn unowned<T, S>(
+        task: T,
+        scheduler: S,
+        id: Id,
+        spawned_at: SpawnLocation,
+    ) -> (UnownedTask<S>, JoinHandle<T::Output>)
     where
         S: Schedule,
         T: Send + Future + 'static,
         T::Output: Send + 'static,
     {
-        let (task, notified, join) = new_task(task, scheduler, id);
+        let (task, notified, join) = new_task(
+            task,
+            scheduler,
+            id,
+            spawned_at,
+        );
 
         // This transfers the ref-count of task and notified into an UnownedTask.
         // This is valid because an UnownedTask holds two ref-counts.
@@ -403,6 +431,24 @@ impl<S: 'static> Task<S> {
         unsafe { Header::get_id(self.raw.header_ptr()) }
     }
 
+    #[cfg(tokio_unstable)]
+    pub(crate) fn spawned_at(&self) -> &'static Location<'static> {
+        // Safety: The header pointer is valid.
+        unsafe { Header::get_spawn_location(self.raw.header_ptr()) }
+    }
+
+    // Explicit `'task` and `'meta` lifetimes are necessary here, as otherwise,
+    // the compiler infers the lifetimes to be the same, and considers the task
+    // to be borrowed for the lifetime of the returned `TaskMeta`.
+    #[cfg(tokio_unstable)]
+    pub(crate) fn task_meta<'task, 'meta>(&'task self) -> crate::runtime::TaskMeta<'meta> {
+        crate::runtime::TaskMeta {
+            id: self.id(),
+            spawned_at: self.spawned_at(),
+            _phantom: PhantomData,
+        }
+    }
+
     cfg_taskdump! {
         /// Notify the task for task dumping.
         ///
@@ -571,3 +617,48 @@ unsafe impl<S> sharded_list::ShardedListItem for Task<S> {
         task_id.0.get() as usize
     }
 }
+
+/// Wrapper around [`std::panic::Location`] that's conditionally compiled out
+/// when `tokio_unstable` is not enabled.
+#[cfg(tokio_unstable)]
+mod spawn_location {
+
+    use std::panic::Location;
+
+    #[derive(Copy, Clone)]
+    pub(crate) struct SpawnLocation(pub &'static Location<'static>);
+
+    impl From<&'static Location<'static>> for SpawnLocation {
+        fn from(location: &'static Location<'static>) -> Self {
+            Self(location)
+        }
+    }
+}
+
+#[cfg(not(tokio_unstable))]
+mod spawn_location {
+    use std::panic::Location;
+
+    #[derive(Copy, Clone)]
+    pub(crate) struct SpawnLocation();
+
+    impl From<&'static Location<'static>> for SpawnLocation {
+        fn from(_: &'static Location<'static>) -> Self {
+            Self()
+        }
+    }
+
+    #[cfg(test)]
+    #[test]
+    fn spawn_location_is_zero_sized() {
+        assert_eq!(std::mem::size_of::<SpawnLocation>(), 0);
+    }
+}
+
+impl SpawnLocation {
+    #[track_caller]
+    #[inline]
+    pub(crate) fn capture() -> Self {
+        Self::from(Location::caller())
+    }
+}

tokio/src/runtime/task/raw.rs

@@ -1,7 +1,8 @@
 use crate::future::Future;
 use crate::runtime::task::core::{Core, Trailer};
 use crate::runtime::task::{Cell, Harness, Header, Id, Schedule, State};
-
+#[cfg(tokio_unstable)]
+use std::panic::Location;
 use std::ptr::NonNull;
 use std::task::{Poll, Waker};
 
@@ -41,6 +42,10 @@ pub(super) struct Vtable {
 
     /// The number of bytes that the `id` field is offset from the header.
     pub(super) id_offset: usize,
+
+    /// The number of bytes that the `spawned_at` field is offset from the header.
+    #[cfg(tokio_unstable)]
+    pub(super) spawn_location_offset: usize,
 }
 
 /// Get the vtable for the requested `T` and `S` generics.
@@ -56,6 +61,8 @@ pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
         trailer_offset: OffsetHelper::<T, S>::TRAILER_OFFSET,
         scheduler_offset: OffsetHelper::<T, S>::SCHEDULER_OFFSET,
         id_offset: OffsetHelper::<T, S>::ID_OFFSET,
+        #[cfg(tokio_unstable)]
+        spawn_location_offset: OffsetHelper::<T, S>::SPAWN_LOCATION_OFFSET,
     }
 }
 
@@ -89,6 +96,16 @@ impl<T: Future, S: Schedule> OffsetHelper<T, S> {
         std::mem::size_of::<S>(),
         std::mem::align_of::<Id>(),
     );
+
+    #[cfg(tokio_unstable)]
+    const SPAWN_LOCATION_OFFSET: usize = get_spawn_location_offset(
+        std::mem::size_of::<Header>(),
+        std::mem::align_of::<Core<T, S>>(),
+        std::mem::size_of::<S>(),
+        std::mem::align_of::<Id>(),
+        std::mem::size_of::<Id>(),
+        std::mem::align_of::<&'static Location<'static>>(),
+    );
 }
 
 /// Compute the offset of the `Trailer` field in `Cell<T, S>` using the
@@ -156,13 +173,50 @@ const fn get_id_offset(
     offset
 }
 
+/// Compute the offset of the `&'static Location<'static>` field in `Cell<T, S>`
+/// using the `#[repr(C)]` algorithm.
+///
+/// Pseudo-code for the `#[repr(C)]` algorithm can be found here:
+/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs>
+#[cfg(tokio_unstable)]
+const fn get_spawn_location_offset(
+    header_size: usize,
+    core_align: usize,
+    scheduler_size: usize,
+    id_align: usize,
+    id_size: usize,
+    spawn_location_align: usize,
+) -> usize {
+    let mut offset = get_id_offset(header_size, core_align, scheduler_size, id_align);
+    offset += id_size;
+
+    let spawn_location_misalign = offset % spawn_location_align;
+    if spawn_location_misalign > 0 {
+        offset += spawn_location_align - spawn_location_misalign;
+    }
+
+    offset
+}
+
 impl RawTask {
-    pub(super) fn new<T, S>(task: T, scheduler: S, id: Id) -> RawTask
+    pub(super) fn new<T, S>(
+        task: T,
+        scheduler: S,
+        id: Id,
+        _spawned_at: super::SpawnLocation,
+    ) -> RawTask
     where
         T: Future,
         S: Schedule,
     {
-        let ptr = Box::into_raw(Cell::<_, S>::new(task, scheduler, State::new(), id));
+        let ptr = Box::into_raw(Cell::<_, S>::new(
+            task,
+            scheduler,
+            State::new(),
+            id,
+            #[cfg(tokio_unstable)]
+            _spawned_at.0,
+        ));
         let ptr = unsafe { NonNull::new_unchecked(ptr.cast()) };
 
         RawTask { ptr }

tokio/src/runtime/task_hooks.rs

@@ -1,4 +1,5 @@
 use std::marker::PhantomData;
+use std::panic::Location;
 
 use super::Config;
 
@@ -23,23 +24,17 @@ impl TaskHooks {
 
     #[cfg(tokio_unstable)]
     #[inline]
-    pub(crate) fn poll_start_callback(&self, id: super::task::Id) {
+    pub(crate) fn poll_start_callback(&self, meta: &TaskMeta<'_>) {
         if let Some(poll_start) = &self.before_poll_callback {
-            (poll_start)(&TaskMeta {
-                id,
-                _phantom: std::marker::PhantomData,
-            })
+            (poll_start)(meta);
         }
     }
 
     #[cfg(tokio_unstable)]
     #[inline]
-    pub(crate) fn poll_stop_callback(&self, id: super::task::Id) {
+    pub(crate) fn poll_stop_callback(&self, meta: &TaskMeta<'_>) {
         if let Some(poll_stop) = &self.after_poll_callback {
-            (poll_stop)(&TaskMeta {
-                id,
-                _phantom: std::marker::PhantomData,
-            })
+            (poll_stop)(meta);
         }
     }
 }
@@ -66,6 +61,8 @@ pub(crate) struct TaskHooks {
 pub struct TaskMeta<'a> {
     /// The opaque ID of the task.
     pub(crate) id: super::task::Id,
+    /// The location where the task was spawned.
+    pub(crate) spawned_at: &'static Location<'static>,
     pub(crate) _phantom: PhantomData<&'a ()>,
 }
 
@@ -75,6 +72,12 @@ impl<'a> TaskMeta<'a> {
     pub fn id(&self) -> super::task::Id {
         self.id
     }
+
+    /// Return the source code location where the task was spawned.
+    #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))]
+    pub fn spawned_at(&self) -> &'static Location<'static> {
+        self.spawned_at
+    }
 }
 
 /// Runs on specific task-related events

tokio/src/runtime/tests/mod.rs

@@ -29,10 +29,11 @@ mod noop_scheduler {
 }
 
 mod unowned_wrapper {
-    use crate::runtime::task::{Id, JoinHandle, Notified};
+    use crate::runtime::task::{Id, JoinHandle, Notified, SpawnLocation};
     use crate::runtime::tests::NoopSchedule;
 
     #[cfg(all(tokio_unstable, feature = "tracing"))]
+    #[track_caller]
     pub(crate) fn unowned<T>(task: T) -> (Notified<NoopSchedule>, JoinHandle<T::Output>)
     where
         T: std::future::Future + Send + 'static,
@@ -41,17 +42,20 @@ mod unowned_wrapper {
         use tracing::Instrument;
         let span = tracing::trace_span!("test_span");
         let task = task.instrument(span);
-        let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next());
+        let (task, handle) =
+            crate::runtime::task::unowned(task, NoopSchedule, Id::next(), SpawnLocation::capture());
         (task.into_notified(), handle)
     }
 
     #[cfg(not(all(tokio_unstable, feature = "tracing")))]
+    #[track_caller]
     pub(crate) fn unowned<T>(task: T) -> (Notified<NoopSchedule>, JoinHandle<T::Output>)
     where
         T: std::future::Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next());
+        let (task, handle) =
+            crate::runtime::task::unowned(task, NoopSchedule, Id::next(), SpawnLocation::capture());
         (task.into_notified(), handle)
     }
 }

tokio/src/runtime/tests/task.rs

@@ -1,10 +1,13 @@
 use crate::runtime::task::{
-    self, unowned, Id, JoinHandle, OwnedTasks, Schedule, Task, TaskHarnessScheduleHooks,
+    self, unowned, Id, JoinHandle, OwnedTasks, Schedule, SpawnLocation, Task,
+    TaskHarnessScheduleHooks,
 };
 use crate::runtime::tests::NoopSchedule;
 
 use std::collections::VecDeque;
 use std::future::Future;
+#[cfg(tokio_unstable)]
+use std::panic::Location;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 
@@ -57,6 +60,7 @@ fn create_drop1() {
         },
         NoopSchedule,
         Id::next(),
+        SpawnLocation::capture(),
     );
     drop(notified);
     handle.assert_not_dropped();
@@ -74,6 +78,7 @@ fn create_drop2() {
         },
         NoopSchedule,
         Id::next(),
+        SpawnLocation::capture(),
     );
     drop(join);
     handle.assert_not_dropped();
@@ -91,6 +96,7 @@ fn drop_abort_handle1() {
         },
         NoopSchedule,
         Id::next(),
+        SpawnLocation::capture(),
     );
     let abort = join.abort_handle();
     drop(join);
@@ -111,6 +117,7 @@ fn drop_abort_handle2() {
         },
         NoopSchedule,
         Id::next(),
+        SpawnLocation::capture(),
     );
     let abort = join.abort_handle();
     drop(notified);
@@ -131,6 +138,7 @@ fn drop_abort_handle_clone() {
         },
         NoopSchedule,
         Id::next(),
+        SpawnLocation::capture(),
     );
     let abort = join.abort_handle();
     let abort_clone = abort.clone();
@@ -155,6 +163,7 @@ fn create_shutdown1() {
         },
         NoopSchedule,
         Id::next(),
+        SpawnLocation::capture(),
     );
     drop(join);
     handle.assert_not_dropped();
@@ -172,6 +181,7 @@ fn create_shutdown2() {
         },
         NoopSchedule,
         Id::next(),
+        SpawnLocation::capture(),
     );
     handle.assert_not_dropped();
     notified.shutdown();
@@ -181,7 +191,7 @@ fn create_shutdown2() {
 
 #[test]
 fn unowned_poll() {
-    let (task, _) = unowned(async {}, NoopSchedule, Id::next());
+    let (task, _) = unowned(async {}, NoopSchedule, Id::next(), SpawnLocation::capture());
     task.run();
 }
 
@@ -385,12 +395,16 @@ struct Core {
 static CURRENT: Mutex<Option<Runtime>> = Mutex::new(None);
 
 impl Runtime {
+    #[track_caller]
     fn spawn<T>(&self, future: T) -> JoinHandle<T::Output>
     where
         T: 'static + Send + Future,
         T::Output: 'static + Send,
     {
-        let (handle, notified) = self.0.owned.bind(future, self.clone(), Id::next());
+        let (handle, notified) =
+            self.0
+                .owned
+                .bind(future, self.clone(), Id::next(), SpawnLocation::capture());
 
         if let Some(notified) = notified {
             self.schedule(notified);

tokio/src/task/local.rs

@@ -3,7 +3,9 @@ use crate::loom::cell::UnsafeCell;
 use crate::loom::sync::{Arc, Mutex};
 #[cfg(tokio_unstable)]
 use crate::runtime;
-use crate::runtime::task::{self, JoinHandle, LocalOwnedTasks, Task, TaskHarnessScheduleHooks};
+use crate::runtime::task::{
+    self, JoinHandle, LocalOwnedTasks, SpawnLocation, Task, TaskHarnessScheduleHooks,
+};
 use crate::runtime::{context, ThreadId, BOX_FUTURE_THRESHOLD};
 use crate::sync::AtomicWaker;
 use crate::util::trace::SpawnMeta;
@@ -1010,10 +1012,12 @@ impl Context {
         // Safety: called from the thread that owns the `LocalSet`
         let (handle, notified) = {
             self.shared.local_state.assert_called_from_owner_thread();
-            self.shared
-                .local_state
-                .owned
-                .bind(future, self.shared.clone(), id)
+            self.shared.local_state.owned.bind(
+                future,
+                self.shared.clone(),
+                id,
+                SpawnLocation::capture(),
+            )
         };
 
         if let Some(notified) = notified {

tokio/tests/task_hooks.rs

@@ -73,3 +73,117 @@ fn terminate_task_hook_fires() {
 
     assert_eq!(TASKS, count.load(Ordering::SeqCst));
 }
+
+/// Test that the correct spawn location is provided to the task hooks on a
+/// current thread runtime.
+#[test]
+fn task_hook_spawn_location_current_thread() {
+    let spawns = Arc::new(AtomicUsize::new(0));
+    let poll_starts = Arc::new(AtomicUsize::new(0));
+    let poll_ends = Arc::new(AtomicUsize::new(0));
+
+    let runtime = Builder::new_current_thread()
+        .on_task_spawn(mk_spawn_location_hook(
+            "(current_thread) on_task_spawn",
+            &spawns,
+        ))
+        .on_before_task_poll(mk_spawn_location_hook(
+            "(current_thread) on_before_task_poll",
+            &poll_starts,
+        ))
+        .on_after_task_poll(mk_spawn_location_hook(
+            "(current_thread) on_after_task_poll",
+            &poll_ends,
+        ))
+        .build()
+        .unwrap();
+
+    let task = runtime.spawn(async move { tokio::task::yield_now().await });
+    runtime.block_on(async move {
+        task.await.unwrap();
+        // tick the runtime a bunch to close out tasks
+        for _ in 0..ITERATIONS {
+            tokio::task::yield_now().await;
+        }
+    });
+
+    assert_eq!(spawns.load(Ordering::SeqCst), 1);
+    let poll_starts = poll_starts.load(Ordering::SeqCst);
+    assert!(poll_starts > 1);
+    assert_eq!(poll_starts, poll_ends.load(Ordering::SeqCst));
+}
+
+/// Test that the correct spawn location is provided to the task hooks on a
+/// multi-thread runtime.
+///
+/// Testing this separately is necessary as the spawn code paths are different
+/// and we should ensure that `#[track_caller]` is passed through correctly
+/// for both runtimes.
+#[cfg_attr(
+    target_os = "wasi",
+    ignore = "WASI does not support multi-threaded runtime"
+)]
+#[test]
+fn task_hook_spawn_location_multi_thread() {
+    let spawns = Arc::new(AtomicUsize::new(0));
+    let poll_starts = Arc::new(AtomicUsize::new(0));
+    let poll_ends = Arc::new(AtomicUsize::new(0));
+
+    let runtime = Builder::new_multi_thread()
+        .on_task_spawn(mk_spawn_location_hook(
+            "(multi_thread) on_task_spawn",
+            &spawns,
+        ))
+        .on_before_task_poll(mk_spawn_location_hook(
+            "(multi_thread) on_before_task_poll",
+            &poll_starts,
+        ))
+        .on_after_task_poll(mk_spawn_location_hook(
+            "(multi_thread) on_after_task_poll",
+            &poll_ends,
+        ))
+        .build()
+        .unwrap();
+
+    let task = runtime.spawn(async move { tokio::task::yield_now().await });
+    runtime.block_on(async move {
+        task.await.unwrap();
+
+        // tick the runtime a bunch to close out tasks
+        for _ in 0..ITERATIONS {
+            tokio::task::yield_now().await;
+        }
+    });
+
+    // Give the runtime to shut down so that we see all the expected calls to
+    // the task hooks.
+    runtime.shutdown_timeout(std::time::Duration::from_secs(60));
+
+    // Note: we "read" the counters using `fetch_add(0, SeqCst)` rather than
+    // `load(SeqCst)` because read-write-modify operations are guaranteed to
+    // observe the latest value, while the load is not.
+    // This avoids a race that may cause test flakiness.
+    assert_eq!(spawns.fetch_add(0, Ordering::SeqCst), 1);
+    let poll_starts = poll_starts.fetch_add(0, Ordering::SeqCst);
+    assert!(poll_starts > 1);
+    assert_eq!(poll_starts, poll_ends.fetch_add(0, Ordering::SeqCst));
+}
+
+fn mk_spawn_location_hook(
+    event: &'static str,
+    count: &Arc<AtomicUsize>,
+) -> impl Fn(&tokio::runtime::TaskMeta<'_>) {
+    let count = Arc::clone(&count);
+    move |data| {
+        eprintln!("{event} ({:?}): {:?}", data.id(), data.spawned_at());
+        // Assert that the spawn location is in this file.
+        // Don't make assertions about line number/column here, as these
+        // may change as new code is added to the test file...
+        assert_eq!(
+            data.spawned_at().file(),
+            file!(),
+            "incorrect spawn location in {event} hook",
+        );
+        count.fetch_add(1, Ordering::SeqCst);
+    }
+}