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    /// The trace context for distributed tracing. This is a subset of the W3C trace context
539    /// which allows stitching together traces with spans from different services.
540    #[derive(Copy, Clone, Default)]
541    pub struct DatadogContext {
542        trace_id: u128,
543        parent_id: u64,
544    }
545
546    impl DatadogContext {
547        /// Parses a context for distributed tracing from W3C trace context headers.
548        ///
549        /// This would be useful in HTTP server middleware.
550        ///
551        /// ```
552        /// # let request = http::Request::builder().body(()).unwrap();
553        /// use tracing_datadog::http::{DatadogContext, DistributedTracingContext};
554        ///
555        /// // Construct a new span.
556        /// let span = tracing::info_span!("http.request");
557        ///
558        /// // Set the context on the span based on request headers.
559        /// span.set_context(DatadogContext::from_w3c_headers(request.headers()));
560        /// ```
561        ///
562        /// An alternative use case is setting the context on the current span, for example
563        /// within `#[instrument]`ed functions.
564        ///
565        /// ```
566        /// # let request = http::Request::builder().body(()).unwrap();
567        /// use tracing_datadog::http::{DatadogContext, DistributedTracingContext};
568        ///
569        /// tracing::Span::current().set_context(DatadogContext::from_w3c_headers(request.headers()));
570        /// ```
571        pub fn from_w3c_headers(headers: &HeaderMap) -> Self {
572            Self::parse_w3c_headers(headers).unwrap_or_default()
573        }
574
575        fn parse_w3c_headers(headers: &HeaderMap) -> Option<Self> {
576            let header = headers.get("traceparent")?.to_str().ok()?;
577
578            let parts: Vec<&str> = header.split('-').collect();
579            if parts.len() != 4 {
580                return None;
581            }
582
583            let Some(0) = u8::from_str_radix(parts[0], 16).ok() else {
584                // Wrong version.
585                return None;
586            };
587
588            let Some(0x01) = u8::from_str_radix(parts[3], 16).ok().map(|n| n & 0x01) else {
589                // Not sampled.
590                return None;
591            };
592
593            let trace_id = u128::from_str_radix(parts[1], 16).ok()?;
594            let parent_id = u64::from_str_radix(parts[2], 16).ok()?;
595
596            Some(Self {
597                trace_id,
598                parent_id,
599            })
600        }
601
602        /// Serializes a context for distributed tracing to W3C trace context headers.
603        ///
604        /// ```
605        /// # use http::Request;
606        /// use tracing_datadog::http::DistributedTracingContext;
607        ///
608        /// // Build the request.
609        /// let mut request = Request::builder().body(()).unwrap();
610        ///
611        /// // Inject distributed tracing headers.
612        /// request.headers_mut().extend(tracing::Span::current().get_context().to_w3c_headers());
613        ///
614        /// // Execute the request.
615        /// // ..
616        /// ```
617        pub fn to_w3c_headers(&self) -> HeaderMap {
618            if self.is_empty() {
619                return Default::default();
620            }
621
622            let header = format!(
623                "{version:02x}-{trace_id:032x}-{parent_id:016x}-{trace_flags:02x}",
624                version = 0,
625                trace_id = self.trace_id,
626                parent_id = self.parent_id,
627                trace_flags = 1,
628            );
629
630            HeaderMap::from_iter([(
631                HeaderName::from_static("traceparent"),
632                header.parse().unwrap(),
633            )])
634        }
635
636        /// Returns `true` if the context is empty, i.e. if it does not contain a trace ID or
637        /// a parent ID.
638        fn is_empty(&self) -> bool {
639            self.trace_id == 0 || self.parent_id == 0
640        }
641    }
642
643    // This function "remembers" the types of the subscriber so that we can downcast to something
644    // aware of them without knowing those types at the call site. Adapted from tracing-error.
645    #[derive(Debug)]
646    pub(crate) struct WithContext(
647        #[allow(clippy::type_complexity)]
648        pub(crate)  fn(&Dispatch, &Id, f: &mut dyn FnMut(&mut DatadogSpan)),
649    );
650
651    impl WithContext {
652        pub(crate) fn with_context(
653            &self,
654            dispatch: &Dispatch,
655            id: &Id,
656            mut f: &mut dyn FnMut(&mut DatadogSpan),
657        ) {
658            self.0(dispatch, id, &mut f);
659        }
660    }
661
662    pub trait DistributedTracingContext {
663        /// Gets the context for distributed tracing from the current span.
664        fn get_context(&self) -> DatadogContext;
665
666        /// Sets the context for distributed tracing on the current span.
667        fn set_context(&self, context: DatadogContext);
668    }
669
670    impl DistributedTracingContext for tracing::Span {
671        fn get_context(&self) -> DatadogContext {
672            let mut ctx = None;
673
674            self.with_subscriber(|(id, subscriber)| {
675                let Some(get_context) = subscriber.downcast_ref::<WithContext>() else {
676                    return;
677                };
678                get_context.with_context(subscriber, id, &mut |dd_span| {
679                    ctx = Some(DatadogContext {
680                        // NB Trace IDs can be 128-bit nowadays, but the 0.4 API still uses 64-bit.
681                        trace_id: dd_span.trace_id as u128,
682                        parent_id: dd_span.parent_id,
683                    })
684                });
685            });
686
687            ctx.unwrap_or_default()
688        }
689
690        fn set_context(&self, context: DatadogContext) {
691            // Avoid setting a null context.
692            if context.is_empty() {
693                return;
694            }
695
696            self.with_subscriber(move |(id, subscriber)| {
697                let Some(get_context) = subscriber.downcast_ref::<WithContext>() else {
698                    return;
699                };
700                get_context.with_context(subscriber, id, &mut |dd_span| {
701                    // NB Trace IDs can be 128-bit nowadays, but the 0.4 API still uses 64-bit.
702                    dd_span.trace_id = context.trace_id as u64;
703                    dd_span.parent_id = context.parent_id;
704                })
705            });
706        }
707    }
708
709    #[cfg(test)]
710    mod tests {
711        use super::*;
712        use crate::DatadogTraceLayer;
713        use rand::random_range;
714        use tracing::info_span;
715        use tracing_subscriber::layer::SubscriberExt;
716
717        #[test]
718        fn w3c_trace_header_round_trip() {
719            let context = DatadogContext {
720                trace_id: random_range(1..=u128::MAX),
721                parent_id: random_range(1..=u64::MAX),
722            };
723
724            let headers = context.to_w3c_headers();
725            let parsed = DatadogContext::from_w3c_headers(&headers);
726
727            assert_eq!(context.trace_id, parsed.trace_id);
728            assert_eq!(context.parent_id, parsed.parent_id);
729        }
730
731        #[test]
732        fn empty_context_doesnt_produce_w3c_trace_header() {
733            assert!(DatadogContext::default().to_w3c_headers().is_empty());
734        }
735
736        #[test]
737        fn w3c_trace_header_with_wrong_version_produces_empty_context() {
738            let headers = HeaderMap::from_iter([(
739                HeaderName::from_static("traceparent"),
740                "01-00000000000000000000000000000001-0000000000000001-01"
741                    .parse()
742                    .unwrap(),
743            )]);
744            let context = DatadogContext::from_w3c_headers(&headers);
745            assert!(context.is_empty());
746        }
747
748        #[test]
749        fn w3c_trace_header_without_sampling_flag_produces_empty_context() {
750            let headers = HeaderMap::from_iter([(
751                HeaderName::from_static("traceparent"),
752                "00-00000000000000000000000000000001-0000000000000001-00"
753                    .parse()
754                    .unwrap(),
755            )]);
756            let context = DatadogContext::from_w3c_headers(&headers);
757            assert!(context.is_empty());
758        }
759
760        #[test]
761        fn span_context_round_trip() {
762            tracing::subscriber::with_default(
763                tracing_subscriber::registry().with(
764                    DatadogTraceLayer::builder()
765                        .service("test-service")
766                        .env("test")
767                        .version("test-version")
768                        .agent_address("localhost:8126")
769                        .build()
770                        .unwrap(),
771                ),
772                || {
773                    let context = DatadogContext {
774                        // Need to limit the size here as we only track 64-bit trace IDs.
775                        trace_id: random_range(1..=u64::MAX) as u128,
776                        parent_id: random_range(1..=u64::MAX),
777                    };
778
779                    let span = info_span!("test");
780
781                    span.set_context(context);
782                    let result = span.get_context();
783
784                    assert_eq!(context.trace_id, result.trace_id);
785                    assert_eq!(context.parent_id, result.parent_id);
786                },
787            );
788        }
789
790        #[test]
791        fn empty_span_context_does_not_erase_ids() {
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                    assert_eq!(result.parent_id, 0);
812                },
813            );
814        }
815    }
816}
817
818#[cfg(test)]
819mod tests {
820    use super::*;
821
822    #[test]
823    fn builder_builds_successfully() {
824        assert!(
825            DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
826                .service("test-service")
827                .env("test")
828                .version("test-version")
829                .agent_address("localhost:8126")
830                .build()
831                .is_ok()
832        );
833    }
834
835    #[test]
836    fn service_is_required() {
837        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
838            .env("test")
839            .version("test-version")
840            .agent_address("localhost:8126")
841            .build();
842        assert!(result.unwrap_err().to_string().contains("service"));
843    }
844
845    #[test]
846    fn env_is_required() {
847        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
848            .service("test-service")
849            .version("test-version")
850            .agent_address("localhost:8126")
851            .build();
852        assert!(result.unwrap_err().to_string().contains("env"));
853    }
854
855    #[test]
856    fn version_is_required() {
857        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
858            .service("test-service")
859            .env("test")
860            .agent_address("localhost:8126")
861            .build();
862        assert!(result.unwrap_err().to_string().contains("version"));
863    }
864
865    #[test]
866    fn agent_address_is_required() {
867        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
868            .service("test-service")
869            .env("test")
870            .version("test-version")
871            .build();
872        assert!(result.unwrap_err().to_string().contains("agent_address"));
873    }
874
875    #[test]
876    fn default_default_tags_include_env_and_version() {
877        let layer: DatadogTraceLayer<tracing_subscriber::Registry> = DatadogTraceLayer::builder()
878            .service("test-service")
879            .env("test")
880            .version("test-version")
881            .agent_address("localhost:8126")
882            .build()
883            .unwrap();
884        let default_tags = &layer.default_tags;
885        assert_eq!(default_tags["env"], "test");
886        assert_eq!(default_tags["version"], "test-version");
887    }
888
889    #[test]
890    fn default_tags_can_be_added() {
891        let layer: DatadogTraceLayer<tracing_subscriber::Registry> = DatadogTraceLayer::builder()
892            .service("test-service")
893            .env("test")
894            .version("test-version")
895            .agent_address("localhost:8126")
896            .default_tag("foo", "bar")
897            .default_tag("baz", "qux")
898            .build()
899            .unwrap();
900        let default_tags = &layer.default_tags;
901        assert_eq!(default_tags["foo"], "bar");
902        assert_eq!(default_tags["baz"], "qux");
903    }
904}
tracing_datadog/lib.rs

tracing_datadog/
lib.rs
