tracing_datadog/
http.rs

1//! Distributed trace context for HTTP.
2//!
3//! This module contains strategies for injecting/extracting distributed trace context into/from
4//! HTTP headers using different formats.
5
6use crate::context::{DatadogContext, Strategy, TraceContextExt};
7use http::{HeaderMap, HeaderName, HeaderValue};
8
9impl TraceContextExt for HeaderMap {}
10
11/// W3C Trace Context headers strategy for [`HeaderMap`].
12///
13/// This uses the `traceparent` header.
14///
15/// ```
16/// # use http::HeaderMap;
17/// use tracing_datadog::{
18///     context::{DatadogContext, TraceContextExt, TracingContextExt},
19///     http::W3CTraceContextHeaders,
20/// };
21///
22/// let mut headers = HeaderMap::new();
23/// let current_span = tracing::Span::current();
24///
25/// // Extract W3C headers and set the context on the current span.
26/// current_span.set_context(headers.extract_trace_context::<W3CTraceContextHeaders>());
27///
28/// // Set the current span's context as W3C headers.
29/// headers.inject_trace_context::<W3CTraceContextHeaders>(current_span.get_context());
30///```
31pub struct W3CTraceContextHeaders;
32
33const W3C_TRACEPARENT_HEADER: HeaderName = HeaderName::from_static("traceparent");
34
35impl Strategy<HeaderMap> for W3CTraceContextHeaders {
36    fn inject(headers: &mut HeaderMap, context: DatadogContext) {
37        if context.is_empty() {
38            return;
39        }
40
41        let header = format!(
42            "{version:02x}-{trace_id:032x}-{parent_id:016x}-{trace_flags:02x}",
43            version = 0,
44            trace_id = context.trace_id,
45            parent_id = context.parent_id,
46            trace_flags = 1,
47        );
48
49        headers.insert(W3C_TRACEPARENT_HEADER, header.parse().unwrap());
50    }
51
52    fn extract(headers: &HeaderMap) -> DatadogContext {
53        move || -> Option<DatadogContext> {
54            let header = headers.get(W3C_TRACEPARENT_HEADER)?.to_str().ok()?;
55
56            let parts: Vec<&str> = header.split('-').collect();
57            if parts.len() != 4 {
58                return None;
59            }
60
61            let Some(0) = u8::from_str_radix(parts[0], 16).ok() else {
62                // Wrong version.
63                return None;
64            };
65
66            let Some(0x01) = u8::from_str_radix(parts[3], 16).ok().map(|n| n & 0x01) else {
67                // Not sampled.
68                return None;
69            };
70
71            let trace_id = u128::from_str_radix(parts[1], 16).ok()?;
72            let parent_id = u64::from_str_radix(parts[2], 16).ok()?;
73
74            Some(DatadogContext {
75                trace_id,
76                parent_id,
77            })
78        }()
79        .unwrap_or_default()
80    }
81}
82
83/// Datadog-specific header strategy for [`HeaderMap`].
84///
85/// This uses the various `x-datadog-` headers.
86///
87/// ```
88/// # use http::HeaderMap;
89/// use tracing_datadog::{
90///     context::{DatadogContext, TraceContextExt, TracingContextExt},
91///     http::DatadogHeaders,
92/// };
93///
94/// let mut headers = HeaderMap::new();
95/// let current_span = tracing::Span::current();
96///
97/// // Extract W3C headers and set the context on the current span.
98/// current_span.set_context(headers.extract_trace_context::<DatadogHeaders>());
99///
100/// // Set the current span's context as W3C headers.
101/// headers.inject_trace_context::<DatadogHeaders>(current_span.get_context());
102///```
103pub struct DatadogHeaders;
104
105const DATADOG_TRACE_ID_HEADER: HeaderName = HeaderName::from_static("x-datadog-trace-id");
106const DATADOG_PARENT_ID_HEADER: HeaderName = HeaderName::from_static("x-datadog-parent-id");
107const DATADOG_SAMPLING_PRIORITY_HEADER: HeaderName =
108    HeaderName::from_static("x-datadog-sampling-priority");
109const DATADOG_TAGS_HEADER: HeaderName = HeaderName::from_static("x-datadog-tags");
110
111impl Strategy<HeaderMap> for DatadogHeaders {
112    fn inject(headers: &mut HeaderMap, context: DatadogContext) {
113        if context.is_empty() {
114            return;
115        }
116
117        let lower_64_bits = context.trace_id as u64;
118        let upper_64_bits = (context.trace_id >> 64) as u64;
119
120        headers.insert(
121            DATADOG_TRACE_ID_HEADER,
122            lower_64_bits.to_string().parse().unwrap(),
123        );
124        headers.insert(
125            DATADOG_PARENT_ID_HEADER,
126            context.parent_id.to_string().parse().unwrap(),
127        );
128        headers.insert(
129            DATADOG_SAMPLING_PRIORITY_HEADER,
130            HeaderValue::from_static("1"),
131        );
132        headers.insert(
133            DATADOG_TAGS_HEADER,
134            format!("_dd.p.tid={upper_64_bits:016x}")
135                .parse()
136                .ok()
137                .unwrap(),
138        );
139    }
140
141    fn extract(headers: &HeaderMap) -> DatadogContext {
142        move || -> Option<DatadogContext> {
143            if headers
144                .get(DATADOG_SAMPLING_PRIORITY_HEADER)?
145                .to_str()
146                .ok()?
147                .parse::<u8>()
148                .ok()?
149                < 1
150            {
151                return None;
152            }
153
154            let lower_64_bits = headers
155                .get(DATADOG_TRACE_ID_HEADER)?
156                .to_str()
157                .ok()?
158                .parse::<u64>()
159                .ok()? as u128;
160            let parent_id = headers
161                .get(DATADOG_PARENT_ID_HEADER)?
162                .to_str()
163                .ok()?
164                .parse()
165                .ok()?;
166
167            let upper_64_bits: u128 = headers
168                .get(DATADOG_TAGS_HEADER)
169                .and_then(|header| {
170                    header.to_str().ok()?.split(',').find_map(|pair| {
171                        pair.strip_prefix("_dd.p.tid=").and_then(|hex_value| {
172                            u64::from_str_radix(hex_value, 16).map(|x| x as u128).ok()
173                        })
174                    })
175                })
176                .unwrap_or_default();
177
178            let trace_id = (upper_64_bits << 64) | lower_64_bits;
179
180            Some(DatadogContext {
181                trace_id,
182                parent_id,
183            })
184        }()
185        .unwrap_or_default()
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use rand::random_range;
193
194    #[test]
195    fn w3c_trace_header_round_trip() {
196        let context = DatadogContext {
197            trace_id: random_range(1..=u128::MAX),
198            parent_id: random_range(1..=u64::MAX),
199        };
200
201        let mut headers = HeaderMap::new();
202        headers.inject_trace_context::<W3CTraceContextHeaders>(context);
203        let parsed = headers.extract_trace_context::<W3CTraceContextHeaders>();
204
205        assert_eq!(context.trace_id, parsed.trace_id);
206        assert_eq!(context.parent_id, parsed.parent_id);
207    }
208
209    #[test]
210    fn empty_context_doesnt_produce_w3c_trace_header() {
211        let mut headers = HeaderMap::new();
212        headers.inject_trace_context::<W3CTraceContextHeaders>(DatadogContext::default());
213        assert!(headers.is_empty());
214    }
215
216    #[test]
217    fn w3c_trace_header_with_wrong_version_produces_empty_context() {
218        let headers = HeaderMap::from_iter([(
219            W3C_TRACEPARENT_HEADER,
220            "01-00000000000000000000000000000001-0000000000000001-01"
221                .parse()
222                .unwrap(),
223        )]);
224        let context = headers.extract_trace_context::<W3CTraceContextHeaders>();
225        assert!(context.is_empty());
226    }
227
228    #[test]
229    fn w3c_trace_header_without_sampling_flag_produces_empty_context() {
230        let headers = HeaderMap::from_iter([(
231            W3C_TRACEPARENT_HEADER,
232            "00-00000000000000000000000000000001-0000000000000001-00"
233                .parse()
234                .unwrap(),
235        )]);
236        let context = headers.extract_trace_context::<W3CTraceContextHeaders>();
237        assert!(context.is_empty());
238    }
239
240    #[test]
241    fn datadog_headers_round_trip() {
242        let context = DatadogContext {
243            // We want to check that the upper 64 bits are preserved.
244            trace_id: random_range((u64::MAX as u128 + 1)..=u128::MAX),
245            parent_id: random_range(1..=u64::MAX),
246        };
247
248        let mut headers = HeaderMap::new();
249        headers.inject_trace_context::<DatadogHeaders>(context);
250        let parsed = headers.extract_trace_context::<DatadogHeaders>();
251
252        assert_eq!(context.trace_id, parsed.trace_id);
253        assert_eq!(context.parent_id, parsed.parent_id);
254    }
255
256    #[test]
257    fn empty_context_doesnt_produce_datadog_headers() {
258        let mut headers = HeaderMap::new();
259        headers.inject_trace_context::<DatadogHeaders>(DatadogContext::default());
260        assert!(headers.is_empty());
261    }
262
263    #[test]
264    fn datadog_headers_without_sampling_produce_empty_context() {
265        let headers = HeaderMap::from_iter([(
266            DATADOG_SAMPLING_PRIORITY_HEADER,
267            HeaderValue::from_static("0"),
268        )]);
269        let context = headers.extract_trace_context::<DatadogHeaders>();
270        assert!(context.is_empty());
271    }
272
273    #[test]
274    fn from_datadog_headers_works_without_tags_header() {
275        let headers = HeaderMap::from_iter([
276            (
277                DATADOG_TRACE_ID_HEADER,
278                HeaderValue::from_static("0000000000000001"),
279            ),
280            (
281                DATADOG_PARENT_ID_HEADER,
282                HeaderValue::from_static("0000000000000001"),
283            ),
284            (
285                DATADOG_SAMPLING_PRIORITY_HEADER,
286                HeaderValue::from_static("1"),
287            ),
288        ]);
289        let context = headers.extract_trace_context::<DatadogHeaders>();
290        assert_eq!(context.trace_id, 0x0000000000000001);
291        assert_eq!(context.parent_id, 0x0000000000000001);
292    }
293
294    #[test]
295    fn from_datadog_header_works_with_other_tags() {
296        let headers = HeaderMap::from_iter([
297            (
298                DATADOG_TRACE_ID_HEADER,
299                HeaderValue::from_static("0000000000000001"),
300            ),
301            (
302                DATADOG_PARENT_ID_HEADER,
303                HeaderValue::from_static("0000000000000001"),
304            ),
305            (
306                DATADOG_SAMPLING_PRIORITY_HEADER,
307                HeaderValue::from_static("1"),
308            ),
309            (
310                DATADOG_TAGS_HEADER,
311                HeaderValue::from_static("other=tags,_dd.p.tid=0000000000000002,more=tags"),
312            ),
313        ]);
314        let context = headers.extract_trace_context::<DatadogHeaders>();
315        assert_eq!(context.trace_id, 0x20000000000000001);
316        assert_eq!(context.parent_id, 0x0000000000000001);
317    }
318}
tracing_datadog/http.rs

tracing_datadog/
http.rs
