Add (feature flagged) tracing spans to spawns

The instrumentation conforms to the standardized naming requirements for `tokio-console` (see tokio-rs/console#41), so they will show up in the console. The `tracing` dependency is an opt-in feature flag. Signed-off-by: Eliza Weisman <[email protected]>
hawkw · Feb 14, 2022 · 1e8da3e · 1e8da3e
1 parent f45eee8
commit 1e8da3e
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 2 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -18,6 +18,9 @@ exclude = ["/ci/*", "/scripts/*", "/.github/*", "/bors.toml"]
 members = ["rayon-demo", "rayon-core"]
 exclude = ["ci"]
 
+[features]
+tracing = ["rayon-core/tracing"]
+
 [dependencies]
 rayon-core = { version = "1.9.1", path = "rayon-core" }
 crossbeam-deque = "0.8.1"

diff --git a/rayon-core/Cargo.toml b/rayon-core/Cargo.toml
@@ -20,6 +20,7 @@ num_cpus = "1.2"
 crossbeam-channel = "0.5.0"
 crossbeam-deque = "0.8.1"
 crossbeam-utils = "0.8.0"
+tracing = { version = "0.1.26", optional = true }
 
 [dev-dependencies]
 rand = "0.8"

diff --git a/rayon-core/src/scope/mod.rs b/rayon-core/src/scope/mod.rs
@@ -534,14 +534,30 @@ impl<'scope> Scope<'scope> {
     /// task spawning.
     ///
     /// [`scope` function]: fn.scope.html
+    #[cfg_attr(feature = "tracing", track_caller)]
     pub fn spawn<BODY>(&self, body: BODY)
     where
         BODY: FnOnce(&Scope<'scope>) + Send + 'scope,
     {
+        // Generate a tracing span for the function.
+        #[cfg(feature = "tracing")]
+        let span = {
+            let location = std::panic::Location::caller();
+            tracing::trace_span!(
+                "runtime.spawn",
+                kind = %"scoped",
+                "fn" = %std::any::type_name::<BODY>(),
+                spawn.location = %format_args!("{}:{}:{}", location.file(), location.line(), location.column()),
+            )
+        };
         self.base.increment();
         unsafe {
             let job_ref = Box::new(HeapJob::new(move || {
-                self.base.execute_job(move || body(self))
+                self.base.execute_job(move || {
+                    #[cfg(feature = "tracing")]
+                    let _s = span.entered();
+                    body(self)
+                })
             }))
             .as_job_ref();
 
@@ -579,10 +595,25 @@ impl<'scope> ScopeFifo<'scope> {
     where
         BODY: FnOnce(&ScopeFifo<'scope>) + Send + 'scope,
     {
+        // Generate a tracing span for the function.
+        #[cfg(feature = "tracing")]
+        let span = {
+            let location = std::panic::Location::caller();
+            tracing::trace_span!(
+                "runtime.spawn",
+                kind = %"scoped-fifo",
+                "fn" = %std::any::type_name::<BODY>(),
+                spawn.location = %format_args!("{}:{}:{}", location.file(), location.line(), location.column()),
+            )
+        };
         self.base.increment();
         unsafe {
             let job_ref = Box::new(HeapJob::new(move || {
-                self.base.execute_job(move || body(self))
+                self.base.execute_job(move || {
+                    #[cfg(feature = "tracing")]
+                    let _s = span.entered();
+                    body(self)
+                })
             }))
             .as_job_ref();
 

diff --git a/rayon-core/src/spawn/mod.rs b/rayon-core/src/spawn/mod.rs
@@ -57,6 +57,7 @@ use std::sync::Arc;
 ///     GLOBAL_COUNTER.fetch_add(1, Ordering::SeqCst);
 /// });
 /// ```
+#[cfg_attr(feature = "tracing", track_caller)]
 pub fn spawn<F>(func: F)
 where
     F: FnOnce() + Send + 'static,
@@ -68,10 +69,27 @@ where
 /// Spawns an asynchronous job in `registry.`
 ///
 /// Unsafe because `registry` must not yet have terminated.
+#[cfg_attr(feature = "tracing", track_caller)]
 pub(super) unsafe fn spawn_in<F>(func: F, registry: &Arc<Registry>)
 where
     F: FnOnce() + Send + 'static,
 {
+    // Generate a tracing span for the function.
+    #[cfg(feature = "tracing")]
+    let func = {
+        let location = std::panic::Location::caller();
+        let span = tracing::trace_span!(
+            "runtime.spawn",
+            kind = %"spawn",
+            "fn" = %std::any::type_name::<F>(),
+            spawn.location = %format_args!("{}:{}:{}", location.file(), location.line(), location.column()),
+        );
+        move || {
+            let _s = span.entered();
+            (func)()
+        }
+    };
+
     // We assert that this does not hold any references (we know
     // this because of the `'static` bound in the inferface);
     // moreover, we assert that the code below is not supposed to
@@ -132,6 +150,7 @@ where
 /// details.
 ///
 /// [ph]: struct.ThreadPoolBuilder.html#method.panic_handler
+#[cfg_attr(feature = "tracing", track_caller)]
 pub fn spawn_fifo<F>(func: F)
 where
     F: FnOnce() + Send + 'static,
@@ -143,10 +162,26 @@ where
 /// Spawns an asynchronous FIFO job in `registry.`
 ///
 /// Unsafe because `registry` must not yet have terminated.
+#[cfg_attr(feature = "tracing", track_caller)]
 pub(super) unsafe fn spawn_fifo_in<F>(func: F, registry: &Arc<Registry>)
 where
     F: FnOnce() + Send + 'static,
 {
+    // Generate a tracing span for the function.
+    #[cfg(feature = "tracing")]
+    let func = {
+        let location = std::panic::Location::caller();
+        let span = tracing::trace_span!(
+            "runtime.spawn",
+            kind = %"fifo",
+            "fn" = %std::any::type_name::<F>(),
+            spawn.location = %format_args!("{}:{}:{}", location.file(), location.line(), location.column()),
+        );
+        move || {
+            let _s = span.entered();
+            (func)()
+        }
+    };
     // We assert that this does not hold any references (we know
     // this because of the `'static` bound in the inferface);
     // moreover, we assert that the code below is not supposed to