use opentelemetry::trace::{SpanId, TraceContextExt, TraceId};
use tracing::{dispatcher::WeakDispatch, Subscriber};
use tracing_opentelemetry::get_otel_context;
use tracing_subscriber::{layer::Context, registry::LookupSpan};
#[derive(serde::Serialize)]
#[cfg_attr(test, derive(Debug, Clone, Copy, serde::Deserialize, PartialEq, Eq))]
pub struct DatadogTraceId(pub(crate) u64);
#[allow(clippy::fallible_impl_from)]
impl From<TraceId> for DatadogTraceId {
fn from(value: TraceId) -> Self {
let bytes = value.to_bytes();
#[allow(clippy::unwrap_used)]
let most_significant_8_bytes = bytes.get(8..16).unwrap();
#[allow(clippy::unwrap_used)]
let bytes_as_sized_slice: [u8; 8] = most_significant_8_bytes.try_into().unwrap();
Self(u64::from_be_bytes(bytes_as_sized_slice))
}
}
#[derive(serde::Serialize)]
#[cfg_attr(test, derive(Debug, Clone, Copy, serde::Deserialize, PartialEq, Eq))]
pub struct DatadogSpanId(pub u64);
impl From<SpanId> for DatadogSpanId {
fn from(value: SpanId) -> Self {
Self(u64::from_be_bytes(value.to_bytes()))
}
}
pub fn read_from_context<S>(
ctx: &Context<'_, S>,
weak_dispatch: Option<&WeakDispatch>,
) -> Option<(DatadogTraceId, DatadogSpanId)>
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
let span_ref = ctx.lookup_current()?;
let dispatch = weak_dispatch?.upgrade()?;
let otel_cx = get_otel_context(&span_ref.id(), &dispatch)?;
let otel_span = otel_cx.span();
let span_context = otel_span.span_context();
let trace_id = span_context.trace_id();
let span_id = span_context.span_id();
if trace_id == TraceId::INVALID && span_id == SpanId::INVALID {
return None;
}
Some((trace_id.into(), span_id.into()))
}