1use 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
12const TIME_FORMAT_DEFAULT: &[time::format_description::FormatItem<'static>] =
14 format_description!("[hour]:[minute]:[second].[subsecond digits:6]");
15
16#[derive(Debug, Default)]
32pub struct PrettyConsoleLayer {
33 format: PrettyFormatOptions,
35}
36
37#[derive(Debug)]
39struct PrettyFormatOptions {
40 pub wrapped: bool,
42 pub oneline: bool,
44 pub time_format: &'static [time::format_description::FormatItem<'static>],
46 pub events_only: bool,
48 pub show_time: bool,
50 pub show_target: bool,
52 pub show_file_info: bool,
54 pub show_span_info: bool,
56 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 pub fn wrapped(mut self, wrapped: bool) -> Self {
79 self.format.wrapped = wrapped;
80 self
81 }
82
83 pub fn oneline(mut self, oneline: bool) -> Self {
85 self.format.oneline = oneline;
86 self
87 }
88
89 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 pub fn events_only(mut self, show: bool) -> Self {
100 self.format.events_only = show;
101 self
102 }
103
104 pub fn show_time(mut self, show: bool) -> Self {
106 self.format.show_time = show;
107 self
108 }
109
110 pub fn show_target(mut self, show: bool) -> Self {
112 self.format.show_target = show;
113 self
114 }
115
116 pub fn show_file_info(mut self, show: bool) -> Self {
118 self.format.show_file_info = show;
119 self
120 }
121
122 pub fn show_span_info(mut self, show: bool) -> Self {
124 self.format.show_span_info = show;
125 self
126 }
127
128 pub fn indent(mut self, indent: usize) -> Self {
130 self.format.indent = indent;
131 self
132 }
133}
134
135#[derive(Debug)]
137struct SpanExtRecord {
138 tree_level: usize,
140 id: u64,
142 name: &'static str,
144 target: String,
146 file: String,
148 line: u32,
150 attrs: HashMap<&'static str, String>,
152 entered: Instant,
154 events: Vec<EventRecord>,
156 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 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 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 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 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 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 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#[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: Option<(usize, u64, String)>,
305}
306
307impl EventRecord {
308 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 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 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 if self.format.wrapped {
436 if let Some(parent_ref) = span_ref.parent() {
437 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 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 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 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 fn output_root_tree(&self, record: &SpanExtRecord) {
512 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}