From c18992ec1d42821935d4dc876272501d16114111 Mon Sep 17 00:00:00 2001
From: "Stainsby, Hayden" <hayden.stainsby@here.com>
Date: Wed, 11 Sep 2024 12:51:54 +0200
Subject: [PATCH] track task spawn location for tokio instrumentation

The tracing instrumentation present in tokio (which feeds into
tokio-console) tracks the spawning location of each task. This can be
useful in understanding what each task is doing, especially if the tasks
aren't named (which is a `tokio_unstable` feature so not usually present
in published crates).

The location tracking isn't as useful when tasks are spawned from some
central location like the `mongodb` crate does, because all tasks have
the same spawn location. However, since the instrumentation uses the
`panic:Location` to determine the spawn location, the `#[track_caller]`
attribute macro can be used to ensure that the actual location is picked
up instead of the common one.

A side effect of using the `#[track_caller]` attribute is that the
single line output in the case of a panic (when `RUST_BACKTRACE` isn't
set) will point to the callsite in the code, and not the common site -
although this is often desirable anyway.

This change adds `#[track_caller]` to the `pub(crate)` function
`crate::runtime::spawn` as well as the method `AsyncJoinHandle::spawn`.
---
 src/runtime.rs             | 1 +
 src/runtime/join_handle.rs | 1 +
 2 files changed, 2 insertions(+)

diff --git a/src/runtime.rs b/src/runtime.rs
index 373bbec85..f76f9e308 100644
--- a/src/runtime.rs
+++ b/src/runtime.rs
@@ -54,6 +54,7 @@ pub(crate) use tls::TlsConfig;
 ///
 /// This must be called from an async block
 /// or function running on a runtime.
+#[track_caller]
 pub(crate) fn spawn<F, O>(fut: F) -> AsyncJoinHandle<O>
 where
     F: Future<Output = O> + Send + 'static,
diff --git a/src/runtime/join_handle.rs b/src/runtime/join_handle.rs
index 7e74d23d7..9bf413833 100644
--- a/src/runtime/join_handle.rs
+++ b/src/runtime/join_handle.rs
@@ -9,6 +9,7 @@ use std::{
 pub(crate) struct AsyncJoinHandle<T>(tokio::task::JoinHandle<T>);
 
 impl<T> AsyncJoinHandle<T> {
+    #[track_caller]
     pub(crate) fn spawn<F>(fut: F) -> Self
     where
         F: Future<Output = T> + Send + 'static,