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 Some(other_span) = ctx.span(follows) else {
159            // The other span might be filtered or closed, so we can't access it.
160            return;
161        };
162
163        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>()
164            && let Some(other_dd_span) = other_span.extensions().get::<DatadogSpan>()
165        {
166            dd_span.span_links.push(SpanLink {
167                trace_id: other_dd_span.trace_id,
168                span_id: other_dd_span.span_id,
169            })
170        }
171    }
172
173    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
174        if !self.logging_enabled {
175            return;
176        }
177
178        let mut fields = {
179            let mut visitor = FieldVisitor::default();
180            event.record(&mut visitor);
181            visitor.finish()
182        };
183
184        fields.extend(
185            ctx.event_scope(event)
186                .into_iter()
187                .flat_map(Scope::from_root)
188                .flat_map(|span| match span.extensions().get::<DatadogSpan>() {
189                    Some(dd_span) => dd_span.meta.clone(),
190                    None => panic!("DatadogSpan extension not found, this is a bug"),
191                }),
192        );
193
194        let message = fields.remove("message").unwrap_or_default();
195
196        let (trace_id, span_id) = ctx
197            .lookup_current()
198            .and_then(|span| {
199                span.extensions()
200                    .get::<DatadogSpan>()
201                    .map(|dd_span| (Some(dd_span.trace_id), Some(dd_span.span_id)))
202            })
203            .unwrap_or_default();
204
205        let log = DatadogLog {
206            timestamp: Zoned::now().timestamp(),
207            level: event.metadata().level().to_owned(),
208            message,
209            trace_id,
210            span_id,
211            fields,
212        };
213
214        let serialized = serde_json::to_string(&log).expect("Failed to serialize log");
215
216        println!("{serialized}");
217    }
218
219    fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
220        let span = ctx.span(id).expect("Span not found, this is a bug");
221        let mut extensions = span.extensions_mut();
222
223        let now = epoch_ns();
224
225        match extensions.get_mut::<DatadogSpan>() {
226            Some(dd_span) if dd_span.start == 0 => dd_span.start = now,
227            _ => {}
228        }
229    }
230
231    fn on_exit(&self, id: &Id, ctx: Context<'_, S>) {
232        let span = ctx.span(id).expect("Span not found, this is a bug");
233        let mut extensions = span.extensions_mut();
234
235        let now = epoch_ns();
236
237        if let Some(dd_span) = extensions.get_mut::<DatadogSpan>() {
238            dd_span.duration = now - dd_span.start
239        }
240    }
241
242    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
243        let span = ctx.span(&id).expect("Span not found, this is a bug");
244        let mut extensions = span.extensions_mut();
245
246        if let Some(mut dd_span) = extensions.remove::<DatadogSpan>() {
247            // Enable trace metrics for select span kinds.
248            if let Some("server" | "client" | "consumer" | "producer") =
249                dd_span.meta.get("span.kind").map(String::as_str)
250            {
251                dd_span.metrics.insert("_dd.measured", 1.0);
252            }
253
254            self.buffer.lock().unwrap().push(dd_span);
255        }
256    }
257
258    // SAFETY: This is safe because the `WithContext` function pointer is valid
259    // for the lifetime of `&self`.
260    #[cfg(feature = "http")]
261    unsafe fn downcast_raw(&self, id: std::any::TypeId) -> Option<*const ()> {
262        match id {
263            id if id == std::any::TypeId::of::<Self>() => Some(self as *const _ as *const ()),
264            id if id == std::any::TypeId::of::<crate::http::WithContext>() => {
265                Some(&self.with_context as *const _ as *const ())
266            }
267            _ => None,
268        }
269    }
270}
271
272/// A builder for [`DatadogTraceLayer`].
273pub struct DatadogTraceLayerBuilder<S> {
274    service: Option<String>,
275    default_tags: HashMap<Cow<'static, str>, String>,
276    agent_address: Option<String>,
277    container_id: Option<String>,
278    logging_enabled: bool,
279    phantom_data: PhantomData<S>,
280}
281
282/// An error that can occur when building a [`DatadogTraceLayer`].
283#[derive(Debug)]
284pub struct BuilderError(&'static str);
285
286impl Display for BuilderError {
287    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
288        f.write_str(self.0)
289    }
290}
291
292impl std::error::Error for BuilderError {}
293
294impl<S> DatadogTraceLayerBuilder<S>
295where
296    S: Subscriber + for<'a> LookupSpan<'a>,
297{
298    /// Sets the `service`. This is required.
299    pub fn service(mut self, service: impl Into<String>) -> Self {
300        self.service = Some(service.into());
301        self
302    }
303
304    /// Sets the `env`. This is required.
305    pub fn env(mut self, env: impl Into<String>) -> Self {
306        self.default_tags.insert("env".into(), env.into());
307        self
308    }
309
310    /// Sets the `version`. This is required.
311    pub fn version(mut self, version: impl Into<String>) -> Self {
312        self.default_tags.insert("version".into(), version.into());
313        self
314    }
315
316    /// Sets the `agent_address`. This is required.
317    pub fn agent_address(mut self, agent_address: impl Into<String>) -> Self {
318        self.agent_address = Some(agent_address.into());
319        self
320    }
321
322    /// Adds a fixed default tag to all spans.
323    ///
324    /// This can be used multiple times for several tags.
325    ///
326    /// Default tags are overridden by tags set explicitly on a span.
327    pub fn default_tag(
328        mut self,
329        key: impl Into<Cow<'static, str>>,
330        value: impl Into<String>,
331    ) -> Self {
332        let _ = self.default_tags.insert(key.into(), value.into());
333        self
334    }
335
336    /// Sets the container ID. This enables infrastructure metrics in APM for supported platforms.
337    pub fn container_id(mut self, container_id: impl Into<String>) -> Self {
338        self.container_id = Some(container_id.into());
339        self
340    }
341
342    /// Enables or disables structured logging with trace correlation to stdout.
343    /// Disabled by default.
344    pub fn enable_logs(mut self, enable_logs: bool) -> Self {
345        self.logging_enabled = enable_logs;
346        self
347    }
348
349    /// Consumes the builder to construct the tracing layer.
350    pub fn build(self) -> Result<DatadogTraceLayer<S>, BuilderError> {
351        let Some(service) = self.service else {
352            return Err(BuilderError("service is required"));
353        };
354        if !self.default_tags.contains_key("env") {
355            return Err(BuilderError("env is required"));
356        };
357        if !self.default_tags.contains_key("version") {
358            return Err(BuilderError("version is required"));
359        };
360        let Some(agent_address) = self.agent_address else {
361            return Err(BuilderError("agent_address is required"));
362        };
363        let container_id = match self.container_id {
364            Some(s) => Some(
365                s.parse::<HeaderValue>()
366                    .map_err(|_| BuilderError("Failed to parse container ID into header"))?,
367            ),
368            _ => None,
369        };
370
371        let buffer = Arc::new(Mutex::new(Vec::new()));
372        let (shutdown, shutdown_rx) = mpsc::channel();
373
374        spawn(crate::export::exporter(
375            agent_address,
376            buffer.clone(),
377            container_id,
378            shutdown_rx,
379        ));
380
381        Ok(DatadogTraceLayer {
382            buffer,
383            service,
384            default_tags: self.default_tags,
385            logging_enabled: self.logging_enabled,
386            #[cfg(feature = "http")]
387            with_context: crate::http::WithContext(DatadogTraceLayer::<S>::get_context),
388            shutdown,
389            _registry: PhantomData,
390        })
391    }
392}
393
394/// Returns the current system time as nanoseconds since 1970.
395fn epoch_ns() -> i64 {
396    SystemTime::now()
397        .duration_since(UNIX_EPOCH)
398        .expect("SystemTime is before UNIX epoch")
399        .as_nanos() as i64
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405
406    #[test]
407    fn builder_builds_successfully() {
408        assert!(
409            DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
410                .service("test-service")
411                .env("test")
412                .version("test-version")
413                .agent_address("localhost:8126")
414                .build()
415                .is_ok()
416        );
417    }
418
419    #[test]
420    fn service_is_required() {
421        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
422            .env("test")
423            .version("test-version")
424            .agent_address("localhost:8126")
425            .build();
426        assert!(result.unwrap_err().to_string().contains("service"));
427    }
428
429    #[test]
430    fn env_is_required() {
431        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
432            .service("test-service")
433            .version("test-version")
434            .agent_address("localhost:8126")
435            .build();
436        assert!(result.unwrap_err().to_string().contains("env"));
437    }
438
439    #[test]
440    fn version_is_required() {
441        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
442            .service("test-service")
443            .env("test")
444            .agent_address("localhost:8126")
445            .build();
446        assert!(result.unwrap_err().to_string().contains("version"));
447    }
448
449    #[test]
450    fn agent_address_is_required() {
451        let result = DatadogTraceLayer::<tracing_subscriber::Registry>::builder()
452            .service("test-service")
453            .env("test")
454            .version("test-version")
455            .build();
456        assert!(result.unwrap_err().to_string().contains("agent_address"));
457    }
458
459    #[test]
460    fn default_default_tags_include_env_and_version() {
461        let layer: DatadogTraceLayer<tracing_subscriber::Registry> = DatadogTraceLayer::builder()
462            .service("test-service")
463            .env("test")
464            .version("test-version")
465            .agent_address("localhost:8126")
466            .build()
467            .unwrap();
468        let default_tags = &layer.default_tags;
469        assert_eq!(default_tags["env"], "test");
470        assert_eq!(default_tags["version"], "test-version");
471    }
472
473    #[test]
474    fn default_tags_can_be_added() {
475        let layer: DatadogTraceLayer<tracing_subscriber::Registry> = DatadogTraceLayer::builder()
476            .service("test-service")
477            .env("test")
478            .version("test-version")
479            .agent_address("localhost:8126")
480            .default_tag("static", "bar")
481            .default_tag(String::from("dynamic"), "qux")
482            .build()
483            .unwrap();
484        let default_tags = &layer.default_tags;
485        assert_eq!(default_tags["static"], "bar");
486        assert_eq!(default_tags["dynamic"], "qux");
487    }
488}
tracing_datadog/layer.rs

tracing_datadog/
layer.rs
