rocket/trace/subscriber/
common.rs

Help
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}
rocket/trace/subscriber/common.rs

rocket/trace/subscriber/
common.rs
