1#![allow(dead_code)]
2use std::collections::HashMap;
9use std::time::{Duration, SystemTime};
10
11#[derive(Debug, Clone, PartialEq, Eq, Hash)]
13pub struct SpanId(String);
14
15impl SpanId {
16 pub fn new(id: impl Into<String>) -> Self {
18 Self(id.into())
19 }
20
21 pub fn as_str(&self) -> &str {
23 &self.0
24 }
25}
26
27impl std::fmt::Display for SpanId {
28 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29 write!(f, "{}", self.0)
30 }
31}
32
33#[derive(Debug, Clone, PartialEq)]
35pub enum SpanStatus {
36 Ok,
38 Warning(String),
40 Error(String),
42 Cancelled,
44 InProgress,
46}
47
48#[derive(Debug, Clone, PartialEq)]
50pub struct SpanAttribute {
51 pub key: String,
53 pub value: String,
55}
56
57impl SpanAttribute {
58 pub fn new(key: impl Into<String>, value: impl Into<String>) -> Self {
60 Self {
61 key: key.into(),
62 value: value.into(),
63 }
64 }
65}
66
67#[derive(Debug, Clone)]
69pub struct TraceSpan {
70 pub id: SpanId,
72 pub parent: Option<SpanId>,
74 pub name: String,
76 pub start_time: SystemTime,
78 pub end_time: Option<SystemTime>,
80 pub duration: Option<Duration>,
82 pub status: SpanStatus,
84 pub attributes: Vec<SpanAttribute>,
86}
87
88impl TraceSpan {
89 pub fn start(id: SpanId, name: impl Into<String>) -> Self {
91 Self {
92 id,
93 parent: None,
94 name: name.into(),
95 start_time: SystemTime::now(),
96 end_time: None,
97 duration: None,
98 status: SpanStatus::InProgress,
99 attributes: Vec::new(),
100 }
101 }
102
103 pub fn start_child(id: SpanId, parent: &SpanId, name: impl Into<String>) -> Self {
105 Self {
106 id,
107 parent: Some(parent.clone()),
108 name: name.into(),
109 start_time: SystemTime::now(),
110 end_time: None,
111 duration: None,
112 status: SpanStatus::InProgress,
113 attributes: Vec::new(),
114 }
115 }
116
117 pub fn finish(&mut self, status: SpanStatus) {
119 let now = SystemTime::now();
120 self.end_time = Some(now);
121 self.duration = now.duration_since(self.start_time).ok();
122 self.status = status;
123 }
124
125 pub fn add_attribute(&mut self, key: impl Into<String>, value: impl Into<String>) {
127 self.attributes.push(SpanAttribute::new(key, value));
128 }
129
130 pub fn is_running(&self) -> bool {
132 self.status == SpanStatus::InProgress
133 }
134
135 pub fn is_root(&self) -> bool {
137 self.parent.is_none()
138 }
139
140 pub fn get_attribute(&self, key: &str) -> Option<&str> {
142 self.attributes
143 .iter()
144 .find(|a| a.key == key)
145 .map(|a| a.value.as_str())
146 }
147}
148
149#[derive(Debug, Clone)]
151pub struct ExecutionTrace {
152 pub trace_id: String,
154 spans: HashMap<String, TraceSpan>,
156 span_order: Vec<String>,
158}
159
160impl ExecutionTrace {
161 pub fn new(trace_id: impl Into<String>) -> Self {
163 Self {
164 trace_id: trace_id.into(),
165 spans: HashMap::new(),
166 span_order: Vec::new(),
167 }
168 }
169
170 pub fn start_span(&mut self, id: SpanId, name: impl Into<String>) -> &mut TraceSpan {
178 let span = TraceSpan::start(id.clone(), name);
179 let key = id.0.clone();
180 self.span_order.push(key.clone());
181 self.spans.entry(key).or_insert(span)
182 }
183
184 pub fn start_child_span(
190 &mut self,
191 id: SpanId,
192 parent: &SpanId,
193 name: impl Into<String>,
194 ) -> &mut TraceSpan {
195 let span = TraceSpan::start_child(id.clone(), parent, name);
196 let key = id.0.clone();
197 self.span_order.push(key.clone());
198 self.spans.entry(key).or_insert(span)
199 }
200
201 pub fn finish_span(&mut self, id: &SpanId, status: SpanStatus) {
203 if let Some(span) = self.spans.get_mut(&id.0) {
204 span.finish(status);
205 }
206 }
207
208 pub fn get_span(&self, id: &SpanId) -> Option<&TraceSpan> {
210 self.spans.get(&id.0)
211 }
212
213 pub fn spans_ordered(&self) -> Vec<&TraceSpan> {
215 self.span_order
216 .iter()
217 .filter_map(|k| self.spans.get(k))
218 .collect()
219 }
220
221 pub fn root_spans(&self) -> Vec<&TraceSpan> {
223 self.spans_ordered()
224 .into_iter()
225 .filter(|s| s.is_root())
226 .collect()
227 }
228
229 pub fn children_of(&self, parent: &SpanId) -> Vec<&TraceSpan> {
231 self.spans
232 .values()
233 .filter(|s| s.parent.as_ref() == Some(parent))
234 .collect()
235 }
236
237 pub fn span_count(&self) -> usize {
239 self.spans.len()
240 }
241
242 pub fn running_count(&self) -> usize {
244 self.spans.values().filter(|s| s.is_running()).count()
245 }
246
247 pub fn error_count(&self) -> usize {
249 self.spans
250 .values()
251 .filter(|s| matches!(s.status, SpanStatus::Error(_)))
252 .count()
253 }
254
255 pub fn total_duration(&self) -> Option<Duration> {
257 let earliest = self.spans.values().map(|s| s.start_time).min()?;
258 let latest = self.spans.values().filter_map(|s| s.end_time).max()?;
259 latest.duration_since(earliest).ok()
260 }
261
262 pub fn summary(&self) -> TraceSummary {
264 TraceSummary {
265 trace_id: self.trace_id.clone(),
266 total_spans: self.span_count(),
267 running: self.running_count(),
268 errors: self.error_count(),
269 total_duration: self.total_duration(),
270 }
271 }
272}
273
274#[derive(Debug, Clone)]
276pub struct TraceSummary {
277 pub trace_id: String,
279 pub total_spans: usize,
281 pub running: usize,
283 pub errors: usize,
285 pub total_duration: Option<Duration>,
287}
288
289impl TraceSummary {
290 pub fn is_complete(&self) -> bool {
292 self.running == 0
293 }
294
295 pub fn has_errors(&self) -> bool {
297 self.errors > 0
298 }
299}
300
301#[cfg(test)]
302mod tests {
303 use super::*;
304
305 #[test]
306 fn test_span_id_display() {
307 let id = SpanId::new("span-1");
308 assert_eq!(id.to_string(), "span-1");
309 assert_eq!(id.as_str(), "span-1");
310 }
311
312 #[test]
313 fn test_span_start_and_finish() {
314 let mut span = TraceSpan::start(SpanId::new("s1"), "test operation");
315 assert!(span.is_running());
316 assert!(span.is_root());
317 assert!(span.end_time.is_none());
318
319 span.finish(SpanStatus::Ok);
320 assert!(!span.is_running());
321 assert!(span.end_time.is_some());
322 assert!(span.duration.is_some());
323 }
324
325 #[test]
326 fn test_child_span() {
327 let parent_id = SpanId::new("parent");
328 let span = TraceSpan::start_child(SpanId::new("child"), &parent_id, "child op");
329 assert!(!span.is_root());
330 assert_eq!(
331 span.parent
332 .as_ref()
333 .expect("should succeed in test")
334 .as_str(),
335 "parent"
336 );
337 }
338
339 #[test]
340 fn test_span_attributes() {
341 let mut span = TraceSpan::start(SpanId::new("s1"), "op");
342 span.add_attribute("task_type", "transcode");
343 span.add_attribute("codec", "h264");
344
345 assert_eq!(span.get_attribute("task_type"), Some("transcode"));
346 assert_eq!(span.get_attribute("codec"), Some("h264"));
347 assert_eq!(span.get_attribute("nonexistent"), None);
348 }
349
350 #[test]
351 fn test_execution_trace_basic() {
352 let mut trace = ExecutionTrace::new("trace-1");
353 trace.start_span(SpanId::new("s1"), "step 1");
354 trace.start_span(SpanId::new("s2"), "step 2");
355
356 assert_eq!(trace.span_count(), 2);
357 assert_eq!(trace.running_count(), 2);
358 assert_eq!(trace.error_count(), 0);
359 }
360
361 #[test]
362 fn test_execution_trace_finish() {
363 let mut trace = ExecutionTrace::new("trace-2");
364 trace.start_span(SpanId::new("s1"), "step 1");
365 trace.finish_span(&SpanId::new("s1"), SpanStatus::Ok);
366
367 assert_eq!(trace.running_count(), 0);
368 let span = trace
369 .get_span(&SpanId::new("s1"))
370 .expect("should succeed in test");
371 assert_eq!(span.status, SpanStatus::Ok);
372 }
373
374 #[test]
375 fn test_execution_trace_errors() {
376 let mut trace = ExecutionTrace::new("trace-3");
377 trace.start_span(SpanId::new("s1"), "good");
378 trace.start_span(SpanId::new("s2"), "bad");
379 trace.finish_span(&SpanId::new("s1"), SpanStatus::Ok);
380 trace.finish_span(&SpanId::new("s2"), SpanStatus::Error("timeout".into()));
381
382 assert_eq!(trace.error_count(), 1);
383 }
384
385 #[test]
386 fn test_root_spans() {
387 let mut trace = ExecutionTrace::new("trace-4");
388 let root_id = SpanId::new("root");
389 trace.start_span(root_id.clone(), "root op");
390 trace.start_child_span(SpanId::new("child"), &root_id, "child op");
391
392 let roots = trace.root_spans();
393 assert_eq!(roots.len(), 1);
394 assert_eq!(roots[0].name, "root op");
395 }
396
397 #[test]
398 fn test_children_of() {
399 let mut trace = ExecutionTrace::new("trace-5");
400 let parent = SpanId::new("p1");
401 trace.start_span(parent.clone(), "parent");
402 trace.start_child_span(SpanId::new("c1"), &parent, "child 1");
403 trace.start_child_span(SpanId::new("c2"), &parent, "child 2");
404
405 let children = trace.children_of(&parent);
406 assert_eq!(children.len(), 2);
407 }
408
409 #[test]
410 fn test_spans_ordered() {
411 let mut trace = ExecutionTrace::new("trace-6");
412 trace.start_span(SpanId::new("first"), "first");
413 trace.start_span(SpanId::new("second"), "second");
414 trace.start_span(SpanId::new("third"), "third");
415
416 let ordered = trace.spans_ordered();
417 assert_eq!(ordered[0].name, "first");
418 assert_eq!(ordered[1].name, "second");
419 assert_eq!(ordered[2].name, "third");
420 }
421
422 #[test]
423 fn test_trace_summary() {
424 let mut trace = ExecutionTrace::new("trace-7");
425 trace.start_span(SpanId::new("s1"), "a");
426 trace.start_span(SpanId::new("s2"), "b");
427 trace.finish_span(&SpanId::new("s1"), SpanStatus::Ok);
428 trace.finish_span(&SpanId::new("s2"), SpanStatus::Error("fail".into()));
429
430 let summary = trace.summary();
431 assert_eq!(summary.total_spans, 2);
432 assert_eq!(summary.running, 0);
433 assert_eq!(summary.errors, 1);
434 assert!(summary.is_complete());
435 assert!(summary.has_errors());
436 }
437
438 #[test]
439 fn test_total_duration() {
440 let mut trace = ExecutionTrace::new("trace-8");
441 trace.start_span(SpanId::new("s1"), "op");
442 std::thread::sleep(Duration::from_millis(5));
444 trace.finish_span(&SpanId::new("s1"), SpanStatus::Ok);
445
446 let dur = trace.total_duration();
447 assert!(dur.is_some());
448 assert!(dur.expect("should succeed in test") >= Duration::from_millis(1));
449 }
450
451 #[test]
452 fn test_span_attribute_struct() {
453 let attr = SpanAttribute::new("key", "value");
454 assert_eq!(attr.key, "key");
455 assert_eq!(attr.value, "value");
456 }
457}