1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172

use crate::CURRENT_THREAD_ID;
use parking_lot::Mutex;
use std::{
    borrow::Cow,
    fmt,
    io::{self, Write as _},
    path::Path,
    sync::Arc,
    time::Instant,
};
use tracing::{
    field::{Field, Visit},
    span, Event, Metadata, Subscriber,
};
use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum EventType {
    Begin,
    Event,
    End,
}

/// A layer to add to a [`tracing_subscriber::Registry`] to output to a chrome
/// trace.
///
/// If you want an easy "set and forget" method of installing this and normal
/// tracing logging, call [`initialize_default_subscriber`].
#[derive(Debug, Clone)]
pub struct ChromeTracingLayer {
    file: Arc<Mutex<std::fs::File>>,
    start_time: Instant,
    process_id: u32,
}

impl ChromeTracingLayer {
    /// Create a trace which outputs to the given file. The file will be cleared if it exits.
    pub fn with_file(file: impl AsRef<Path>) -> io::Result<Self> {
        std::fs::File::create(file).map(|mut file| {
            writeln!(file, "[").unwrap();
            ChromeTracingLayer {
                file: Arc::new(Mutex::new(file)),
                start_time: Instant::now(),
                process_id: std::process::id(),
            }
        })
    }

    fn write_event(
        &self,
        mut fields: Option<EventVisitor>,
        metadata: &'static Metadata<'static>,
        event_type: EventType,
    ) {
        if let Some(EventVisitor { trace: false, .. }) = fields {
            return;
        }

        let current_time = Instant::now();

        let diff = current_time - self.start_time;
        let diff_in_us = diff.as_micros();

        let event_type_str = match event_type {
            EventType::Begin => "B",
            EventType::Event => "i",
            EventType::End => "E",
        };

        let instant_scope = if event_type == EventType::Event {
            r#","s": "p""#
        } else {
            ""
        };

        let name = match event_type {
            EventType::Event => fields
                .as_mut()
                .and_then(|fields| fields.message.take())
                .map_or_else(
                    || {
                        let name = metadata.name();
                        // The default name for events has a path in it, filter paths only on windows
                        if cfg!(target_os = "windows") {
                            Cow::Owned(name.replace("\\", "\\\\"))
                        } else {
                            Cow::Borrowed(name)
                        }
                    },
                    Cow::from,
                ),
            EventType::Begin | EventType::End => Cow::Borrowed(metadata.name()),
        };

        let category = fields
            .as_mut()
            .and_then(|fields| fields.category.take())
            .map_or_else(|| Cow::Borrowed("trace"), Cow::from);

        let mut file = self.file.lock();
        writeln!(
            file,
            r#"{{ "name": "{}", "cat": "{}", "ph": "{}", "ts": {}, "pid": {}, "tid": {} {} }},"#,
            name,
            category,
            event_type_str,
            diff_in_us,
            self.process_id,
            CURRENT_THREAD_ID.with(|v| *v),
            instant_scope,
        )
        .unwrap();
    }
}

impl Drop for ChromeTracingLayer {
    fn drop(&mut self) {
        let mut file = self.file.lock();
        writeln!(file, "]").unwrap();
        file.flush().unwrap();
    }
}

#[derive(Debug, Default)]
struct EventVisitor {
    message: Option<String>,
    category: Option<String>,
    trace: bool,
}

impl Visit for EventVisitor {
    fn record_bool(&mut self, field: &Field, value: bool) {
        match field.name() {
            "trace" => self.trace = value,
            _ => {}
        }
    }

    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
        match field.name() {
            "message" => self.message = Some(format!("{:?}", value)),
            "category" => self.category = Some(format!("{:?}", value)),
            _ => {}
        }
    }
}

impl<S> Layer<S> for ChromeTracingLayer
where
    S: Subscriber + for<'span> LookupSpan<'span>,
{
    fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
        let mut event_visitor = EventVisitor::default();
        event.record(&mut event_visitor);

        self.write_event(Some(event_visitor), event.metadata(), EventType::Event);
    }

    fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
        let span = ctx.span(id).unwrap();
        self.write_event(None, span.metadata(), EventType::Begin);
    }

    fn on_exit(&self, id: &span::Id, ctx: Context<'_, S>) {
        if std::thread::panicking() {
            return;
        }

        let span = ctx.span(id).unwrap();
        self.write_event(None, span.metadata(), EventType::End);
    }
}