tracing_datadog/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use jiff::{Timestamp, Zoned};
4use reqwest::header::{self, HeaderMap, HeaderName, HeaderValue};
5use rmp_serde::Serializer as MpSerializer;
6use serde::{Serialize, Serializer, ser::SerializeMap};
7use std::{
8    collections::HashMap,
9    fmt::{Debug, Display, Formatter},
10    marker::PhantomData,
11    ops::DerefMut,
12    sync::{Arc, Mutex, mpsc},
13    thread::{sleep, spawn},
14    time::{Duration, SystemTime, UNIX_EPOCH},
15};
16use tracing_core::{
17    Event, Field, Level, Subscriber,
18    field::Visit,
19    span::{Attributes, Id, Record},
20};
21use tracing_subscriber::{
22    Layer,
23    layer::Context,
24    registry::{LookupSpan, Scope},
25};
26
27/// A [`Layer`] that sends traces to Datadog.
28///
29/// ```
30/// # use tracing_subscriber::prelude::*;
31/// # use tracing_datadog::DatadogTraceLayer;
32/// tracing_subscriber::registry()
33///    .with(
34///        DatadogTraceLayer::builder()
35///            .service("my-service")
36///            .env("production")
37///            .version("git sha")
38///            .agent_address("localhost:8126")
39///            .build()
40///            .expect("failed to build DatadogTraceLayer"),
41///    )
42///    .init();
43/// ```
44#[derive(Debug)]
45pub struct DatadogTraceLayer<S> {
46    buffer: Arc<Mutex<Vec<DatadogSpan>>>,
47    service: String,
48    default_tags: HashMap<String, String>,
49    logging_enabled: bool,
50    #[cfg(feature = "http")]
51    with_context: http::WithContext,
52    shutdown: mpsc::Sender<()>,
53    _registry: PhantomData<S>,
54}
55
56impl<S> DatadogTraceLayer<S>
57where
58    S: Subscriber + for<'a> LookupSpan<'a>,
59{
60    /// Creates a builder to construct a [`DatadogTraceLayer`].
61    pub fn builder() -> DatadogTraceLayerBuilder<S> {
62        DatadogTraceLayerBuilder {
63            service: None,
64            default_tags: HashMap::new(),
65            agent_address: None,
66            container_id: None,
67            logging_enabled: false,
68            phantom_data: Default::default(),
69        }
70    }
71
72    #[cfg(feature = "http")]
73    fn get_context(
74        dispatch: &tracing_core::Dispatch,
75        id: &Id,
76        f: &mut dyn FnMut(&mut DatadogSpan),
77    ) {
78        let subscriber = dispatch
79            .downcast_ref::<S>()
80            .expect("Subscriber did not downcast to expected type, this is a bug");
81        let span = subscriber.span(id).expect("Span not found, this is a bug");
82
83        let mut extensions = span.extensions_mut();
84        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>() {
85            f(dd_span);
86        }
87    }
88}
89
90impl<S> Drop for DatadogTraceLayer<S> {
91    fn drop(&mut self) {
92        let _ = self.shutdown.send(());
93    }
94}
95
96impl<S> Layer<S> for DatadogTraceLayer<S>
97where
98    S: Subscriber + for<'a> LookupSpan<'a>,
99{
100    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
101        let span = ctx.span(id).expect("Span not found, this is a bug");
102        let mut extensions = span.extensions_mut();
103
104        let trace_id = span
105            .parent()
106            .map(|parent| {
107                parent
108                    .extensions()
109                    .get::<DatadogSpan>()
110                    .expect("Parent span didn't have a DatadogSpan extension, this is a bug")
111                    .trace_id
112            })
113            .unwrap_or(rand::random_range(1..=u64::MAX));
114
115        debug_assert!(trace_id != 0, "Trace ID is zero, this is a bug");
116
117        let mut dd_span = DatadogSpan {
118            name: span.name().to_string(),
119            service: self.service.clone(),
120            r#type: "internal".into(),
121            span_id: span.id().into_u64(),
122            start: epoch_ns(),
123            parent_id: span
124                .parent()
125                .map(|parent| parent.id().into_u64())
126                .unwrap_or_default(),
127            trace_id,
128            meta: self.default_tags.clone(),
129            ..Default::default()
130        };
131
132        attrs.record(&mut SpanAttributeVisitor::new(&mut dd_span));
133
134        extensions.insert(dd_span);
135    }
136
137    fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
138        let span = ctx.span(id).expect("Span not found, this is a bug");
139        let mut extensions = span.extensions_mut();
140
141        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>() {
142            values.record(&mut SpanAttributeVisitor::new(dd_span));
143        }
144    }
145
146    fn on_follows_from(&self, id: &Id, follows: &Id, ctx: Context<'_, S>) {
147        let span = ctx.span(id).expect("Span not found, this is a bug");
148        let mut extensions = span.extensions_mut();
149
150        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>() {
151            dd_span.parent_id = follows.into_u64();
152        }
153    }
154
155    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
156        if !self.logging_enabled {
157            return;
158        }
159
160        let mut fields = {
161            let mut visitor = FieldVisitor::default();
162            event.record(&mut visitor);
163            visitor.fields
164        };
165
166        fields.extend(
167            ctx.event_scope(event)
168                .into_iter()
169                .flat_map(Scope::from_root)
170                .flat_map(|span| match span.extensions().get::<DatadogSpan>() {
171                    Some(dd_span) => dd_span.meta.clone(),
172                    None => panic!("DatadogSpan extension not found, this is a bug"),
173                }),
174        );
175
176        let message = fields.remove("message").unwrap_or_default();
177
178        let (trace_id, span_id) = ctx
179            .lookup_current()
180            .and_then(|span| {
181                span.extensions()
182                    .get::<DatadogSpan>()
183                    .map(|dd_span| (Some(dd_span.trace_id), Some(dd_span.span_id)))
184            })
185            .unwrap_or_default();
186
187        let log = DatadogLog {
188            timestamp: Zoned::now().timestamp(),
189            level: event.metadata().level().to_owned(),
190            message,
191            trace_id,
192            span_id,
193            fields,
194        };
195
196        let serialized = serde_json::to_string(&log).expect("Failed to serialize log");
197
198        println!("{serialized}");
199    }
200
201    fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
202        let span = ctx.span(id).expect("Span not found, this is a bug");
203        let mut extensions = span.extensions_mut();
204
205        let now = epoch_ns();
206
207        match extensions.get_mut::<DatadogSpan>() {
208            Some(dd_span) if dd_span.start == 0 => dd_span.start = now,
209            _ => {}
210        }
211    }
212
213    fn on_exit(&self, id: &Id, ctx: Context<'_, S>) {
214        let span = ctx.span(id).expect("Span not found, this is a bug");
215        let mut extensions = span.extensions_mut();
216
217        let now = epoch_ns();
218
219        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>() {
220            dd_span.duration = now - dd_span.start
221        }
222    }
223
224    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
225        let span = ctx.span(&id).expect("Span not found, this is a bug");
226        let mut extensions = span.extensions_mut();
227
228        if let Some(dd_span) = extensions.remove::<DatadogSpan>() {
229            self.buffer.lock().unwrap().push(dd_span);
230        }
231    }
232
233    // SAFETY: This is safe because the `WithContext` function pointer is valid
234    // for the lifetime of `&self`.
235    #[cfg(feature = "http")]
236    unsafe fn downcast_raw(&self, id: std::any::TypeId) -> Option<*const ()> {
237        match id {
238            id if id == std::any::TypeId::of::<Self>() => Some(self as *const _ as *const ()),
239            id if id == std::any::TypeId::of::<http::WithContext>() => {
240                Some(&self.with_context as *const _ as *const ())
241            }
242            _ => None,
243        }
244    }
245}
246
247/// A builder for [`DatadogTraceLayer`].
248pub struct DatadogTraceLayerBuilder<S> {
249    service: Option<String>,
250    default_tags: HashMap<String, String>,
251    agent_address: Option<String>,
252    container_id: Option<String>,
253    logging_enabled: bool,
254    phantom_data: PhantomData<S>,
255}
256
257/// An error that can occur when building a [`DatadogTraceLayer`].
258#[derive(Debug)]
259pub struct BuilderError(&'static str);
260
261impl Display for BuilderError {
262    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
263        f.write_str(self.0)
264    }
265}
266
267impl std::error::Error for BuilderError {}
268
269const DATADOG_LANGUAGE_HEADER: HeaderName = HeaderName::from_static("datadog-meta-lang");
270const DATADOG_TRACER_VERSION_HEADER: HeaderName =
271    HeaderName::from_static("datadog-meta-tracer-version");
272const DATADOG_CONTAINER_ID_HEADER: HeaderName = HeaderName::from_static("datadog-container-id");
273
274impl<S> DatadogTraceLayerBuilder<S>
275where
276    S: Subscriber + for<'a> LookupSpan<'a>,
277{
278    /// Sets the `service`. This is required.
279    pub fn service(mut self, service: impl Into<String>) -> Self {
280        self.service = Some(service.into());
281        self
282    }
283
284    /// Sets the `env`. This is required.
285    pub fn env(mut self, env: impl Into<String>) -> Self {
286        self.default_tags.insert("env".into(), env.into());
287        self
288    }
289
290    /// Sets the `version`. This is required.
291    pub fn version(mut self, version: impl Into<String>) -> Self {
292        self.default_tags.insert("version".into(), version.into());
293        self
294    }
295
296    /// Sets the `agent_address`. This is required.
297    pub fn agent_address(mut self, agent_address: impl Into<String>) -> Self {
298        self.agent_address = Some(agent_address.into());
299        self
300    }
301
302    /// Adds a fixed default tag to all spans.
303    ///
304    /// This can be used multiple times for several tags.
305    ///
306    /// Default tags are overridden by tags set explicitly on a span.
307    pub fn default_tag(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
308        let _ = self.default_tags.insert(key.into(), value.into());
309        self
310    }
311
312    /// Sets the container ID. This enables infrastructure metrics in APM for supported platforms.
313    pub fn container_id(mut self, container_id: impl Into<String>) -> Self {
314        self.container_id = Some(container_id.into());
315        self
316    }
317
318    /// Enables or disables structured logging with trace correlation to stdout.
319    /// Disabled by default.
320    pub fn enable_logs(mut self, enable_logs: bool) -> Self {
321        self.logging_enabled = enable_logs;
322        self
323    }
324
325    /// Consumes the builder to construct the tracing layer.
326    pub fn build(self) -> Result<DatadogTraceLayer<S>, BuilderError> {
327        let Some(service) = self.service else {
328            return Err(BuilderError("service is required"));
329        };
330        if !self.default_tags.contains_key("env") {
331            return Err(BuilderError("env is required"));
332        };
333        if !self.default_tags.contains_key("version") {
334            return Err(BuilderError("version is required"));
335        };
336        let Some(agent_address) = self.agent_address else {
337            return Err(BuilderError("agent_address is required"));
338        };
339        let container_id = match self.container_id {
340            Some(s) => Some(
341                s.parse::<HeaderValue>()
342                    .map_err(|_| BuilderError("Failed to parse container ID into header"))?,
343            ),
344            _ => None,
345        };
346
347        let buffer = Arc::new(Mutex::new(Vec::new()));
348        let exporter_buffer = buffer.clone();
349        let url = format!("http://{}/v0.4/traces", agent_address);
350        let (tx, rx) = mpsc::channel();
351
352        spawn(move || {
353            let client = {
354                let mut default_headers = HeaderMap::from_iter([(
355                    DATADOG_LANGUAGE_HEADER,
356                    HeaderValue::from_static("rust"),
357                )]);
358
359                if let Some(container_id) = container_id {
360                    default_headers.insert(DATADOG_CONTAINER_ID_HEADER, container_id);
361                };
362
363                reqwest::blocking::Client::builder()
364                    .default_headers(default_headers)
365                    .retry(reqwest::retry::for_host(agent_address).max_retries_per_request(2))
366                    .build()
367                    .expect("Failed to build reqwest client")
368            };
369            let mut spans = Vec::new();
370
371            loop {
372                if rx.try_recv().is_ok() {
373                    break;
374                }
375
376                sleep(Duration::from_secs(1));
377
378                std::mem::swap(&mut spans, exporter_buffer.lock().unwrap().deref_mut());
379
380                if spans.is_empty() {
381                    continue;
382                }
383
384                let mut body = vec![0b10010001];
385                let _ = spans
386                    .serialize(&mut MpSerializer::new(&mut body).with_struct_map())
387                    .inspect_err(|error| println!("Error serializing spans: {error:?}"));
388
389                spans.clear();
390
391                let _ = client
392                    .post(&url)
393                    .header(DATADOG_TRACER_VERSION_HEADER, "v1.27.0")
394                    .header(header::CONTENT_TYPE, "application/msgpack")
395                    .body(body)
396                    .send()
397                    .inspect_err(|error| println!("Error exporting spans: {error:?}"));
398            }
399        });
400
401        Ok(DatadogTraceLayer {
402            buffer,
403            service,
404            default_tags: self.default_tags,
405            logging_enabled: self.logging_enabled,
406            #[cfg(feature = "http")]
407            with_context: http::WithContext(DatadogTraceLayer::<S>::get_context),
408            shutdown: tx,
409            _registry: PhantomData,
410        })
411    }
412}
413
414/// Returns the current system time as nanoseconds since 1970.
415fn epoch_ns() -> i64 {
416    SystemTime::now()
417        .duration_since(UNIX_EPOCH)
418        .expect("SystemTime is before UNIX epoch")
419        .as_nanos() as i64
420}
421
422/// The v0.4 Datadog trace API format for spans. This is what we write to MessagePack.
423#[derive(Default, Debug, Serialize)]
424struct DatadogSpan {
425    trace_id: u64,
426    span_id: u64,
427    parent_id: u64,
428    start: i64,
429    duration: i64,
430    /// This is what maps to the operation in Datadog.
431    name: String,
432    service: String,
433    r#type: String,
434    resource: String,
435    meta: HashMap<String, String>,
436    error_code: i32,
437}
438
439/// A visitor that converts tracing span attributes to a [`DatadogSpan`].
440struct SpanAttributeVisitor<'a> {
441    dd_span: &'a mut DatadogSpan,
442}
443
444impl<'a> SpanAttributeVisitor<'a> {
445    fn new(dd_span: &'a mut DatadogSpan) -> Self {
446        Self { dd_span }
447    }
448}
449
450impl<'a> Visit for SpanAttributeVisitor<'a> {
451    fn record_str(&mut self, field: &Field, value: &str) {
452        // Strings are broken out because their debug representation includes quotation marks.
453        match field.name() {
454            "service" => self.dd_span.service = value.to_string(),
455            "span.type" => self.dd_span.r#type = value.to_string(),
456            "operation" => self.dd_span.name = value.to_string(),
457            "resource" => self.dd_span.resource = value.to_string(),
458            name => {
459                self.dd_span
460                    .meta
461                    .insert(name.to_string(), value.to_string());
462            }
463        };
464    }
465
466    fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
467        match field.name() {
468            "service" => self.dd_span.service = format!("{value:?}"),
469            "span.type" => self.dd_span.r#type = format!("{value:?}"),
470            "operation" => self.dd_span.name = format!("{value:?}"),
471            "resource" => self.dd_span.resource = format!("{value:?}"),
472            name => {
473                self.dd_span
474                    .meta
475                    .insert(name.to_string(), format!("{value:?}"));
476            }
477        };
478    }
479}
480
481/// The Datadog structure log format. This is what we write to JSON.
482struct DatadogLog {
483    timestamp: Timestamp,
484    level: Level,
485    message: String,
486    trace_id: Option<u64>,
487    span_id: Option<u64>,
488    fields: HashMap<String, String>,
489}
490
491impl Serialize for DatadogLog {
492    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
493    where
494        S: Serializer,
495    {
496        let mut map = serializer.serialize_map(None)?;
497        map.serialize_entry("timestamp", &self.timestamp)?;
498        map.serialize_entry("level", &self.level.as_str())?;
499        map.serialize_entry("message", &self.message)?;
500        if let Some(trace_id) = &self.trace_id {
501            map.serialize_entry("dd.trace_id", &trace_id)?;
502        }
503        if let Some(span_id) = &self.span_id {
504            map.serialize_entry("dd.span_id", &span_id)?;
505        }
506        for (key, value) in &self.fields {
507            map.serialize_entry(&format!("fields.{key}"), value)?;
508        }
509        map.end()
510    }
511}
512
513/// A visitor that collects tracing attributes into a map.
514#[derive(Default)]
515struct FieldVisitor {
516    fields: HashMap<String, String>,
517}
518
519impl Visit for FieldVisitor {
520    fn record_str(&mut self, field: &Field, value: &str) {
521        self.fields
522            .insert(field.name().to_string(), value.to_string());
523    }
524
525    fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
526        self.fields
527            .insert(field.name().to_string(), format!("{value:?}"));
528    }
529}
530
531#[cfg(feature = "http")]
532#[doc = "Functionality for working with distributed tracing HTTP headers"]
533pub mod http {
534    use crate::DatadogSpan;
535    use http::{HeaderMap, HeaderName};
536    use tracing_core::{Dispatch, span::Id};
537
538    const W3C_TRACEPARENT_HEADER: HeaderName = HeaderName::from_static("traceparent");
539
540    /// The trace context for distributed tracing. This is a subset of the W3C trace context
541    /// which allows stitching together traces with spans from different services.
542    #[derive(Copy, Clone, Default)]
543    pub struct DatadogContext {
544        trace_id: u128,
545        parent_id: u64,
546    }
547
548    impl DatadogContext {
549        /// Parses a context for distributed tracing from W3C trace context headers.
550        ///
551        /// This would be useful in HTTP server middleware.
552        ///
553        /// ```
554        /// # let request = http::Request::builder().body(()).unwrap();
555        /// use tracing_datadog::http::{DatadogContext, DistributedTracingContext};
556        ///
557        /// // Construct a new span.
558        /// let span = tracing::info_span!("http.request");
559        ///
560        /// // Set the context on the span based on request headers.
561        /// span.set_context(DatadogContext::from_w3c_headers(request.headers()));
562        /// ```
563        ///
564        /// An alternative use case is setting the context on the current span, for example
565        /// within `#[instrument]`ed functions.
566        ///
567        /// ```
568        /// # let request = http::Request::builder().body(()).unwrap();
569        /// use tracing_datadog::http::{DatadogContext, DistributedTracingContext};
570        ///
571        /// tracing::Span::current().set_context(DatadogContext::from_w3c_headers(request.headers()));
572        /// ```
573        pub fn from_w3c_headers(headers: &HeaderMap) -> Self {
574            Self::parse_w3c_headers(headers).unwrap_or_default()
575        }
576
577        fn parse_w3c_headers(headers: &HeaderMap) -> Option<Self> {
578            let header = headers.get(W3C_TRACEPARENT_HEADER)?.to_str().ok()?;
579
580            let parts: Vec<&str> = header.split('-').collect();
581            if parts.len() != 4 {
582                return None;
583            }
584
585            let Some(0) = u8::from_str_radix(parts[0], 16).ok() else {
586                // Wrong version.
587                return None;
588            };
589
590            let Some(0x01) = u8::from_str_radix(parts[3], 16).ok().map(|n| n & 0x01) else {
591                // Not sampled.
592                return None;
593            };
594
595            let trace_id = u128::from_str_radix(parts[1], 16).ok()?;
596            let parent_id = u64::from_str_radix(parts[2], 16).ok()?;
597
598            Some(Self {
599                trace_id,
600                parent_id,
601            })
602        }
603
604        /// Serializes a context for distributed tracing to W3C trace context headers.
605        ///
606        /// ```
607        /// # use http::Request;
608        /// use tracing_datadog::http::DistributedTracingContext;
609        ///
610        /// // Build the request.
611        /// let mut request = Request::builder().body(()).unwrap();
612        ///
613        /// // Inject distributed tracing headers.
614        /// request.headers_mut().extend(tracing::Span::current().get_context().to_w3c_headers());
615        ///
616        /// // Execute the request.
617        /// // ..
618        /// ```
619        pub fn to_w3c_headers(&self) -> HeaderMap {
620            if self.is_empty() {
621                return Default::default();
622            }
623
624            let header = format!(
625                "{version:02x}-{trace_id:032x}-{parent_id:016x}-{trace_flags:02x}",
626                version = 0,
627                trace_id = self.trace_id,
628                parent_id = self.parent_id,
629                trace_flags = 1,
630            );
631
632            HeaderMap::from_iter([(W3C_TRACEPARENT_HEADER, header.parse().unwrap())])
633        }
634
635        /// Returns `true` if the context is empty, i.e. if it does not contain a trace ID or
636        /// a parent ID.
637        fn is_empty(&self) -> bool {
638            self.trace_id == 0 || self.parent_id == 0
639        }
640    }
641
642    // This function "remembers" the types of the subscriber so that we can downcast to something
643    // aware of them without knowing those types at the call site. Adapted from tracing-error.
644    #[derive(Debug)]
645    pub(crate) struct WithContext(
646        #[allow(clippy::type_complexity)]
647        pub(crate)  fn(&Dispatch, &Id, f: &mut dyn FnMut(&mut DatadogSpan)),
648    );
649
650    impl WithContext {
651        pub(crate) fn with_context(
652            &self,
653            dispatch: &Dispatch,
654            id: &Id,
655            mut f: &mut dyn FnMut(&mut DatadogSpan),
656        ) {
657            self.0(dispatch, id, &mut f);
658        }
659    }
660
661    pub trait DistributedTracingContext {
662        /// Gets the context for distributed tracing from the current span.
663        fn get_context(&self) -> DatadogContext;
664
665        /// Sets the context for distributed tracing on the current span.
666        fn set_context(&self, context: DatadogContext);
667    }
668
669    impl DistributedTracingContext for tracing::Span {
670        fn get_context(&self) -> DatadogContext {
671            let mut ctx = None;
672
673            self.with_subscriber(|(id, subscriber)| {
674                let Some(get_context) = subscriber.downcast_ref::<WithContext>() else {
675                    return;
676                };
677                get_context.with_context(subscriber, id, &mut |dd_span| {
678                    ctx = Some(DatadogContext {
679                        // NB Trace IDs can be 128-bit nowadays, but the 0.4 API still uses 64-bit.
680                        trace_id: dd_span.trace_id as u128,
681                        parent_id: dd_span.span_id,
682                    })
683                });
684            });
685
686            ctx.unwrap_or_default()
687        }
688
689        fn set_context(&self, context: DatadogContext) {
690            // Avoid setting a null context.
691            if context.is_empty() {
692                return;
693            }
694
695            self.with_subscriber(move |(id, subscriber)| {
696                let Some(get_context) = subscriber.downcast_ref::<WithContext>() else {
697                    return;
698                };
699                get_context.with_context(subscriber, id, &mut |dd_span| {
700                    // NB Trace IDs can be 128-bit nowadays, but the 0.4 API still uses 64-bit.
701                    dd_span.trace_id = context.trace_id as u64;
702                    dd_span.parent_id = context.parent_id;
703                })
704            });
705        }
706    }
707
708    #[cfg(test)]
709    mod tests {
710        use super::*;
711        use crate::DatadogTraceLayer;
712        use rand::random_range;
713        use tracing::info_span;
714        use tracing_subscriber::layer::SubscriberExt;
715
716        #[test]
717        fn w3c_trace_header_round_trip() {
718            let context = DatadogContext {
719                trace_id: random_range(1..=u128::MAX),
720                parent_id: random_range(1..=u64::MAX),
721            };
722
723            let headers = context.to_w3c_headers();
724            let parsed = DatadogContext::from_w3c_headers(&headers);
725
726            assert_eq!(context.trace_id, parsed.trace_id);
727            assert_eq!(context.parent_id, parsed.parent_id);
728        }
729
730        #[test]
731        fn empty_context_doesnt_produce_w3c_trace_header() {
732            assert!(DatadogContext::default().to_w3c_headers().is_empty());
733        }
734
735        #[test]
736        fn w3c_trace_header_with_wrong_version_produces_empty_context() {
737            let headers = HeaderMap::from_iter([(
738                HeaderName::from_static("traceparent"),
739                "01-00000000000000000000000000000001-0000000000000001-01"
740                    .parse()
741                    .unwrap(),
742            )]);
743            let context = DatadogContext::from_w3c_headers(&headers);
744            assert!(context.is_empty());
745        }
746
747        #[test]
748        fn w3c_trace_header_without_sampling_flag_produces_empty_context() {
749            let headers = HeaderMap::from_iter([(
750                HeaderName::from_static("traceparent"),
751                "00-00000000000000000000000000000001-0000000000000001-00"
752                    .parse()
753                    .unwrap(),
754            )]);
755            let context = DatadogContext::from_w3c_headers(&headers);
756            assert!(context.is_empty());
757        }
758
759        #[test]
760        fn span_context_round_trip() {
761            tracing::subscriber::with_default(
762                tracing_subscriber::registry().with(
763                    DatadogTraceLayer::builder()
764                        .service("test-service")
765                        .env("test")
766                        .version("test-version")
767                        .agent_address("localhost:8126")
768                        .build()
769                        .unwrap(),
770                ),
771                || {
772                    let context = DatadogContext {
773                        // Need to limit the size here as we only track 64-bit trace IDs.
774                        trace_id: random_range(1..=u64::MAX) as u128,
775                        parent_id: random_range(1..=u64::MAX),
776                    };
777
778                    let span = info_span!("test");
779
780                    span.set_context(context);
781                    let result = span.get_context();
782
783                    assert_eq!(context.trace_id, result.trace_id);
784                    // NB Parent ID is asymmetrical, this span's ID becomes the next span's parent ID.
785                    assert_eq!(span.id().unwrap().into_u64(), result.parent_id);
786                },
787            );
788        }
789
790        #[test]
791        fn empty_span_context_does_not_erase_trace_id() {
792            tracing::subscriber::with_default(
793                tracing_subscriber::registry().with(
794                    DatadogTraceLayer::builder()
795                        .service("test-service")
796                        .env("test")
797                        .version("test-version")
798                        .agent_address("localhost:8126")
799                        .build()
800                        .unwrap(),
801                ),
802                || {
803                    let context = DatadogContext::default();
804
805                    let span = info_span!("test");
806
807                    span.set_context(context);
808                    let result = span.get_context();
809
810                    assert_ne!(result.trace_id, 0);
811                },
812            );
813        }
814    }
815}
816
817#[cfg(test)]
818mod tests {
819    use super::*;
820
821    #[test]
822    fn builder_builds_successfully() {
823        assert!(
824            DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
825                .service("test-service")
826                .env("test")
827                .version("test-version")
828                .agent_address("localhost:8126")
829                .build()
830                .is_ok()
831        );
832    }
833
834    #[test]
835    fn service_is_required() {
836        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
837            .env("test")
838            .version("test-version")
839            .agent_address("localhost:8126")
840            .build();
841        assert!(result.unwrap_err().to_string().contains("service"));
842    }
843
844    #[test]
845    fn env_is_required() {
846        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
847            .service("test-service")
848            .version("test-version")
849            .agent_address("localhost:8126")
850            .build();
851        assert!(result.unwrap_err().to_string().contains("env"));
852    }
853
854    #[test]
855    fn version_is_required() {
856        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
857            .service("test-service")
858            .env("test")
859            .agent_address("localhost:8126")
860            .build();
861        assert!(result.unwrap_err().to_string().contains("version"));
862    }
863
864    #[test]
865    fn agent_address_is_required() {
866        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
867            .service("test-service")
868            .env("test")
869            .version("test-version")
870            .build();
871        assert!(result.unwrap_err().to_string().contains("agent_address"));
872    }
873
874    #[test]
875    fn default_default_tags_include_env_and_version() {
876        let layer: DatadogTraceLayer<tracing_subscriber::Registry> = DatadogTraceLayer::builder()
877            .service("test-service")
878            .env("test")
879            .version("test-version")
880            .agent_address("localhost:8126")
881            .build()
882            .unwrap();
883        let default_tags = &layer.default_tags;
884        assert_eq!(default_tags["env"], "test");
885        assert_eq!(default_tags["version"], "test-version");
886    }
887
888    #[test]
889    fn default_tags_can_be_added() {
890        let layer: DatadogTraceLayer<tracing_subscriber::Registry> = DatadogTraceLayer::builder()
891            .service("test-service")
892            .env("test")
893            .version("test-version")
894            .agent_address("localhost:8126")
895            .default_tag("foo", "bar")
896            .default_tag("baz", "qux")
897            .build()
898            .unwrap();
899        let default_tags = &layer.default_tags;
900        assert_eq!(default_tags["foo"], "bar");
901        assert_eq!(default_tags["baz"], "qux");
902    }
903}
tracing_datadog/lib.rs

tracing_datadog/
lib.rs
