1use std::{
2 collections::VecDeque,
3 sync::{
4 Mutex,
5 atomic::{AtomicBool, Ordering},
6 },
7 time::{Duration, Instant},
8};
9
10#[derive(Clone, Debug, serde::Serialize)]
13pub struct ExecutionTrace {
14 pub events: Vec<TraceEvent>,
16 pub dropped: u64,
18}
19#[derive(Clone, Debug, serde::Serialize)]
21pub struct TraceEvent {
22 pub start_us: f64,
24 pub duration_us: f64,
26 pub lane: usize,
28 pub copied_bytes: usize,
30 pub failed: bool,
32}
33impl ExecutionTrace {
34 pub fn to_json(&self) -> crate::Result<String> {
36 let events = self
37 .events
38 .iter()
39 .map(|event| {
40 serde_json::json!({
41 "name": "native region (host-observed)", "cat": "hrx.host", "ph": "X",
42 "ts": event.start_us, "dur": event.duration_us, "pid": 0, "tid": event.lane,
43 "args": {"copied_bytes": event.copied_bytes, "failed": event.failed},
44 })
45 })
46 .collect::<Vec<_>>();
47 serde_json::to_string(&serde_json::json!({
48 "traceEvents": events, "displayTimeUnit": "ms", "dropped_events": self.dropped,
49 "measurement": "host-observed synchronized region latency; not device timestamps",
50 }))
51 .map_err(|error| crate::Error::Message(error.to_string()))
52 }
53}
54struct Capture {
55 epoch: Instant,
56 capacity: usize,
57 events: VecDeque<TraceEvent>,
58 dropped: u64,
59}
60#[derive(Default)]
61pub(super) struct Tracer {
62 enabled: AtomicBool,
63 capture: Mutex<Option<Capture>>,
64}
65impl Tracer {
66 pub fn start(&self, capacity: usize) -> crate::Result<()> {
67 if capacity == 0 {
68 return Err(crate::Error::Message(
69 "trace capacity must be nonzero".into(),
70 ));
71 }
72 let mut capture = self.capture.lock().unwrap_or_else(|e| e.into_inner());
73 if capture.is_some() {
74 return Err(crate::Error::Busy("a trace is already running".into()));
75 }
76 *capture = Some(Capture {
77 epoch: Instant::now(),
78 capacity,
79 events: VecDeque::new(),
80 dropped: 0,
81 });
82 self.enabled.store(true, Ordering::Release);
83 Ok(())
84 }
85 pub fn finish(&self) -> Option<ExecutionTrace> {
86 let mut capture = self.capture.lock().unwrap_or_else(|e| e.into_inner());
87 self.enabled.store(false, Ordering::Release);
88 capture.take().map(|capture| ExecutionTrace {
89 events: capture.events.into(),
90 dropped: capture.dropped,
91 })
92 }
93 pub fn record(
94 &self,
95 start: Instant,
96 elapsed: Duration,
97 lane: usize,
98 copied_bytes: usize,
99 failed: bool,
100 ) {
101 if !self.enabled.load(Ordering::Acquire) {
102 return;
103 }
104 let mut capture = self.capture.lock().unwrap_or_else(|e| e.into_inner());
105 let Some(capture) = capture.as_mut() else {
106 return;
107 };
108 let Some(offset) = start.checked_duration_since(capture.epoch) else {
109 return;
110 };
111 if capture.events.len() == capture.capacity {
112 capture.events.pop_front();
113 capture.dropped += 1;
114 }
115 capture.events.push_back(TraceEvent {
116 start_us: offset.as_secs_f64() * 1e6,
117 duration_us: elapsed.as_secs_f64() * 1e6,
118 lane,
119 copied_bytes,
120 failed,
121 });
122 }
123}