rust_analyzer/tracing/
hprof.rs

1//! Consumer of `tracing` data, which prints a hierarchical profile.
2//!
3//! Based on <https://github.com/davidbarsky/tracing-tree>, but does less, while
4//! actually printing timings for spans by default. The code here is vendored from
5//! <https://github.com/matklad/tracing-span-tree>.
6//!
7//! Usage:
8//!
9//! ```ignore
10//! # use tracing_subscriber::Registry;
11//! let layer = hprof::SpanTree::default();
12//! Registry::default().with(layer).init();
13//! ```
14//!
15//! Example output:
16//!
17//! ```text
18//! 8.37ms           top_level
19//!   1.09ms           middle
20//!     1.06ms           leaf
21//!   1.06ms           middle
22//!   3.12ms           middle
23//!     1.06ms           leaf
24//!   3.06ms           middle
25//! ```
26//!
27//! Same data, but with `.aggregate(true)`:
28//!
29//! ```text
30//! 8.39ms           top_level
31//!  8.35ms    4      middle
32//!    2.13ms    2      leaf
33//! ```
34
35use std::{
36    fmt::Write,
37    marker::PhantomData,
38    mem,
39    time::{Duration, Instant},
40};
41
42use rustc_hash::FxHashSet;
43use tracing::{
44    Event, Id, Level, Subscriber,
45    field::{Field, Visit},
46    span::Attributes,
47};
48use tracing_subscriber::{
49    Layer, Registry, filter,
50    layer::{Context, SubscriberExt},
51    registry::LookupSpan,
52};
53
54pub fn init(spec: &str) -> tracing::subscriber::DefaultGuard {
55    let subscriber = Registry::default().with(SpanTree::new(spec));
56    tracing::subscriber::set_default(subscriber)
57}
58
59#[derive(Debug)]
60pub(crate) struct SpanTree<S> {
61    aggregate: bool,
62    write_filter: WriteFilter,
63    _inner: PhantomData<fn(S)>,
64}
65
66impl<S> SpanTree<S>
67where
68    S: Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
69{
70    pub(crate) fn new(spec: &str) -> impl Layer<S> {
71        let (write_filter, allowed_names) = WriteFilter::from_spec(spec);
72
73        // this filter the first pass for `tracing`: these are all the "profiling" spans, but things like
74        // span depth or duration are not filtered here: that only occurs at write time.
75        let profile_filter = filter::filter_fn(move |metadata| {
76            let allowed = match &allowed_names {
77                Some(names) => names.contains(metadata.name()),
78                None => true,
79            };
80
81            allowed
82                && metadata.is_span()
83                && metadata.level() >= &Level::INFO
84                && !metadata.target().starts_with("salsa")
85                && metadata.name() != "compute_exhaustiveness_and_usefulness"
86                && !metadata.target().starts_with("chalk")
87        });
88
89        Self { aggregate: true, write_filter, _inner: PhantomData }.with_filter(profile_filter)
90    }
91}
92
93struct Data {
94    start: Instant,
95    children: Vec<Node>,
96    fields: String,
97}
98
99impl Data {
100    fn new(attrs: &Attributes<'_>) -> Self {
101        let mut data = Self { start: Instant::now(), children: Vec::new(), fields: String::new() };
102
103        let mut visitor = DataVisitor { string: &mut data.fields };
104        attrs.record(&mut visitor);
105        data
106    }
107
108    fn into_node(self, name: &'static str) -> Node {
109        Node {
110            name,
111            fields: self.fields,
112            count: 1,
113            duration: self.start.elapsed(),
114            children: self.children,
115        }
116    }
117}
118
119pub struct DataVisitor<'a> {
120    string: &'a mut String,
121}
122
123impl Visit for DataVisitor<'_> {
124    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
125        write!(self.string, "{} = {:?} ", field.name(), value).unwrap();
126    }
127}
128
129impl<S> Layer<S> for SpanTree<S>
130where
131    S: Subscriber + for<'span> LookupSpan<'span>,
132{
133    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
134        let span = ctx.span(id).unwrap();
135
136        let data = Data::new(attrs);
137        span.extensions_mut().insert(data);
138    }
139
140    fn on_event(&self, _event: &Event<'_>, _ctx: Context<'_, S>) {}
141
142    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
143        let span = ctx.span(&id).unwrap();
144        let data = span.extensions_mut().remove::<Data>().unwrap();
145        let mut node = data.into_node(span.name());
146
147        match span.parent() {
148            Some(parent_span) => {
149                parent_span.extensions_mut().get_mut::<Data>().unwrap().children.push(node);
150            }
151            None => {
152                if self.aggregate {
153                    node.aggregate()
154                }
155                node.print(&self.write_filter)
156            }
157        }
158    }
159}
160
161#[derive(Default)]
162struct Node {
163    name: &'static str,
164    fields: String,
165    count: u32,
166    duration: Duration,
167    children: Vec<Node>,
168}
169
170impl Node {
171    fn print(&self, filter: &WriteFilter) {
172        self.go(0, filter)
173    }
174
175    #[allow(clippy::print_stderr)]
176    fn go(&self, level: usize, filter: &WriteFilter) {
177        if self.duration > filter.longer_than && level < filter.depth {
178            let duration = ms(self.duration);
179            let current_indent = level * 2;
180
181            let mut out = String::new();
182            let _ = write!(out, "{:current_indent$}   {duration} {:<6}", "", self.name);
183
184            if !self.fields.is_empty() {
185                let _ = write!(out, " @ {}", self.fields);
186            }
187
188            if self.count > 1 {
189                let _ = write!(out, " ({} calls)", self.count);
190            }
191
192            eprintln!("{out}");
193
194            for child in &self.children {
195                child.go(level + 1, filter)
196            }
197        }
198    }
199
200    fn aggregate(&mut self) {
201        if self.children.is_empty() {
202            return;
203        }
204
205        self.children.sort_by_key(|it| it.name);
206        let mut idx = 0;
207        for i in 1..self.children.len() {
208            if self.children[idx].name == self.children[i].name {
209                let child = mem::take(&mut self.children[i]);
210                self.children[idx].duration += child.duration;
211                self.children[idx].count += child.count;
212                self.children[idx].children.extend(child.children);
213            } else {
214                idx += 1;
215                assert!(idx <= i);
216                self.children.swap(idx, i);
217            }
218        }
219        self.children.truncate(idx + 1);
220        for child in &mut self.children {
221            child.aggregate()
222        }
223    }
224}
225
226#[derive(Default, Clone, Debug)]
227pub(crate) struct WriteFilter {
228    depth: usize,
229    longer_than: Duration,
230}
231
232impl WriteFilter {
233    pub(crate) fn from_spec(mut spec: &str) -> (WriteFilter, Option<FxHashSet<String>>) {
234        let longer_than = if let Some(idx) = spec.rfind('>') {
235            let longer_than = spec[idx + 1..].parse().expect("invalid profile longer_than");
236            spec = &spec[..idx];
237            Duration::from_millis(longer_than)
238        } else {
239            Duration::new(0, 0)
240        };
241
242        let depth = if let Some(idx) = spec.rfind('@') {
243            let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
244            spec = &spec[..idx];
245            depth
246        } else {
247            999
248        };
249        let allowed = if spec == "*" {
250            None
251        } else {
252            Some(FxHashSet::from_iter(spec.split('|').map(String::from)))
253        };
254        (WriteFilter { depth, longer_than }, allowed)
255    }
256}
257
258#[allow(non_camel_case_types)]
259struct ms(Duration);
260
261impl std::fmt::Display for ms {
262    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
263        let n = self.0.as_millis();
264        write!(f, "{n:5}ms")
265    }
266}
rust_analyzer/tracing/hprof.rs

rust_analyzer/tracing/
hprof.rs
