tower_http_tracing/
datadog.rs1use core::fmt;
6
7pub use tracing_datadog;
8use tracing_datadog::context::{self, DatadogContext, Strategy};
9
10pub const W3C_TRACEPARENT_NAME: http::HeaderName = http::HeaderName::from_static("traceparent");
12const W3C_VERSION: u8 = 0;
13const TRACE_SAMPLED_FLAG: u8 = 0x01;
14
15pub 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
35impl 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 if *trace_id == 0 || *parent_id == 0 {
43 return;
44 }
45 let mut out = BytesWriter(bytes::BytesMut::new());
46
47 let _ = write!(out, "{W3C_VERSION:02x}-{trace_id:032x}-{parent_id:016x}-{TRACE_SAMPLED_FLAG:02x}");
49
50 let value = unsafe {
51 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 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 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)]
91pub 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)]
101pub 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)]
108pub fn on_response_error(_span: &tracing::Span, _error: &impl std::error::Error) {
110}