From 24114ed74fbd46cf07a0cea6039f2c80c403e530 Mon Sep 17 00:00:00 2001
From: David Steiner <david_j_steiner@yahoo.co.nz>
Date: Fri, 15 Aug 2025 15:49:02 +0200
Subject: [PATCH] Retrieve the trace ID and span ID from the otel context to
 ensure it's populated

---
 Cargo.toml       |  2 +-
 src/formatter.rs | 79 +++++++++---------------------------------------
 2 files changed, 15 insertions(+), 66 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index cb802e9..468f1b2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ddtrace"
-version = "0.2.0"
+version = "0.2.1"
 authors = ["David Steiner <david_j_steiner@yahoo.co.nz", "Fergus Strangways-Dixon <fergusdixon101@gmail.com>"]
 edition = "2021"
 license = "MIT"
diff --git a/src/formatter.rs b/src/formatter.rs
index 37fb532..c672042 100644
--- a/src/formatter.rs
+++ b/src/formatter.rs
@@ -8,47 +8,15 @@
 use std::io;
 
 use chrono::Utc;
-use opentelemetry::trace::{SpanId, TraceId};
+use opentelemetry::trace::TraceContextExt;
 use serde::ser::{SerializeMap, Serializer as _};
-use serde::Serialize;
-use tracing::{Event, Subscriber};
-use tracing_opentelemetry::OtelData;
+use tracing::{Event, Span, Subscriber};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 use tracing_serde::fields::AsMap;
 use tracing_serde::AsSerde;
 use tracing_subscriber::fmt::format::Writer;
 use tracing_subscriber::fmt::{FmtContext, FormatEvent, FormatFields};
-use tracing_subscriber::registry::{LookupSpan, SpanRef};
-
-#[derive(Serialize)]
-struct DatadogId(u64);
-
-struct TraceInfo {
-    trace_id: DatadogId,
-    span_id: DatadogId,
-}
-
-impl From<TraceId> for DatadogId {
-    fn from(value: TraceId) -> Self {
-        let bytes = &value.to_bytes()[size_of::<u64>()..size_of::<u128>()];
-        Self(u64::from_be_bytes(bytes.try_into().unwrap()))
-    }
-}
-
-impl From<SpanId> for DatadogId {
-    fn from(value: SpanId) -> Self {
-        Self(u64::from_be_bytes(value.to_bytes()))
-    }
-}
-
-fn lookup_trace_info<S>(span_ref: &SpanRef<S>) -> Option<TraceInfo>
-where
-    S: Subscriber + for<'a> LookupSpan<'a>,
-{
-    span_ref.extensions().get::<OtelData>().map(|o| TraceInfo {
-        trace_id: o.builder.trace_id.unwrap_or(TraceId::INVALID).into(),
-        span_id: o.builder.span_id.unwrap_or(SpanId::INVALID).into(),
-    })
-}
+use tracing_subscriber::registry::LookupSpan;
 
 // mostly stolen from here: https://github.com/tokio-rs/tracing/issues/1531
 pub struct DatadogFormatter;
@@ -60,7 +28,7 @@ where
 {
     fn format_event(
         &self,
-        ctx: &FmtContext<'_, S, N>,
+        _ctx: &FmtContext<'_, S, N>,
         mut writer: Writer<'_>,
         event: &Event<'_>,
     ) -> std::fmt::Result
@@ -77,12 +45,15 @@ where
             serializer.serialize_entry("fields", &event.field_map())?;
             serializer.serialize_entry("target", meta.target())?;
 
-            if let Some(ref span_ref) = ctx.lookup_current() {
-                if let Some(trace_info) = lookup_trace_info(span_ref) {
-                    serializer.serialize_entry("dd.span_id", &trace_info.span_id)?;
-                    serializer.serialize_entry("dd.trace_id", &trace_info.trace_id)?;
-                }
-            }
+            let otel_context = Span::current().context();
+            let span = otel_context.span();
+            let span_context = span.span_context();
+
+            let span_id = span_context.span_id();
+            let trace_id = span_context.trace_id();
+
+            serializer.serialize_entry("dd.span_id", &format!("{span_id}"))?;
+            serializer.serialize_entry("dd.trace_id", &format!("{trace_id}"))?;
 
             serializer.end()
         };
@@ -115,25 +86,3 @@ impl<'a> io::Write for WriteAdaptor<'a> {
         Ok(())
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use crate::formatter::DatadogId;
-    use opentelemetry::trace::{SpanId, TraceId};
-
-    #[test]
-    fn test_trace_id_converted_to_datadog_id() {
-        let trace_id = TraceId::from_hex("2de7888d8f42abc9c7ba048b78f7a9fb").unwrap();
-        let datadog_id: DatadogId = trace_id.into();
-
-        assert_eq!(datadog_id.0, 14391820556292303355);
-    }
-
-    #[test]
-    fn test_span_id_converted_to_datadog_id() {
-        let span_id = SpanId::from_hex("58406520a0066491").unwrap();
-        let datadog_id: DatadogId = span_id.into();
-
-        assert_eq!(datadog_id.0, 6359193864645272721);
-    }
-}