wgpu_subscriber/
chrome.rs

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/// A layer to add to a [`tracing_subscriber::Registry`] to output to a chrome
25/// trace.
26///
27/// If you want an easy "set and forget" method of installing this and normal
28/// tracing logging, call [`initialize_default_subscriber`].
29#[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    /// Create a trace which outputs to the given file. The file will be cleared if it exits.
38    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                        // The default name for events has a path in it, filter paths only on windows
84                        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}