Move tracing setup to the `tracing` module

3 files changed, 304 insertions, 298 deletions

diff --git a/src/bootstrap/src/bin/main.rs b/src/bootstrap/src/bin/main.rs
index aa062e14202..0dc3cad0cd8 100644
--- a/src/bootstrap/src/bin/main.rs
+++ b/src/bootstrap/src/bin/main.rs
@@ -26,7 +26,7 @@ fn is_tracing_enabled() -> bool {
 
 fn main() {
     #[cfg(feature = "tracing")]
-    let guard = setup_tracing(is_profiling_enabled());
+    let guard = bootstrap::setup_tracing(is_profiling_enabled());
 
     let start_time = Instant::now();
 
@@ -247,300 +247,3 @@ fn check_version(config: &Config) -> Option<String> {
 
     Some(msg)
 }
-
-// # Note on `tracing` usage in bootstrap
-//
-// Due to the conditional compilation via the `tracing` cargo feature, this means that `tracing`
-// usages in bootstrap need to be also gated behind the `tracing` feature:
-//
-// - `tracing` macros with log levels (`trace!`, `debug!`, `warn!`, `info`, `error`) should not be
-//   used *directly*. You should use the wrapped `tracing` macros which gate the actual invocations
-//   behind `feature = "tracing"`.
-// - `tracing`'s `#[instrument(..)]` macro will need to be gated like `#![cfg_attr(feature =
-//   "tracing", instrument(..))]`.
-#[cfg(feature = "tracing")]
-fn setup_tracing(profiling_enabled: bool) -> Option<TracingGuard> {
-    use std::fmt::Debug;
-    use std::fs::File;
-    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};
-    use tracing_subscriber::layer::{Context, SubscriberExt};
-    use tracing_subscriber::registry::{LookupSpan, SpanRef};
-    use tracing_subscriber::{EnvFilter, Layer};
-
-    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) {
-            let formatted = format!("{value:?}");
-            match field.name() {
-                "message" => {
-                    self.message = Some(formatted);
-                }
-                name => {
-                    self.fields.push((name, formatted));
-                }
-            }
-        }
-    }
-
-    #[derive(Copy, Clone)]
-    enum SpanAction {
-        Enter,
-    }
-
-    /// Holds the name of a step, stored in `tracing_subscriber`'s extensions.
-    struct StepNameExtension(String);
-
-    #[derive(Default)]
-    struct TracingPrinter {
-        indent: AtomicU32,
-        span_values: std::sync::Mutex<std::collections::HashMap<tracing::Id, FieldValues>>,
-    }
-
-    impl TracingPrinter {
-        fn format_header<W: Write>(
-            &self,
-            writer: &mut W,
-            time: DateTime<Utc>,
-            level: &Level,
-        ) -> std::io::Result<()> {
-            // Use a fixed-width timestamp without date, that shouldn't be very important
-            let timestamp = time.format("%H:%M:%S.%3f");
-            write!(writer, "{timestamp} ")?;
-            // Make sure that levels are aligned to the same number of characters, in order not to
-            // break the layout
-            write!(writer, "{level:>5} ")?;
-            write!(writer, "{}", " ".repeat(self.indent.load(Ordering::Relaxed) as usize))
-        }
-
-        fn write_event<W: Write>(&self, writer: &mut W, event: &Event<'_>) -> std::io::Result<()> {
-            let now = Utc::now();
-
-            self.format_header(writer, now, event.metadata().level())?;
-
-            let mut field_values = FieldValues::default();
-            event.record(&mut field_values);
-
-            if let Some(msg) = &field_values.message {
-                write!(writer, "{msg}")?;
-            }
-
-            if !field_values.fields.is_empty() {
-                if field_values.message.is_some() {
-                    write!(writer, " ")?;
-                }
-                write!(writer, "[")?;
-                for (index, (name, value)) in field_values.fields.iter().enumerate() {
-                    write!(writer, "{name} = {value}")?;
-                    if index < field_values.fields.len() - 1 {
-                        write!(writer, ", ")?;
-                    }
-                }
-                write!(writer, "]")?;
-            }
-            write_location(writer, event.metadata())?;
-            writeln!(writer)?;
-            Ok(())
-        }
-
-        fn write_span<W: Write, S>(
-            &self,
-            writer: &mut W,
-            span: SpanRef<'_, S>,
-            field_values: Option<&FieldValues>,
-            action: SpanAction,
-        ) -> std::io::Result<()>
-        where
-            S: for<'lookup> LookupSpan<'lookup>,
-        {
-            let now = Utc::now();
-
-            self.format_header(writer, now, span.metadata().level())?;
-            match action {
-                SpanAction::Enter => {
-                    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_location(writer, span.metadata())?;
-            writeln!(writer)?;
-            Ok(())
-        }
-    }
-
-    fn write_location<W: Write>(
-        writer: &mut W,
-        metadata: &'static tracing::Metadata<'static>,
-    ) -> std::io::Result<()> {
-        use std::path::{Path, PathBuf};
-
-        if let Some(filename) = metadata.file() {
-            // Keep only the module name and file name to make it shorter
-            let filename: PathBuf = Path::new(filename)
-                .components()
-                // Take last two path components
-                .rev()
-                .take(2)
-                .collect::<Vec<_>>()
-                .into_iter()
-                .rev()
-                .collect();
-
-            write!(writer, " ({}", filename.display())?;
-            if let Some(line) = metadata.line() {
-                write!(writer, ":{line}")?;
-            }
-            write!(writer, ")")?;
-        }
-        Ok(())
-    }
-
-    impl<S> Layer<S> for TracingPrinter
-    where
-        S: Subscriber,
-        S: for<'lookup> LookupSpan<'lookup>,
-    {
-        fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
-            let mut writer = std::io::stderr().lock();
-            self.write_event(&mut writer, event).unwrap();
-        }
-
-        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);
-        }
-
-        fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
-            if let Some(span) = ctx.span(id) {
-                let mut writer = std::io::stderr().lock();
-                let values = self.span_values.lock().unwrap();
-                let values = values.get(id);
-                self.write_span(&mut writer, span, values, SpanAction::Enter).unwrap();
-            }
-            self.indent.fetch_add(1, Ordering::Relaxed);
-        }
-
-        fn on_exit(&self, _id: &Id, _ctx: Context<'_, S>) {
-            self.indent.fetch_sub(1, Ordering::Relaxed);
-        }
-    }
-
-    let registry = tracing_subscriber::registry().with(filter).with(TracingPrinter::default());
-
-    let guard = if profiling_enabled {
-        // When we're creating this layer, we do not yet know the location of the tracing output
-        // directory, because it is stored in the output directory determined after Config is parsed,
-        // but we already want to make tracing calls during (and before) config parsing.
-        // So we store the output into a temporary file, and then move it to the tracing directory
-        // before bootstrap ends.
-        let tempdir = tempfile::TempDir::new().expect("Cannot create temporary directory");
-        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)
-            .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();
-        Some(TracingGuard { guard, _tempdir: tempdir, chrome_tracing_path })
-    } else {
-        tracing::subscriber::set_global_default(registry).unwrap();
-        None
-    };
-
-    guard
-}
-
-#[cfg(feature = "tracing")]
-struct TracingGuard {
-    guard: tracing_chrome::FlushGuard,
-    _tempdir: tempfile::TempDir,
-    chrome_tracing_path: std::path::PathBuf,
-}
-
-#[cfg(feature = "tracing")]
-impl TracingGuard {
-    fn copy_to_dir(self, dir: &std::path::Path) {
-        drop(self.guard);
-        std::fs::rename(&self.chrome_tracing_path, dir.join("chrome-trace.json")).unwrap();
-    }
-}
diff --git a/src/bootstrap/src/lib.rs b/src/bootstrap/src/lib.rs
index d980c91d8eb..ba3d3fb46e2 100644
--- a/src/bootstrap/src/lib.rs
+++ b/src/bootstrap/src/lib.rs
@@ -54,6 +54,8 @@ pub use utils::change_tracker::{
     CONFIG_CHANGE_HISTORY, find_recent_config_change_ids, human_readable_changes,
 };
 pub use utils::helpers::{PanicTracker, symlink_dir};
