Skip to main content

oximedia_workflow/
execution_trace.rs

1#![allow(dead_code)]
2//! Execution tracing and span tracking for workflow debugging.
3//!
4//! Records detailed timing and status information for every step in a
5//! workflow execution, enabling post-mortem analysis, performance
6//! profiling, and audit trails.
7
8use std::collections::HashMap;
9use std::time::{Duration, SystemTime};
10
11/// Unique identifier for a trace span.
12#[derive(Debug, Clone, PartialEq, Eq, Hash)]
13pub struct SpanId(String);
14
15impl SpanId {
16    /// Create a new span ID from a string.
17    pub fn new(id: impl Into<String>) -> Self {
18        Self(id.into())
19    }
20
21    /// Return the string representation.
22    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/// The outcome of a traced span.
34#[derive(Debug, Clone, PartialEq)]
35pub enum SpanStatus {
36    /// Span completed successfully.
37    Ok,
38    /// Span completed with a warning.
39    Warning(String),
40    /// Span failed with an error.
41    Error(String),
42    /// Span was cancelled before completion.
43    Cancelled,
44    /// Span is still running.
45    InProgress,
46}
47
48/// A key-value attribute attached to a span.
49#[derive(Debug, Clone, PartialEq)]
50pub struct SpanAttribute {
51    /// Attribute key.
52    pub key: String,
53    /// Attribute value.
54    pub value: String,
55}
56
57impl SpanAttribute {
58    /// Create a new attribute.
59    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/// A single span in an execution trace.
68#[derive(Debug, Clone)]
69pub struct TraceSpan {
70    /// Unique span identifier.
71    pub id: SpanId,
72    /// Optional parent span (for nesting).
73    pub parent: Option<SpanId>,
74    /// Human-readable name of the operation.
75    pub name: String,
76    /// When the span started.
77    pub start_time: SystemTime,
78    /// When the span ended (if finished).
79    pub end_time: Option<SystemTime>,
80    /// Duration (calculated when ended).
81    pub duration: Option<Duration>,
82    /// Current status of the span.
83    pub status: SpanStatus,
84    /// Arbitrary key-value attributes.
85    pub attributes: Vec<SpanAttribute>,
86}
87
88impl TraceSpan {
89    /// Create a new span that starts now.
90    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    /// Create a new child span.
104    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    /// End the span with the given status.
118    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    /// Add an attribute to the span.
126    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    /// Check if this span is still running.
131    pub fn is_running(&self) -> bool {
132        self.status == SpanStatus::InProgress
133    }
134
135    /// Check if this span has a parent.
136    pub fn is_root(&self) -> bool {
137        self.parent.is_none()
138    }
139
140    /// Get the attribute value for a given key.
141    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/// An execution trace that collects spans for a single workflow run.
150#[derive(Debug, Clone)]
151pub struct ExecutionTrace {
152    /// Unique trace identifier (typically the workflow run ID).
153    pub trace_id: String,
154    /// All spans in this trace, indexed by span ID.
155    spans: HashMap<String, TraceSpan>,
156    /// Order in which spans were created.
157    span_order: Vec<String>,
158}
159
160impl ExecutionTrace {
161    /// Create a new execution trace.
162    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    /// Start a new root span.
171    ///
172    /// # Panics
173    ///
174    /// This method will not panic; if the internal map entry is somehow
175    /// missing after insertion the returned reference is obtained via
176    /// `Entry::or_insert_with`, guaranteeing it exists.
177    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    /// Start a child span under a parent.
185    ///
186    /// # Panics
187    ///
188    /// Same guarantee as [`start_span`](Self::start_span).
189    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    /// Finish a span with the given status.
202    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    /// Get a span by ID.
209    pub fn get_span(&self, id: &SpanId) -> Option<&TraceSpan> {
210        self.spans.get(&id.0)
211    }
212
213    /// Get all spans in creation order.
214    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    /// Get only root spans (those without a parent).
222    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    /// Get all children of a given span.
230    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    /// Return how many spans are in this trace.
238    pub fn span_count(&self) -> usize {
239        self.spans.len()
240    }
241
242    /// Return how many spans are still running.
243    pub fn running_count(&self) -> usize {
244        self.spans.values().filter(|s| s.is_running()).count()
245    }
246
247    /// Return how many spans failed.
248    pub fn error_count(&self) -> usize {
249        self.spans
250            .values()
251            .filter(|s| matches!(s.status, SpanStatus::Error(_)))
252            .count()
253    }
254
255    /// Compute total trace duration (from earliest start to latest end).
256    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    /// Generate a simple summary of the trace.
263    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/// A summary of an execution trace.
275#[derive(Debug, Clone)]
276pub struct TraceSummary {
277    /// Trace identifier.
278    pub trace_id: String,
279    /// Total number of spans.
280    pub total_spans: usize,
281    /// Number of spans still running.
282    pub running: usize,
283    /// Number of spans that errored.
284    pub errors: usize,
285    /// Total wall-clock duration of the trace.
286    pub total_duration: Option<Duration>,
287}
288
289impl TraceSummary {
290    /// Check if all spans have completed (none running).
291    pub fn is_complete(&self) -> bool {
292        self.running == 0
293    }
294
295    /// Check if the trace had any errors.
296    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        // Small sleep so duration is measurable.
443        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}