rocket_community/trace/subscriber/
common.rs

1use std::cell::Cell;
2use std::fmt;
3
4use tracing::field::Field;
5use tracing::{Level, Metadata};
6use tracing_subscriber::field::RecordFields;
7use tracing_subscriber::filter;
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>(
80        &self,
81        init_prefix: &'a dyn fmt::Display,
82        cont_prefix: &'a dyn fmt::Display,
83        meta: &'a Metadata<'_>,
84        data: F,
85    ) -> impl fmt::Display + 'a {
86        let style = self.style(meta);
87        Formatter(move |f| {
88            let fields = meta.fields();
89            let message = fields.field("message");
90            if let Some(message_field) = &message {
91                data.record_display(|field: &Field, value: &dyn fmt::Display| {
92                    if field != message_field {
93                        return;
94                    }
95
96                    for (i, line) in value.to_string().lines().enumerate() {
97                        let line = line.paint(style);
98                        if i == 0 {
99                            let _ = writeln!(f, "{init_prefix}{line}");
100                        } else {
101                            let _ = writeln!(f, "{cont_prefix}{line}");
102                        }
103                    }
104                });
105            }
106
107            Ok(())
108        })
109    }
110
111    pub(crate) fn compact_fields<'a, F: RecordFields + 'a>(
112        &self,
113        meta: &'a Metadata<'_>,
114        data: F,
115    ) -> impl fmt::Display + 'a {
116        let key_style = self.style(meta).bold();
117        let val_style = self.style(meta).primary();
118
119        Formatter(move |f| {
120            let mut printed = false;
121            data.record_display(|field: &Field, val: &dyn fmt::Display| {
122                let key = field.name();
123                if key != "message" {
124                    if printed {
125                        let _ = write!(f, " ");
126                    }
127                    let _ = write!(f, "{}: {}", key.paint(key_style), val.paint(val_style));
128                    printed = true;
129                }
130            });
131
132            Ok(())
133        })
134    }
135
136    pub(crate) fn print<F: RecordFields>(
137        &self,
138        prefix: &dyn fmt::Display,
139        cont_prefix: &dyn fmt::Display,
140        m: &Metadata<'_>,
141        data: F,
142    ) {
143        if self.has_message(m) {
144            let message = self.message(prefix, cont_prefix, m, &data);
145            if self.has_data_fields(m) {
146                println!("{message}{cont_prefix}{}", self.compact_fields(m, &data));
147            } else {
148                print!("{message}");
149            }
150        } else if self.has_data_fields(m) {
151            println!("{prefix}{}", self.compact_fields(m, &data));
152        }
153    }
154}
rocket_community/trace/subscriber/common.rs

rocket_community/trace/subscriber/
common.rs
