add task IDs to tracing spans

This commit adds the runtime generated task IDs to the task `tracing`
spans, when tracing is enabled. Unfortunately, this means the IDs have
to be generated much higher up the callstack so they can be added to the
tracing spans, and then passed down to the actual task machinery. But, I
think this is worth it to be able to record that data.

Signed-off-by: Eliza Weisman <[email protected]>

tokio/src/runtime/basic_scheduler.rs

@@ -365,12 +365,12 @@ impl Context {
 
 impl Spawner {
     /// Spawns a future onto the basic scheduler
-    pub(crate) fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
+    pub(crate) fn spawn<F>(&self, future: F, id: super::task::Id) -> JoinHandle<F::Output>
     where
         F: crate::future::Future + Send + 'static,
         F::Output: Send + 'static,
     {
-        let (handle, notified) = self.shared.owned.bind(future, self.shared.clone());
+        let (handle, notified) = self.shared.owned.bind(future, self.shared.clone(), id);
 
         if let Some(notified) = notified {
             self.shared.schedule(notified);

tokio/src/runtime/handle.rs

@@ -166,9 +166,10 @@ impl Handle {
         F: Future + Send + 'static,
         F::Output: Send + 'static,
     {
+        let id = crate::runtime::task::Id::next();
         #[cfg(all(tokio_unstable, feature = "tracing"))]
-        let future = crate::util::trace::task(future, "task", None);
-        self.spawner.spawn(future)
+        let future = crate::util::trace::task(future, "task", None, id.as_usize());
+        self.spawner.spawn(future, id)
     }
 
     /// Runs the provided function on an executor dedicated to blocking.
@@ -251,7 +252,7 @@ impl Handle {
         R: Send + 'static,
     {
         let fut = BlockingTask::new(func);
-
+        let id = super::task::Id::next();
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let fut = {
             use tracing::Instrument;
@@ -261,6 +262,7 @@ impl Handle {
                 "runtime.spawn",
                 kind = %"blocking",
                 task.name = %name.unwrap_or_default(),
+                task.id = id.as_usize(),
                 "fn" = %std::any::type_name::<F>(),
                 spawn.location = %format_args!("{}:{}:{}", location.file(), location.line(), location.column()),
             );
@@ -270,7 +272,7 @@ impl Handle {
         #[cfg(not(all(tokio_unstable, feature = "tracing")))]
         let _ = name;
 
-        let (task, handle) = task::unowned(fut, NoopSchedule);
+        let (task, handle) = task::unowned(fut, NoopSchedule, id);
         let spawned = self
             .blocking_spawner
             .spawn(blocking::Task::new(task, is_mandatory), self);

tokio/src/runtime/spawner.rs

@@ -1,5 +1,6 @@
 use crate::future::Future;
 use crate::runtime::basic_scheduler;
+use crate::runtime::task::Id;
 use crate::task::JoinHandle;
 
 cfg_rt_multi_thread! {
@@ -23,15 +24,15 @@ impl Spawner {
         }
     }
 
-    pub(crate) fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
+    pub(crate) fn spawn<F>(&self, future: F, id: Id) -> JoinHandle<F::Output>
     where
         F: Future + Send + 'static,
         F::Output: Send + 'static,
     {
         match self {
-            Spawner::Basic(spawner) => spawner.spawn(future),
+            Spawner::Basic(spawner) => spawner.spawn(future, id),
             #[cfg(feature = "rt-multi-thread")]
-            Spawner::ThreadPool(spawner) => spawner.spawn(future),
+            Spawner::ThreadPool(spawner) => spawner.spawn(future, id),
         }
     }
 }

tokio/src/runtime/task/list.rs

@@ -84,13 +84,14 @@ impl<S: 'static> OwnedTasks<S> {
         &self,
         task: T,
         scheduler: S,
+        id: super::Id,
     ) -> (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);
+        let (task, notified, join) = super::new_task(task, scheduler, id);
 
         unsafe {
             // safety: We just created the task, so we have exclusive access
@@ -187,13 +188,14 @@ impl<S: 'static> LocalOwnedTasks<S> {
         &self,
         task: T,
         scheduler: S,
+        id: super::Id,
     ) -> (JoinHandle<T::Output>, Option<Notified<S>>)
     where
         S: Schedule,
         T: Future + 'static,
         T::Output: 'static,
     {
-        let (task, notified, join) = super::new_task(task, scheduler);
+        let (task, notified, join) = super::new_task(task, scheduler, id);
 
         unsafe {
             // safety: We just created the task, so we have exclusive access

tokio/src/runtime/task/mod.rs

@@ -271,14 +271,14 @@ cfg_rt! {
     /// notification.
     fn new_task<T, S>(
         task: T,
-        scheduler: S
+        scheduler: S,
+        id: Id,
     ) -> (Task<S>, Notified<S>, JoinHandle<T::Output>)
     where
         S: Schedule,
         T: Future + 'static,
         T::Output: 'static,
     {
-        let id = Id::next();
         let raw = RawTask::new::<T, S>(task, scheduler, id.clone());
         let task = Task {
             raw,
@@ -297,13 +297,13 @@ 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) -> (UnownedTask<S>, JoinHandle<T::Output>)
+    pub(crate) fn unowned<T, S>(task: T, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>)
     where
         S: Schedule,
         T: Send + Future + 'static,
         T::Output: Send + 'static,
     {
-        let (task, notified, join) = new_task(task, scheduler);
+        let (task, notified, join) = new_task(task, scheduler, id);
 
         // This transfers the ref-count of task and notified into an UnownedTask.
         // This is valid because an UnownedTask holds two ref-counts.
@@ -480,9 +480,13 @@ impl fmt::Display for Id {
 }
 
 impl Id {
-    fn next() -> Self {
+    pub(crate) fn next() -> Self {
         use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
         static NEXT_ID: AtomicUsize = AtomicUsize::new(1);
         Self(NEXT_ID.fetch_add(1, Relaxed))
     }
+
+    pub(crate) fn as_usize(&self) -> usize {
+        self.0
+    }
 }

tokio/src/runtime/thread_pool/mod.rs

@@ -9,7 +9,7 @@ pub(crate) use worker::Launch;
 pub(crate) use worker::block_in_place;
 
 use crate::loom::sync::Arc;
-use crate::runtime::task::JoinHandle;
+use crate::runtime::task::{self, JoinHandle};
 use crate::runtime::{Callback, Parker};
 
 use std::fmt;
@@ -90,12 +90,12 @@ impl Drop for ThreadPool {
 
 impl Spawner {
     /// Spawns a future onto the thread pool
-    pub(crate) fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
+    pub(crate) fn spawn<F>(&self, future: F, id: task::Id) -> JoinHandle<F::Output>
     where
         F: crate::future::Future + Send + 'static,
         F::Output: Send + 'static,
     {
-        worker::Shared::bind_new_task(&self.shared, future)
+        worker::Shared::bind_new_task(&self.shared, future, id)
     }
 
     pub(crate) fn shutdown(&mut self) {

tokio/src/runtime/thread_pool/worker.rs

@@ -715,12 +715,16 @@ impl task::Schedule for Arc<Shared> {
 }
 
 impl Shared {
-    pub(super) fn bind_new_task<T>(me: &Arc<Self>, future: T) -> JoinHandle<T::Output>
+    pub(super) fn bind_new_task<T>(
+        me: &Arc<Self>,
+        future: T,
+        id: crate::runtime::task::Id,
+    ) -> JoinHandle<T::Output>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        let (handle, notified) = me.owned.bind(future, me.clone());
+        let (handle, notified) = me.owned.bind(future, me.clone(), id);
 
         if let Some(notified) = notified {
             me.schedule(notified, false);

tokio/src/task/local.rs

@@ -301,12 +301,13 @@ cfg_rt! {
     where F: Future + 'static,
           F::Output: 'static
     {
-        let future = crate::util::trace::task(future, "local", name);
+        let id = crate::runtime::task::Id::next();
+        let future = crate::util::trace::task(future, "local", name, id.as_usize());
         CURRENT.with(|maybe_cx| {
             let cx = maybe_cx
                 .expect("`spawn_local` called from outside of a `task::LocalSet`");
 
-            let (handle, notified) = cx.owned.bind(future, cx.shared.clone());
+            let (handle, notified) = cx.owned.bind(future, cx.shared.clone(), id);
 
             if let Some(notified) = notified {
                 cx.shared.schedule(notified);
@@ -385,9 +386,13 @@ impl LocalSet {
         F: Future + 'static,
         F::Output: 'static,
     {
-        let future = crate::util::trace::task(future, "local", None);
+        let id = crate::runtime::task::Id::next();
+        let future = crate::util::trace::task(future, "local", None, id.as_usize());
 
-        let (handle, notified) = self.context.owned.bind(future, self.context.shared.clone());
+        let (handle, notified) = self
+            .context
+            .owned
+            .bind(future, self.context.shared.clone(), id);
 
         if let Some(notified) = notified {
             self.context.shared.schedule(notified);

tokio/src/task/spawn.rs

@@ -142,8 +142,10 @@ cfg_rt! {
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        let spawn_handle = crate::runtime::context::spawn_handle().expect(CONTEXT_MISSING_ERROR);
-        let task = crate::util::trace::task(future, "task", name);
-        spawn_handle.spawn(task)
+        use crate::runtime::{task, context};
+        let id = task::Id::next();
+        let spawn_handle = context::spawn_handle().expect(CONTEXT_MISSING_ERROR);
+        let task = crate::util::trace::task(future, "task", name, id.as_usize());
+        spawn_handle.spawn(task, id)
     }
 }

tokio/src/util/trace.rs

@@ -10,14 +10,15 @@ cfg_trace! {
 
         #[inline]
         #[track_caller]
-        pub(crate) fn task<F>(task: F, kind: &'static str, name: Option<&str>) -> Instrumented<F> {
+        pub(crate) fn task<F>(task: F, kind: &'static str, name: Option<&str>, id: usize) -> Instrumented<F> {
             use tracing::instrument::Instrument;
             let location = std::panic::Location::caller();
             let span = tracing::trace_span!(
                 target: "tokio::task",
                 "runtime.spawn",
                 %kind,
                 task.name = %name.unwrap_or_default(),
+                task.id = id,
                 loc.file = location.file(),
                 loc.line = location.line(),
                 loc.col = location.column(),
@@ -91,7 +92,7 @@ cfg_time! {
 cfg_not_trace! {
     cfg_rt! {
         #[inline]
-        pub(crate) fn task<F>(task: F, _: &'static str, _name: Option<&str>) -> F {
+        pub(crate) fn task<F>(task: F, _: &'static str, _name: Option<&str>, _: usize) -> F {
             // nop
             task
         }