1use 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#[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#[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
133pub 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 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