tracing_span_tree/
lib.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.
5//!
6//! Usage:
7//!
8//! ```rust
9//! tracing_span_tree::span_tree()
10//!     .aggregate(true)
11//!     .enable();
12//! ```
13//!
14//! Example output:
15//!
16//! ```text
17//! 8.37ms           top_level
18//!   1.09ms           middle
19//!     1.06ms           leaf
20//!   1.06ms           middle
21//!   3.12ms           middle
22//!     1.06ms           leaf
23//!   3.06ms           middle
24//! ```
25//!
26//! Same data, but with `.aggregate(true)`:
27//!
28//! ```text
29//! 8.39ms           top_level
30//!  8.35ms    4      middle
31//!    2.13ms    2      leaf
32//! ```
33
34use std::{
35    fmt, mem,
36    time::{Duration, Instant},
37};
38
39use tracing::{
40    debug,
41    field::{Field, Visit},
42    span::Attributes,
43    Event, Id, Subscriber,
44};
45use tracing_subscriber::{
46    layer::Context,
47    prelude::*,
48    registry::{LookupSpan, Registry},
49    Layer,
50};
51
52pub fn span_tree() -> SpanTree {
53    SpanTree::default()
54}
55
56#[derive(Default)]
57pub struct SpanTree {
58    aggregate: bool,
59}
60
61impl SpanTree {
62    /// Merge identical sibling spans together.
63    pub fn aggregate(self, yes: bool) -> SpanTree {
64        SpanTree { aggregate: yes, ..self }
65    }
66    /// Set as a global subscriber
67    pub fn enable(self) {
68        let subscriber = Registry::default().with(self);
69        tracing::subscriber::set_global_default(subscriber)
70            .unwrap_or_else(|_| debug!("Global subscriber is already set"));
71    }
72}
73
74struct Data {
75    start: Instant,
76    children: Vec<Node>,
77}
78
79impl Data {
80    fn new(attrs: &Attributes<'_>) -> Self {
81        let mut span = Self { start: Instant::now(), children: Vec::new() };
82        attrs.record(&mut span);
83        span
84    }
85    fn into_node(self, name: &'static str) -> Node {
86        Node { name, count: 1, duration: self.start.elapsed(), children: self.children }
87    }
88}
89
90impl Visit for Data {
91    fn record_debug(&mut self, _field: &Field, _value: &dyn fmt::Debug) {}
92}
93
94impl<S> Layer<S> for SpanTree
95where
96    S: Subscriber + for<'span> LookupSpan<'span> + fmt::Debug,
97{
98    fn on_new_span(&self, attrs: &Attributes, id: &Id, ctx: Context<S>) {
99        let span = ctx.span(id).unwrap();
100
101        let data = Data::new(attrs);
102        span.extensions_mut().insert(data);
103    }
104
105    fn on_event(&self, _event: &Event<'_>, _ctx: Context<S>) {}
106
107    fn on_close(&self, id: Id, ctx: Context<S>) {
108        let span = ctx.span(&id).unwrap();
109        let data = span.extensions_mut().remove::<Data>().unwrap();
110        let mut node = data.into_node(span.name());
111
112        match span.parent() {
113            Some(parent_span) => {
114                parent_span.extensions_mut().get_mut::<Data>().unwrap().children.push(node);
115            }
116            None => {
117                if self.aggregate {
118                    node.aggregate()
119                }
120                node.print()
121            }
122        }
123    }
124}
125
126#[derive(Default)]
127struct Node {
128    name: &'static str,
129    count: u32,
130    duration: Duration,
131    children: Vec<Node>,
132}
133
134impl Node {
135    fn print(&self) {
136        self.go(0)
137    }
138    fn go(&self, level: usize) {
139        let bold = "\u{001b}[1m";
140        let reset = "\u{001b}[0m";
141
142        let duration = format!("{:3.2?}", self.duration);
143        let count = if self.count > 1 { self.count.to_string() } else { String::new() };
144        eprintln!(
145            "{:width$}  {:<9} {:<6} {bold}{}{reset}",
146            "",
147            duration,
148            count,
149            self.name,
150            bold = bold,
151            reset = reset,
152            width = level * 2
153        );
154        for child in &self.children {
155            child.go(level + 1)
156        }
157        if level == 0 {
158            eprintln!()
159        }
160    }
161
162    fn aggregate(&mut self) {
163        if self.children.is_empty() {
164            return;
165        }
166
167        self.children.sort_by_key(|it| it.name);
168        let mut idx = 0;
169        for i in 1..self.children.len() {
170            if self.children[idx].name == self.children[i].name {
171                let child = mem::take(&mut self.children[i]);
172                self.children[idx].duration += child.duration;
173                self.children[idx].count += child.count;
174                self.children[idx].children.extend(child.children);
175            } else {
176                idx += 1;
177                assert!(idx <= i);
178                self.children.swap(idx, i);
179            }
180        }
181        self.children.truncate(idx + 1);
182        for child in &mut self.children {
183            child.aggregate()
184        }
185    }
186}
tracing_span_tree/lib.rs

tracing_span_tree/
lib.rs
