Skip to main content

forge_runtime/gateway/
tracing.rs

1use uuid::Uuid;
2
3/// Header name for trace ID.
4pub const TRACE_ID_HEADER: &str = "x-trace-id";
5/// Header name for request ID.
6pub const REQUEST_ID_HEADER: &str = "x-request-id";
7/// Header name for parent span ID.
8pub const SPAN_ID_HEADER: &str = "x-span-id";
9
10/// Request tracing state.
11#[derive(Debug, Clone)]
12pub struct TracingState {
13    /// Unique trace ID for distributed tracing.
14    pub trace_id: String,
15    /// Unique request ID.
16    pub request_id: String,
17    /// Parent span ID (if propagated).
18    pub parent_span_id: Option<String>,
19    /// When the request started.
20    pub start_time: std::time::Instant,
21}
22
23impl TracingState {
24    /// Create a new tracing state.
25    pub fn new() -> Self {
26        Self {
27            trace_id: Uuid::new_v4().to_string(),
28            request_id: Uuid::new_v4().to_string(),
29            parent_span_id: None,
30            start_time: std::time::Instant::now(),
31        }
32    }
33
34    /// Create with an existing trace ID (for propagation).
35    pub fn with_trace_id(trace_id: String) -> Self {
36        Self {
37            trace_id,
38            request_id: Uuid::new_v4().to_string(),
39            parent_span_id: None,
40            start_time: std::time::Instant::now(),
41        }
42    }
43
44    /// Set parent span ID.
45    pub fn with_parent_span(mut self, span_id: String) -> Self {
46        self.parent_span_id = Some(span_id);
47        self
48    }
49
50    /// Get elapsed time since request start.
51    pub fn elapsed(&self) -> std::time::Duration {
52        self.start_time.elapsed()
53    }
54}
55
56impl Default for TracingState {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62/// Tracing middleware marker type.
63#[derive(Debug, Clone)]
64pub struct TracingMiddleware;
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[test]
71    fn test_tracing_state_new() {
72        let state = TracingState::new();
73        assert!(!state.trace_id.is_empty());
74        assert!(!state.request_id.is_empty());
75        assert!(state.parent_span_id.is_none());
76    }
77
78    #[test]
79    fn test_tracing_state_with_trace_id() {
80        let state = TracingState::with_trace_id("trace-123".to_string());
81        assert_eq!(state.trace_id, "trace-123");
82        assert!(!state.request_id.is_empty());
83    }
84
85    #[test]
86    fn test_tracing_state_with_parent_span() {
87        let state = TracingState::new().with_parent_span("span-456".to_string());
88        assert_eq!(state.parent_span_id, Some("span-456".to_string()));
89    }
90
91    #[test]
92    fn test_tracing_state_elapsed() {
93        let state = TracingState::new();
94        std::thread::sleep(std::time::Duration::from_millis(10));
95        let elapsed = state.elapsed();
96        assert!(elapsed.as_millis() >= 10);
97    }
98}