+#[cfg(feature = "tracing")]
+pub use utils::tracing::setup_tracing;
 
 use crate::core::build_steps::vendor::VENDOR_DIR;
 
diff --git a/src/bootstrap/src/utils/tracing.rs b/src/bootstrap/src/utils/tracing.rs
index 4e596e35060..68063131d15 100644
--- a/src/bootstrap/src/utils/tracing.rs
+++ b/src/bootstrap/src/utils/tracing.rs
@@ -62,3 +62,304 @@ macro_rules! trace_cmd {
         }
     };
 }
+
+// # Note on `tracing` usage in bootstrap
+//
+// Due to the conditional compilation via the `tracing` cargo feature, this means that `tracing`
+// usages in bootstrap need to be also gated behind the `tracing` feature:
+//
+// - `tracing` macros with log levels (`trace!`, `debug!`, `warn!`, `info`, `error`) should not be
+//   used *directly*. You should use the wrapped `tracing` macros which gate the actual invocations
+//   behind `feature = "tracing"`.
+// - `tracing`'s `#[instrument(..)]` macro will need to be gated like `#![cfg_attr(feature =
+//   "tracing", instrument(..))]`.
+#[cfg(feature = "tracing")]
+mod inner {
+    use std::fmt::Debug;
+    use std::fs::File;
+    use std::io::Write;
+    use std::sync::atomic::Ordering;
+
+    use chrono::{DateTime, Utc};
+    use tracing::field::{Field, Visit};
+    use tracing::{Event, Id, Level, Subscriber};
+    use tracing_subscriber::layer::{Context, SubscriberExt};
+    use tracing_subscriber::registry::{LookupSpan, SpanRef};
+    use tracing_subscriber::{EnvFilter, Layer};
+
+    use crate::STEP_NAME_TARGET;
+
+    pub fn setup_tracing(profiling_enabled: bool) -> Option<TracingGuard> {
+        let filter = EnvFilter::from_env("BOOTSTRAP_TRACING");
+
+        let registry = tracing_subscriber::registry().with(filter).with(TracingPrinter::default());
+
+        let guard = if profiling_enabled {
+            // When we're creating this layer, we do not yet know the location of the tracing output
+            // directory, because it is stored in the output directory determined after Config is parsed,
+            // but we already want to make tracing calls during (and before) config parsing.
+            // So we store the output into a temporary file, and then move it to the tracing directory
+            // before bootstrap ends.
+            let tempdir = tempfile::TempDir::new().expect("Cannot create temporary directory");
+            let chrome_tracing_path = tempdir.path().join("bootstrap-trace.json");
+            let file = std::io::BufWriter::new(File::create(&chrome_tracing_path).unwrap());
+
+            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();
+            Some(TracingGuard { guard, _tempdir: tempdir, chrome_tracing_path })
+        } else {
+            tracing::subscriber::set_global_default(registry).unwrap();
+            None
+        };
+
+        guard
+    }
+
+    pub struct TracingGuard {
+        guard: tracing_chrome::FlushGuard,
+        _tempdir: tempfile::TempDir,
+        chrome_tracing_path: std::path::PathBuf,
+    }
+
+    impl TracingGuard {
+        pub fn copy_to_dir(self, dir: &std::path::Path) {
+            drop(self.guard);
+            std::fs::rename(&self.chrome_tracing_path, dir.join("chrome-trace.json")).unwrap();
+        }
+    }
+
+    /// 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) {
+            let formatted = format!("{value:?}");
+            match field.name() {
+                "message" => {
+                    self.message = Some(formatted);
+                }
+                name => {
+                    self.fields.push((name, formatted));
+                }
+            }
+        }
+    }
+
+    #[derive(Copy, Clone)]
+    enum SpanAction {
+        Enter,
+    }
+
+    /// Holds the name of a step, stored in `tracing_subscriber`'s extensions.
+    struct StepNameExtension(String);
+
+    #[derive(Default)]
+    struct TracingPrinter {
+        indent: std::sync::atomic::AtomicU32,
+        span_values: std::sync::Mutex<std::collections::HashMap<tracing::Id, FieldValues>>,
+    }
+
+    impl TracingPrinter {
+        fn format_header<W: Write>(
+            &self,
+            writer: &mut W,
+            time: DateTime<Utc>,
+            level: &Level,
+        ) -> std::io::Result<()> {
+            // Use a fixed-width timestamp without date, that shouldn't be very important
+            let timestamp = time.format("%H:%M:%S.%3f");
+            write!(writer, "{timestamp} ")?;
+            // Make sure that levels are aligned to the same number of characters, in order not to
+            // break the layout
+            write!(writer, "{level:>5} ")?;
+            write!(writer, "{}", " ".repeat(self.indent.load(Ordering::Relaxed) as usize))
+        }
+
+        fn write_event<W: Write>(&self, writer: &mut W, event: &Event<'_>) -> std::io::Result<()> {
+            let now = Utc::now();
+
+            self.format_header(writer, now, event.metadata().level())?;
+
+            let mut field_values = FieldValues::default();
+            event.record(&mut field_values);
+
+            if let Some(msg) = &field_values.message {
+                write!(writer, "{msg}")?;
+            }
+
+            if !field_values.fields.is_empty() {
+                if field_values.message.is_some() {
+                    write!(writer, " ")?;
+                }
+                write!(writer, "[")?;
+                for (index, (name, value)) in field_values.fields.iter().enumerate() {
+                    write!(writer, "{name} = {value}")?;
+                    if index < field_values.fields.len() - 1 {
+                        write!(writer, ", ")?;
+                    }
+                }
+                write!(writer, "]")?;
+            }
+            write_location(writer, event.metadata())?;
+            writeln!(writer)?;
+            Ok(())
+        }
+
+        fn write_span<W: Write, S>(
+            &self,
+            writer: &mut W,
+            span: SpanRef<'_, S>,
+            field_values: Option<&FieldValues>,
+            action: SpanAction,
+        ) -> std::io::Result<()>
+        where
+            S: for<'lookup> LookupSpan<'lookup>,
+        {
+            let now = Utc::now();
+
+            self.format_header(writer, now, span.metadata().level())?;
+            match action {
+                SpanAction::Enter => {
+                    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_location(writer, span.metadata())?;
+            writeln!(writer)?;
+            Ok(())
+        }
+    }
+
+    fn write_location<W: Write>(
+        writer: &mut W,
+        metadata: &'static tracing::Metadata<'static>,
+    ) -> std::io::Result<()> {
+        use std::path::{Path, PathBuf};
+
+        if let Some(filename) = metadata.file() {
+            // Keep only the module name and file name to make it shorter
+            let filename: PathBuf = Path::new(filename)
+                .components()
+                // Take last two path components
+                .rev()
+                .take(2)
+                .collect::<Vec<_>>()
+                .into_iter()
+                .rev()
+                .collect();
+
+            write!(writer, " ({}", filename.display())?;
+            if let Some(line) = metadata.line() {
+                write!(writer, ":{line}")?;
+            }
+            write!(writer, ")")?;
+        }
+        Ok(())
+    }
+
+    impl<S> Layer<S> for TracingPrinter
+    where
+        S: Subscriber,
+        S: for<'lookup> LookupSpan<'lookup>,
+    {
+        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);
+        }
+
+        fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
+            let mut writer = std::io::stderr().lock();
+            self.write_event(&mut writer, event).unwrap();
+        }
+
+        fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
+            if let Some(span) = ctx.span(id) {
+                let mut writer = std::io::stderr().lock();
+                let values = self.span_values.lock().unwrap();
+                let values = values.get(id);
+                self.write_span(&mut writer, span, values, SpanAction::Enter).unwrap();
+            }
+            self.indent.fetch_add(1, Ordering::Relaxed);
+        }
+
+        fn on_exit(&self, _id: &Id, _ctx: Context<'_, S>) {
+            self.indent.fetch_sub(1, Ordering::Relaxed);
+        }
+    }
+}
+
+#[cfg(feature = "tracing")]
+pub use inner::setup_tracing;