tower_http_tracing/
datadog.rs

1//! datadog integration
2//!
3//! This is alternative to using specialized OTEL exporter: <https://crates.io/crates/opentelemetry-datadog>
4
5use core::fmt;
6
7pub use tracing_datadog;
8use tracing_datadog::context::{self, DatadogContext, Strategy};
9
10///W3C header name
11pub const W3C_TRACEPARENT_NAME: http::HeaderName = http::HeaderName::from_static("traceparent");
12const W3C_VERSION: u8 = 0;
13const TRACE_SAMPLED_FLAG: u8 = 0x01;
14
15///Propagation strategy for W3C header
16pub struct Propagation;
17
18struct BytesWriter(bytes::BytesMut);
19
20impl BytesWriter {
21    #[inline(always)]
22    fn finish(self) -> bytes::Bytes {
23        self.0.freeze()
24    }
25}
26
27impl fmt::Write for BytesWriter {
28    #[inline(always)]
29    fn write_str(&mut self, input: &str) -> fmt::Result {
30        self.0.extend_from_slice(input.as_bytes());
31        Ok(())
32    }
33}
34
35//A more efficient re-implementation of tracing_datadog's default impl
36impl Strategy<http::HeaderMap> for Propagation {
37    fn inject(headers: &mut http::HeaderMap, context: DatadogContext) {
38        use fmt::Write;
39
40        let DatadogContext { trace_id, parent_id } = &context;
41        //Make DatadogContext::is_empty() public
42        if *trace_id == 0 || *parent_id == 0 {
43            return;
44        }
45        let mut out = BytesWriter(bytes::BytesMut::new());
46
47        //Cannot fail, will panic on OOM
48        let _ = write!(out, "{W3C_VERSION:02x}-{trace_id:032x}-{parent_id:016x}-{TRACE_SAMPLED_FLAG:02x}");
49
50        let value = unsafe {
51            //BytesWriter is guaranteed to only write via `fmt::Write` so all content is valid utf-8
52            http::HeaderValue::from_maybe_shared_unchecked(out.finish())
53        };
54        headers.insert(W3C_TRACEPARENT_NAME, value);
55    }
56
57    fn extract(headers: &http::HeaderMap) -> DatadogContext {
58        let header = match headers.get(W3C_TRACEPARENT_NAME).and_then(|value| value.to_str().ok()) {
59            Some(header) => header,
60            None => return DatadogContext::default(),
61        };
62
63        let mut parts = header.split('-');
64
65        match (parts.next(), parts.next(), parts.next(), parts.next()) {
66            (Some(version), Some(trace_id), Some(parent_id), Some(trace_flag)) => {
67                if u8::from_str_radix(version, 16).map(|version| version != W3C_VERSION).unwrap_or(true) {
68                    //Cannot recognize version, then we cannot parse it
69                    return DatadogContext::default();
70                }
71
72                if u8::from_str_radix(trace_flag, 16).map(|flag| flag & TRACE_SAMPLED_FLAG != TRACE_SAMPLED_FLAG).unwrap_or(true) {
73                    //Not sampled = no propagation
74                    return DatadogContext::default();
75                }
76
77                match (u128::from_str_radix(trace_id, 16), u64::from_str_radix(parent_id, 16)) {
78                    (Ok(trace_id), Ok(parent_id)) => DatadogContext {
79                        trace_id,
80                        parent_id,
81                    },
82                    _ => DatadogContext::default()
83                }
84            },
85            _ => DatadogContext::default(),
86        }
87    }
88}
89
90#[inline(always)]
91///Extracts datadog context from `request` propagating it as `span`'s parent
92///
93///Datadog format is somewhat different from opentelemetry:
94///<https://docs.datadoghq.com/tracing/trace_collection/trace_context_propagation/?tab=java#datadog-format>
95pub fn on_request<T>(span: &tracing::Span, request: &http::Request<T>) {
96    let context = Propagation::extract(request.headers());
97    context::TracingContextExt::set_context(span, context);
98}
99
100#[inline(always)]
101///Propagates context into response's headers
102pub fn on_response_ok<T>(span: &tracing::Span, response: &mut http::Response<T>) {
103    let context = tracing_datadog::context::TracingContextExt::get_context(span);
104    Propagation::inject(response.headers_mut(), context);
105}
106
107#[inline(always)]
108///No `error` propagation is done aside from default one
109pub fn on_response_error(_span: &tracing::Span, _error: &impl std::error::Error) {
110}