From 3e890cc0171ddb210acdcfec831b7c7bcbb0d2d9 Mon Sep 17 00:00:00 2001 From: Eliza Weisman Date: Mon, 30 Jun 2025 11:13:42 -0700 Subject: [PATCH] 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 | 7 +- .../runtime/scheduler/current_thread/mod.rs | 23 +++- tokio/src/runtime/scheduler/mod.rs | 2 + .../runtime/scheduler/multi_thread/handle.rs | 10 +- .../runtime/scheduler/multi_thread/worker.rs | 18 ++- tokio/src/runtime/task/core.rs | 69 ++++++++++- tokio/src/runtime/task/harness.rs | 3 + tokio/src/runtime/task/list.rs | 11 +- tokio/src/runtime/task/mod.rs | 101 +++++++++++++++- tokio/src/runtime/task/raw.rs | 60 ++++++++- tokio/src/runtime/task_hooks.rs | 23 ++-- tokio/src/runtime/tests/mod.rs | 10 +- tokio/src/runtime/tests/task.rs | 20 ++- tokio/src/task/local.rs | 14 ++- tokio/tests/task_hooks.rs | 114 ++++++++++++++++++ 15 files changed, 437 insertions(+), 48 deletions(-) diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 23180dc52..04edcac1d 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -379,7 +379,12 @@ impl Spawner { let fut = blocking_task::>(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) diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index 72e12fae8..c12308bf5 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/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( me: &Arc, 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( me: &Arc, 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; diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs index 8241b57c1..7e2cb91ed 100644 --- a/tokio/src/runtime/scheduler/mod.rs +++ b/tokio/src/runtime/scheduler/mod.rs @@ -117,6 +117,7 @@ cfg_rt! { } } + #[track_caller] pub(crate) fn spawn(&self, future: F, id: Id) -> JoinHandle 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(&self, future: F, id: Id) -> JoinHandle where F: Future + 'static, diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index 47299659f..7ac525067 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/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(me: &Arc, future: F, id: task::Id) -> JoinHandle where F: crate::future::Future + Send + 'static, @@ -49,15 +51,21 @@ impl Handle { self.close(); } + #[track_caller] pub(super) fn bind_new_task(me: &Arc, future: T, id: task::Id) -> JoinHandle 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(), }); diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 838c15840..8a9719b9b 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -568,7 +568,7 @@ impl Context { fn run_task(&self, task: Notified, mut core: Box) -> 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); } }) } diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index 5d3ca0e00..e91e8be40 100644 --- a/tokio/src/runtime/task/core.rs +++ b/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 { /// 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, } @@ -208,7 +217,13 @@ pub(super) enum Stage { impl Cell { /// 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> { + pub(super) fn new( + future: T, + scheduler: S, + state: State, + task_id: Id, + #[cfg(tokio_unstable)] spawned_at: &'static Location<'static>, + ) -> Box> { // Separated into a non-generic function to reduce LLVM codegen fn new_header( state: State, @@ -242,13 +257,21 @@ impl Cell { 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(header: &Header, trailer: &Trailer, scheduler: &S, task_id: &Id) { + unsafe fn check( + 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 Cell { 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 Cell { &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
, + ) -> NonNull<&'static Location<'static>> { + let offset = me.as_ref().vtable.spawn_location_offset; + let spawned_at = me + .as_ptr() + .cast::() + .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
) -> &'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 diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 9bf73b74f..543d2e527 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/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(), }) })); diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs index 54bfc01aa..908ce07ec 100644 --- a/tokio/src/runtime/task/list.rs +++ b/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 OwnedTasks { task: T, scheduler: S, id: super::Id, + spawned_at: SpawnLocation, ) -> (JoinHandle, Option>) 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 OwnedTasks { task: T, scheduler: S, id: super::Id, + spawned_at: SpawnLocation, ) -> (JoinHandle, Option>) 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 LocalOwnedTasks { task: T, scheduler: S, id: super::Id, + spawned_at: SpawnLocation, ) -> (JoinHandle, Option>) 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 diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index 7d314c3b1..e6b1a0804 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/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 Sync for Task {} #[repr(transparent)] pub(crate) struct Notified(Task); +impl Notified { + #[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 Send for Notified {} @@ -258,8 +269,9 @@ pub(crate) struct LocalNotified { impl LocalNotified { #[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, Notified, JoinHandle) where S: Schedule, T: Future + 'static, T::Output: 'static, { - let raw = RawTask::new::(task, scheduler, id); + let raw = RawTask::new::( + 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(task: T, scheduler: S, id: Id) -> (UnownedTask, JoinHandle) + pub(crate) fn unowned( + task: T, + scheduler: S, + id: Id, + spawned_at: SpawnLocation, + ) -> (UnownedTask, JoinHandle) 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 Task { 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 sharded_list::ShardedListItem for Task { 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::(), 0); + } +} + +impl SpawnLocation { + #[track_caller] + #[inline] + pub(crate) fn capture() -> Self { + Self::from(Location::caller()) + } +} diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs index 6699551f3..e9a378022 100644 --- a/tokio/src/runtime/task/raw.rs +++ b/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() -> &'static Vtable { trailer_offset: OffsetHelper::::TRAILER_OFFSET, scheduler_offset: OffsetHelper::::SCHEDULER_OFFSET, id_offset: OffsetHelper::::ID_OFFSET, + #[cfg(tokio_unstable)] + spawn_location_offset: OffsetHelper::::SPAWN_LOCATION_OFFSET, } } @@ -89,6 +96,16 @@ impl OffsetHelper { std::mem::size_of::(), std::mem::align_of::(), ); + + #[cfg(tokio_unstable)] + const SPAWN_LOCATION_OFFSET: usize = get_spawn_location_offset( + std::mem::size_of::
(), + std::mem::align_of::>(), + std::mem::size_of::(), + std::mem::align_of::(), + std::mem::size_of::(), + std::mem::align_of::<&'static Location<'static>>(), + ); } /// Compute the offset of the `Trailer` field in `Cell` using the @@ -156,13 +173,50 @@ const fn get_id_offset( offset } +/// Compute the offset of the `&'static Location<'static>` field in `Cell` +/// using the `#[repr(C)]` algorithm. +/// +/// Pseudo-code for the `#[repr(C)]` algorithm can be found here: +/// +#[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(task: T, scheduler: S, id: Id) -> RawTask + pub(super) fn new( + 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 } diff --git a/tokio/src/runtime/task_hooks.rs b/tokio/src/runtime/task_hooks.rs index 13865ed51..bc470302e 100644 --- a/tokio/src/runtime/task_hooks.rs +++ b/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 diff --git a/tokio/src/runtime/tests/mod.rs b/tokio/src/runtime/tests/mod.rs index 6fcf8a2ec..5471ffc2e 100644 --- a/tokio/src/runtime/tests/mod.rs +++ b/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(task: T) -> (Notified, JoinHandle) 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(task: T) -> (Notified, JoinHandle) 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) } } diff --git a/tokio/src/runtime/tests/task.rs b/tokio/src/runtime/tests/task.rs index ea48b8e51..7a10ac4a3 100644 --- a/tokio/src/runtime/tests/task.rs +++ b/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> = Mutex::new(None); impl Runtime { + #[track_caller] fn spawn(&self, future: T) -> JoinHandle 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); diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index f00836730..11b877fd1 100644 --- a/tokio/src/task/local.rs +++ b/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 { diff --git a/tokio/tests/task_hooks.rs b/tokio/tests/task_hooks.rs index 185b9126c..2c8c4717a 100644 --- a/tokio/tests/task_hooks.rs +++ b/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, +) -> 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); + } +}