From a04557059620436d4509ab51d0544d1d5b141796 Mon Sep 17 00:00:00 2001
From: hal8174 <henning.lindemann@pm.me>
Date: Mon, 24 Feb 2025 22:14:36 +0100
Subject: [PATCH] refactor code

---
 src/lib.rs | 46 +++++++++++++++++++++++++++-------------------
 1 file changed, 27 insertions(+), 19 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0562ac2..b36af3a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -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();
     }
 }