1use crate::tape::{
2 FieldValueOwned, Instruction, InstructionSet, SpanRecords, TapeMachine, ValueOwned,
3};
4use chrono::{DateTime, Utc};
5use nu_ansi_term::{Color, Style};
6use std::borrow::Cow;
7use std::fmt::Write;
8use std::num::NonZeroU64;
9use std::{collections::HashMap, io};
10use tracing::Level;
11
12pub struct Printer<W> {
13 out: W,
14 color: bool,
15 span: HashMap<NonZeroU64, SpanRecords>,
16 new_records: Option<(NonZeroU64, SpanRecords)>,
17 new_event: Option<NewEvent>,
18}
19impl<W> Printer<W>
20where
21 W: io::Write + Send + 'static,
22{
23 pub fn new(out: W, color: bool) -> Self {
24 Self {
25 out,
26 color,
27 span: Default::default(),
28 new_records: None,
29 new_event: None,
30 }
31 }
32
33 fn get_span(&self, span: NonZeroU64) -> Cow<SpanRecords> {
34 match self.span.get(&span) {
35 Some(span) => Cow::Borrowed(span),
36 None => Cow::Owned(SpanRecords::lost(span)),
37 }
38 }
39
40 fn take_span(&mut self, span: NonZeroU64) -> SpanRecords {
41 match self.span.remove(&span) {
42 Some(records) => records,
43 None => SpanRecords::lost(span),
44 }
45 }
46
47 fn span_iter<'a, F>(&'a self, span: NonZeroU64, f: &mut F)
48 where
49 F: FnMut(Cow<'a, SpanRecords>),
50 {
51 let records = self.get_span(span);
52 if let Some(parent) = records.parent {
53 self.span_iter(parent, f);
54 }
55 f(records);
56 }
57
58 fn span_from_root(&self, span: NonZeroU64) -> Vec<Cow<SpanRecords>> {
59 let mut r = Vec::new();
60 self.span_iter(span, &mut |records| {
61 r.push(records);
62 });
63 r
64 }
65}
66impl<W> TapeMachine<InstructionSet> for Printer<W>
67where
68 W: io::Write + Send + 'static,
69{
70 fn needs_restart(&mut self) -> bool {
71 false
72 }
73
74 fn handle(&mut self, instruction: Instruction) {
75 match instruction {
76 Instruction::Restart => {
77 self.new_event = None;
78 self.new_records = None;
79 }
80 Instruction::NewSpan { parent, span, name } => {
81 assert!(self.new_records.is_none());
82 self.new_records = Some((
83 span,
84 SpanRecords {
85 parent,
86 name: name.to_owned(),
87 records: Default::default(),
88 },
89 ));
90 }
91 Instruction::FinishedSpan | Instruction::FinishedRecord => {
92 let new = self.new_records.take().unwrap();
93 self.span.insert(new.0, new.1);
94 }
95 Instruction::NewRecord(id) => {
96 assert!(self.new_records.is_none());
97 self.new_records = Some((id, self.take_span(id)));
98 }
99 Instruction::StartEvent {
100 time,
101 span,
102 target,
103 priority,
104 } => {
105 assert!(self.new_event.is_none());
106 self.new_event = Some(NewEvent {
107 time,
108 span,
109 target: target.to_owned(),
110 priority,
111 records: Default::default(),
112 });
113 }
114 Instruction::FinishedEvent => {
115 let new_event = self.new_event.take().unwrap();
116 let spans = new_event
117 .span
118 .map(|span| self.span_from_root(span))
119 .unwrap_or_default();
120
121 let line = new_event.to_line(self.color, &spans);
122
123 let _ = self.out.write_all(line.as_bytes());
124 let _ = self.out.write_all(b"\n");
125 let _ = self.out.flush();
126 }
127 Instruction::AddValue(field_value) => {
128 match (&mut self.new_records, &mut self.new_event) {
129 (Some(new_records), None) => {
130 new_records.1.records.push(field_value.to_owned());
131 }
132 (None, Some(new_event)) => {
133 new_event.records.push(field_value.to_owned());
134 }
135 _ => panic!(),
136 }
137 }
138 Instruction::DeleteSpan(id) => {
139 self.span.remove(&id);
140 }
141 }
142 }
143}
144
145pub struct NewEvent {
146 pub time: DateTime<Utc>,
147 pub span: Option<NonZeroU64>,
148 pub target: String,
149 pub priority: Level,
150 pub records: Vec<FieldValueOwned>,
151}
152impl NewEvent {
153 pub fn to_line(&self, color: bool, spans: &[Cow<SpanRecords>]) -> String {
154 let mut line = String::new();
155 self.write_line(color, spans, &mut line);
156 line
157 }
158
159 pub fn write_line<W>(&self, color: bool, spans: &[Cow<SpanRecords>], line: &mut W)
160 where
161 W: Write,
162 {
163 let dimmed = color.then(|| Style::new().dimmed());
164 let bold = color.then(|| Style::new().bold());
165 let level_color = color.then(|| Self::level_style(self.priority));
166 let field_style = color.then(|| Style::new().italic());
167
168 Self::with_style(dimmed, line, |line| write!(line, "{:?}", self.time)).unwrap();
169 Self::with_style(level_color, line, |line| {
170 write!(line, " {}", Self::level_padded(self.priority))
171 })
172 .unwrap();
173
174 for (idx, span) in spans.iter().enumerate() {
175 if idx == 0 {
176 write!(line, " ").unwrap();
177 }
178
179 let name = &span.name;
180
181 Self::with_style(bold, line, |line| write!(line, "{name}{{")).unwrap();
182
183 for (idx, record) in span.records.iter().enumerate() {
184 if idx > 0 {
185 write!(line, " ").unwrap();
186 }
187 Self::write_record(record, field_style, false, line).unwrap();
188 }
189 write!(line, "}}").unwrap();
190 Self::with_style(dimmed, line, |line| write!(line, ":")).unwrap();
191 }
192
193 Self::with_style(dimmed, line, |line| write!(line, " {}:", self.target)).unwrap();
194
195 for record in self.records.iter() {
196 write!(line, " ").unwrap();
197 Self::write_record(record, field_style, true, line).unwrap();
198 }
199 }
200
201 fn level_style(level: Level) -> Style {
202 match level {
203 Level::TRACE => Color::Purple,
204 Level::DEBUG => Color::Blue,
205 Level::INFO => Color::Green,
206 Level::WARN => Color::Yellow,
207 Level::ERROR => Color::Red,
208 }
209 .normal()
210 }
211
212 fn level_padded(level: Level) -> &'static str {
213 match level {
214 Level::TRACE => "TRACE",
215 Level::DEBUG => "DEBUG",
216 Level::INFO => " INFO",
217 Level::WARN => " WARN",
218 Level::ERROR => "ERROR",
219 }
220 }
221
222 fn write_record<W>(
223 record: &FieldValueOwned,
224 field_style: Option<Style>,
225 with_message: bool,
226 out: &mut W,
227 ) -> std::fmt::Result
228 where
229 W: Write,
230 {
231 let name = &record.name;
232
233 if name == "message" && with_message {
234 if let ValueOwned::Debug(str) = &record.value {
235 return write!(out, "{}", str);
236 }
237 }
238
239 Self::with_style(field_style, out, |out| write!(out, "{name}"))?;
240
241 write!(out, "=")?;
242 Self::write_value(&record.value, out)
243 }
244
245 fn write_value<W>(value: &ValueOwned, out: &mut W) -> std::fmt::Result
246 where
247 W: Write,
248 {
249 match value {
250 ValueOwned::Debug(str) => write!(out, "{str}"),
251 ValueOwned::String(str) => write!(out, "{str:?}"),
252 ValueOwned::Float(value) => write!(out, "{value}"),
253 ValueOwned::Integer(value) => write!(out, "{value}"),
254 ValueOwned::Unsigned(value) => write!(out, "{value}"),
255 ValueOwned::Bool(value) => write!(out, "{value}"),
256 ValueOwned::ByteArray(items) => {
257 for &char in items.iter() {
258 write!(out, "{char:02x}")?;
259 }
260 Ok(())
261 }
262 }
263 }
264
265 fn with_style<W, F>(style: Option<Style>, out: &mut W, f: F) -> std::fmt::Result
266 where
267 W: Write,
268 F: FnOnce(&mut W) -> std::fmt::Result,
269 {
270 match style {
271 Some(style) => {
272 write!(out, "{}", style.prefix())?;
273 f(out)?;
274 write!(out, "{}", style.suffix())?;
275 Ok(())
276 }
277 None => f(out),
278 }
279 }
280}
281
282#[cfg(test)]
283pub mod tests {
284 use super::*;
285
286 #[test]
287 fn print_debug() {
288 let event = NewEvent {
289 time: Default::default(),
290 span: None,
291 target: "target".to_string(),
292 priority: Level::INFO,
293 records: vec![
294 FieldValueOwned {
295 name: "dbg".to_string(),
296 value: ValueOwned::Debug("thing".to_string()),
297 },
298 FieldValueOwned {
299 name: "str".to_string(),
300 value: ValueOwned::String("thing".to_string()),
301 },
302 ],
303 };
304
305 assert_eq!(
306 event.to_line(false, &[]),
307 r#"1970-01-01T00:00:00Z INFO target: dbg=thing str="thing""#
308 );
309 }
310
311 #[test]
312 fn log_levels_ident() {
313 for (priority, str) in [
314 (Level::ERROR, "ERROR"),
315 (Level::WARN, " WARN"),
316 (Level::INFO, " INFO"),
317 (Level::DEBUG, "DEBUG"),
318 (Level::TRACE, "TRACE"),
319 ] {
320 let event = NewEvent {
321 time: Default::default(),
322 span: None,
323 target: "target".to_string(),
324 priority,
325 records: Default::default(),
326 };
327
328 assert_eq!(
329 event.to_line(false, Default::default()),
330 format!("1970-01-01T00:00:00Z {str} target:")
331 )
332 }
333 }
334
335 #[test]
336 fn message_field_name_is_omitted() {
337 let event = NewEvent {
338 time: Default::default(),
339 span: None,
340 target: "target".to_string(),
341 priority: Level::INFO,
342 records: vec![FieldValueOwned {
343 name: "message".to_string(),
344 value: ValueOwned::Debug("a log".to_string()),
345 }],
346 };
347
348 assert_eq!(
349 event.to_line(false, Default::default()),
350 "1970-01-01T00:00:00Z INFO target: a log"
351 )
352 }
353
354 #[test]
355 fn span_print() {
356 let event = NewEvent {
357 time: Default::default(),
358 span: None,
359 target: "target".to_string(),
360 priority: Level::INFO,
361 records: Default::default(),
362 };
363
364 let spans = [
365 SpanRecords {
366 parent: None,
367 name: "record".to_string(),
368 records: vec![
369 FieldValueOwned {
370 name: "message".to_string(),
371 value: ValueOwned::String("a log".to_string()),
372 },
373 FieldValueOwned {
374 name: "a".to_string(),
375 value: ValueOwned::Debug("b".to_string()),
376 },
377 ],
378 },
379 SpanRecords {
380 parent: None,
381 name: "second".to_string(),
382 records: Default::default(),
383 },
384 ];
385 let spans = spans.iter().map(Cow::Borrowed).collect::<Vec<_>>();
386
387 assert_eq!(
388 event.to_line(false, &spans),
389 r#"1970-01-01T00:00:00Z INFO record{message="a log" a=b}:second{}: target:"#
390 );
391 }
392}