Create a span for each executed step

3 files changed, 108 insertions, 28 deletions

diff --git a/src/bootstrap/src/bin/main.rs b/src/bootstrap/src/bin/main.rs
index a98fc447e89..bdf85a50afd 100644
--- a/src/bootstrap/src/bin/main.rs
+++ b/src/bootstrap/src/bin/main.rs
@@ -268,6 +268,7 @@ fn setup_tracing(profiling_enabled: bool) -> Option<TracingGuard> {
     use std::io::BufWriter;
     use std::sync::atomic::{AtomicU32, Ordering};
 
+    use bootstrap::STEP_NAME_TARGET;
     use chrono::{DateTime, Utc};
     use tracing::field::{Field, Visit};
     use tracing::{Event, Id, Level, Subscriber};
@@ -277,18 +278,40 @@ fn setup_tracing(profiling_enabled: bool) -> Option<TracingGuard> {
 
     let filter = EnvFilter::from_env("BOOTSTRAP_TRACING");
 
+    /// Visitor that extracts both known and unknown field values from events and spans.
     #[derive(Default)]
     struct FieldValues {
+        /// Main event message
         message: Option<String>,
+        /// Name of a recorded psna
+        step_name: Option<String>,
+        /// The rest of arbitrary event/span fields
         fields: Vec<(&'static str, String)>,
     }
 
     impl Visit for FieldValues {
+        /// Record fields if possible using `record_str`, to avoid rendering simple strings with
+        /// their `Debug` representation, which adds extra quotes.
+        fn record_str(&mut self, field: &Field, value: &str) {
+            match field.name() {
+                "step_name" => {
+                    self.step_name = Some(value.to_string());
+                }
+                name => {
+                    self.fields.push((name, value.to_string()));
+                }
+            }
+        }
+
         fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
-            if field.name() == "message" {
-                self.message = Some(format!("{value:?}"));
-            } else {
-                self.fields.push((field.name(), format!("{value:?}")));
+            let formatted = format!("{value:?}");
+            match field.name() {
+                "message" => {
+                    self.message = Some(formatted);
+                }
+                name => {
+                    self.fields.push((name, formatted));
+                }
             }
         }
     }
@@ -298,6 +321,9 @@ fn setup_tracing(profiling_enabled: bool) -> Option<TracingGuard> {
         Enter,
     }
 
+    /// Holds the name of a step, stored in `tracing_subscriber`'s extensions.
+    struct StepNameExtension(String);
+
     #[derive(Default)]
     struct TracingPrinter {
         indent: AtomicU32,
@@ -369,17 +395,31 @@ fn setup_tracing(profiling_enabled: bool) -> Option<TracingGuard> {
                 }
             }
 
-            write!(writer, "{}", span.name())?;
-            if let Some(values) = field_values.filter(|v| !v.fields.is_empty()) {
-                write!(writer, " [")?;
-                for (index, (name, value)) in values.fields.iter().enumerate() {
-                    write!(writer, "{name} = {value}")?;
-                    if index < values.fields.len() - 1 {
-                        write!(writer, ", ")?;
+            // We handle steps specially. We instrument them dynamically in `Builder::ensure`,
+            // and we want to have custom name for each step span. But tracing doesn't allow setting
+            // dynamic span names. So we detect step spans here and override their name.
+            if span.metadata().target() == STEP_NAME_TARGET {
+                let name = field_values.and_then(|v| v.step_name.as_deref()).unwrap_or(span.name());
+                write!(writer, "{name}")?;
+
+                // There should be only one more field called `args`
+                if let Some(values) = field_values {
+                    let field = &values.fields[0];
+                    write!(writer, " {{{}}}", field.1)?;
+                }
+            } else {
+                write!(writer, "{}", span.name())?;
+                if let Some(values) = field_values.filter(|v| !v.fields.is_empty()) {
+                    write!(writer, " [")?;
+                    for (index, (name, value)) in values.fields.iter().enumerate() {
+                        write!(writer, "{name} = {value}")?;
+                        if index < values.fields.len() - 1 {
+                            write!(writer, ", ")?;
+                        }
                     }
+                    write!(writer, "]")?;
                 }
-                write!(writer, "]")?;
-            }
+            };
 
             write_location(writer, span.metadata())?;
             writeln!(writer)?;
@@ -424,18 +464,18 @@ fn setup_tracing(profiling_enabled: bool) -> Option<TracingGuard> {
             self.write_event(&mut writer, event).unwrap();
         }
 
-        fn on_new_span(
-            &self,
-            attrs: &tracing::span::Attributes<'_>,
-            id: &Id,
-            _ctx: Context<'_, S>,
-        ) {
+        fn on_new_span(&self, attrs: &tracing::span::Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
             // Record value of span fields
             // Note that we do not implement changing values of span fields after they are created.
             // For that we would also need to implement the `on_record` method
-
             let mut field_values = FieldValues::default();
             attrs.record(&mut field_values);
+
+            // We need to propagate the actual name of the span to the Chrome layer below, because
+            // it cannot access field values. We do that through extensions.
+            if let Some(step_name) = field_values.step_name.clone() {
+                ctx.span(id).unwrap().extensions_mut().insert(StepNameExtension(step_name));
+            }
             self.span_values.lock().unwrap().insert(id.clone(), field_values);
         }
 
@@ -466,8 +506,21 @@ fn setup_tracing(profiling_enabled: bool) -> Option<TracingGuard> {
         let chrome_tracing_path = tempdir.path().join("bootstrap-trace.json");
         let file = BufWriter::new(File::create(&chrome_tracing_path).unwrap());
 
-        let chrome_layer =
-            tracing_chrome::ChromeLayerBuilder::new().writer(file).include_args(true);
+        let chrome_layer = tracing_chrome::ChromeLayerBuilder::new()
+            .writer(file)
+            .include_args(true)
+            .name_fn(Box::new(|event_or_span| match event_or_span {
+                tracing_chrome::EventOrSpan::Event(e) => e.metadata().name().to_string(),
+                tracing_chrome::EventOrSpan::Span(s) => {
+                    if s.metadata().target() == STEP_NAME_TARGET
+                        && let Some(extension) = s.extensions().get::<StepNameExtension>()
+                    {
+                        extension.0.clone()
+                    } else {
+                        s.metadata().name().to_string()
+                    }
+                }
+            }));
         let (chrome_layer, guard) = chrome_layer.build();
 
         tracing::subscriber::set_global_default(registry.with(chrome_layer)).unwrap();
diff --git a/src/bootstrap/src/core/builder/mod.rs b/src/bootstrap/src/core/builder/mod.rs
index fc454ebe819..0eae2f58a68 100644
--- a/src/bootstrap/src/core/builder/mod.rs
+++ b/src/bootstrap/src/core/builder/mod.rs
@@ -947,6 +947,9 @@ impl Step for Libdir {
     }
 }
 
+#[cfg(feature = "tracing")]
+pub const STEP_NAME_TARGET: &str = "STEP";
+
 impl<'a> Builder<'a> {
     fn get_step_descriptions(kind: Kind) -> Vec<StepDescription> {
         macro_rules! describe {
@@ -1709,6 +1712,20 @@ You have to build a stage1 compiler for `{}` first, and then use it to build a s
             let zero = Duration::new(0, 0);
             let parent = self.time_spent_on_dependencies.replace(zero);
 
+            #[cfg(feature = "tracing")]
+            let _span = {
+                // Keep the target and field names synchronized with `setup_tracing`.
+                let span = tracing::info_span!(
+                    target: STEP_NAME_TARGET,
+                    // We cannot use a dynamic name here, so instead we record the actual step name
+                    // in the step_name field.
+                    "step",
+                    step_name = step_name::<S>(),
+                    args = step_debug_args(&step)
+                );
+                span.entered()
+            };
+
             let out = step.clone().run(self);
             let dur = start.elapsed();
             let deps = self.time_spent_on_dependencies.replace(parent + dur);
@@ -1805,15 +1822,23 @@ You have to build a stage1 compiler for `{}` first, and then use it to build a s
     }
 }
 
-fn pretty_print_step<S: Step>(step: &S) -> String {
-    let step_dbg_repr = format!("{step:?}");
-    let brace_index = step_dbg_repr.find('{').unwrap_or(0);
-
+/// Return qualified step name, e.g. `compile::Rustc`.
+fn step_name<S: Step>() -> String {
     // Normalize step type path to only keep the module and the type name
     let path = type_name::<S>().rsplit("::").take(2).collect::<Vec<_>>();
-    let type_string = path.into_iter().rev().collect::<Vec<_>>().join("::");
+    path.into_iter().rev().collect::<Vec<_>>().join("::")
+}
+
+/// Renders `step` using its `Debug` implementation and extract the field arguments out of it.
+fn step_debug_args<S: Step>(step: &S) -> String {
+    let step_dbg_repr = format!("{step:?}");
+    let brace_start = step_dbg_repr.find('{').unwrap_or(0);
+    let brace_end = step_dbg_repr.rfind('}').unwrap_or(step_dbg_repr.len());
+    step_dbg_repr[brace_start + 1..brace_end - 1].trim().to_string()
+}
 
-    format!("{type_string} {}", &step_dbg_repr[brace_index..])
+fn pretty_print_step<S: Step>(step: &S) -> String {
+    format!("{} {{ {} }}", step_name::<S>(), step_debug_args(step))
 }
 
 impl<'a> AsRef<ExecutionContext> for Builder<'a> {
diff --git a/src/bootstrap/src/lib.rs b/src/bootstrap/src/lib.rs
index 7d6c0658f47..d980c91d8eb 100644
--- a/src/bootstrap/src/lib.rs
+++ b/src/bootstrap/src/lib.rs
@@ -43,6 +43,8 @@ mod core;
 mod utils;
 
 pub use core::builder::PathSet;
+#[cfg(feature = "tracing")]
+pub use core::builder::STEP_NAME_TARGET;
 pub use core::config::flags::{Flags, Subcommand};
 pub use core::config::{ChangeId, Config};