tracing_chrome/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use tracing_core::{field::Field, span, Event, Subscriber};
4use tracing_subscriber::{
5    layer::Context,
6    registry::{LookupSpan, SpanRef},
7    Layer,
8};
9
10use serde_json::{json, Value as JsonValue};
11use std::{
12    marker::PhantomData,
13    path::Path,
14    sync::{
15        atomic::{AtomicUsize, Ordering},
16        Arc, Mutex,
17    },
18};
19
20use std::io::{BufWriter, Write};
21use std::sync::mpsc;
22use std::sync::mpsc::Sender;
23use std::{
24    cell::{Cell, RefCell},
25    thread::JoinHandle,
26};
27
28thread_local! {
29    static OUT: RefCell<Option<Sender<Message>>> = const { RefCell::new(None) };
30    static TID: RefCell<Option<usize>> = const { RefCell::new(None) };
31}
32
33type NameFn<S> = Box<dyn Fn(&EventOrSpan<'_, '_, S>) -> String + Send + Sync>;
34type Object = serde_json::Map<String, JsonValue>;
35
36/// A [`Layer`](tracing_subscriber::Layer) that writes a Chrome trace file.
37pub struct ChromeLayer<S>
38where
39    S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
40{
41    out: Arc<Mutex<Sender<Message>>>,
42    start: std::time::Instant,
43    max_tid: AtomicUsize,
44    include_args: bool,
45    include_locations: bool,
46    trace_style: TraceStyle,
47    name_fn: Option<NameFn<S>>,
48    cat_fn: Option<NameFn<S>>,
49    _inner: PhantomData<S>,
50}
51
52/// A builder for [`ChromeLayer`](crate::ChromeLayer).
53#[derive(Default)]
54pub struct ChromeLayerBuilder<S>
55where
56    S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
57{
58    out_writer: Option<Box<dyn Write + Send>>,
59    name_fn: Option<NameFn<S>>,
60    cat_fn: Option<NameFn<S>>,
61    include_args: bool,
62    include_locations: bool,
63    trace_style: TraceStyle,
64    _inner: PhantomData<S>,
65}
66
67/// Decides how traces will be recorded.
68#[derive(Default)]
69pub enum TraceStyle {
70    /// Traces will be recorded as a group of threads.
71    /// In this style, spans should be entered and exited on the same thread.
72    #[default]
73    Threaded,
74
75    /// Traces will recorded as a group of asynchronous operations.
76    Async,
77}
78
79impl<S> ChromeLayerBuilder<S>
80where
81    S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
82{
83    pub fn new() -> Self {
84        ChromeLayerBuilder {
85            out_writer: None,
86            name_fn: None,
87            cat_fn: None,
88            include_args: false,
89            include_locations: true,
90            trace_style: TraceStyle::Threaded,
91            _inner: PhantomData,
92        }
93    }
94
95    /// Set the file to which to output the trace.
96    ///
97    /// Defaults to `./trace-{unix epoch in micros}.json`.
98    ///
99    /// # Panics
100    ///
101    /// If `file` could not be opened/created. To handle errors,
102    /// open a file and pass it to [`writer`](crate::ChromeLayerBuilder::writer) instead.
103    pub fn file<P: AsRef<Path>>(self, file: P) -> Self {
104        self.writer(std::fs::File::create(file).expect("Failed to create trace file."))
105    }
106
107    /// Supply an arbitrary writer to which to write trace contents.
108    ///
109    /// # Examples
110    ///
111    /// ```rust
112    /// # use tracing_chrome::ChromeLayerBuilder;
113    /// # use tracing_subscriber::prelude::*;
114    /// let (layer, guard) = ChromeLayerBuilder::new().writer(std::io::sink()).build();
115    /// # tracing_subscriber::registry().with(layer).init();
116    /// ```
117    pub fn writer<W: Write + Send + 'static>(mut self, writer: W) -> Self {
118        self.out_writer = Some(Box::new(writer));
119        self
120    }
121
122    /// Include arguments in each trace entry.
123    ///
124    /// Defaults to `false`.
125    ///
126    /// Includes the arguments used when creating a span/event
127    /// in the "args" section of the trace entry.
128    pub fn include_args(mut self, include: bool) -> Self {
129        self.include_args = include;
130        self
131    }
132
133    /// Include file+line with each trace entry.
134    ///
135    /// Defaults to `true`.
136    ///
137    /// This can add quite a bit of data to the output so turning
138    /// it off might be helpful when collecting larger traces.
139    pub fn include_locations(mut self, include: bool) -> Self {
140        self.include_locations = include;
141        self
142    }
143
144    /// Sets the style used when recording trace events.
145    ///
146    /// See [`TraceStyle`](crate::TraceStyle) for details.
147    pub fn trace_style(mut self, style: TraceStyle) -> Self {
148        self.trace_style = style;
149        self
150    }
151
152    /// Allows supplying a function that derives a name from
153    /// an Event or Span. The result is used as the "name" field
154    /// on trace entries.
155    ///
156    /// # Example
157    /// ```
158    /// use tracing_chrome::{ChromeLayerBuilder, EventOrSpan};
159    /// use tracing_subscriber::{registry::Registry, prelude::*};
160    ///
161    /// let (chrome_layer, _guard) = ChromeLayerBuilder::new().name_fn(Box::new(|event_or_span| {
162    ///     match event_or_span {
163    ///         EventOrSpan::Event(ev) => { ev.metadata().name().into() },
164    ///         EventOrSpan::Span(_s) => { "span".into() },
165    ///     }
166    /// })).build();
167    /// tracing_subscriber::registry().with(chrome_layer).init()
168    /// ```
169    pub fn name_fn(mut self, name_fn: NameFn<S>) -> Self {
170        self.name_fn = Some(name_fn);
171        self
172    }
173
174    /// Allows supplying a function that derives a category from
175    /// an Event or Span. The result is used as the "cat" field on
176    /// trace entries.
177    ///
178    /// # Example
179    /// ```
180    /// use tracing_chrome::{ChromeLayerBuilder, EventOrSpan};
181    /// use tracing_subscriber::{registry::Registry, prelude::*};
182    ///
183    /// let (chrome_layer, _guard) = ChromeLayerBuilder::new().category_fn(Box::new(|_| {
184    ///     "my_module".into()
185    /// })).build();
186    /// tracing_subscriber::registry().with(chrome_layer).init()
187    /// ```
188    pub fn category_fn(mut self, cat_fn: NameFn<S>) -> Self {
189        self.cat_fn = Some(cat_fn);
190        self
191    }
192
193    /// Creates a [`ChromeLayer`](crate::ChromeLayer) and associated [`FlushGuard`](crate::FlushGuard).
194    ///
195    /// # Panics
196    ///
197    /// If no file or writer was specified and the default trace file could not be opened/created.
198    pub fn build(self) -> (ChromeLayer<S>, FlushGuard) {
199        ChromeLayer::new(self)
200    }
201}
202
203/// This guard will signal the thread writing the trace file to stop and join it when dropped.
204pub struct FlushGuard {
205    sender: Sender<Message>,
206    handle: Cell<Option<JoinHandle<()>>>,
207}
208
209impl FlushGuard {
210    /// Signals the trace writing thread to flush to disk.
211    pub fn flush(&self) {
212        if let Some(handle) = self.handle.take() {
213            let _ignored = self.sender.send(Message::Flush);
214            self.handle.set(Some(handle));
215        }
216    }
217
218    /// Finishes the current trace and starts a new one.
219    ///
220    /// If a [`Write`](std::io::Write) implementation is supplied,
221    /// the new trace is written to it. Otherwise, the new trace
222    /// goes to `./trace-{unix epoc in micros}.json`.
223    pub fn start_new(&self, writer: Option<Box<dyn Write + Send>>) {
224        if let Some(handle) = self.handle.take() {
225            let _ignored = self.sender.send(Message::StartNew(writer));
226            self.handle.set(Some(handle));
227        }
228    }
229}
230
231impl Drop for FlushGuard {
232    fn drop(&mut self) {
233        if let Some(handle) = self.handle.take() {
234            let _ignored = self.sender.send(Message::Drop);
235            if handle.join().is_err() {
236                eprintln!("tracing_chrome: Trace writing thread panicked.");
237            }
238        }
239    }
240}
241
242struct Callsite {
243    tid: usize,
244    name: String,
245    target: String,
246    file: Option<&'static str>,
247    line: Option<u32>,
248    args: Option<Arc<Object>>,
249}
250
251enum Message {
252    Enter(f64, Callsite, Option<u64>),
253    Event(f64, Callsite),
254    Exit(f64, Callsite, Option<u64>),
255    NewThread(usize, String),
256    Flush,
257    Drop,
258    StartNew(Option<Box<dyn Write + Send>>),
259}
260
261/// Represents either an [`Event`](tracing_core::Event) or [`SpanRef`](tracing_subscriber::registry::SpanRef).
262pub enum EventOrSpan<'a, 'b, S>
263where
264    S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
265{
266    Event(&'a Event<'b>),
267    Span(&'a SpanRef<'b, S>),
268}
269
270fn create_default_writer() -> Box<dyn Write + Send> {
271    Box::new(
272        std::fs::File::create(format!(
273            "./trace-{}.json",
274            std::time::SystemTime::UNIX_EPOCH
275                .elapsed()
276                .unwrap()
277                .as_micros()
278        ))
279        .expect("Failed to create trace file."),
280    )
281}
282
283impl<S> ChromeLayer<S>
284where
285    S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
286{
287    fn new(mut builder: ChromeLayerBuilder<S>) -> (ChromeLayer<S>, FlushGuard) {
288        let (tx, rx) = mpsc::channel();
289        OUT.with(|val| val.replace(Some(tx.clone())));
290
291        let out_writer = builder
292            .out_writer
293            .unwrap_or_else(|| create_default_writer());
294
295        let handle = std::thread::spawn(move || {
296            let mut write = BufWriter::new(out_writer);
297            write.write_all(b"[\n").unwrap();
298
299            let mut has_started = false;
300            let mut thread_names: Vec<(usize, String)> = Vec::new();
301            for msg in rx {
302                if let Message::Flush = &msg {
303                    write.flush().unwrap();
304                    continue;
305                } else if let Message::Drop = &msg {
306                    break;
307                } else if let Message::StartNew(writer) = msg {
308                    // Finish off current file
309                    write.write_all(b"\n]").unwrap();
310                    write.flush().unwrap();
311
312                    // Get or create new writer
313                    let out_writer = writer.unwrap_or_else(|| create_default_writer());
314                    write = BufWriter::new(out_writer);
315                    write.write_all(b"[\n").unwrap();
316                    has_started = false;
317
318                    // Write saved thread names
319                    for (tid, name) in thread_names.iter() {
320                        let entry = json!({
321                            "ph": "M",
322                            "pid": 1,
323                            "name": "thread_name",
324                            "tid": *tid,
325                            "args": {
326                                "name": name,
327                            },
328                        });
329
330                        if has_started {
331                            write.write_all(b",\n").unwrap();
332                        }
333                        serde_json::to_writer(&mut write, &entry).unwrap();
334                        has_started = true;
335                    }
336                    continue;
337                }
338
339                let (ph, ts, callsite, id) = match &msg {
340                    Message::Enter(ts, callsite, None) => ("B", Some(ts), Some(callsite), None),
341                    Message::Enter(ts, callsite, Some(root_id)) => {
342                        ("b", Some(ts), Some(callsite), Some(root_id))
343                    }
344                    Message::Event(ts, callsite) => ("i", Some(ts), Some(callsite), None),
345                    Message::Exit(ts, callsite, None) => ("E", Some(ts), Some(callsite), None),
346                    Message::Exit(ts, callsite, Some(root_id)) => {
347                        ("e", Some(ts), Some(callsite), Some(root_id))
348                    }
349                    Message::NewThread(_tid, _name) => ("M", None, None, None),
350                    Message::Flush | Message::Drop | Message::StartNew(_) => {
351                        panic!("Was supposed to break by now.")
352                    }
353                };
354                let mut entry = json!({
355                    "ph": ph,
356                    "pid": 1,
357                });
358
359                if let Message::NewThread(tid, name) = msg {
360                    thread_names.push((tid, name.clone()));
361                    entry["name"] = "thread_name".into();
362                    entry["tid"] = tid.into();
363                    entry["args"] = json!({ "name": name });
364                } else {
365                    let ts = ts.unwrap();
366                    let callsite = callsite.unwrap();
367                    entry["ts"] = (*ts).into();
368                    entry["name"] = callsite.name.clone().into();
369                    entry["cat"] = callsite.target.clone().into();
370                    entry["tid"] = callsite.tid.into();
371
372                    if let Some(&id) = id {
373                        entry["id"] = id.into();
374                    }
375
376                    if ph == "i" {
377                        entry["s"] = "t".into();
378                    }
379
380                    if let (Some(file), Some(line)) = (callsite.file, callsite.line) {
381                        entry[".file"] = file.into();
382                        entry[".line"] = line.into();
383                    }
384
385                    if let Some(call_args) = &callsite.args {
386                        if !call_args.is_empty() {
387                            entry["args"] = (**call_args).clone().into();
388                        }
389                    }
390                }
391
392                if has_started {
393                    write.write_all(b",\n").unwrap();
394                }
395                serde_json::to_writer(&mut write, &entry).unwrap();
396                has_started = true;
397            }
398
399            write.write_all(b"\n]").unwrap();
400            write.flush().unwrap();
401        });
402
403        let guard = FlushGuard {
404            sender: tx.clone(),
405            handle: Cell::new(Some(handle)),
406        };
407        let layer = ChromeLayer {
408            out: Arc::new(Mutex::new(tx)),
409            start: std::time::Instant::now(),
410            max_tid: AtomicUsize::new(0),
411            name_fn: builder.name_fn.take(),
412            cat_fn: builder.cat_fn.take(),
413            include_args: builder.include_args,
414            include_locations: builder.include_locations,
415            trace_style: builder.trace_style,
416            _inner: PhantomData,
417        };
418
419        (layer, guard)
420    }
421
422    fn get_tid(&self) -> (usize, bool) {
423        TID.with(|value| {
424            let tid = *value.borrow();
425            match tid {
426                Some(tid) => (tid, false),
427                None => {
428                    let tid = self.max_tid.fetch_add(1, Ordering::SeqCst);
429                    value.replace(Some(tid));
430                    (tid, true)
431                }
432            }
433        })
434    }
435
436    fn get_callsite(&self, data: EventOrSpan<S>) -> Callsite {
437        let (tid, new_thread) = self.get_tid();
438        let name = self.name_fn.as_ref().map(|name_fn| name_fn(&data));
439        let target = self.cat_fn.as_ref().map(|cat_fn| cat_fn(&data));
440        let meta = match data {
441            EventOrSpan::Event(e) => e.metadata(),
442            EventOrSpan::Span(s) => s.metadata(),
443        };
444        let args = match data {
445            EventOrSpan::Event(e) => {
446                if self.include_args {
447                    let mut args = Object::new();
448                    e.record(&mut JsonVisitor { object: &mut args });
449                    Some(Arc::new(args))
450                } else {
451                    None
452                }
453            }
454            EventOrSpan::Span(s) => s
455                .extensions()
456                .get::<ArgsWrapper>()
457                .map(|e| &e.args)
458                .cloned(),
459        };
460        let name = name.unwrap_or_else(|| meta.name().into());
461        let target = target.unwrap_or_else(|| meta.target().into());
462        let (file, line) = if self.include_locations {
463            (meta.file(), meta.line())
464        } else {
465            (None, None)
466        };
467
468        if new_thread {
469            let name = match std::thread::current().name() {
470                Some(name) => name.to_owned(),
471                None => tid.to_string(),
472            };
473            self.send_message(Message::NewThread(tid, name));
474        }
475
476        Callsite {
477            tid,
478            name,
479            target,
480            file,
481            line,
482            args,
483        }
484    }
485
486    fn get_root_id(span: SpanRef<S>) -> u64 {
487        span.scope()
488            .from_root()
489            .take(1)
490            .next()
491            .unwrap_or(span)
492            .id()
493            .into_u64()
494    }
495
496    fn enter_span(&self, span: SpanRef<S>, ts: f64) {
497        let callsite = self.get_callsite(EventOrSpan::Span(&span));
498        let root_id = match self.trace_style {
499            TraceStyle::Async => Some(ChromeLayer::get_root_id(span)),
500            _ => None,
501        };
502        self.send_message(Message::Enter(ts, callsite, root_id));
503    }
504
505    fn exit_span(&self, span: SpanRef<S>, ts: f64) {
506        let callsite = self.get_callsite(EventOrSpan::Span(&span));
507        let root_id = match self.trace_style {
508            TraceStyle::Async => Some(ChromeLayer::get_root_id(span)),
509            _ => None,
510        };
511        self.send_message(Message::Exit(ts, callsite, root_id));
512    }
513
514    fn get_ts(&self) -> f64 {
515        self.start.elapsed().as_nanos() as f64 / 1000.0
516    }
517
518    fn send_message(&self, message: Message) {
519        OUT.with(move |val| {
520            if val.borrow().is_some() {
521                let _ignored = val.borrow().as_ref().unwrap().send(message);
522            } else {
523                let out = self.out.lock().unwrap().clone();
524                let _ignored = out.send(message);
525                val.replace(Some(out));
526            }
527        });
528    }
529}
530
531impl<S> Layer<S> for ChromeLayer<S>
532where
533    S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
534{
535    fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
536        if let TraceStyle::Async = self.trace_style {
537            return;
538        }
539
540        let ts = self.get_ts();
541        self.enter_span(ctx.span(id).expect("Span not found."), ts);
542    }
543
544    fn on_record(&self, id: &span::Id, values: &span::Record<'_>, ctx: Context<'_, S>) {
545        if self.include_args {
546            let span = ctx.span(id).unwrap();
547            let mut exts = span.extensions_mut();
548
549            let args = exts.get_mut::<ArgsWrapper>();
550
551            if let Some(args) = args {
552                let args = Arc::make_mut(&mut args.args);
553                values.record(&mut JsonVisitor { object: args });
554            }
555        }
556    }
557
558    fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
559        let ts = self.get_ts();
560        let callsite = self.get_callsite(EventOrSpan::Event(event));
561        self.send_message(Message::Event(ts, callsite));
562    }
563
564    fn on_exit(&self, id: &span::Id, ctx: Context<'_, S>) {
565        if let TraceStyle::Async = self.trace_style {
566            return;
567        }
568        let ts = self.get_ts();
569        self.exit_span(ctx.span(id).expect("Span not found."), ts);
570    }
571
572    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: Context<'_, S>) {
573        if self.include_args {
574            let mut args = Object::new();
575            attrs.record(&mut JsonVisitor { object: &mut args });
576            ctx.span(id).unwrap().extensions_mut().insert(ArgsWrapper {
577                args: Arc::new(args),
578            });
579        }
580        if let TraceStyle::Threaded = self.trace_style {
581            return;
582        }
583
584        let ts = self.get_ts();
585        self.enter_span(ctx.span(id).expect("Span not found."), ts);
586    }
587
588    fn on_close(&self, id: span::Id, ctx: Context<'_, S>) {
589        if let TraceStyle::Threaded = self.trace_style {
590            return;
591        }
592
593        let ts = self.get_ts();
594        self.exit_span(ctx.span(&id).expect("Span not found."), ts);
595    }
596}
597
598struct JsonVisitor<'a> {
599    object: &'a mut Object,
600}
601
602impl<'a> tracing_subscriber::field::Visit for JsonVisitor<'a> {
603    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
604        self.object
605            .insert(field.name().to_owned(), format!("{value:?}").into());
606    }
607}
608
609struct ArgsWrapper {
610    args: Arc<Object>,
611}
tracing_chrome/lib.rs

tracing_chrome/
lib.rs
