Auto merge of #18227 - davidbarsky:davidbarsky/push-lmntvwvznyyx, r=davidbarsky

internal: add json `tracing` Layer for profiling startup

On `buck2/integrations/rust-project`, this results in the following being printed:

```json
{"name":"discover_command","elapsed_ms":18703}
{"name":"parallel_prime_caches","elapsed_ms":0}
{"name":"vfs_load","elapsed_ms":5895}
{"name":"vfs_load","elapsed_ms":547}
{"name":"parallel_prime_caches","elapsed_ms":23}
{"name":"parallel_prime_caches","elapsed_ms":84}
{"name":"parallel_prime_caches","elapsed_ms":5819}
```

crates/hir/src/lib.rs

@@ -556,7 +556,7 @@ impl Module {
         acc: &mut Vec<AnyDiagnostic>,
         style_lints: bool,
     ) {
-        let _p = tracing::info_span!("Module::diagnostics", name = ?self.name(db)).entered();
+        let _p = tracing::info_span!("diagnostics", name = ?self.name(db)).entered();
         let edition = db.crate_graph()[self.id.krate()].edition;
         let def_map = self.id.def_map(db.upcast());
         for diag in def_map.diagnostics() {

crates/rust-analyzer/src/bin/main.rs

@@ -137,6 +137,7 @@ fn setup_logging(log_file_flag: Option<PathBuf>) -> anyhow::Result<()> {
         filter: env::var("RA_LOG").ok().unwrap_or_else(|| "error".to_owned()),
         chalk_filter: env::var("CHALK_DEBUG").ok(),
         profile_filter: env::var("RA_PROFILE").ok(),
+        json_profile_filter: std::env::var("RA_PROFILE_JSON").ok(),
     }
     .init()?;
 

crates/rust-analyzer/src/lib.rs

@@ -34,6 +34,7 @@ mod handlers {
 
 pub mod tracing {
     pub mod config;
+    pub mod json;
     pub use config::Config;
     pub mod hprof;
 }

crates/rust-analyzer/src/tracing/config.rs

@@ -1,17 +1,20 @@
 //! Simple logger that logs either to stderr or to a file, using `tracing_subscriber`
 //! filter syntax and `tracing_appender` for non blocking output.
 
-use std::io;
+use std::io::{self};
 
 use anyhow::Context;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::{
-    filter::Targets, fmt::MakeWriter, layer::SubscriberExt, util::SubscriberInitExt, Layer,
-    Registry,
+    filter::{filter_fn, Targets},
+    fmt::MakeWriter,
+    layer::SubscriberExt,
+    Layer, Registry,
 };
 use tracing_tree::HierarchicalLayer;
 
 use crate::tracing::hprof;
+use crate::tracing::json;
 
 #[derive(Debug)]
 pub struct Config<T> {
@@ -34,6 +37,12 @@ pub struct Config<T> {
     /// env RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10
     /// ```
     pub profile_filter: Option<String>,
+
+    /// Filtering syntax, set in a shell:
+    /// ```
+    /// env RA_PROFILE_JSON=foo|bar|baz
+    /// ```
+    pub json_profile_filter: Option<String>,
 }
 
 impl<T> Config<T>
@@ -41,7 +50,7 @@ where
     T: for<'writer> MakeWriter<'writer> + Send + Sync + 'static,
 {
     pub fn init(self) -> anyhow::Result<()> {
-        let filter: Targets = self
+        let targets_filter: Targets = self
             .filter
             .parse()
             .with_context(|| format!("invalid log filter: `{}`", self.filter))?;
@@ -52,10 +61,10 @@ where
             .with_target(false)
             .with_ansi(false)
             .with_writer(writer)
-            .with_filter(filter);
+            .with_filter(targets_filter);
 
-        let mut chalk_layer = None;
-        if let Some(chalk_filter) = self.chalk_filter {
+        let chalk_layer = match self.chalk_filter {
+            Some(chalk_filter) => {
                 let level: LevelFilter =
                     chalk_filter.parse().with_context(|| "invalid chalk log filter")?;
 
@@ -63,19 +72,47 @@ where
                     .with_target("chalk_solve", level)
                     .with_target("chalk_ir", level)
                     .with_target("chalk_recursive", level);
-            chalk_layer = Some(
+                // TODO: remove `.with_filter(LevelFilter::OFF)` on the `None` branch.
                 HierarchicalLayer::default()
                     .with_indent_lines(true)
                     .with_ansi(false)
                     .with_indent_amount(2)
                     .with_writer(io::stderr)
-                    .with_filter(chalk_filter),
-            );
+                    .with_filter(chalk_filter)
+                    .boxed()
+            }
+            None => None::<HierarchicalLayer>.with_filter(LevelFilter::OFF).boxed(),
         };
 
-        let profiler_layer = self.profile_filter.map(|spec| hprof::layer(&spec));
+        // TODO: remove `.with_filter(LevelFilter::OFF)` on the `None` branch.
+        let profiler_layer = match self.profile_filter {
+            Some(spec) => Some(hprof::SpanTree::new(&spec)).with_filter(LevelFilter::INFO),
+            None => None.with_filter(LevelFilter::OFF),
+        };
 
-        Registry::default().with(ra_fmt_layer).with(chalk_layer).with(profiler_layer).try_init()?;
+        let json_profiler_layer = match self.json_profile_filter {
+            Some(spec) => {
+                let filter = json::JsonFilter::from_spec(&spec);
+                let filter = filter_fn(move |metadata| {
+                    let allowed = match &filter.allowed_names {
+                        Some(names) => names.contains(metadata.name()),
+                        None => true,
+                    };
+
+                    allowed && metadata.is_span()
+                });
+                Some(json::TimingLayer::new(std::io::stderr).with_filter(filter))
+            }
+            None => None,
+        };
+
+        let subscriber = Registry::default()
+            .with(ra_fmt_layer)
+            .with(json_profiler_layer)
+            .with(profiler_layer)
+            .with(chalk_layer);
+
+        tracing::subscriber::set_global_default(subscriber)?;
 
         Ok(())
     }

crates/rust-analyzer/src/tracing/hprof.rs

@@ -33,6 +33,7 @@
 
 use std::{
     fmt::Write,
+    marker::PhantomData,
     mem,
     time::{Duration, Instant},
 };
@@ -50,17 +51,23 @@ use tracing_subscriber::{
     Layer, Registry,
 };
 
-use crate::tracing::hprof;
-
 pub fn init(spec: &str) -> tracing::subscriber::DefaultGuard {
-    let subscriber = Registry::default().with(layer(spec));
+    let subscriber = Registry::default().with(SpanTree::new(spec));
     tracing::subscriber::set_default(subscriber)
 }
 
-pub fn layer<S>(spec: &str) -> impl Layer<S>
+#[derive(Debug)]
+pub(crate) struct SpanTree<S> {
+    aggregate: bool,
+    write_filter: WriteFilter,
+    _inner: PhantomData<fn(S)>,
+}
+
+impl<S> SpanTree<S>
 where
     S: Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
 {
+    pub(crate) fn new(spec: &str) -> impl Layer<S> {
         let (write_filter, allowed_names) = WriteFilter::from_spec(spec);
 
         // this filter the first pass for `tracing`: these are all the "profiling" spans, but things like
@@ -79,24 +86,7 @@ where
                 && !metadata.target().starts_with("chalk")
         });
 
-    hprof::SpanTree::default().aggregate(true).spec_filter(write_filter).with_filter(profile_filter)
-}
-
-#[derive(Default, Debug)]
-pub(crate) struct SpanTree {
-    aggregate: bool,
-    write_filter: WriteFilter,
-}
-
-impl SpanTree {
-    /// Merge identical sibling spans together.
-    pub(crate) fn aggregate(self, yes: bool) -> SpanTree {
-        SpanTree { aggregate: yes, ..self }
-    }
-
-    /// Add a write-time filter for span duration or tree depth.
-    pub(crate) fn spec_filter(self, write_filter: WriteFilter) -> SpanTree {
-        SpanTree { write_filter, ..self }
+        Self { aggregate: true, write_filter, _inner: PhantomData }.with_filter(profile_filter)
     }
 }
 
@@ -136,7 +126,7 @@ impl<'a> Visit for DataVisitor<'a> {
     }
 }
 
-impl<S> Layer<S> for SpanTree
+impl<S> Layer<S> for SpanTree<S>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {

crates/rust-analyzer/src/tracing/json.rs

@@ -0,0 +1,90 @@
+//! A [tracing_subscriber::layer::Layer] that exports new-line delinated JSON.
+//!
+//! Usage:
+//!
+//! ```rust
+//! let layer = json::TimingLayer::new(std::io::stderr);
+//! Registry::default().with(layer).init();
+//! ```
+
+use std::{io::Write as _, marker::PhantomData, time::Instant};
+
+use ide_db::FxHashSet;
+use tracing::{
+    span::{Attributes, Id},
+    Event, Subscriber,
+};
+use tracing_subscriber::{fmt::MakeWriter, layer::Context, registry::LookupSpan, Layer};
+
+struct JsonData {
+    name: &'static str,
+    start: std::time::Instant,
+}
+
+impl JsonData {
+    fn new(name: &'static str) -> Self {
+        Self { name, start: Instant::now() }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct TimingLayer<S, W> {
+    writer: W,
+    _inner: PhantomData<fn(S)>,
+}
+
+impl<S, W> TimingLayer<S, W> {
+    pub(crate) fn new(writer: W) -> Self {
+        Self { writer, _inner: PhantomData }
+    }
+}
+
+impl<S, W> Layer<S> for TimingLayer<S, W>
+where
+    S: Subscriber + for<'span> LookupSpan<'span>,
+    W: for<'writer> MakeWriter<'writer> + Send + Sync + 'static,
+{
+    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
+        let span = ctx.span(id).unwrap();
+
+        let data = JsonData::new(attrs.metadata().name());
+        span.extensions_mut().insert(data);
+    }
+
+    fn on_event(&self, _event: &Event<'_>, _ctx: Context<'_, S>) {}
+
+    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
+        #[derive(serde::Serialize)]
+        struct JsonDataInner {
+            name: &'static str,
+            elapsed_ms: u128,
+        }
+
+        let span = ctx.span(&id).unwrap();
+        let Some(data) = span.extensions_mut().remove::<JsonData>() else {
+            return;
+        };
+
+        let data = JsonDataInner { name: data.name, elapsed_ms: data.start.elapsed().as_millis() };
+        let mut out = serde_json::to_string(&data).expect("Unable to serialize data");
+        out.push('\n');
+        self.writer.make_writer().write_all(out.as_bytes()).expect("Unable to write data");
+    }
+}
+
+#[derive(Default, Clone, Debug)]
+pub(crate) struct JsonFilter {
+    pub(crate) allowed_names: Option<FxHashSet<String>>,
+}
+
+impl JsonFilter {
+    pub(crate) fn from_spec(spec: &str) -> Self {
+        let allowed_names = if spec == "*" {
+            None
+        } else {
+            Some(FxHashSet::from_iter(spec.split('|').map(String::from)))
+        };
+
+        Self { allowed_names }
+    }
+}

crates/rust-analyzer/tests/slow-tests/support.rs

@@ -103,6 +103,7 @@ impl Project<'_> {
                 filter: std::env::var("RA_LOG").ok().unwrap_or_else(|| "error".to_owned()),
                 chalk_filter: std::env::var("CHALK_DEBUG").ok(),
                 profile_filter: std::env::var("RA_PROFILE").ok(),
+                json_profile_filter: std::env::var("RA_PROFILE_JSON").ok(),
             };
         });
 

docs/dev/README.md

@@ -178,7 +178,15 @@ RA_PROFILE=foo|bar|baz   // enabled only selected entries
 RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10 ms
 ```
 
-In particular, I have `export RA_PROFILE='*>10'` in my shell profile.
+Some rust-analyzer contributors have `export RA_PROFILE='*>10'` in my shell profile.
+
+For machine-readable JSON output, we have the `RA_PROFILE_JSON` env variable. We support
+filtering only by span name:
+
+```
+RA_PROFILE=* // dump everything
+RA_PROFILE_JSON="vfs_load|parallel_prime_caches|discover_command" // dump selected spans
+```
 
 We also have a "counting" profiler which counts number of instances of popular structs.
 It is enabled by `RA_COUNT=1`.