tracing-perfetto/src/lib.rs

363 lines
12 KiB
Rust

use std::{
collections::HashMap,
sync::{atomic::AtomicU64, Mutex},
thread::ThreadId,
time::Instant,
};
mod perfetto {
include!(concat!(env!("OUT_DIR"), "/perfetto.trace.rs"));
}
use perfetto::DebugAnnotation;
use prost::Message;
use tracing::{field::Visit, span, Dispatch, Level, Metadata, Subscriber};
pub fn install_perfetto_subscriber() {
let subscriber = PerfettoSubscriber::new();
let dispatcher = Dispatch::new(subscriber);
tracing::dispatcher::set_global_default(dispatcher).unwrap();
}
pub struct PerfettoSubscriber {
id_counter: AtomicU64,
start_time: Instant,
inner: Mutex<InnerPerfettoSubscriber>,
}
struct SpanInfo {
debug_annotations: Vec<DebugAnnotation>,
level: Level,
name: String,
}
struct InnerPerfettoSubscriber {
buffer: Vec<perfetto::TracePacket>,
spans: HashMap<u64, SpanInfo>,
threads: HashMap<ThreadId, u64>,
}
impl InnerPerfettoSubscriber {
fn new(buffer: Vec<perfetto::TracePacket>) -> Self {
Self {
buffer,
spans: HashMap::new(),
threads: HashMap::new(),
}
}
fn get_track_id(&mut self) -> u64 {
let current = std::thread::current();
let nextfree = self.threads.len() as u64;
match self.threads.entry(current.id()) {
std::collections::hash_map::Entry::Occupied(occupied_entry) => *occupied_entry.get(),
std::collections::hash_map::Entry::Vacant(vacant_entry) => {
vacant_entry.insert(nextfree);
self.buffer.push(perfetto::TracePacket {
timestamp: None,
data: Some(perfetto::trace_packet::Data::TrackDescriptor(
perfetto::TrackDescriptor {
uuid: Some(nextfree),
parent_uuid: None,
static_or_dynamic_name: None,
process: None,
thread: Some(perfetto::ThreadDescriptor {
pid: Some(1234),
thread_name: current.name().map(|s| s.to_owned()),
tid: Some(nextfree as i32),
}),
},
)),
optional_trusted_packet_sequence_id: Some(
perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
3903809,
),
),
});
nextfree
}
}
}
}
struct DebugAnnotationVisitor {
data: Vec<DebugAnnotation>,
}
impl DebugAnnotationVisitor {
fn new_with_metadata(metadata: &'static Metadata<'static>) -> Self {
Self {
data: vec![
DebugAnnotation {
name_field: Some(perfetto::debug_annotation::NameField::Name(String::from(
"module_path",
))),
value: metadata
.module_path()
.map(|m| perfetto::debug_annotation::Value::StringValue(m.to_owned())),
},
DebugAnnotation {
name_field: Some(perfetto::debug_annotation::NameField::Name(String::from(
"location",
))),
value: Some(perfetto::debug_annotation::Value::StringValue(format!(
"{}: {}",
metadata.file().unwrap_or(""),
metadata.line().unwrap_or(0)
))),
},
],
}
}
fn new() -> Self {
Self { data: Vec::new() }
}
}
impl Visit for DebugAnnotationVisitor {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
self.data.push(DebugAnnotation {
name_field: Some(perfetto::debug_annotation::NameField::Name(
field.name().to_owned(),
)),
value: Some(perfetto::debug_annotation::Value::StringValue(format!(
"{value:?}"
))),
});
}
fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
self.data.push(DebugAnnotation {
name_field: Some(perfetto::debug_annotation::NameField::Name(
field.name().to_owned(),
)),
value: Some(perfetto::debug_annotation::Value::IntValue(value)),
});
}
fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
self.data.push(DebugAnnotation {
name_field: Some(perfetto::debug_annotation::NameField::Name(
field.name().to_owned(),
)),
value: Some(perfetto::debug_annotation::Value::UintValue(value)),
});
}
fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
self.data.push(DebugAnnotation {
name_field: Some(perfetto::debug_annotation::NameField::Name(
field.name().to_owned(),
)),
value: Some(perfetto::debug_annotation::Value::DoubleValue(value)),
});
}
fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
self.data.push(DebugAnnotation {
name_field: Some(perfetto::debug_annotation::NameField::Name(
field.name().to_owned(),
)),
value: Some(perfetto::debug_annotation::Value::BoolValue(value)),
});
}
fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
self.data.push(DebugAnnotation {
name_field: Some(perfetto::debug_annotation::NameField::Name(
field.name().to_owned(),
)),
value: Some(perfetto::debug_annotation::Value::StringValue(
value.to_owned(),
)),
});
}
}
impl PerfettoSubscriber {
pub fn new() -> Self {
let mut packets = Vec::new();
packets.push(perfetto::TracePacket {
timestamp: None,
data: Some(perfetto::trace_packet::Data::TrackDescriptor(
perfetto::TrackDescriptor {
uuid: Some(894893984),
parent_uuid: None,
static_or_dynamic_name: None,
process: Some(perfetto::ProcessDescriptor {
pid: Some(1234),
process_name: Some(String::from("process")),
}),
thread: None,
},
)),
optional_trusted_packet_sequence_id: Some(
perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
3903809,
),
),
});
Self {
id_counter: AtomicU64::new(1),
inner: Mutex::new(InnerPerfettoSubscriber::new(packets)),
start_time: Instant::now(),
}
}
fn flush(&self) {
let data = self.inner.lock().unwrap();
let trace = perfetto::Trace {
packet: data.buffer.clone(),
};
let buf = trace.encode_to_vec();
let _ = std::fs::write("out.trace", buf);
}
}
impl Subscriber for PerfettoSubscriber {
fn enabled(&self, metadata: &tracing::Metadata<'_>) -> bool {
true
}
fn new_span(&self, span: &span::Attributes<'_>) -> span::Id {
// println!("new_span {:?}", span);
let id = self
.id_counter
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let mut debug_annotations = DebugAnnotationVisitor::new_with_metadata(span.metadata());
span.record(&mut debug_annotations);
let mut inner = self.inner.lock().unwrap();
inner.spans.insert(
id,
SpanInfo {
debug_annotations: debug_annotations.data,
level: *span.metadata().level(),
name: span.metadata().name().to_owned(),
},
);
span::Id::from_u64(id)
}
fn record(&self, span: &span::Id, values: &span::Record<'_>) {
// println!("record {:?} {:?}", span, values);
let mut debug_annotations = DebugAnnotationVisitor::new();
values.record(&mut debug_annotations);
let mut inner = self.inner.lock().unwrap();
inner
.spans
.get_mut(&span.into_u64())
.unwrap()
.debug_annotations
.append(&mut debug_annotations.data);
}
fn record_follows_from(&self, span: &span::Id, follows: &span::Id) {
// println!("record_follows_from {:?} {:?}", span, follows);
}
fn event(&self, event: &tracing::Event<'_>) {
// println!("event {:?}", event);
{
let time = self.start_time.elapsed().as_nanos() as u64;
let mut inner = self.inner.lock().unwrap();
let mut debug_annotations = DebugAnnotationVisitor::new_with_metadata(event.metadata());
event.record(&mut debug_annotations);
let track_id = inner.get_track_id();
inner.buffer.push(perfetto::TracePacket {
timestamp: Some(time),
data: Some(perfetto::trace_packet::Data::TrackEvent(
perfetto::TrackEvent {
track_uuid: Some(track_id),
categories: vec![event.metadata().level().as_str().to_owned()],
category_iids: vec![],
name_field: Some(perfetto::track_event::NameField::Name(String::from(event.metadata().name()))),
r#type: Some(perfetto::track_event::Type::Instant as i32),
debug_annotations: debug_annotations.data,
},
)),
optional_trusted_packet_sequence_id: Some(
perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
3903809,
),
),
});
}
self.flush();
}
fn enter(&self, span: &span::Id) {
// println!("enter {:?}", span);
{
let time = self.start_time.elapsed().as_nanos() as u64;
let mut inner = self.inner.lock().unwrap();
let debug_annotations = inner.spans[&span.into_u64()].debug_annotations.clone();
let level = inner.spans[&span.into_u64()].level;
let name = inner.spans[&span.into_u64()].name.clone();
let track_id = inner.get_track_id();
inner.buffer.push(perfetto::TracePacket {
timestamp: Some(time),
data: Some(perfetto::trace_packet::Data::TrackEvent(
perfetto::TrackEvent {
track_uuid: Some(track_id),
categories: vec![level.as_str().to_owned()],
category_iids: vec![],
name_field: Some(perfetto::track_event::NameField::Name(name)),
r#type: Some(perfetto::track_event::Type::SliceBegin as i32),
debug_annotations,
},
)),
optional_trusted_packet_sequence_id: Some(
perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
3903809,
),
),
});
}
self.flush();
}
fn exit(&self, span: &span::Id) {
// println!("exit {:?}", span);
{
let time = self.start_time.elapsed().as_nanos() as u64;
let mut inner = self.inner.lock().unwrap();
let debug_annotations = inner.spans[&span.into_u64()].debug_annotations.clone();
let level = inner.spans[&span.into_u64()].level;
let name = inner.spans[&span.into_u64()].name.clone();
let track_id = inner.get_track_id();
inner.buffer.push(perfetto::TracePacket {
timestamp: Some(time),
data: Some(perfetto::trace_packet::Data::TrackEvent(
perfetto::TrackEvent {
track_uuid: Some(track_id),
categories: vec![level.as_str().to_owned()],
category_iids: vec![],
name_field: Some(perfetto::track_event::NameField::Name(name)),
r#type: Some(perfetto::track_event::Type::SliceEnd as i32),
debug_annotations,
},
)),
optional_trusted_packet_sequence_id: Some(
perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
3903809,
),
),
});
}
self.flush();
}
}