forge_core/observability/
trace.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5/// Trace ID for distributed tracing.
6#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
7pub struct TraceId(String);
8
9impl TraceId {
10    /// Create a new random trace ID.
11    pub fn new() -> Self {
12        Self(uuid::Uuid::new_v4().to_string().replace('-', ""))
13    }
14
15    /// Create from a string.
16    pub fn from_string(s: impl Into<String>) -> Self {
17        Self(s.into())
18    }
19
20    /// Get the trace ID as a string.
21    pub fn as_str(&self) -> &str {
22        &self.0
23    }
24}
25
26impl Default for TraceId {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl std::fmt::Display for TraceId {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        write!(f, "{}", self.0)
35    }
36}
37
38/// Span ID within a trace.
39#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
40pub struct SpanId(String);
41
42impl SpanId {
43    /// Create a new random span ID.
44    pub fn new() -> Self {
45        Self(uuid::Uuid::new_v4().to_string().replace('-', "")[..16].to_string())
46    }
47
48    /// Create from a string.
49    pub fn from_string(s: impl Into<String>) -> Self {
50        Self(s.into())
51    }
52
53    /// Get the span ID as a string.
54    pub fn as_str(&self) -> &str {
55        &self.0
56    }
57}
58
59impl Default for SpanId {
60    fn default() -> Self {
61        Self::new()
62    }
63}
64
65impl std::fmt::Display for SpanId {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        write!(f, "{}", self.0)
68    }
69}
70
71/// Span context for propagation.
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct SpanContext {
74    /// Trace ID.
75    pub trace_id: TraceId,
76    /// Span ID.
77    pub span_id: SpanId,
78    /// Parent span ID if any.
79    pub parent_span_id: Option<SpanId>,
80    /// Trace flags (e.g., sampled).
81    pub trace_flags: u8,
82}
83
84impl SpanContext {
85    /// Create a new root context.
86    pub fn new_root() -> Self {
87        Self {
88            trace_id: TraceId::new(),
89            span_id: SpanId::new(),
90            parent_span_id: None,
91            trace_flags: 0x01, // sampled
92        }
93    }
94
95    /// Create a child context.
96    pub fn child(&self) -> Self {
97        Self {
98            trace_id: self.trace_id.clone(),
99            span_id: SpanId::new(),
100            parent_span_id: Some(self.span_id.clone()),
101            trace_flags: self.trace_flags,
102        }
103    }
104
105    /// Check if the trace is sampled.
106    pub fn is_sampled(&self) -> bool {
107        self.trace_flags & 0x01 != 0
108    }
109
110    /// Create a W3C traceparent header value.
111    pub fn to_traceparent(&self) -> String {
112        format!(
113            "00-{}-{}-{:02x}",
114            self.trace_id, self.span_id, self.trace_flags
115        )
116    }
117
118    /// Parse from W3C traceparent header.
119    pub fn from_traceparent(traceparent: &str) -> Option<Self> {
120        let parts: Vec<&str> = traceparent.split('-').collect();
121        if parts.len() != 4 || parts[0] != "00" {
122            return None;
123        }
124
125        let trace_id = TraceId::from_string(parts[1]);
126        let span_id = SpanId::from_string(parts[2]);
127        let trace_flags = u8::from_str_radix(parts[3], 16).ok()?;
128
129        Some(Self {
130            trace_id,
131            span_id,
132            parent_span_id: None,
133            trace_flags,
134        })
135    }
136}
137
138/// Span kind indicating the relationship.
139#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
140#[serde(rename_all = "snake_case")]
141pub enum SpanKind {
142    /// Internal operation.
143    #[default]
144    Internal,
145    /// Server handling a request.
146    Server,
147    /// Client making a request.
148    Client,
149    /// Producer sending a message.
150    Producer,
151    /// Consumer receiving a message.
152    Consumer,
153}
154
155impl std::fmt::Display for SpanKind {
156    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157        match self {
158            Self::Internal => write!(f, "internal"),
159            Self::Server => write!(f, "server"),
160            Self::Client => write!(f, "client"),
161            Self::Producer => write!(f, "producer"),
162            Self::Consumer => write!(f, "consumer"),
163        }
164    }
165}
166
167/// Span status.
168#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
169#[serde(rename_all = "snake_case")]
170pub enum SpanStatus {
171    /// Unset status.
172    #[default]
173    Unset,
174    /// Operation completed successfully.
175    Ok,
176    /// Operation failed with an error.
177    Error,
178}
179
180impl std::fmt::Display for SpanStatus {
181    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
182        match self {
183            Self::Unset => write!(f, "unset"),
184            Self::Ok => write!(f, "ok"),
185            Self::Error => write!(f, "error"),
186        }
187    }
188}
189
190/// A trace span.
191#[derive(Debug, Clone, Serialize, Deserialize)]
192pub struct Span {
193    /// Span context.
194    pub context: SpanContext,
195    /// Span name.
196    pub name: String,
197    /// Span kind.
198    pub kind: SpanKind,
199    /// Span status.
200    pub status: SpanStatus,
201    /// Status message (for errors).
202    pub status_message: Option<String>,
203    /// Start time.
204    pub start_time: chrono::DateTime<chrono::Utc>,
205    /// End time.
206    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
207    /// Attributes.
208    pub attributes: HashMap<String, serde_json::Value>,
209    /// Events within the span.
210    pub events: Vec<SpanEvent>,
211    /// Node ID that generated this span.
212    pub node_id: Option<uuid::Uuid>,
213}
214
215impl Span {
216    /// Create a new span.
217    pub fn new(name: impl Into<String>) -> Self {
218        Self {
219            context: SpanContext::new_root(),
220            name: name.into(),
221            kind: SpanKind::Internal,
222            status: SpanStatus::Unset,
223            status_message: None,
224            start_time: chrono::Utc::now(),
225            end_time: None,
226            attributes: HashMap::new(),
227            events: Vec::new(),
228            node_id: None,
229        }
230    }
231
232    /// Create a child span.
233    pub fn child(&self, name: impl Into<String>) -> Self {
234        Self {
235            context: self.context.child(),
236            name: name.into(),
237            kind: SpanKind::Internal,
238            status: SpanStatus::Unset,
239            status_message: None,
240            start_time: chrono::Utc::now(),
241            end_time: None,
242            attributes: HashMap::new(),
243            events: Vec::new(),
244            node_id: self.node_id,
245        }
246    }
247
248    /// Set the span kind.
249    pub fn with_kind(mut self, kind: SpanKind) -> Self {
250        self.kind = kind;
251        self
252    }
253
254    /// Add an attribute.
255    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Serialize) -> Self {
256        if let Ok(v) = serde_json::to_value(value) {
257            self.attributes.insert(key.into(), v);
258        }
259        self
260    }
261
262    /// Set the node ID.
263    pub fn with_node_id(mut self, node_id: uuid::Uuid) -> Self {
264        self.node_id = Some(node_id);
265        self
266    }
267
268    /// Set from a parent context.
269    pub fn with_parent(mut self, parent: &SpanContext) -> Self {
270        self.context = parent.child();
271        self
272    }
273
274    /// Add an event.
275    pub fn add_event(&mut self, name: impl Into<String>) {
276        self.events.push(SpanEvent {
277            name: name.into(),
278            timestamp: chrono::Utc::now(),
279            attributes: HashMap::new(),
280        });
281    }
282
283    /// End the span successfully.
284    pub fn end_ok(&mut self) {
285        self.status = SpanStatus::Ok;
286        self.end_time = Some(chrono::Utc::now());
287    }
288
289    /// End the span with an error.
290    pub fn end_error(&mut self, message: impl Into<String>) {
291        self.status = SpanStatus::Error;
292        self.status_message = Some(message.into());
293        self.end_time = Some(chrono::Utc::now());
294    }
295
296    /// Get the duration in milliseconds.
297    pub fn duration_ms(&self) -> Option<f64> {
298        self.end_time
299            .map(|end| (end - self.start_time).num_milliseconds() as f64)
300    }
301
302    /// Check if the span is complete.
303    pub fn is_complete(&self) -> bool {
304        self.end_time.is_some()
305    }
306}
307
308/// An event within a span.
309#[derive(Debug, Clone, Serialize, Deserialize)]
310pub struct SpanEvent {
311    /// Event name.
312    pub name: String,
313    /// Event timestamp.
314    pub timestamp: chrono::DateTime<chrono::Utc>,
315    /// Event attributes.
316    pub attributes: HashMap<String, serde_json::Value>,
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[test]
324    fn test_trace_id_generation() {
325        let id1 = TraceId::new();
326        let id2 = TraceId::new();
327        assert_ne!(id1.as_str(), id2.as_str());
328        assert!(!id1.as_str().contains('-'));
329    }
330
331    #[test]
332    fn test_span_context_root() {
333        let ctx = SpanContext::new_root();
334        assert!(ctx.parent_span_id.is_none());
335        assert!(ctx.is_sampled());
336    }
337
338    #[test]
339    fn test_span_context_child() {
340        let parent = SpanContext::new_root();
341        let child = parent.child();
342
343        assert_eq!(child.trace_id, parent.trace_id);
344        assert_ne!(child.span_id, parent.span_id);
345        assert_eq!(child.parent_span_id, Some(parent.span_id));
346    }
347
348    #[test]
349    fn test_traceparent_roundtrip() {
350        let ctx = SpanContext::new_root();
351        let header = ctx.to_traceparent();
352        let parsed = SpanContext::from_traceparent(&header).unwrap();
353
354        assert_eq!(parsed.trace_id, ctx.trace_id);
355        assert_eq!(parsed.span_id, ctx.span_id);
356        assert_eq!(parsed.trace_flags, ctx.trace_flags);
357    }
358
359    #[test]
360    fn test_span_lifecycle() {
361        let mut span = Span::new("test_operation")
362            .with_kind(SpanKind::Server)
363            .with_attribute("http.method", "GET");
364
365        assert!(!span.is_complete());
366        assert!(span.duration_ms().is_none());
367
368        span.add_event("started processing");
369        span.end_ok();
370
371        assert!(span.is_complete());
372        assert!(span.duration_ms().is_some());
373        assert_eq!(span.status, SpanStatus::Ok);
374    }
375
376    #[test]
377    fn test_span_error() {
378        let mut span = Span::new("failing_operation");
379        span.end_error("Something went wrong");
380
381        assert_eq!(span.status, SpanStatus::Error);
382        assert_eq!(
383            span.status_message,
384            Some("Something went wrong".to_string())
385        );
386    }
387
388    #[test]
389    fn test_child_span() {
390        let parent = Span::new("parent");
391        let child = parent.child("child");
392
393        assert_eq!(child.context.trace_id, parent.context.trace_id);
394        assert_eq!(
395            child.context.parent_span_id,
396            Some(parent.context.span_id.clone())
397        );
398    }
399}
forge_core/observability/trace.rs

forge_core/observability/
trace.rs
