tracing_perfetto/
lib.rs

1#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/README.md"))]
2#![forbid(unsafe_code)]
3
4use bytes::BytesMut;
5use prost::Message;
6use std::io::Write;
7use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
8use tracing::field::Field;
9use tracing::field::Visit;
10use tracing::span;
11use tracing::Event;
12use tracing::Id;
13use tracing::Subscriber;
14use tracing_subscriber::fmt::MakeWriter;
15use tracing_subscriber::layer::Context;
16use tracing_subscriber::registry::LookupSpan;
17use tracing_subscriber::Layer;
18
19#[path = "perfetto.protos.rs"]
20#[allow(clippy::all)]
21#[rustfmt::skip]
22mod idl;
23
24thread_local! {
25    static THREAD_TRACK_UUID: AtomicU64 = AtomicU64::new(rand::random::<u64>());
26    static THREAD_DESCRIPTOR_SENT: AtomicBool = const { AtomicBool::new(false) };
27}
28
29// This is thread safe, since duplicated descriptor will be combined into one by perfetto.
30static PROCESS_DESCRIPTOR_SENT: AtomicBool = AtomicBool::new(false);
31
32/// A `Layer` that records span as perfetto's
33/// `TYPE_SLICE_BEGIN`/`TYPE_SLICE_END`, and event as `TYPE_INSTANT`.
34///
35/// `PerfettoLayer` will output the records as encoded [protobuf messages](https://github.com/google/perfetto).
36pub struct PerfettoLayer<W = fn() -> std::io::Stdout> {
37    sequence_id: SequenceId,
38    track_uuid: TrackUuid,
39    writer: W,
40    config: Config,
41}
42
43/// Writes encoded records into provided instance.
44///
45/// This is implemented for types implements [`MakeWriter`].
46pub trait PerfettoWriter {
47    fn write_log(&self, buf: BytesMut) -> std::io::Result<()>;
48}
49
50impl<W: for<'writer> MakeWriter<'writer> + 'static> PerfettoWriter for W {
51    fn write_log(&self, buf: BytesMut) -> std::io::Result<()> {
52        self.make_writer().write_all(&buf)
53    }
54}
55
56#[derive(Default)]
57struct Config {
58    debug_annotations: bool,
59    filter: Option<fn(&str) -> bool>,
60}
61
62impl<W: PerfettoWriter> PerfettoLayer<W> {
63    pub fn new(writer: W) -> Self {
64        Self {
65            sequence_id: SequenceId::new(rand::random()),
66            track_uuid: TrackUuid::new(rand::random()),
67            writer,
68            config: Config::default(),
69        }
70    }
71
72    /// Configures whether or not spans/events shoulde be recored with their metadata and fields.
73    pub fn with_debug_annotations(mut self, value: bool) -> Self {
74        self.config.debug_annotations = value;
75        self
76    }
77
78    /// Configures whether or not spans/events be recored based on the occurrence of a field name.
79    ///
80    /// Sometimes, not all the events/spans should be treated as perfetto trace, you can append a
81    /// field to indicate that this even/span should be captured into trace:
82    ///
83    /// ```rust
84    /// use tracing_perfetto::PerfettoLayer;
85    /// use tracing_subscriber::{layer::SubscriberExt, Registry, prelude::*};
86    ///
87    /// let layer = PerfettoLayer::new(std::fs::File::open("/tmp/test.pftrace").unwrap())
88    ///                 .with_filter_by_marker(|field_name| field_name == "perfetto");
89    /// tracing_subscriber::registry().with(layer).init();
90    ///
91    /// // this event will be record, as it contains a `perfetto` field
92    /// tracing::info!(perfetto = true, my_bool = true);
93    ///
94    /// // this span will be record, as it contains a `perfetto` field
95    /// #[tracing::instrument(fields(perfetto = true))]
96    /// fn to_instr() {
97    ///
98    ///   // this event will be ignored
99    ///   tracing::info!(my_bool = true);
100    /// }
101    /// ```
102    pub fn with_filter_by_marker(mut self, filter: fn(&str) -> bool) -> Self {
103        self.config.filter = Some(filter);
104        self
105    }
106
107    fn thread_descriptor(&self) -> Option<idl::TracePacket> {
108        let thread_first_frame_sent =
109            THREAD_DESCRIPTOR_SENT.with(|v| v.fetch_or(true, Ordering::SeqCst));
110        if thread_first_frame_sent {
111            return None;
112        }
113        let thread_track_uuid = THREAD_TRACK_UUID.with(|id| id.load(Ordering::Relaxed));
114        let mut packet = idl::TracePacket::default();
115        let thread = create_thread_descriptor().into();
116        let track_desc = create_track_descriptor(
117            thread_track_uuid.into(),
118            None,
119            std::thread::current().name(),
120            None,
121            thread,
122            None,
123        );
124        packet.data = Some(idl::trace_packet::Data::TrackDescriptor(track_desc));
125        Some(packet)
126    }
127
128    fn process_descriptor(&self) -> Option<idl::TracePacket> {
129        let process_first_frame_sent = PROCESS_DESCRIPTOR_SENT.fetch_or(true, Ordering::SeqCst);
130        if process_first_frame_sent {
131            return None;
132        }
133        let mut packet = idl::TracePacket::default();
134        let process = create_process_descriptor().into();
135        let track_desc = create_track_descriptor(
136            self.track_uuid.get().into(),
137            None,
138            None::<&str>,
139            process,
140            None,
141            None,
142        );
143        packet.data = Some(idl::trace_packet::Data::TrackDescriptor(track_desc));
144        Some(packet)
145    }
146
147    fn write_log(&self, mut log: idl::Trace) {
148        let mut buf = BytesMut::new();
149
150        if let Some(p) = self.process_descriptor() {
151            log.packet.insert(0, p);
152        }
153        if let Some(t) = self.thread_descriptor() {
154            log.packet.insert(1, t);
155        }
156
157        let Ok(_) = log.encode(&mut buf) else {
158            return;
159        };
160        _ = self.writer.write_log(buf);
161    }
162}
163
164struct SequenceId(u64);
165
166impl SequenceId {
167    fn new(n: u64) -> Self {
168        Self(n)
169    }
170
171    fn get(&self) -> u64 {
172        self.0
173    }
174}
175
176struct TrackUuid(u64);
177
178impl TrackUuid {
179    fn new(n: u64) -> Self {
180        Self(n)
181    }
182
183    fn get(&self) -> u64 {
184        self.0
185    }
186}
187
188struct PerfettoVisitor {
189    perfetto: bool,
190    filter: fn(&str) -> bool,
191}
192
193impl PerfettoVisitor {
194    fn new(filter: fn(&str) -> bool) -> PerfettoVisitor {
195        Self {
196            filter,
197            perfetto: false,
198        }
199    }
200}
201
202impl Visit for PerfettoVisitor {
203    fn record_debug(&mut self, field: &Field, _value: &dyn std::fmt::Debug) {
204        if (self.filter)(field.name()) {
205            self.perfetto = true;
206        }
207    }
208}
209
210impl<W, S: Subscriber> Layer<S> for PerfettoLayer<W>
211where
212    S: for<'a> LookupSpan<'a>,
213    W: for<'writer> MakeWriter<'writer> + 'static,
214{
215    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: Context<'_, S>) {
216        let Some(span) = ctx.span(id) else {
217            return;
218        };
219
220        let enabled = self
221            .config
222            .filter
223            .map(|f| {
224                let mut visitor = PerfettoVisitor::new(f);
225                attrs.record(&mut visitor);
226                visitor.perfetto
227            })
228            .unwrap_or(true);
229
230        if !enabled {
231            return;
232        }
233
234        let mut debug_annotations = DebugAnnotations::default();
235        if self.config.debug_annotations {
236            attrs.record(&mut debug_annotations);
237        }
238
239        let mut packet = idl::TracePacket::default();
240        let thread_track_uuid = THREAD_TRACK_UUID.with(|id| id.load(Ordering::Relaxed));
241        let event = create_event(
242            thread_track_uuid,
243            Some(span.name()),
244            span.metadata().file().zip(span.metadata().line()),
245            debug_annotations,
246            Some(idl::track_event::Type::SliceBegin),
247        );
248        packet.data = Some(idl::trace_packet::Data::TrackEvent(event));
249        packet.timestamp = chrono::Local::now().timestamp_nanos_opt().map(|t| t as _);
250        packet.trusted_pid = Some(std::process::id() as _);
251        packet.optional_trusted_packet_sequence_id = Some(
252            idl::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
253                self.sequence_id.get() as _,
254            ),
255        );
256        span.extensions_mut().insert(idl::Trace {
257            packet: vec![packet],
258        });
259    }
260
261    fn on_record(&self, span: &span::Id, values: &span::Record<'_>, ctx: Context<'_, S>) {
262        let Some(span) = ctx.span(span) else {
263            return;
264        };
265        if let Some(f) = self.config.filter {
266            if !PerfettoVisitor::new(f).perfetto {
267                return;
268            }
269        }
270        if let Some(extension) = span.extensions_mut().get_mut::<idl::Trace>() {
271            if let Some(idl::trace_packet::Data::TrackEvent(ref mut event)) =
272                &mut extension.packet[0].data
273            {
274                let mut debug_annotations = DebugAnnotations::default();
275                values.record(&mut debug_annotations);
276                event
277                    .debug_annotations
278                    .append(&mut debug_annotations.annotations);
279            }
280        };
281    }
282
283    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
284        let enabled = self
285            .config
286            .filter
287            .map(|f| {
288                let mut visitor = PerfettoVisitor::new(f);
289                event.record(&mut visitor);
290                visitor.perfetto
291            })
292            .unwrap_or(true);
293
294        if !enabled {
295            return;
296        }
297
298        let metadata = event.metadata();
299        let location = metadata.file().zip(metadata.line());
300
301        let mut debug_annotations = DebugAnnotations::default();
302
303        if self.config.debug_annotations {
304            event.record(&mut debug_annotations);
305        }
306
307        let track_event = THREAD_TRACK_UUID.with(|id| {
308            create_event(
309                id.load(Ordering::Relaxed),
310                Some(metadata.name()),
311                location,
312                debug_annotations,
313                Some(idl::track_event::Type::Instant),
314            )
315        });
316        let mut packet = idl::TracePacket::default();
317        packet.data = Some(idl::trace_packet::Data::TrackEvent(track_event));
318        packet.trusted_pid = Some(std::process::id() as _);
319        packet.timestamp = chrono::Local::now().timestamp_nanos_opt().map(|t| t as _);
320        packet.optional_trusted_packet_sequence_id = Some(
321            idl::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
322                self.sequence_id.get() as _,
323            ),
324        );
325
326        if let Some(span) = ctx.event_span(event) {
327            if let Some(trace) = span.extensions_mut().get_mut::<idl::Trace>() {
328                trace.packet.push(packet);
329                return;
330            }
331        }
332        let trace = idl::Trace {
333            packet: vec![packet],
334        };
335        self.write_log(trace);
336    }
337
338    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
339        let Some(span) = ctx.span(&id) else {
340            return;
341        };
342
343        let Some(mut trace) = span.extensions_mut().remove::<idl::Trace>() else {
344            return;
345        };
346
347        let debug_annotations = DebugAnnotations::default();
348
349        let mut packet = idl::TracePacket::default();
350        let meta = span.metadata();
351        let event = THREAD_TRACK_UUID.with(|id| {
352            create_event(
353                id.load(Ordering::Relaxed),
354                Some(meta.name()),
355                meta.file().zip(meta.line()),
356                debug_annotations,
357                Some(idl::track_event::Type::SliceEnd),
358            )
359        });
360        packet.data = Some(idl::trace_packet::Data::TrackEvent(event));
361        packet.timestamp = chrono::Local::now().timestamp_nanos_opt().map(|t| t as _);
362        packet.trusted_pid = Some(std::process::id() as _);
363        packet.optional_trusted_packet_sequence_id = Some(
364            idl::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
365                self.sequence_id.get() as _,
366            ),
367        );
368        trace.packet.push(packet);
369
370        self.write_log(trace);
371    }
372}
373
374fn create_thread_descriptor() -> idl::ThreadDescriptor {
375    let mut thread = idl::ThreadDescriptor::default();
376    thread.pid = Some(std::process::id() as _);
377    thread.tid = Some(thread_id::get() as _);
378    thread.thread_name = std::thread::current().name().map(|n| n.to_string());
379    thread
380}
381
382fn create_process_descriptor() -> idl::ProcessDescriptor {
383    let mut process = idl::ProcessDescriptor::default();
384    process.pid = Some(std::process::id() as _);
385    process
386}
387
388fn create_track_descriptor(
389    uuid: Option<u64>,
390    parent_uuid: Option<u64>,
391    name: Option<impl AsRef<str>>,
392    process: Option<idl::ProcessDescriptor>,
393    thread: Option<idl::ThreadDescriptor>,
394    counter: Option<idl::CounterDescriptor>,
395) -> idl::TrackDescriptor {
396    let mut desc = idl::TrackDescriptor::default();
397    desc.uuid = uuid;
398    desc.parent_uuid = parent_uuid;
399    desc.static_or_dynamic_name = name
400        .map(|s| s.as_ref().to_string())
401        .map(idl::track_descriptor::StaticOrDynamicName::Name);
402    desc.process = process;
403    desc.thread = thread;
404    desc.counter = counter;
405    desc
406}
407
408fn create_event(
409    track_uuid: u64,
410    name: Option<&str>,
411    location: Option<(&str, u32)>,
412    debug_annotations: DebugAnnotations,
413    r#type: Option<idl::track_event::Type>,
414) -> idl::TrackEvent {
415    let mut event = idl::TrackEvent::default();
416    event.track_uuid = Some(track_uuid);
417    event.categories = vec!["".to_string()];
418    if let Some(name) = name {
419        event.name_field = Some(idl::track_event::NameField::Name(name.to_string()));
420    }
421    if let Some(t) = r#type {
422        event.set_type(t);
423    }
424    if !debug_annotations.annotations.is_empty() {
425        event.debug_annotations = debug_annotations.annotations;
426    }
427    if let Some((file, line)) = location {
428        let mut source_location = idl::SourceLocation::default();
429        source_location.file_name = Some(file.to_owned());
430        source_location.line_number = Some(line);
431        let location = idl::track_event::SourceLocationField::SourceLocation(source_location);
432        event.source_location_field = Some(location);
433    }
434    event
435}
436
437#[derive(Default)]
438struct DebugAnnotations {
439    annotations: Vec<idl::DebugAnnotation>,
440}
441
442macro_rules! impl_record {
443    ($method:ident, $type:ty, $value_variant:ident) => {
444        fn $method(&mut self, field: &Field, value: $type) {
445            let mut annotation = idl::DebugAnnotation::default();
446            annotation.name_field = Some(idl::debug_annotation::NameField::Name(
447                field.name().to_string(),
448            ));
449            annotation.value = Some(idl::debug_annotation::Value::$value_variant(value.into()));
450            self.annotations.push(annotation);
451        }
452    };
453    ($method:ident, $type:ty, $value_variant:ident, $conversion:expr) => {
454        fn $method(&mut self, field: &Field, value: $type) {
455            let mut annotation = idl::DebugAnnotation::default();
456            annotation.name_field = Some(idl::debug_annotation::NameField::Name(
457                field.name().to_string(),
458            ));
459            annotation.value = Some(idl::debug_annotation::Value::$value_variant($conversion(
460                value,
461            )));
462            self.annotations.push(annotation);
463        }
464    };
465}
466
467impl Visit for DebugAnnotations {
468    impl_record!(record_bool, bool, BoolValue);
469    impl_record!(record_str, &str, StringValue, String::from);
470    impl_record!(record_f64, f64, DoubleValue);
471    impl_record!(record_i64, i64, IntValue);
472    impl_record!(record_i128, i128, StringValue, |v: i128| v.to_string());
473    impl_record!(record_u128, u128, StringValue, |v: u128| v.to_string());
474    impl_record!(record_u64, u64, IntValue, |v: u64| v as i64);
475
476    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
477        let mut annotation = idl::DebugAnnotation::default();
478        annotation.name_field = Some(idl::debug_annotation::NameField::Name(
479            field.name().to_string(),
480        ));
481        annotation.value = Some(idl::debug_annotation::Value::StringValue(format!(
482            "{value:?}"
483        )));
484        self.annotations.push(annotation);
485    }
486
487    fn record_error(&mut self, field: &Field, value: &(dyn std::error::Error + 'static)) {
488        let mut annotation = idl::DebugAnnotation::default();
489        annotation.name_field = Some(idl::debug_annotation::NameField::Name(
490            field.name().to_string(),
491        ));
492        annotation.value = Some(idl::debug_annotation::Value::StringValue(format!(
493            "{value}"
494        )));
495        self.annotations.push(annotation);
496    }
497}
tracing_perfetto/lib.rs

tracing_perfetto/
lib.rs
