avx_async/
tracing.rs

1//! Distributed tracing support for Industry 4.0
2//!
3//! Provides context propagation and span tracking across async boundaries
4
5use std::sync::Arc;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::time::Instant;
8use std::collections::HashMap;
9
10static TRACE_ID_COUNTER: AtomicU64 = AtomicU64::new(1);
11static SPAN_ID_COUNTER: AtomicU64 = AtomicU64::new(1);
12
13/// Trace context for distributed tracing
14#[derive(Clone, Debug)]
15pub struct TraceContext {
16    pub trace_id: u64,
17    pub span_id: u64,
18    pub parent_span_id: Option<u64>,
19    pub service_name: String,
20    pub attributes: HashMap<String, String>,
21}
22
23impl TraceContext {
24    pub fn new(service_name: impl Into<String>) -> Self {
25        Self {
26            trace_id: TRACE_ID_COUNTER.fetch_add(1, Ordering::Relaxed),
27            span_id: SPAN_ID_COUNTER.fetch_add(1, Ordering::Relaxed),
28            parent_span_id: None,
29            service_name: service_name.into(),
30            attributes: HashMap::new(),
31        }
32    }
33
34    pub fn child_span(&self, operation: impl Into<String>) -> Span {
35        let span_id = SPAN_ID_COUNTER.fetch_add(1, Ordering::Relaxed);
36        Span {
37            trace_id: self.trace_id,
38            span_id,
39            parent_span_id: Some(self.span_id),
40            operation: operation.into(),
41            start_time: Instant::now(),
42            end_time: None,
43            attributes: HashMap::new(),
44            events: Vec::new(),
45        }
46    }
47
48    pub fn set_attribute(&mut self, key: impl Into<String>, value: impl Into<String>) {
49        self.attributes.insert(key.into(), value.into());
50    }
51}
52
53/// A span representing a unit of work
54#[derive(Debug)]
55pub struct Span {
56    pub trace_id: u64,
57    pub span_id: u64,
58    pub parent_span_id: Option<u64>,
59    pub operation: String,
60    pub start_time: Instant,
61    pub end_time: Option<Instant>,
62    pub attributes: HashMap<String, String>,
63    pub events: Vec<SpanEvent>,
64}
65
66#[derive(Debug, Clone)]
67pub struct SpanEvent {
68    pub name: String,
69    pub timestamp: Instant,
70    pub attributes: HashMap<String, String>,
71}
72
73impl Span {
74    pub fn set_attribute(&mut self, key: impl Into<String>, value: impl Into<String>) {
75        self.attributes.insert(key.into(), value.into());
76    }
77
78    pub fn add_event(&mut self, name: impl Into<String>) {
79        self.events.push(SpanEvent {
80            name: name.into(),
81            timestamp: Instant::now(),
82            attributes: HashMap::new(),
83        });
84    }
85
86    pub fn add_event_with_attributes(
87        &mut self,
88        name: impl Into<String>,
89        attributes: HashMap<String, String>,
90    ) {
91        self.events.push(SpanEvent {
92            name: name.into(),
93            timestamp: Instant::now(),
94            attributes,
95        });
96    }
97
98    pub fn end(mut self) -> CompletedSpan {
99        self.end_time = Some(Instant::now());
100        CompletedSpan {
101            trace_id: self.trace_id,
102            span_id: self.span_id,
103            parent_span_id: self.parent_span_id,
104            operation: self.operation,
105            duration: self.end_time.unwrap() - self.start_time,
106            attributes: self.attributes,
107            events: self.events,
108        }
109    }
110}
111
112#[derive(Debug, Clone)]
113pub struct CompletedSpan {
114    pub trace_id: u64,
115    pub span_id: u64,
116    pub parent_span_id: Option<u64>,
117    pub operation: String,
118    pub duration: std::time::Duration,
119    pub attributes: HashMap<String, String>,
120    pub events: Vec<SpanEvent>,
121}
122
123impl std::fmt::Display for CompletedSpan {
124    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125        write!(
126            f,
127            "Span[trace={:016x} span={:016x} parent={:?}] {} {:?}",
128            self.trace_id, self.span_id, self.parent_span_id, self.operation, self.duration
129        )
130    }
131}
132
133/// Tracer for collecting spans
134pub struct Tracer {
135    spans: Arc<std::sync::Mutex<Vec<CompletedSpan>>>,
136}
137
138impl Tracer {
139    pub fn new() -> Self {
140        Self {
141            spans: Arc::new(std::sync::Mutex::new(Vec::new())),
142        }
143    }
144
145    pub fn record(&self, span: CompletedSpan) {
146        let mut spans = self.spans.lock().unwrap();
147        spans.push(span);
148    }
149
150    pub fn get_spans(&self) -> Vec<CompletedSpan> {
151        let spans = self.spans.lock().unwrap();
152        spans.clone()
153    }
154
155    pub fn clear(&self) {
156        let mut spans = self.spans.lock().unwrap();
157        spans.clear();
158    }
159
160    /// Export spans in Jaeger format
161    pub fn to_jaeger_json(&self) -> String {
162        let spans = self.get_spans();
163        let mut json = String::from("[\n");
164
165        for (i, span) in spans.iter().enumerate() {
166            if i > 0 {
167                json.push_str(",\n");
168            }
169            json.push_str(&format!(
170                "  {{\n    \"traceId\": \"{:016x}\",\n    \"spanId\": \"{:016x}\",\n    \
171                 \"operationName\": \"{}\",\n    \"duration\": {},\n    \
172                 \"startTime\": 0\n  }}",
173                span.trace_id,
174                span.span_id,
175                span.operation,
176                span.duration.as_micros()
177            ));
178        }
179
180        json.push_str("\n]");
181        json
182    }
183}
184
185impl Default for Tracer {
186    fn default() -> Self {
187        Self::new()
188    }
189}
190
191impl Clone for Tracer {
192    fn clone(&self) -> Self {
193        Self {
194            spans: Arc::clone(&self.spans),
195        }
196    }
197}
198
199
200
201
202