refactor code

This commit is contained in:
hal8174 2025-02-24 22:14:36 +01:00
parent 3a4822a7ab
commit a045570596

View file

@ -1,6 +1,7 @@
use std::{
collections::HashMap,
sync::{atomic::AtomicU64, Mutex},
io::Write,
sync::{Mutex, atomic::AtomicU64},
thread::ThreadId,
time::Instant,
};
@ -11,7 +12,7 @@ mod perfetto {
use perfetto::DebugAnnotation;
use prost::Message;
use tracing::{field::Visit, span, Dispatch, Level, Metadata, Subscriber};
use tracing::{Dispatch, Level, Metadata, Subscriber, field::Visit, span};
pub fn install_perfetto_subscriber() {
let subscriber = PerfettoSubscriber::new();
@ -80,6 +81,26 @@ impl InnerPerfettoSubscriber {
}
}
}
fn push(&mut self, packet: perfetto::TracePacket) {
self.buffer.push(packet);
self.flush();
}
fn flush(&mut self) {
let trace = perfetto::Trace {
packet: std::mem::take(&mut self.buffer),
};
let buf = trace.encode_to_vec();
let mut file = std::fs::File::options()
.append(true)
.open("out.trace")
.unwrap();
let _ = file.write(&buf);
}
}
struct DebugAnnotationVisitor {
@ -179,6 +200,7 @@ impl Visit for DebugAnnotationVisitor {
impl PerfettoSubscriber {
pub fn new() -> Self {
std::fs::File::create("out.trace").unwrap();
let mut packets = Vec::new();
packets.push(perfetto::TracePacket {
timestamp: None,
@ -206,17 +228,6 @@ impl PerfettoSubscriber {
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 {
@ -273,7 +284,7 @@ impl Subscriber for PerfettoSubscriber {
event.record(&mut debug_annotations);
let track_id = inner.get_track_id();
inner.buffer.push(perfetto::TracePacket {
inner.push(perfetto::TracePacket {
timestamp: Some(time),
data: Some(perfetto::trace_packet::Data::TrackEvent(
perfetto::TrackEvent {
@ -292,7 +303,6 @@ impl Subscriber for PerfettoSubscriber {
),
});
}
self.flush();
}
fn enter(&self, span: &span::Id) {
@ -306,7 +316,7 @@ impl Subscriber for PerfettoSubscriber {
let name = inner.spans[&span.into_u64()].name.clone();
let track_id = inner.get_track_id();
inner.buffer.push(perfetto::TracePacket {
inner.push(perfetto::TracePacket {
timestamp: Some(time),
data: Some(perfetto::trace_packet::Data::TrackEvent(
perfetto::TrackEvent {
@ -325,7 +335,6 @@ impl Subscriber for PerfettoSubscriber {
),
});
}
self.flush();
}
fn exit(&self, span: &span::Id) {
@ -339,7 +348,7 @@ impl Subscriber for PerfettoSubscriber {
let name = inner.spans[&span.into_u64()].name.clone();
let track_id = inner.get_track_id();
inner.buffer.push(perfetto::TracePacket {
inner.push(perfetto::TracePacket {
timestamp: Some(time),
data: Some(perfetto::trace_packet::Data::TrackEvent(
perfetto::TrackEvent {
@ -358,6 +367,5 @@ impl Subscriber for PerfettoSubscriber {
),
});
}
self.flush();
}
}