Merge pull request #4406 from Stypox/tracing

Add tracing_chrome under "tracing" feature

Author: RalfJung

src/tools/miri/CONTRIBUTING.md

@@ -158,6 +158,15 @@ compiler that has `debug=true` set in `bootstrap.toml`.
 You can set `MIRI_BACKTRACE=1` to get a backtrace of where an
 evaluation error was originally raised.
 
+#### Tracing
+
+You can generate a Chrome trace file from a Miri execution by passing `--features=tracing` during the
+build and then setting `MIRI_TRACING=1` when running Miri. This will generate a `.json` file that
+you can visualize in [Perfetto](https://ui.perfetto.dev/). For example:
+
+```sh
+MIRI_TRACING=1 ./miri run --features=tracing tests/pass/hello.rs
+```
 
 ### UI testing
 

src/tools/miri/Cargo.lock

@@ -627,6 +627,7 @@ dependencies = [
  "regex",
  "rustc_version",
  "serde",
+ "serde_json",
  "smallvec",
  "tempfile",
  "tikv-jemalloc-sys",

src/tools/miri/Cargo.toml

@@ -27,6 +27,7 @@ chrono = { version = "0.4.38", default-features = false }
 chrono-tz = "0.10"
 directories = "6"
 bitflags = "2.6"
+serde_json = { version = "1.0", optional = true }
 
 # Copied from `compiler/rustc/Cargo.toml`.
 # But only for some targets, it fails for others. Rustc configures this in its CI, but we can't
@@ -67,6 +68,7 @@ default = ["stack-cache"]
 genmc = []
 stack-cache = []
 stack-cache-consistency-check = ["stack-cache"]
+tracing = ["serde_json"]
 
 [lints.rust.unexpected_cfgs]
 level = "warn"

src/tools/miri/rustfmt.toml

@@ -8,3 +8,8 @@ imports_granularity = "Module"
 force_multiline_blocks = true
 match_arm_blocks = false
 match_arm_leading_pipes = "Preserve"
+
+ignore = [
+    # This file is copy-pasted from the tracing_chrome crate and should remain like the original.
+    "src/bin/log/tracing_chrome.rs"
+]

src/tools/miri/src/bin/log/mod.rs

@@ -0,0 +1,2 @@
+pub mod setup;
+mod tracing_chrome;

src/tools/miri/src/bin/log/setup.rs

@@ -0,0 +1,126 @@
+use std::env::{self, VarError};
+use std::str::FromStr;
+use std::sync::{Mutex, OnceLock};
+
+use rustc_middle::ty::TyCtxt;
+use rustc_session::{CtfeBacktrace, EarlyDiagCtxt};
+
+/// The tracing layer from `tracing-chrome` starts a thread in the background that saves data to
+/// file and closes the file when stopped. If the thread is not stopped properly, the file will be
+/// missing end terminators (`]` for JSON arrays) and other data may also not be flushed. Therefore
+/// we need to keep a guard that, when [Drop]ped, will send a signal to stop the thread. Make sure
+/// to manually drop this guard using [deinit_loggers], if you are exiting the program with
+/// [std::process::exit]!
+#[must_use]
+struct TracingGuard {
+    #[cfg(feature = "tracing")]
+    _chrome: super::tracing_chrome::FlushGuard,
+    _no_construct: (),
+}
+
+// This ensures TracingGuard is always a drop-type, even when the `_chrome` field is disabled.
+impl Drop for TracingGuard {
+    fn drop(&mut self) {}
+}
+
+fn rustc_logger_config() -> rustc_log::LoggerConfig {
+    // Start with the usual env vars.
+    let mut cfg = rustc_log::LoggerConfig::from_env("RUSTC_LOG");
+
+    // Overwrite if MIRI_LOG is set.
+    if let Ok(var) = env::var("MIRI_LOG") {
+        // MIRI_LOG serves as default for RUSTC_LOG, if that is not set.
+        if matches!(cfg.filter, Err(VarError::NotPresent)) {
+            // We try to be a bit clever here: if `MIRI_LOG` is just a single level
+            // used for everything, we only apply it to the parts of rustc that are
+            // CTFE-related. Otherwise, we use it verbatim for `RUSTC_LOG`.
+            // This way, if you set `MIRI_LOG=trace`, you get only the right parts of
+            // rustc traced, but you can also do `MIRI_LOG=miri=trace,rustc_const_eval::interpret=debug`.
+            if tracing::Level::from_str(&var).is_ok() {
+                cfg.filter = Ok(format!(
+                    "rustc_middle::mir::interpret={var},rustc_const_eval::interpret={var},miri={var}"
+                ));
+            } else {
+                cfg.filter = Ok(var);
+            }
+        }
+    }
+
+    cfg
+}
+
+/// The global logger can only be set once per process, so track whether that already happened and
+/// keep a [TracingGuard] so it can be [Drop]ped later using [deinit_loggers].
+static LOGGER_INITED: OnceLock<Mutex<Option<TracingGuard>>> = OnceLock::new();
+
+fn init_logger_once(early_dcx: &EarlyDiagCtxt) {
+    // If the logger is not yet initialized, initialize it.
+    LOGGER_INITED.get_or_init(|| {
+        let guard = if env::var_os("MIRI_TRACING").is_some() {
+            #[cfg(not(feature = "tracing"))]
+            {
+                crate::fatal_error!(
+                    "fatal error: cannot enable MIRI_TRACING since Miri was not built with the \"tracing\" feature"
+                );
+            }
+
+            #[cfg(feature = "tracing")]
+            {
+                let (chrome_layer, chrome_guard) =
+                    super::tracing_chrome::ChromeLayerBuilder::new().include_args(true).build();
+                rustc_driver::init_logger_with_additional_layer(
+                    early_dcx,
+                    rustc_logger_config(),
+                    || {
+                        tracing_subscriber::layer::SubscriberExt::with(
+                            tracing_subscriber::Registry::default(),
+                            chrome_layer,
+                        )
+                    },
+                );
+
+                Some(TracingGuard { _chrome: chrome_guard, _no_construct: () })
+            }
+        } else {
+            // initialize the logger without any tracing enabled
+            rustc_driver::init_logger(early_dcx, rustc_logger_config());
+            None
+        };
+        Mutex::new(guard)
+    });
+}
+
+pub fn init_early_loggers(early_dcx: &EarlyDiagCtxt) {
+    // We only initialize `rustc` if the env var is set (so the user asked for it).
+    // If it is not set, we avoid initializing now so that we can initialize later with our custom
+    // settings, and *not* log anything for what happens before `miri` starts interpreting.
+    if env::var_os("RUSTC_LOG").is_some() {
+        init_logger_once(early_dcx);
+    }
+}
+
+pub fn init_late_loggers(early_dcx: &EarlyDiagCtxt, tcx: TyCtxt<'_>) {
+    // If the logger is not yet initialized, initialize it.
+    init_logger_once(early_dcx);
+
+    // If `MIRI_BACKTRACE` is set and `RUSTC_CTFE_BACKTRACE` is not, set `RUSTC_CTFE_BACKTRACE`.
+    // Do this late, so we ideally only apply this to Miri's errors.
+    if let Some(val) = env::var_os("MIRI_BACKTRACE") {
+        let ctfe_backtrace = match &*val.to_string_lossy() {
+            "immediate" => CtfeBacktrace::Immediate,
+            "0" => CtfeBacktrace::Disabled,
+            _ => CtfeBacktrace::Capture,
+        };
+        *tcx.sess.ctfe_backtrace.borrow_mut() = ctfe_backtrace;
+    }
+}
+
+/// Must be called before the program terminates to ensure the trace file is closed correctly. Not
+/// doing so will result in invalid trace files. Also see [TracingGuard].
+pub fn deinit_loggers() {
+    if let Some(guard) = LOGGER_INITED.get()
+        && let Ok(mut guard) = guard.lock()
+    {
+        std::mem::drop(guard.take());
+    }
+}