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