Line data Source code
1 : use once_cell::sync::Lazy;
2 : use std::cell::RefCell;
3 : use std::sync::atomic::{AtomicU64, Ordering};
4 : use tokio::task_local;
5 : use tracing::Span;
6 : use uuid::Uuid;
7 :
8 : #[allow(dead_code)]
9 : pub const TRACE_ID_KEY: &str = "trace_id";
10 : #[allow(dead_code)]
11 : pub const SPAN_ID_KEY: &str = "span_id";
12 : #[allow(dead_code)]
13 : pub const PARENT_SPAN_ID_KEY: &str = "parent_span_id";
14 :
15 : /// The tracing identifiers attached to a proxied operation
16 : #[derive(Clone, Debug)]
17 : pub struct TraceContext {
18 : /// The id of the trace, a random UUID for new traces
19 : pub trace_id: String,
20 : /// The id of the current span
21 : pub span_id: String,
22 : /// The id of the parent span, `None` for a root span
23 : pub parent_span_id: Option<String>,
24 : }
25 :
26 : #[allow(dead_code)]
27 1 : static NEXT_SPAN_ID: Lazy<AtomicU64> = Lazy::new(|| AtomicU64::new(1));
28 :
29 : task_local! {
30 : static TRACE_CONTEXT: TraceContext;
31 : }
32 :
33 : // Thread-local storage for accessing trace context from synchronous code
34 : thread_local! {
35 : static THREAD_TRACE_CONTEXT: RefCell<Option<TraceContext>> = RefCell::new(None);
36 : }
37 :
38 : impl TraceContext {
39 : #[allow(dead_code)]
40 4 : pub fn new_trace() -> Self {
41 4 : Self {
42 4 : trace_id: Uuid::new_v4().to_string(),
43 4 : span_id: NEXT_SPAN_ID.fetch_add(1, Ordering::SeqCst).to_string(),
44 4 : parent_span_id: None,
45 4 : }
46 4 : }
47 :
48 : #[allow(dead_code)]
49 5 : pub fn from_headers(
50 5 : trace_id: String,
51 5 : span_id: String,
52 5 : parent_span_id: Option<String>,
53 5 : ) -> Self {
54 5 : Self {
55 5 : trace_id,
56 5 : span_id,
57 5 : parent_span_id,
58 5 : }
59 5 : }
60 :
61 : #[allow(dead_code)]
62 3 : pub fn create_span(&self, name: &str) -> Span {
63 3 : tracing::span!(
64 3 : tracing::Level::INFO,
65 : "traced_operation",
66 : name = name,
67 : trace_id = %self.trace_id,
68 : span_id = %self.span_id,
69 0 : parent_span_id = self.parent_span_id.as_deref().unwrap_or("")
70 : )
71 3 : }
72 :
73 : /// Return trace context fields for callers that use the `log` crate.
74 : /// Returns (trace_id, span_id, parent_span_id) for manual KV addition
75 : #[allow(dead_code)]
76 2 : pub fn trace_fields(&self) -> (&str, &str, &str) {
77 2 : (
78 2 : &self.trace_id,
79 2 : &self.span_id,
80 2 : self.parent_span_id.as_deref().unwrap_or(""),
81 2 : )
82 2 : }
83 : }
84 :
85 : #[allow(dead_code)]
86 2 : pub async fn enter_trace_context<F, Fut, R>(ctx: TraceContext, f: F) -> R
87 2 : where
88 2 : F: FnOnce() -> Fut,
89 2 : Fut: std::future::Future<Output = R>,
90 2 : {
91 2 : let span = ctx.create_span("traced_operation");
92 :
93 : // Store in thread-local storage
94 2 : THREAD_TRACE_CONTEXT.with(|tls| {
95 2 : *tls.borrow_mut() = Some(ctx.clone());
96 2 : });
97 :
98 2 : let result = TRACE_CONTEXT.scope(ctx, span.in_scope(|| f())).await;
99 :
100 : // Clear thread-local storage
101 2 : THREAD_TRACE_CONTEXT.with(|tls| {
102 2 : *tls.borrow_mut() = None;
103 2 : });
104 :
105 2 : result
106 2 : }
107 :
108 : #[allow(dead_code)]
109 2 : pub fn current_trace_context() -> Option<TraceContext> {
110 : // Try task_local first (for async contexts)
111 2 : if let Ok(ctx) = TRACE_CONTEXT.try_with(|ctx| ctx.clone()) {
112 2 : return Some(ctx);
113 0 : }
114 :
115 : // Fall back to thread_local (for sync contexts)
116 0 : THREAD_TRACE_CONTEXT.with(|tls| tls.borrow().clone())
117 2 : }
118 :
119 : #[cfg(test)]
120 : mod tests {
121 : use super::*;
122 :
123 : #[test]
124 1 : fn test_new_trace_generates_uuid_trace_id() {
125 1 : let ctx = TraceContext::new_trace();
126 :
127 : // UUID v4 format: 8-4-4-4-12 hex digits
128 1 : assert_eq!(ctx.trace_id.len(), 36);
129 1 : assert!(ctx.trace_id.contains('-'));
130 1 : assert!(ctx.parent_span_id.is_none());
131 1 : }
132 :
133 : #[test]
134 1 : fn test_span_ids_increment() {
135 1 : let ctx1 = TraceContext::new_trace();
136 1 : let ctx2 = TraceContext::new_trace();
137 :
138 1 : let span1: u64 = ctx1.span_id.parse().unwrap();
139 1 : let span2: u64 = ctx2.span_id.parse().unwrap();
140 :
141 1 : assert!(span2 > span1, "Span IDs should increment");
142 1 : }
143 :
144 : #[test]
145 1 : fn test_from_headers() {
146 1 : let trace_id = "test-trace-123".to_string();
147 1 : let span_id = "456".to_string();
148 1 : let parent_span_id = Some("789".to_string());
149 :
150 1 : let ctx = TraceContext::from_headers(
151 1 : trace_id.clone(),
152 1 : span_id.clone(),
153 1 : parent_span_id.clone(),
154 : );
155 :
156 1 : assert_eq!(ctx.trace_id, trace_id);
157 1 : assert_eq!(ctx.span_id, span_id);
158 1 : assert_eq!(ctx.parent_span_id, parent_span_id);
159 1 : }
160 :
161 : #[test]
162 1 : fn test_trace_fields() {
163 1 : let ctx = TraceContext::from_headers(
164 1 : "trace-abc".to_string(),
165 1 : "123".to_string(),
166 1 : Some("456".to_string()),
167 : );
168 :
169 1 : let (trace_id, span_id, parent_span_id) = ctx.trace_fields();
170 :
171 1 : assert_eq!(trace_id, "trace-abc");
172 1 : assert_eq!(span_id, "123");
173 1 : assert_eq!(parent_span_id, "456");
174 1 : }
175 :
176 : #[test]
177 1 : fn test_trace_fields_no_parent() {
178 1 : let ctx = TraceContext::from_headers(
179 1 : "trace-xyz".to_string(),
180 1 : "999".to_string(),
181 1 : None,
182 : );
183 :
184 1 : let (trace_id, span_id, parent_span_id) = ctx.trace_fields();
185 :
186 1 : assert_eq!(trace_id, "trace-xyz");
187 1 : assert_eq!(span_id, "999");
188 1 : assert_eq!(parent_span_id, ""); // Empty string when no parent
189 1 : }
190 :
191 : #[test]
192 1 : fn test_create_span() {
193 1 : let ctx = TraceContext::new_trace();
194 1 : let span = ctx.create_span("test_operation");
195 :
196 : // Verify span is created (doesn't panic)
197 : // Note: span may be disabled if no tracing subscriber is set up
198 1 : drop(span); // Just verify it was created successfully
199 1 : }
200 :
201 : #[tokio::test]
202 1 : async fn test_task_local_isolation() {
203 1 : let ctx1 = TraceContext::from_headers("trace-1".to_string(), "1".to_string(), None);
204 1 : let ctx2 = TraceContext::from_headers("trace-2".to_string(), "2".to_string(), None);
205 :
206 1 : let handle1 = tokio::spawn(async move {
207 1 : enter_trace_context(ctx1.clone(), || async {
208 1 : let current = current_trace_context();
209 1 : assert!(current.is_some());
210 1 : assert_eq!(current.unwrap().trace_id, "trace-1");
211 2 : }).await
212 1 : });
213 :
214 1 : let handle2 = tokio::spawn(async move {
215 1 : enter_trace_context(ctx2.clone(), || async {
216 1 : let current = current_trace_context();
217 1 : assert!(current.is_some());
218 1 : assert_eq!(current.unwrap().trace_id, "trace-2");
219 2 : }).await
220 1 : });
221 :
222 1 : handle1.await.unwrap();
223 1 : handle2.await.unwrap();
224 1 : }
225 : }
|
