1use crate::CURRENT_THREAD_ID;
2use parking_lot::Mutex;
3use std::{
4 borrow::Cow,
5 fmt,
6 io::{self, Write as _},
7 path::Path,
8 sync::Arc,
9 time::Instant,
10};
11use tracing::{
12 field::{Field, Visit},
13 span, Event, Metadata, Subscriber,
14};
15use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};
16
17#[derive(Debug, Copy, Clone, Eq, PartialEq)]
18enum EventType {
19 Begin,
20 Event,
21 End,
22}
23
24#[derive(Debug, Clone)]
30pub struct ChromeTracingLayer {
31 file: Arc<Mutex<std::fs::File>>,
32 start_time: Instant,
33 process_id: u32,
34}
35
36impl ChromeTracingLayer {
37 pub fn with_file(file: impl AsRef<Path>) -> io::Result<Self> {
39 std::fs::File::create(file).map(|mut file| {
40 writeln!(file, "[").unwrap();
41 ChromeTracingLayer {
42 file: Arc::new(Mutex::new(file)),
43 start_time: Instant::now(),
44 process_id: std::process::id(),
45 }
46 })
47 }
48
49 fn write_event(
50 &self,
51 mut fields: Option<EventVisitor>,
52 metadata: &'static Metadata<'static>,
53 event_type: EventType,
54 ) {
55 if let Some(EventVisitor { trace: false, .. }) = fields {
56 return;
57 }
58
59 let current_time = Instant::now();
60
61 let diff = current_time - self.start_time;
62 let diff_in_us = diff.as_micros();
63
64 let event_type_str = match event_type {
65 EventType::Begin => "B",
66 EventType::Event => "i",
67 EventType::End => "E",
68 };
69
70 let instant_scope = if event_type == EventType::Event {
71 r#","s": "p""#
72 } else {
73 ""
74 };
75
76 let name = match event_type {
77 EventType::Event => fields
78 .as_mut()
79 .and_then(|fields| fields.message.take())
80 .map_or_else(
81 || {
82 let name = metadata.name();
83 if cfg!(target_os = "windows") {
85 Cow::Owned(name.replace("\\", "\\\\"))
86 } else {
87 Cow::Borrowed(name)
88 }
89 },
90 Cow::from,
91 ),
92 EventType::Begin | EventType::End => Cow::Borrowed(metadata.name()),
93 };
94
95 let category = fields
96 .as_mut()
97 .and_then(|fields| fields.category.take())
98 .map_or_else(|| Cow::Borrowed("trace"), Cow::from);
99
100 let mut file = self.file.lock();
101 writeln!(
102 file,
103 r#"{{ "name": "{}", "cat": "{}", "ph": "{}", "ts": {}, "pid": {}, "tid": {} {} }},"#,
104 name,
105 category,
106 event_type_str,
107 diff_in_us,
108 self.process_id,
109 CURRENT_THREAD_ID.with(|v| *v),
110 instant_scope,
111 )
112 .unwrap();
113 }
114}
115
116impl Drop for ChromeTracingLayer {
117 fn drop(&mut self) {
118 let mut file = self.file.lock();
119 writeln!(file, "]").unwrap();
120 file.flush().unwrap();
121 }
122}
123
124#[derive(Debug, Default)]
125struct EventVisitor {
126 message: Option<String>,
127 category: Option<String>,
128 trace: bool,
129}
130
131impl Visit for EventVisitor {
132 fn record_bool(&mut self, field: &Field, value: bool) {
133 match field.name() {
134 "trace" => self.trace = value,
135 _ => {}
136 }
137 }
138
139 fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
140 match field.name() {
141 "message" => self.message = Some(format!("{:?}", value)),
142 "category" => self.category = Some(format!("{:?}", value)),
143 _ => {}
144 }
145 }
146}
147
148impl<S> Layer<S> for ChromeTracingLayer
149where
150 S: Subscriber + for<'span> LookupSpan<'span>,
151{
152 fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
153 let mut event_visitor = EventVisitor::default();
154 event.record(&mut event_visitor);
155
156 self.write_event(Some(event_visitor), event.metadata(), EventType::Event);
157 }
158
159 fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
160 let span = ctx.span(id).unwrap();
161 self.write_event(None, span.metadata(), EventType::Begin);
162 }
163
164 fn on_exit(&self, id: &span::Id, ctx: Context<'_, S>) {
165 if std::thread::panicking() {
166 return;
167 }
168
169 let span = ctx.span(id).unwrap();
170 self.write_event(None, span.metadata(), EventType::End);
171 }
172}