tracing_datadog/
layer.rs

1use crate::{
2    log::{DatadogLog, FieldVisitor},
3    span::{DatadogSpan, SpanAttributeVisitor, SpanLink},
4};
5use jiff::Zoned;
6use reqwest::header::HeaderValue;
7use std::{
8    borrow::Cow,
9    collections::HashMap,
10    fmt::{Display, Formatter},
11    marker::PhantomData,
12    sync::{Arc, Mutex, mpsc},
13    thread::spawn,
14    time::{SystemTime, UNIX_EPOCH},
15};
16use tracing_core::{
17    Event, Subscriber,
18    span::{Attributes, Id, Record},
19};
20use tracing_subscriber::{
21    Layer,
22    layer::Context,
23    registry::{LookupSpan, Scope},
24};
25
26/// A [`Layer`] that sends traces to Datadog.
27///
28/// ```
29/// # use tracing_subscriber::prelude::*;
30/// # use tracing_datadog::DatadogTraceLayer;
31/// tracing_subscriber::registry()
32///    .with(
33///        DatadogTraceLayer::builder()
34///            .service("my-service")
35///            .env("production")
36///            .version("git sha")
37///            .agent_address("localhost:8126")
38///            .build()
39///            .expect("failed to build DatadogTraceLayer"),
40///    )
41///    .init();
42/// ```
43#[derive(Debug)]
44pub struct DatadogTraceLayer<S> {
45    buffer: Arc<Mutex<Vec<DatadogSpan>>>,
46    service: String,
47    default_tags: HashMap<Cow<'static, str>, String>,
48    logging_enabled: bool,
49    #[cfg(feature = "http")]
50    with_context: crate::http::WithContext,
51    shutdown: mpsc::Sender<()>,
52    _registry: PhantomData<S>,
53}
54
55impl<S> DatadogTraceLayer<S>
56where
57    S: Subscriber + for<'a> LookupSpan<'a>,
58{
59    /// Creates a builder to construct a [`DatadogTraceLayer`].
60    pub fn builder() -> DatadogTraceLayerBuilder<S> {
61        DatadogTraceLayerBuilder {
62            service: None,
63            default_tags: HashMap::from_iter([("span.kind".into(), "internal".to_string())]),
64            agent_address: None,
65            container_id: None,
66            logging_enabled: false,
67            phantom_data: Default::default(),
68        }
69    }
70
71    #[cfg(feature = "http")]
72    fn get_context(
73        dispatch: &tracing_core::Dispatch,
74        id: &Id,
75        f: &mut dyn FnMut(&mut DatadogSpan),
76    ) {
77        let subscriber = dispatch
78            .downcast_ref::<S>()
79            .expect("Subscriber did not downcast to expected type, this is a bug");
80        let span = subscriber.span(id).expect("Span not found, this is a bug");
81
82        let mut extensions = span.extensions_mut();
83        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>() {
84            f(dd_span);
85        }
86    }
87}
88
89impl<S> Drop for DatadogTraceLayer<S> {
90    fn drop(&mut self) {
91        let _ = self.shutdown.send(());
92    }
93}
94
95impl<S> Layer<S> for DatadogTraceLayer<S>
96where
97    S: Subscriber + for<'a> LookupSpan<'a>,
98{
99    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
100        let span = ctx.span(id).expect("Span not found, this is a bug");
101        let mut extensions = span.extensions_mut();
102
103        let trace_id = span
104            .parent()
105            .map(|parent| {
106                parent
107                    .extensions()
108                    .get::<DatadogSpan>()
109                    .expect("Parent span didn't have a DatadogSpan extension, this is a bug")
110                    .trace_id
111            })
112            .unwrap_or(rand::random_range(1..=u64::MAX));
113
114        debug_assert!(trace_id != 0, "Trace ID is zero, this is a bug");
115
116        let mut dd_span = DatadogSpan {
117            name: span.name().to_string(),
118            service: self.service.clone(),
119            r#type: "custom".into(),
120            span_id: span.id().into_u64(),
121            start: epoch_ns(),
122            parent_id: span
123                .parent()
124                .map(|parent| parent.id().into_u64())
125                .unwrap_or_default(),
126            trace_id,
127            meta: self.default_tags.clone(),
128            metrics: {
129                let mut m = HashMap::new();
130                if span.parent().is_none() {
131                    // Special tag to mark the service entry span.
132                    m.insert("_dd.top_level", 1.0);
133                }
134                m.insert("_sampling_priority_v1", 1.0);
135                m
136            },
137            ..Default::default()
138        };
139
140        attrs.record(&mut SpanAttributeVisitor::new(&mut dd_span));
141
142        extensions.insert(dd_span);
143    }
144
145    fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
146        let span = ctx.span(id).expect("Span not found, this is a bug");
147        let mut extensions = span.extensions_mut();
148
149        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>() {
150            values.record(&mut SpanAttributeVisitor::new(dd_span));
151        }
152    }
153
154    fn on_follows_from(&self, id: &Id, follows: &Id, ctx: Context<'_, S>) {
155        let span = ctx.span(id).expect("Span not found, this is a bug");
156        let mut extensions = span.extensions_mut();
157
158        let other_span = ctx.span(follows).expect("Span not found, this is a bug");
159
160        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>()
161            && let Some(other_dd_span) = other_span.extensions().get::<DatadogSpan>()
162        {
163            dd_span.span_links.push(SpanLink {
164                trace_id: other_dd_span.trace_id,
165                span_id: other_dd_span.span_id,
166            })
167        }
168    }
169
170    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
171        if !self.logging_enabled {
172            return;
173        }
174
175        let mut fields = {
176            let mut visitor = FieldVisitor::default();
177            event.record(&mut visitor);
178            visitor.finish()
179        };
180
181        fields.extend(
182            ctx.event_scope(event)
183                .into_iter()
184                .flat_map(Scope::from_root)
185                .flat_map(|span| match span.extensions().get::<DatadogSpan>() {
186                    Some(dd_span) => dd_span.meta.clone(),
187                    None => panic!("DatadogSpan extension not found, this is a bug"),
188                }),
189        );
190
191        let message = fields.remove("message").unwrap_or_default();
192
193        let (trace_id, span_id) = ctx
194            .lookup_current()
195            .and_then(|span| {
196                span.extensions()
197                    .get::<DatadogSpan>()
198                    .map(|dd_span| (Some(dd_span.trace_id), Some(dd_span.span_id)))
199            })
200            .unwrap_or_default();
201
202        let log = DatadogLog {
203            timestamp: Zoned::now().timestamp(),
204            level: event.metadata().level().to_owned(),
205            message,
206            trace_id,
207            span_id,
208            fields,
209        };
210
211        let serialized = serde_json::to_string(&log).expect("Failed to serialize log");
212
213        println!("{serialized}");
214    }
215
216    fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
217        let span = ctx.span(id).expect("Span not found, this is a bug");
218        let mut extensions = span.extensions_mut();
219
220        let now = epoch_ns();
221
222        match extensions.get_mut::<DatadogSpan>() {
223            Some(dd_span) if dd_span.start == 0 => dd_span.start = now,
224            _ => {}
225        }
226    }
227
228    fn on_exit(&self, id: &Id, ctx: Context<'_, S>) {
229        let span = ctx.span(id).expect("Span not found, this is a bug");
230        let mut extensions = span.extensions_mut();
231
232        let now = epoch_ns();
233
234        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>() {
235            dd_span.duration = now - dd_span.start
236        }
237    }
238
239    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
240        let span = ctx.span(&id).expect("Span not found, this is a bug");
241        let mut extensions = span.extensions_mut();
242
243        if let Some(mut dd_span) = extensions.remove::<DatadogSpan>() {
244            // Enable trace metrics for select span kinds.
245            if let Some("server" | "client" | "consumer" | "producer") =
246                dd_span.meta.get("span.kind").map(String::as_str)
247            {
248                dd_span.metrics.insert("_dd.measured", 1.0);
249            }
250
251            self.buffer.lock().unwrap().push(dd_span);
252        }
253    }
254
255    // SAFETY: This is safe because the `WithContext` function pointer is valid
256    // for the lifetime of `&self`.
257    #[cfg(feature = "http")]
258    unsafe fn downcast_raw(&self, id: std::any::TypeId) -> Option<*const ()> {
259        match id {
260            id if id == std::any::TypeId::of::<Self>() => Some(self as *const _ as *const ()),
261            id if id == std::any::TypeId::of::<crate::http::WithContext>() => {
262                Some(&self.with_context as *const _ as *const ())
263            }
264            _ => None,
265        }
266    }
267}
268
269/// A builder for [`DatadogTraceLayer`].
270pub struct DatadogTraceLayerBuilder<S> {
271    service: Option<String>,
272    default_tags: HashMap<Cow<'static, str>, String>,
273    agent_address: Option<String>,
274    container_id: Option<String>,
275    logging_enabled: bool,
276    phantom_data: PhantomData<S>,
277}
278
279/// An error that can occur when building a [`DatadogTraceLayer`].
280#[derive(Debug)]
281pub struct BuilderError(&'static str);
282
283impl Display for BuilderError {
284    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
285        f.write_str(self.0)
286    }
287}
288
289impl std::error::Error for BuilderError {}
290
291impl<S> DatadogTraceLayerBuilder<S>
292where
293    S: Subscriber + for<'a> LookupSpan<'a>,
294{
295    /// Sets the `service`. This is required.
296    pub fn service(mut self, service: impl Into<String>) -> Self {
297        self.service = Some(service.into());
298        self
299    }
300
301    /// Sets the `env`. This is required.
302    pub fn env(mut self, env: impl Into<String>) -> Self {
303        self.default_tags.insert("env".into(), env.into());
304        self
305    }
306
307    /// Sets the `version`. This is required.
308    pub fn version(mut self, version: impl Into<String>) -> Self {
309        self.default_tags.insert("version".into(), version.into());
310        self
311    }
312
313    /// Sets the `agent_address`. This is required.
314    pub fn agent_address(mut self, agent_address: impl Into<String>) -> Self {
315        self.agent_address = Some(agent_address.into());
316        self
317    }
318
319    /// Adds a fixed default tag to all spans.
320    ///
321    /// This can be used multiple times for several tags.
322    ///
323    /// Default tags are overridden by tags set explicitly on a span.
324    pub fn default_tag(
325        mut self,
326        key: impl Into<Cow<'static, str>>,
327        value: impl Into<String>,
328    ) -> Self {
329        let _ = self.default_tags.insert(key.into(), value.into());
330        self
331    }
332
333    /// Sets the container ID. This enables infrastructure metrics in APM for supported platforms.
334    pub fn container_id(mut self, container_id: impl Into<String>) -> Self {
335        self.container_id = Some(container_id.into());
336        self
337    }
338
339    /// Enables or disables structured logging with trace correlation to stdout.
340    /// Disabled by default.
341    pub fn enable_logs(mut self, enable_logs: bool) -> Self {
342        self.logging_enabled = enable_logs;
343        self
344    }
345
346    /// Consumes the builder to construct the tracing layer.
347    pub fn build(self) -> Result<DatadogTraceLayer<S>, BuilderError> {
348        let Some(service) = self.service else {
349            return Err(BuilderError("service is required"));
350        };
351        if !self.default_tags.contains_key("env") {
352            return Err(BuilderError("env is required"));
353        };
354        if !self.default_tags.contains_key("version") {
355            return Err(BuilderError("version is required"));
356        };
357        let Some(agent_address) = self.agent_address else {
358            return Err(BuilderError("agent_address is required"));
359        };
360        let container_id = match self.container_id {
361            Some(s) => Some(
362                s.parse::<HeaderValue>()
363                    .map_err(|_| BuilderError("Failed to parse container ID into header"))?,
364            ),
365            _ => None,
366        };
367
368        let buffer = Arc::new(Mutex::new(Vec::new()));
369        let (shutdown, shutdown_rx) = mpsc::channel();
370
371        spawn(crate::export::exporter(
372            agent_address,
373            buffer.clone(),
374            container_id,
375            shutdown_rx,
376        ));
377
378        Ok(DatadogTraceLayer {
379            buffer,
380            service,
381            default_tags: self.default_tags,
382            logging_enabled: self.logging_enabled,
383            #[cfg(feature = "http")]
384            with_context: crate::http::WithContext(DatadogTraceLayer::<S>::get_context),
385            shutdown,
386            _registry: PhantomData,
387        })
388    }
389}
390
391/// Returns the current system time as nanoseconds since 1970.
392fn epoch_ns() -> i64 {
393    SystemTime::now()
394        .duration_since(UNIX_EPOCH)
395        .expect("SystemTime is before UNIX epoch")
396        .as_nanos() as i64
397}
398
399#[cfg(test)]
400mod tests {
401    use super::*;
402
403    #[test]
404    fn builder_builds_successfully() {
405        assert!(
406            DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
407                .service("test-service")
408                .env("test")
409                .version("test-version")
410                .agent_address("localhost:8126")
411                .build()
412                .is_ok()
413        );
414    }
415
416    #[test]
417    fn service_is_required() {
418        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
419            .env("test")
420            .version("test-version")
421            .agent_address("localhost:8126")
422            .build();
423        assert!(result.unwrap_err().to_string().contains("service"));
424    }
425
426    #[test]
427    fn env_is_required() {
428        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
429            .service("test-service")
430            .version("test-version")
431            .agent_address("localhost:8126")
432            .build();
433        assert!(result.unwrap_err().to_string().contains("env"));
434    }
435
436    #[test]
437    fn version_is_required() {
438        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
439            .service("test-service")
440            .env("test")
441            .agent_address("localhost:8126")
442            .build();
443        assert!(result.unwrap_err().to_string().contains("version"));
444    }
445
446    #[test]
447    fn agent_address_is_required() {
448        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
449            .service("test-service")
450            .env("test")
451            .version("test-version")
452            .build();
453        assert!(result.unwrap_err().to_string().contains("agent_address"));
454    }
455
456    #[test]
457    fn default_default_tags_include_env_and_version() {
458        let layer: DatadogTraceLayer<tracing_subscriber::Registry> = DatadogTraceLayer::builder()
459            .service("test-service")
460            .env("test")
461            .version("test-version")
462            .agent_address("localhost:8126")
463            .build()
464            .unwrap();
465        let default_tags = &layer.default_tags;
466        assert_eq!(default_tags["env"], "test");
467        assert_eq!(default_tags["version"], "test-version");
468    }
469
470    #[test]
471    fn default_tags_can_be_added() {
472        let layer: DatadogTraceLayer<tracing_subscriber::Registry> = DatadogTraceLayer::builder()
473            .service("test-service")
474            .env("test")
475            .version("test-version")
476            .agent_address("localhost:8126")
477            .default_tag("static", "bar")
478            .default_tag(String::from("dynamic"), "qux")
479            .build()
480            .unwrap();
481        let default_tags = &layer.default_tags;
482        assert_eq!(default_tags["static"], "bar");
483        assert_eq!(default_tags["dynamic"], "qux");
484    }
485}
tracing_datadog/layer.rs

tracing_datadog/
layer.rs
