rocket_community/trace/subscriber/
compact.rs

1use std::fmt;
2use std::num::NonZeroU64;
3use std::time::Instant;
4
5use tracing::span::{Attributes, Id, Record};
6use tracing::{Event, Level, Metadata, Subscriber};
7use tracing_subscriber::field::RecordFields;
8use tracing_subscriber::layer::{Context, Layer};
9use tracing_subscriber::registry::LookupSpan;
10
11use time::OffsetDateTime;
12use yansi::{Paint, Painted};
13
14use super::RecordDisplay;
15use crate::http::{Status, StatusClass};
16use crate::trace::subscriber::{Data, RocketFmt};
17use crate::util::Formatter;
18
19#[derive(Debug, Default, Copy, Clone)]
20pub struct Compact {
21    /// The `tracing::Span::Id` of the request we're in, if any.
22    request: Option<NonZeroU64>,
23}
24
25#[derive(Debug)]
26pub struct RequestData {
27    start: Instant,
28    fields: Data,
29    item: Option<(String, String)>,
30}
31
32impl RequestData {
33    pub fn new<T: RecordFields>(attrs: T) -> Self {
34        Self {
35            start: Instant::now(),
36            fields: Data::new(attrs),
37            item: None,
38        }
39    }
40}
41
42impl RocketFmt<Compact> {
43    fn request_span_id(&self) -> Option<Id> {
44        self.state().request.map(Id::from_non_zero_u64)
45    }
46
47    fn timestamp_for(&self, datetime: OffsetDateTime) -> impl fmt::Display {
48        Formatter(move |f| {
49            let (date, time) = (datetime.date(), datetime.time());
50            let (year, month, day) = (date.year(), date.month() as u8, date.day());
51            let (h, m, s, l) = (
52                time.hour(),
53                time.minute(),
54                time.second(),
55                time.millisecond(),
56            );
57            write!(
58                f,
59                "{year:04}-{month:02}-{day:02}T{h:02}:{m:02}:{s:02}.{l:03}Z"
60            )
61        })
62    }
63
64    fn in_debug(&self) -> bool {
65        self.level.is_some_and(|l| l >= Level::DEBUG)
66    }
67
68    fn prefix<'a>(&self, meta: &'a Metadata<'_>) -> impl fmt::Display + 'a {
69        let style = self.style(meta);
70        let name = if meta.name().starts_with("event ") {
71            meta.target()
72        } else {
73            meta.name()
74        };
75
76        let pad = self.level.map_or(0, |lvl| lvl.as_str().len());
77        let timestamp = self.timestamp_for(OffsetDateTime::now_utc());
78        Formatter(move |f| {
79            write!(
80                f,
81                "{} {:>pad$} {} ",
82                timestamp.paint(style).primary().dim(),
83                meta.level().paint(style),
84                name.paint(style).primary()
85            )
86        })
87    }
88
89    fn chevron(&self, meta: &Metadata<'_>) -> Painted<&'static str> {
90        "›".paint(self.style(meta)).bold()
91    }
92
93    fn print_compact<F: RecordFields>(&self, m: &Metadata<'_>, data: F) {
94        let style = self.style(m);
95        let prefix = self.prefix(m);
96        let chevron = self.chevron(m);
97        let init_prefix = Formatter(|f| write!(f, "{prefix}{chevron} "));
98        let cont_prefix = Formatter(|f| write!(f, "{prefix}{} ", "+".paint(style).dim()));
99        self.print(&init_prefix, &cont_prefix, m, data);
100    }
101}
102
103impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for RocketFmt<Compact> {
104    fn enabled(&self, metadata: &Metadata<'_>, _: Context<'_, S>) -> bool {
105        self.filter
106            .would_enable(metadata.target(), metadata.level())
107            && (self.in_debug()
108                || self.request_span_id().is_none()
109                || metadata.name() == "request"
110                || metadata.name() == "response")
111    }
112
113    fn on_event(&self, event: &Event<'_>, ctxt: Context<'_, S>) {
114        if let Some(id) = self.request_span_id() {
115            let name = event.metadata().name();
116            if name == "response" {
117                let req_span = ctxt.span(&id).expect("on_event: req does not exist");
118                let mut exts = req_span.extensions_mut();
119                let data = exts.get_mut::<RequestData>().unwrap();
120                event.record(&mut data.fields);
121            } else if name == "catcher" || name == "route" {
122                let req_span = ctxt.span(&id).expect("on_event: req does not exist");
123                let mut exts = req_span.extensions_mut();
124                let data = exts.get_mut::<RequestData>().unwrap();
125                data.item = event.find_map_display("name", |v| (name.into(), v.to_string()))
126            }
127
128            if !self.in_debug() {
129                return;
130            }
131        }
132
133        self.print_compact(event.metadata(), event);
134    }
135
136    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctxt: Context<'_, S>) {
137        let span = ctxt.span(id).expect("new_span: span does not exist");
138
139        if span.name() == "request" {
140            let data = RequestData::new(attrs);
141            span.extensions_mut().replace(data);
142
143            if !self.in_debug() {
144                return;
145            }
146        }
147
148        if self.state().request.is_none() {
149            self.print_compact(span.metadata(), attrs);
150        }
151    }
152
153    fn on_record(&self, id: &Id, values: &Record<'_>, ctxt: Context<'_, S>) {
154        let span = ctxt.span(id).expect("record: span does not exist");
155        if self.request_span_id().as_ref() == Some(id) {
156            let mut extensions = span.extensions_mut();
157            match extensions.get_mut::<RequestData>() {
158                Some(data) => values.record(&mut data.fields),
159                None => span.extensions_mut().insert(RequestData::new(values)),
160            }
161        }
162
163        if self.in_debug() {
164            println!(
165                "{}{} {}",
166                self.prefix(span.metadata()),
167                self.chevron(span.metadata()),
168                self.compact_fields(span.metadata(), values)
169            );
170        }
171    }
172
173    fn on_enter(&self, id: &Id, ctxt: Context<'_, S>) {
174        let span = ctxt.span(id).expect("new_span: span does not exist");
175        if span.name() == "request" {
176            self.update_state(|state| state.request = Some(id.into_non_zero_u64()));
177        }
178    }
179
180    fn on_exit(&self, id: &Id, ctxt: Context<'_, S>) {
181        let span = ctxt.span(id).expect("new_span: span does not exist");
182        if span.name() == "request" {
183            self.update_state(|state| state.request = None);
184        }
185    }
186
187    fn on_close(&self, id: Id, ctxt: Context<'_, S>) {
188        let span = ctxt.span(&id).expect("new_span: span does not exist");
189        if span.name() == "request" {
190            let extensions = span.extensions();
191            let data = extensions.get::<RequestData>().unwrap();
192
193            let elapsed = data.start.elapsed();
194            let datetime = OffsetDateTime::now_utc() - elapsed;
195            let timestamp = self.timestamp_for(datetime);
196
197            let s = self.style(span.metadata());
198            let prefix = self.prefix(span.metadata());
199            let chevron = self.chevron(span.metadata());
200            let arrow = "→".paint(s.primary().bright());
201
202            let status_class = data.fields["status"]
203                .parse()
204                .ok()
205                .and_then(Status::from_code)
206                .map(|status| status.class());
207
208            let status_style = match status_class {
209                Some(StatusClass::Informational) => s,
210                Some(StatusClass::Success) => s.green(),
211                Some(StatusClass::Redirection) => s.magenta(),
212                Some(StatusClass::ClientError) => s.yellow(),
213                Some(StatusClass::ServerError) => s.red(),
214                Some(StatusClass::Unknown) => s.cyan(),
215                None => s.primary(),
216            };
217
218            let autohandle = Formatter(|f| match data.fields.get("autohandled") {
219                Some("true") => write!(f, " {} {}", "via".paint(s.dim()), "GET".paint(s)),
220                _ => Ok(()),
221            });
222
223            let item = Formatter(|f| match &data.item {
224                Some((kind, name)) => {
225                    write!(f, "{} {} {arrow} ", kind.paint(s), name.paint(s.bold()),)
226                }
227                None => Ok(()),
228            });
229
230            println!(
231                "{prefix}{chevron} ({} {}ms) {}{autohandle} {} {arrow} {item}{}",
232                timestamp.paint(s).primary().dim(),
233                elapsed.as_millis(),
234                &data.fields["method"].paint(s),
235                &data.fields["uri"],
236                &data.fields["status"].paint(status_style),
237            );
238        }
239    }
240}