From 052355f0649bf54ee1850060a636d045eb7e2bf1 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Mon, 16 May 2022 09:36:24 -0700
Subject: [PATCH] task: add `#[track_caller]` to `JoinSet`/`JoinMap` (#4697)

## Motivation

Currently, the various spawning methods on `tokio::task::JoinSet` and
`tokio_util::task::JoinMap` lack `#[track_caller]` attributes, so the
`tracing` spans generated for tasks spawned on `JoinSet`s/`JoinMap`s
will have the `JoinSet` spawning method as their spawn location, rather
than the user code that called that method.

## Solution

This PR fixes that by...adding a bunch of `#[track_caller]` attributes.

## Future Work

In the future, we may also want to consider adding additional data to
the task span indicating that the task is spawned on a `JoinSet`...

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 tokio-util/src/task/join_map.rs | 4 ++++
 tokio/src/task/join_set.rs      | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/tokio-util/src/task/join_map.rs b/tokio-util/src/task/join_map.rs
index 41c82f448ab..299c92faed3 100644
--- a/tokio-util/src/task/join_map.rs
+++ b/tokio-util/src/task/join_map.rs
@@ -285,6 +285,7 @@ where
     /// This method panics if called outside of a Tokio runtime.
     ///
     /// [`join_one`]: Self::join_one
+    #[track_caller]
     pub fn spawn<F>(&mut self, key: K, task: F)
     where
         F: Future<Output = V>,
@@ -304,6 +305,7 @@ where
     /// *not* return a cancelled [`JoinError`] for that task.
     ///
     /// [`join_one`]: Self::join_one
+    #[track_caller]
     pub fn spawn_on<F>(&mut self, key: K, task: F, handle: &Handle)
     where
         F: Future<Output = V>,
@@ -328,6 +330,7 @@ where
     ///
     /// [`LocalSet`]: tokio::task::LocalSet
     /// [`join_one`]: Self::join_one
+    #[track_caller]
     pub fn spawn_local<F>(&mut self, key: K, task: F)
     where
         F: Future<Output = V>,
@@ -347,6 +350,7 @@ where
     ///
     /// [`LocalSet`]: tokio::task::LocalSet
     /// [`join_one`]: Self::join_one
+    #[track_caller]
     pub fn spawn_local_on<F>(&mut self, key: K, task: F, local_set: &LocalSet)
     where
         F: Future<Output = V>,
diff --git a/tokio/src/task/join_set.rs b/tokio/src/task/join_set.rs
index d036fcc3cda..42f55a034cc 100644
--- a/tokio/src/task/join_set.rs
+++ b/tokio/src/task/join_set.rs
@@ -81,6 +81,7 @@ impl<T: 'static> JoinSet<T> {
     /// This method panics if called outside of a Tokio runtime.
     ///
     /// [`AbortHandle`]: crate::task::AbortHandle
+    #[track_caller]
     pub fn spawn<F>(&mut self, task: F) -> AbortHandle
     where
         F: Future<Output = T>,
@@ -95,6 +96,7 @@ impl<T: 'static> JoinSet<T> {
     /// cancel the task.
     ///
     /// [`AbortHandle`]: crate::task::AbortHandle
+    #[track_caller]
     pub fn spawn_on<F>(&mut self, task: F, handle: &Handle) -> AbortHandle
     where
         F: Future<Output = T>,
@@ -114,6 +116,7 @@ impl<T: 'static> JoinSet<T> {
     ///
     /// [`LocalSet`]: crate::task::LocalSet
     /// [`AbortHandle`]: crate::task::AbortHandle
+    #[track_caller]
     pub fn spawn_local<F>(&mut self, task: F) -> AbortHandle
     where
         F: Future<Output = T>,
@@ -128,6 +131,7 @@ impl<T: 'static> JoinSet<T> {
     ///
     /// [`LocalSet`]: crate::task::LocalSet
     /// [`AbortHandle`]: crate::task::AbortHandle
+    #[track_caller]
     pub fn spawn_local_on<F>(&mut self, task: F, local_set: &LocalSet) -> AbortHandle
     where
         F: Future<Output = T>,