tracing_ext/sub/
pretty.rs

1//! A pretty tracing layer for console printing
2
3use std::{collections::HashMap, io::Write, time::Instant};
4
5use colored::Colorize;
6use time::macros::format_description;
7use tracing::Level;
8use tracing_subscriber::registry::SpanRef;
9
10use super::{EventVisitor, SpanExtension};
11
12/// Default time format
13const TIME_FORMAT_DEFAULT: &[time::format_description::FormatItem<'static>] =
14    format_description!("[hour]:[minute]:[second].[subsecond digits:6]");
15
16/// A tracing layer with pretty print to the console
17///
18/// ```
19///  use tracing_ext::sub::PrettyConsoleLayer;
20///
21///  let pretty_layer = PrettyConsoleLayer::default()
22///     .wrapped(true)
23///     .oneline(false)
24///     .events_only(true)
25///     .show_time(true)
26///     .show_target(true)
27///     .show_file_info(true)
28///     .show_span_info(true)
29///     .indent(6);
30/// ```
31#[derive(Debug, Default)]
32pub struct PrettyConsoleLayer {
33    /// Format
34    format: PrettyFormatOptions,
35}
36
37/// Formatting options (for spans and events)
38#[derive(Debug)]
39struct PrettyFormatOptions {
40    /// Defines if the display is wrapped
41    pub wrapped: bool,
42    /// If true, spans and events are printed in 1 line
43    pub oneline: bool,
44    /// Time format
45    pub time_format: &'static [time::format_description::FormatItem<'static>],
46    /// The span is shown (enter and exit info)
47    pub events_only: bool,
48    /// The timestanp is shown
49    pub show_time: bool,
50    /// The target is shown
51    pub show_target: bool,
52    /// The file info is shown
53    pub show_file_info: bool,
54    /// Shows the span info
55    pub show_span_info: bool,
56    /// Indentation (x spaces) - invalid if the `oneline` option is set
57    pub indent: usize,
58}
59
60impl Default for PrettyFormatOptions {
61    fn default() -> Self {
62        Self {
63            wrapped: false,
64            oneline: false,
65            time_format: TIME_FORMAT_DEFAULT,
66            events_only: false,
67            show_time: true,
68            show_target: true,
69            show_file_info: true,
70            show_span_info: true,
71            indent: 6,
72        }
73    }
74}
75
76impl PrettyConsoleLayer {
77    /// Sets the kind is wrapped
78    pub fn wrapped(mut self, wrapped: bool) -> Self {
79        self.format.wrapped = wrapped;
80        self
81    }
82
83    /// Shows each span and event as 1 line
84    pub fn oneline(mut self, oneline: bool) -> Self {
85        self.format.oneline = oneline;
86        self
87    }
88
89    /// Sets the time format
90    pub fn time_format(
91        mut self,
92        format: &'static [time::format_description::FormatItem<'static>],
93    ) -> Self {
94        self.format.time_format = format;
95        self
96    }
97
98    /// Sets if only the events are shown
99    pub fn events_only(mut self, show: bool) -> Self {
100        self.format.events_only = show;
101        self
102    }
103
104    /// Sets if the time is shown
105    pub fn show_time(mut self, show: bool) -> Self {
106        self.format.show_time = show;
107        self
108    }
109
110    /// Sets if the target is shown
111    pub fn show_target(mut self, show: bool) -> Self {
112        self.format.show_target = show;
113        self
114    }
115
116    /// Sets if the file info is shown
117    pub fn show_file_info(mut self, show: bool) -> Self {
118        self.format.show_file_info = show;
119        self
120    }
121
122    /// Sets if the span inline info is shown
123    pub fn show_span_info(mut self, show: bool) -> Self {
124        self.format.show_span_info = show;
125        self
126    }
127
128    /// Sets the indentation (in x spaces)
129    pub fn indent(mut self, indent: usize) -> Self {
130        self.format.indent = indent;
131        self
132    }
133}
134
135/// A span extension for the span record
136#[derive(Debug)]
137struct SpanExtRecord {
138    /// Level within the tree
139    tree_level: usize,
140    /// Span ID
141    id: u64,
142    /// Span name
143    name: &'static str,
144    /// Span target
145    target: String,
146    /// File
147    file: String,
148    /// Line
149    line: u32,
150    /// Span attributes
151    attrs: HashMap<&'static str, String>,
152    /// Entered time
153    entered: Instant,
154    /// Events within the span
155    events: Vec<EventRecord>,
156    // children
157    children: Vec<SpanExtRecord>,
158}
159
160impl tracing::field::Visit for SpanExtRecord {
161    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
162        let value = format!("{value:?}");
163        self.attrs.insert(field.name(), value);
164    }
165}
166
167impl SpanExtension for SpanExtRecord {}
168
169impl SpanExtRecord {
170    /// Instantiates from a [SpanRef]
171    ///
172    /// NB: attributes are not collected yet
173    fn new_from_span_ref<S>(span_ref: &SpanRef<S>) -> Self
174    where
175        S: for<'b> tracing_subscriber::registry::LookupSpan<'b>,
176    {
177        let tree_level = if let Some(parent) = span_ref.parent() {
178            let extensions = parent.extensions();
179            let tree_level = extensions.get::<Self>().unwrap().tree_level;
180            tree_level + 1
181        } else {
182            0
183        };
184
185        Self {
186            tree_level,
187            id: span_ref.id().into_u64(),
188            name: span_ref.name(),
189            target: span_ref.metadata().target().to_string(),
190            file: span_ref.metadata().file().unwrap_or("").to_string(),
191            line: span_ref.metadata().line().unwrap_or(0),
192            attrs: HashMap::new(),
193            entered: Instant::now(),
194            events: Vec::new(),
195            children: Vec::new(),
196        }
197    }
198
199    /// Serializes the span entry
200    fn serialize_span_entry(&self, opts: &PrettyFormatOptions) -> Vec<u8> {
201        if opts.events_only {
202            return vec![];
203        }
204
205        let mut buf: Vec<u8> = vec![];
206
207        let tree_indent = if opts.wrapped {
208            self.tree_level * opts.indent
209        } else {
210            0
211        };
212        let tree_indent_str = " ".repeat(tree_indent);
213        write!(buf, "{}", tree_indent_str).unwrap();
214
215        if !opts.wrapped {
216            write!(buf, "{:w$}", format!("-->"), w = opts.indent).unwrap();
217        }
218        write!(buf, "{}", format!("{{{}}}", self.name).magenta()).unwrap();
219
220        let field_indent = tree_indent + opts.indent;
221        let field_indent_str = " ".repeat(field_indent);
222        let field_new_line = if opts.oneline {
223            " ".to_string()
224        } else {
225            format!("\n{field_indent_str}")
226        };
227
228        if opts.show_time {
229            let time_str = time::OffsetDateTime::now_utc()
230                .format(opts.time_format)
231                .expect("invalid datetime");
232            let line = format!("{}: {}", "time".italic(), time_str);
233            write!(buf, "{field_new_line}{}", line.dimmed()).unwrap();
234        };
235
236        // span info
237        if opts.show_span_info {
238            let span_id = format!("{}: {}", "span.id".italic(), self.id);
239            write!(buf, "{field_new_line}{}", span_id.dimmed()).unwrap();
240        }
241
242        if opts.show_target {
243            let target = format!("{}: {}", "target".italic(), self.target);
244            write!(buf, "{field_new_line}{}", target.dimmed()).unwrap();
245        }
246
247        if opts.show_file_info {
248            let target = format!("{}: {}:{}", "file".italic(), self.file, self.line);
249            write!(buf, "{field_new_line}{}", target.dimmed()).unwrap();
250        }
251
252        // span attributes
253        for (k, v) in &self.attrs {
254            write!(buf, "{field_new_line}{}={}", k.to_string().italic(), v).unwrap();
255        }
256
257        buf
258    }
259
260    /// Serializes the span exit
261    fn serialize_span_exit(&self, opts: &PrettyFormatOptions) -> Vec<u8> {
262        if opts.events_only {
263            return vec![];
264        }
265
266        let mut buf: Vec<u8> = vec![];
267
268        let tree_indent = if opts.wrapped {
269            self.tree_level * opts.indent
270        } else {
271            0
272        };
273        let tree_indent_str = " ".repeat(tree_indent);
274        write!(buf, "{}", tree_indent_str).unwrap();
275
276        if !opts.wrapped {
277            write!(buf, "{:w$}", format!("<--"), w = opts.indent).unwrap();
278        }
279        write!(buf, "{}", format!("!{{{}}}", self.name).magenta()).unwrap();
280
281        // span info
282        if opts.show_span_info {
283            let span_id = format!("({}={})", "id".italic(), self.id);
284            write!(buf, " {}", span_id.dimmed()).unwrap();
285        }
286
287        let duration_us = self.entered.elapsed().as_micros();
288        write!(buf, " {}", format!("{duration_us}us").dimmed()).unwrap();
289
290        buf
291    }
292}
293
294/// An event record
295#[derive(Debug)]
296struct EventRecord {
297    level: Level,
298    target: String,
299    file: String,
300    line: u32,
301    message: String,
302    meta_fields: HashMap<&'static str, String>,
303    /// Span info (tree level, id, name)
304    span: Option<(usize, u64, String)>,
305}
306
307impl EventRecord {
308    /// Serializes an event
309    fn serialize(&self, opts: &PrettyFormatOptions) -> Vec<u8> {
310        let mut buf: Vec<u8> = vec![];
311
312        let tree_indent = if opts.wrapped {
313            let tree_level = self.span.as_ref().map(|(l, _, _)| *l).unwrap_or(0);
314            tree_level * opts.indent
315        } else {
316            0
317        };
318        let tree_indent_str = " ".repeat(tree_indent);
319        write!(buf, "{}", tree_indent_str).unwrap();
320
321        let level_str = match self.level {
322            tracing::Level::TRACE => format!("{:w$}", "TRACE", w = opts.indent).magenta(),
323            tracing::Level::DEBUG => format!("{:w$}", "DEBUG", w = opts.indent).blue(),
324            tracing::Level::INFO => format!("{:w$}", "INFO", w = opts.indent).green(),
325            tracing::Level::WARN => format!("{:w$}", "WARN", w = opts.indent).yellow(),
326            tracing::Level::ERROR => format!("{:w$}", "ERROR", w = opts.indent).red(),
327        };
328        write!(buf, "{}", level_str).unwrap();
329        write!(buf, "{}", self.message).unwrap();
330
331        let field_indent = tree_indent + opts.indent;
332        let field_indent_str = " ".repeat(field_indent);
333        let field_new_line = if opts.oneline {
334            " ".to_string()
335        } else {
336            format!("\n{field_indent_str}")
337        };
338
339        if opts.show_time {
340            let time_str = time::OffsetDateTime::now_utc()
341                .format(opts.time_format)
342                .expect("invalid datetime");
343            let line = format!("{}: {}", "time".italic(), time_str);
344            write!(buf, "{field_new_line}{}", line.dimmed()).unwrap();
345        };
346
347        // event context
348        if opts.show_span_info {
349            if let Some((_, id, name)) = &self.span {
350                let span_id = format!("{}: {}", "span.id".italic(), id);
351                write!(buf, "{field_new_line}{}", span_id.dimmed()).unwrap();
352
353                let span_name = format!(
354                    "{field_new_line}{}{} {}",
355                    "span.name".italic().dimmed(),
356                    ":".dimmed(),
357                    name.truecolor(191, 160, 217)
358                );
359                write!(buf, "{}", span_name.dimmed()).unwrap();
360            }
361        }
362
363        if opts.show_target {
364            let target = format!("{}: {}", "target".italic(), self.target);
365            write!(buf, "{field_new_line}{}", target.dimmed()).unwrap();
366        }
367
368        if opts.show_file_info {
369            let target = format!("{}: {}:{}", "file".italic(), self.file, self.line);
370            write!(buf, "{field_new_line}{}", target.dimmed()).unwrap();
371        }
372
373        // event fields
374        for (k, v) in &self.meta_fields {
375            write!(buf, "{field_new_line}{}={}", k.to_string().italic(), v).unwrap();
376        }
377
378        buf
379    }
380}
381
382impl<S> tracing_subscriber::Layer<S> for PrettyConsoleLayer
383where
384    S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
385{
386    fn on_new_span(
387        &self,
388        attrs: &tracing::span::Attributes<'_>,
389        id: &tracing::span::Id,
390        ctx: tracing_subscriber::layer::Context<'_, S>,
391    ) {
392        let span_ref = ctx.span(id).expect("span not found");
393        let record = SpanExtRecord::new_from_span_ref(&span_ref);
394        SpanExtRecord::register_value(record, &span_ref);
395        SpanExtRecord::record_attrs(&span_ref, attrs);
396    }
397
398    fn on_enter(&self, id: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
399        let span_ref = ctx.span(id).expect("span not found");
400
401        let mut extensions = span_ref.extensions_mut();
402        let record = extensions
403            .get_mut::<SpanExtRecord>()
404            .expect("Extension not initialized");
405
406        if !self.format.wrapped {
407            let buf = record.serialize_span_entry(&self.format);
408            if !buf.is_empty() {
409                eprintln!("{}", std::str::from_utf8(&buf).unwrap());
410            }
411        }
412    }
413
414    fn on_exit(&self, id: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
415        let span_ref = ctx.span(id).expect("span not found");
416
417        let mut extensions = span_ref.extensions_mut();
418        let record = extensions
419            .get_mut::<SpanExtRecord>()
420            .expect("Extension not initialized");
421
422        if !self.format.wrapped {
423            let buf = record.serialize_span_exit(&self.format);
424            if !buf.is_empty() {
425                eprintln!("{}", std::str::from_utf8(&buf).unwrap());
426            }
427        }
428    }
429
430    fn on_close(&self, id: tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
431        let span_ref = ctx.span(&id).expect("span not found");
432
433        // When wrapping, if the span has a parent, we record it as a child of the parent.
434        // If it is the root, the span tree is outputted
435        if self.format.wrapped {
436            if let Some(parent_ref) = span_ref.parent() {
437                // => the span has a parent and hence it is recorded on the parent
438                let mut parent_extensions = parent_ref.extensions_mut();
439                let parent_record = parent_extensions
440                    .get_mut::<SpanExtRecord>()
441                    .expect("Extension not initialized");
442
443                let mut extensions = span_ref.extensions_mut();
444                let record = extensions
445                    .remove::<SpanExtRecord>()
446                    .expect("Extension not initialized");
447
448                parent_record.children.push(record);
449            } else {
450                // => root of span tree => print
451                let mut extensions = span_ref.extensions_mut();
452                let record = extensions
453                    .remove::<SpanExtRecord>()
454                    .expect("Extension not initialized");
455                self.output_root_tree(&record);
456            }
457        }
458    }
459
460    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
461        let visitor = EventVisitor::record_event(event);
462
463        let evt_record = EventRecord {
464            level: *event.metadata().level(),
465            target: event.metadata().target().to_string(),
466            file: event.metadata().file().unwrap_or("").to_string(),
467            line: event.metadata().line().unwrap_or(0),
468            message: visitor.message().to_string(),
469            meta_fields: visitor
470                .meta_fields()
471                .iter()
472                .map(|(k, v)| (*k, v.to_string()))
473                .collect(),
474            span: ctx.current_span().id().map(|id| {
475                let parent_ref = ctx.span(id).expect("span not found");
476                let mut extensions = parent_ref.extensions_mut();
477                let span_record = extensions
478                    .get_mut::<SpanExtRecord>()
479                    .expect("Extension not initialized");
480                (
481                    span_record.tree_level + 1,
482                    id.into_u64(),
483                    ctx.current_span().metadata().unwrap().name().to_string(),
484                )
485            }),
486        };
487
488        // we print the event is we print by chronological order, or if the event is at the root
489        match (self.format.wrapped, ctx.current_span().id().is_some()) {
490            (false, _) | (true, false) => {
491                let buf = evt_record.serialize(&self.format);
492                eprintln!("{}", std::str::from_utf8(&buf).unwrap());
493            }
494            _ => {
495                // NB: push the events to the span record if everything is printed at the end
496                let curr_span = ctx.current_span();
497                let parent_span_id = curr_span.id().unwrap();
498                let span_ref = ctx.span(parent_span_id).expect("span not found");
499                let mut extensions = span_ref.extensions_mut();
500                let span_record = extensions
501                    .get_mut::<SpanExtRecord>()
502                    .expect("Extension not initialized");
503                span_record.events.push(evt_record);
504            }
505        }
506    }
507}
508
509impl PrettyConsoleLayer {
510    /// Outputs a tree of spans from the root
511    fn output_root_tree(&self, record: &SpanExtRecord) {
512        // eprintln!("ENTER SPAN {}", record.id);
513        let buf = record.serialize_span_entry(&self.format);
514        if !buf.is_empty() {
515            eprintln!("{}", std::str::from_utf8(&buf).unwrap());
516        }
517
518        for event in &record.events {
519            let buf = event.serialize(&self.format);
520            if !buf.is_empty() {
521                println!("{}", std::str::from_utf8(&buf).unwrap());
522            }
523        }
524
525        for child in &record.children {
526            self.output_root_tree(child);
527        }
528
529        let buf = record.serialize_span_exit(&self.format);
530        if !buf.is_empty() {
531            eprintln!("{}", std::str::from_utf8(&buf).unwrap());
532        }
533    }
534}