rocket/trace/subscriber/
common.rs
1use std::fmt;
2use std::cell::Cell;
3
4use tracing::field::Field;
5use tracing::{Level, Metadata};
6use tracing_subscriber::filter;
7use tracing_subscriber::field::RecordFields;
8
9use thread_local::ThreadLocal;
10use yansi::{Condition, Paint, Style};
11
12use crate::config::CliColors;
13use crate::trace::subscriber::RecordDisplay;
14use crate::util::Formatter;
15
16mod private {
17 pub trait FmtKind: Send + Sync + 'static { }
18
19 impl FmtKind for crate::trace::subscriber::Pretty { }
20 impl FmtKind for crate::trace::subscriber::Compact { }
21}
22
23#[derive(Default)]
24pub struct RocketFmt<K: private::FmtKind> {
25 state: ThreadLocal<Cell<K>>,
26 pub(crate) level: Option<Level>,
27 pub(crate) filter: filter::Targets,
28 pub(crate) style: Style,
29}
30
31impl<K: private::FmtKind + Default + Copy> RocketFmt<K> {
32 pub(crate) fn state(&self) -> K {
33 self.state.get_or_default().get()
34 }
35
36 pub(crate) fn update_state<F: FnOnce(&mut K)>(&self, update: F) {
37 let cell = self.state.get_or_default();
38 let mut old = cell.get();
39 update(&mut old);
40 cell.set(old);
41 }
42}
43
44impl<K: private::FmtKind> RocketFmt<K> {
45 pub fn new(workers: usize, cli_colors: CliColors, level: Option<Level>) -> Self {
46 Self {
47 state: ThreadLocal::with_capacity(workers),
48 level,
49 filter: filter::Targets::new()
50 .with_default(level)
51 .with_target("rustls", level.filter(|&l| l == Level::TRACE))
52 .with_target("hyper", level.filter(|&l| l == Level::TRACE)),
53 style: match cli_colors {
54 CliColors::Always => Style::new().whenever(Condition::ALWAYS),
55 CliColors::Auto => Style::new().whenever(Condition::DEFAULT),
56 CliColors::Never => Style::new().whenever(Condition::NEVER),
57 }
58 }
59 }
60
61 pub fn style(&self, metadata: &Metadata<'_>) -> Style {
62 match *metadata.level() {
63 Level::ERROR => self.style.red(),
64 Level::WARN => self.style.yellow(),
65 Level::INFO => self.style.blue(),
66 Level::DEBUG => self.style.green(),
67 Level::TRACE => self.style.magenta(),
68 }
69 }
70
71 pub(crate) fn has_message(&self, meta: &Metadata<'_>) -> bool {
72 meta.fields().field("message").is_some()
73 }
74
75 pub(crate) fn has_data_fields(&self, meta: &Metadata<'_>) -> bool {
76 meta.fields().iter().any(|f| f.name() != "message")
77 }
78
79 pub(crate) fn message<'a, F: RecordFields + 'a>(&self,
80 init_prefix: &'a dyn fmt::Display,
81 cont_prefix: &'a dyn fmt::Display,
82 meta: &'a Metadata<'_>,
83 data: F
84 ) -> impl fmt::Display + 'a {
85 let style = self.style(meta);
86 Formatter(move |f| {
87 let fields = meta.fields();
88 let message = fields.field("message");
89 if let Some(message_field) = &message {
90 data.record_display(|field: &Field, value: &dyn fmt::Display| {
91 if field != message_field {
92 return;
93 }
94
95 for (i, line) in value.to_string().lines().enumerate() {
96 let line = line.paint(style);
97 if i == 0 {
98 let _ = writeln!(f, "{init_prefix}{line}");
99 } else {
100 let _ = writeln!(f, "{cont_prefix}{line}");
101 }
102 }
103 });
104 }
105
106 Ok(())
107 })
108 }
109
110 pub(crate) fn compact_fields<'a, F: RecordFields + 'a>(
111 &self,
112 meta: &'a Metadata<'_>,
113 data: F
114 ) -> impl fmt::Display + 'a {
115 let key_style = self.style(meta).bold();
116 let val_style = self.style(meta).primary();
117
118 Formatter(move |f| {
119 let mut printed = false;
120 data.record_display(|field: &Field, val: &dyn fmt::Display| {
121 let key = field.name();
122 if key != "message" {
123 if printed { let _ = write!(f, " "); }
124 let _ = write!(f, "{}: {}", key.paint(key_style), val.paint(val_style));
125 printed = true;
126 }
127 });
128
129 Ok(())
130 })
131 }
132
133 pub(crate) fn print<F: RecordFields>(
134 &self,
135 prefix: &dyn fmt::Display,
136 cont_prefix: &dyn fmt::Display,
137 m: &Metadata<'_>,
138 data: F
139 ) {
140 if self.has_message(m) {
141 let message = self.message(prefix, cont_prefix, m, &data);
142 if self.has_data_fields(m) {
143 println!("{message}{cont_prefix}{}", self.compact_fields(m, &data));
144 } else {
145 print!("{message}");
146 }
147 } else if self.has_data_fields(m) {
148 println!("{prefix}{}", self.compact_fields(m, &data));
149 }
150 }
151}