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