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
36pub 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#[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#[derive(Default)]
69pub enum TraceStyle {
70 #[default]
73 Threaded,
74
75 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 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 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 pub fn include_args(mut self, include: bool) -> Self {
129 self.include_args = include;
130 self
131 }
132
133 pub fn include_locations(mut self, include: bool) -> Self {
140 self.include_locations = include;
141 self
142 }
143
144 pub fn trace_style(mut self, style: TraceStyle) -> Self {
148 self.trace_style = style;
149 self
150 }
151
152 pub fn name_fn(mut self, name_fn: NameFn<S>) -> Self {
170 self.name_fn = Some(name_fn);
171 self
172 }
173
174 pub fn category_fn(mut self, cat_fn: NameFn<S>) -> Self {
189 self.cat_fn = Some(cat_fn);
190 self
191 }
192
193 pub fn build(self) -> (ChromeLayer<S>, FlushGuard) {
199 ChromeLayer::new(self)
200 }
201}
202
203pub struct FlushGuard {
205 sender: Sender<Message>,
206 handle: Cell<Option<JoinHandle<()>>>,
207}
208
209impl FlushGuard {
210 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 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
261pub 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 write.write_all(b"\n]").unwrap();
310 write.flush().unwrap();
311
312 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 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}