From 130e7183dd74c18bc99c5a7e84210e43fa955d8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Tue, 1 Apr 2025 17:25:03 +0200 Subject: [PATCH 01/12] feat: OpenTelemetry context activation POC Right now this is completely focused on the Spans but there could be a feature that only propagates the context as well. --- Cargo.toml | 6 +- src/layer.rs | 663 ++++++++++++++++++++++++++---------------------- src/lib.rs | 11 +- src/span_ext.rs | 73 ++++-- src/stack.rs | 98 +++++++ src/tracer.rs | 70 ++--- 6 files changed, 567 insertions(+), 354 deletions(-) create mode 100644 src/stack.rs diff --git a/Cargo.toml b/Cargo.toml index b1aa1f3..ce7c278 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,9 +16,11 @@ edition = "2021" rust-version = "1.75.0" [features] -default = ["tracing-log", "metrics"] +default = ["tracing-log", "metrics", "activate_context"] # Enables support for exporting OpenTelemetry metrics metrics = ["opentelemetry/metrics","opentelemetry_sdk/metrics", "smallvec"] +# Enables support for activating an OpenTelemetry Context on Span entry +activate_context = [] [dependencies] opentelemetry = { version = "0.30.0", default-features = false, features = ["trace"] } @@ -37,7 +39,7 @@ lazy_static = { version = "1.0.2", optional = true } async-trait = "0.1.56" criterion = { version = "0.5.1", default-features = false, features = ["html_reports"] } opentelemetry = { version = "0.30.0", features = ["trace", "metrics"] } -opentelemetry_sdk = { version = "0.30.0", default-features = false, features = ["trace", "rt-tokio", "experimental_metrics_custom_reader"] } +opentelemetry_sdk = { version = "0.30.0", default-features = false, features = ["trace", "rt-tokio", "experimental_metrics_custom_reader", "testing"] } opentelemetry-stdout = { version = "0.30.0", features = ["trace", "metrics"] } opentelemetry-otlp = { version = "0.30.0", features = ["metrics", "grpc-tonic"] } opentelemetry-semantic-conventions = { version = "0.30.0", features = ["semconv_experimental"] } diff --git a/src/layer.rs b/src/layer.rs index 4ff1f84..0303a93 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -1,8 +1,14 @@ -use crate::{OtelData, PreSampledTracer}; +#[cfg(feature = "activate_context")] +use crate::stack::IdValueStack; +use crate::OtelData; +#[cfg(feature = "activate_context")] +use opentelemetry::ContextGuard; use opentelemetry::{ trace::{self as otel, noop, SpanBuilder, SpanKind, Status, TraceContextExt}, Context as OtelContext, Key, KeyValue, StringValue, Value, }; +#[cfg(feature = "activate_context")] +use std::cell::RefCell; use std::fmt; use std::marker; use std::thread; @@ -82,7 +88,7 @@ where // See https://github.com/tokio-rs/tracing/blob/4dad420ee1d4607bad79270c1520673fa6266a3d/tracing-error/src/layer.rs pub(crate) struct WithContext( #[allow(clippy::type_complexity)] - fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut OtelData, &dyn PreSampledTracer)), + fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut OtelData)), ); impl WithContext { @@ -92,7 +98,7 @@ impl WithContext { &self, dispatch: &tracing::Dispatch, id: &span::Id, - mut f: impl FnMut(&mut OtelData, &dyn PreSampledTracer), + mut f: impl FnMut(&mut OtelData), ) { (self.0)(dispatch, id, &mut f) } @@ -126,7 +132,16 @@ struct SpanBuilderUpdates { } impl SpanBuilderUpdates { - fn update(self, span_builder: &mut SpanBuilder) { + fn update(self, span_builder: &mut Option) -> Option { + if let Some(builder) = span_builder.as_mut() { + self.apply(builder); + None + } else { + Some(self) + } + } + + fn apply(self, span_builder: &mut SpanBuilder) { let Self { name, span_kind, @@ -151,6 +166,19 @@ impl SpanBuilderUpdates { } } } + + fn update_span(self, span: &opentelemetry::trace::SpanRef<'_>) { + let Self { + status, attributes, .. + } = self; + + if let Some(status) = status { + span.set_status(status); + } + if let Some(attributes) = attributes { + span.set_attributes(attributes); + } + } } struct SpanEventVisitor<'a, 'b> { @@ -519,10 +547,19 @@ impl field::Visit for SpanAttributeVisitor<'_> { } } +pub trait LayerTracer: otel::Tracer {} +impl LayerTracer for T +where + T: otel::Tracer, + T::Span: Send + Sync, +{ +} + impl OpenTelemetryLayer where S: Subscriber + for<'span> LookupSpan<'span>, - T: otel::Tracer + PreSampledTracer + 'static, + T: LayerTracer + 'static, + T::Span: Send + Sync, { /// Set the [`Tracer`] that this layer will use to produce and track /// OpenTelemetry [`Span`]s. @@ -612,7 +649,8 @@ where /// ``` pub fn with_tracer(self, tracer: Tracer) -> OpenTelemetryLayer where - Tracer: otel::Tracer + PreSampledTracer + 'static, + Tracer: LayerTracer + 'static, + Tracer::Span: Send + Sync, { OpenTelemetryLayer { tracer, @@ -787,47 +825,50 @@ where // In these case, we prefer to emit a smaller span tree instead of panicking. if let Some(span) = ctx.span(parent) { let mut extensions = span.extensions_mut(); - return extensions - .get_mut::() - .map(|builder| self.tracer.sampled_context(builder)) - .unwrap_or_default(); + if let Some(otel_data) = extensions.get_mut::() { + // If the parent span has a span builder the parent span should be started + // so we get a proper context with the parent span. + return self.with_started_cx(otel_data, &|cx| cx.clone()); + } } } - // Else if the span is inferred from context, look up any available current span. if attrs.is_contextual() { - ctx.lookup_current() - .and_then(|span| { - let mut extensions = span.extensions_mut(); - extensions - .get_mut::() - .map(|builder| self.tracer.sampled_context(builder)) - }) - .unwrap_or_else(OtelContext::current) - // Explicit root spans should have no parent context. + #[cfg(feature = "activate_context")] + // If the span is contextual and we are using the activate_context feature, + // we should use the current OTel context + { + OtelContext::current() + } + #[cfg(not(feature = "activate_context"))] + // If the span is contextual and we are not using the activate_context feature, + // we should use the current tracing context + { + ctx.lookup_current() + .and_then(|span| { + let mut extensions = span.extensions_mut(); + extensions + .get_mut::() + .map(|data| self.with_started_cx(data, &|cx| cx.clone())) + }) + .unwrap_or_else(|| OtelContext::current()) + } } else { - OtelContext::new() + OtelContext::default() } } - fn get_context( - dispatch: &tracing::Dispatch, - id: &span::Id, - f: &mut dyn FnMut(&mut OtelData, &dyn PreSampledTracer), - ) { + fn get_context(dispatch: &tracing::Dispatch, id: &span::Id, f: &mut dyn FnMut(&mut OtelData)) { let subscriber = dispatch .downcast_ref::() .expect("subscriber should downcast to expected type; this is a bug!"); let span = subscriber .span(id) .expect("registry should have a span for the current ID"); - let layer = dispatch - .downcast_ref::>() - .expect("layer should downcast to expected type; this is a bug!"); let mut extensions = span.extensions_mut(); - if let Some(builder) = extensions.get_mut::() { - f(builder, &layer.tracer); + if let Some(otel_data) = extensions.get_mut::() { + f(otel_data); } } @@ -844,6 +885,18 @@ where } extra_attrs } + + fn start_cx(&self, otel_data: &mut OtelData) { + if let Some(builder) = otel_data.builder.take() { + let span = builder.start_with_context(&self.tracer, &otel_data.parent_cx); + otel_data.parent_cx = otel_data.parent_cx.with_span(span); + } + } + + fn with_started_cx(&self, otel_data: &mut OtelData, f: &dyn Fn(&OtelContext) -> U) -> U { + self.start_cx(otel_data); + f(&otel_data.parent_cx) + } } thread_local! { @@ -859,10 +912,19 @@ thread_local! { }; } +#[cfg(feature = "activate_context")] +thread_local! { + static GUARD_STACK: RefCell = RefCell::new(IdContextGuardStack::new()); +} + +#[cfg(feature = "activate_context")] +type IdContextGuardStack = IdValueStack; + impl Layer for OpenTelemetryLayer where S: Subscriber + for<'span> LookupSpan<'span>, - T: otel::Tracer + PreSampledTracer + 'static, + T: LayerTracer + 'static, + T::Span: Send + Sync, { /// Creates an [OpenTelemetry `Span`] for the corresponding [tracing `Span`]. /// @@ -880,14 +942,7 @@ where let mut builder = self .tracer .span_builder(attrs.metadata().name()) - .with_start_time(crate::time::now()) - // Eagerly assign span id so children have stable parent id - .with_span_id(self.tracer.new_span_id()); - - // Record new trace id if there is no active parent span - if !parent_cx.has_active_span() { - builder.trace_id = Some(self.tracer.new_trace_id()); - } + .with_start_time(crate::time::now()); let builder_attrs = builder.attributes.get_or_insert(Vec::with_capacity( attrs.fields().len() + self.extra_span_attrs(), @@ -930,11 +985,17 @@ where sem_conv_config: self.sem_conv_config, }); + let mut builder = Some(builder); updates.update(&mut builder); - extensions.insert(OtelData { builder, parent_cx }); + extensions.insert(OtelData { + builder, + parent_cx, + ..Default::default() + }); } fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) { + #[cfg(not(feature = "activate_context"))] if !self.tracked_inactivity { return; } @@ -942,6 +1003,16 @@ where let span = ctx.span(id).expect("Span not found, this is a bug"); let mut extensions = span.extensions_mut(); + #[cfg(feature = "activate_context")] + { + if let Some(otel_data) = extensions.get_mut::() { + self.with_started_cx(otel_data, &|cx| { + let guard = cx.clone().attach(); + GUARD_STACK.with(|stack| stack.borrow_mut().push(id.clone(), guard)); + }); + } + } + if let Some(timings) = extensions.get_mut::() { if timings.entered_count == 0 { let now = Instant::now(); @@ -957,7 +1028,9 @@ where let mut extensions = span.extensions_mut(); if let Some(otel_data) = extensions.get_mut::() { - otel_data.builder.end_time = Some(crate::time::now()); + otel_data.end_time = Some(crate::time::now()); + #[cfg(feature = "activate_context")] + GUARD_STACK.with(|stack| stack.borrow_mut().pop(id)); } if !self.tracked_inactivity { @@ -985,15 +1058,17 @@ where sem_conv_config: self.sem_conv_config, }); let mut extensions = span.extensions_mut(); - if let Some(data) = extensions.get_mut::() { - updates.update(&mut data.builder); + if let Some(otel_data) = extensions.get_mut::() { + if let Some(updates) = updates.update(&mut otel_data.builder) { + updates.update_span(&otel_data.parent_cx.span()); + } } } fn on_follows_from(&self, id: &Id, follows: &Id, ctx: Context) { let span = ctx.span(id).expect("Span not found, this is a bug"); let mut extensions = span.extensions_mut(); - let data = extensions + let _data = extensions .get_mut::() .expect("Missing otel data span extensions"); @@ -1002,22 +1077,26 @@ where // uses the same reasoning as `parent_context` above. if let Some(follows_span) = ctx.span(follows) { let mut follows_extensions = follows_span.extensions_mut(); - let follows_data = follows_extensions + let _follows_data = follows_extensions .get_mut::() .expect("Missing otel data span extensions"); - - let follows_context = self - .tracer - .sampled_context(follows_data) - .span() - .span_context() - .clone(); - let follows_link = otel::Link::with_context(follows_context); - if let Some(ref mut links) = data.builder.links { - links.push(follows_link); - } else { - data.builder.links = Some(vec![follows_link]); - } + // TODO:ban There are no tests that check this code :( + // TODO:ban if the follows span has a span builder the follows span should be _started_ here + // let follows_link = self.with_started_cx(follows_data, &|cx| { + // otel::Link::with_context(cx.span().span_context().clone()) + // }); + // let follows_context = self + // .tracer + // .sampled_context(follows_data) + // .span() + // .span_context() + // .clone(); + // let follows_link = otel::Link::with_context(follows_context); + // if let Some(ref mut links) = data.builder.links { + // links.push(follows_link); + // } else { + // data.builder.links = Some(vec![follows_link]); + // } } } @@ -1088,19 +1167,21 @@ where } let mut extensions = span.extensions_mut(); - let otel_data = extensions.get_mut::(); - if let Some(otel_data) = otel_data { - let builder = &mut otel_data.builder; + if let Some(otel_data) = extensions.get_mut::() { + self.start_cx(otel_data); + let span = otel_data.parent_cx.span(); - if builder.status == otel::Status::Unset - && *meta.level() == tracing_core::Level::ERROR - { - builder.status = otel::Status::error("") + // TODO:ban fix this with accessor in SpanRef that can check the span status + // if builder.status == otel::Status::Unset + // && *meta.level() == tracing_core::Level::ERROR + // There is no test that checks this behavior + if *meta.level() == tracing_core::Level::ERROR { + span.set_status(otel::Status::error("")); } if let Some(builder_updates) = builder_updates { - builder_updates.update(builder); + builder_updates.update_span(&span); } if self.location { @@ -1134,11 +1215,7 @@ where } } - if let Some(ref mut events) = builder.events { - events.push(otel_event); - } else { - builder.events = Some(vec![otel_event]); - } + span.add_event(otel_event.name, otel_event.attributes); } }; } @@ -1159,24 +1236,35 @@ where }; if let Some(OtelData { - mut builder, + builder, parent_cx, + end_time, }) = otel_data { + let cx = if let Some(builder) = builder { + let span = builder.start_with_context(&self.tracer, &parent_cx); + parent_cx.with_span(span) + } else { + parent_cx + }; + + let span = cx.span(); // Append busy/idle timings when enabled. if let Some(timings) = timings { let busy_ns = Key::new("busy_ns"); let idle_ns = Key::new("idle_ns"); - let attributes = builder - .attributes - .get_or_insert_with(|| Vec::with_capacity(2)); + let mut attributes = Vec::with_capacity(2); attributes.push(KeyValue::new(busy_ns, timings.busy)); attributes.push(KeyValue::new(idle_ns, timings.idle)); + span.set_attributes(attributes); } - // Build and start span, drop span to export - builder.start_with_context(&self.tracer, &parent_cx); + if let Some(end_time) = end_time { + span.end_with_timestamp(end_time); + } else { + span.end(); + } } } @@ -1223,62 +1311,58 @@ fn thread_id_integer(id: thread::ThreadId) -> u64 { #[cfg(test)] mod tests { use super::*; - use opentelemetry::trace::{SpanContext, TraceFlags}; - use std::{ - collections::HashMap, - error::Error, - fmt::Display, - sync::{Arc, Mutex}, - time::SystemTime, - }; + use opentelemetry::trace::{SpanContext, TraceFlags, TracerProvider}; + use opentelemetry_sdk::trace::SpanExporter; + use std::{collections::HashMap, error::Error, fmt::Display, time::SystemTime}; use tracing_subscriber::prelude::*; #[derive(Debug, Clone)] - struct TestTracer(Arc>>); - impl otel::Tracer for TestTracer { - type Span = noop::NoopSpan; - fn start_with_context(&self, _name: T, _context: &OtelContext) -> Self::Span - where - T: Into>, - { - noop::NoopSpan::DEFAULT + struct TestTracer { + tracer: opentelemetry_sdk::trace::Tracer, + exporter: opentelemetry_sdk::trace::InMemorySpanExporter, + } + + impl TestTracer { + fn spans(&mut self) -> Vec { + self.exporter + .force_flush() + .expect("problems flushing spans"); + self.exporter + .get_finished_spans() + .expect("problems recording spans") } - fn span_builder(&self, name: T) -> otel::SpanBuilder - where - T: Into>, - { - otel::SpanBuilder::from_name(name) + + fn with_data(&mut self, f: impl FnOnce(&opentelemetry_sdk::trace::SpanData) -> T) -> T { + let spans = self.spans(); + f(spans.first().expect("no spans recorded")) } - fn build_with_context( - &self, - builder: otel::SpanBuilder, - parent_cx: &OtelContext, - ) -> Self::Span { - *self.0.lock().unwrap() = Some(OtelData { - builder, - parent_cx: parent_cx.clone(), - }); - noop::NoopSpan::DEFAULT + + fn attributes(&mut self) -> HashMap { + self.with_data(|data| { + data.attributes + .iter() + .map(|kv| (kv.key.to_string(), kv.value.clone())) + .collect() + }) } } - impl PreSampledTracer for TestTracer { - fn sampled_context(&self, _builder: &mut crate::OtelData) -> OtelContext { - OtelContext::new() - } - fn new_trace_id(&self) -> otel::TraceId { - otel::TraceId::INVALID - } - fn new_span_id(&self) -> otel::SpanId { - otel::SpanId::INVALID + impl Default for TestTracer { + fn default() -> Self { + let exporter = opentelemetry_sdk::trace::InMemorySpanExporter::default(); + let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + let tracer = provider.tracer("test-tracer"); + Self { tracer, exporter } } } - impl TestTracer { - fn with_data(&self, f: impl FnOnce(&OtelData) -> T) -> T { - let lock = self.0.lock().unwrap(); - let data = lock.as_ref().expect("no span data has been recorded yet"); - f(data) + impl opentelemetry::trace::Tracer for TestTracer { + type Span = opentelemetry_sdk::trace::Span; + + fn build_with_context(&self, builder: SpanBuilder, parent_cx: &OtelContext) -> Self::Span { + self.tracer.build_with_context(builder, parent_cx) } } @@ -1338,51 +1422,46 @@ mod tests { #[test] fn dynamic_span_names() { let dynamic_name = "GET http://example.com".to_string(); - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("static_name", otel.name = dynamic_name.as_str()); }); - let recorded_name = tracer - .0 - .lock() - .unwrap() - .as_ref() - .map(|b| b.builder.name.clone()); - assert_eq!(recorded_name, Some(dynamic_name.into())) + let recorded_name = tracer.spans().first().unwrap().name.clone(); + assert_eq!(recorded_name, dynamic_name.as_str()) } #[test] fn span_kind() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("request", otel.kind = "server"); }); - let recorded_kind = tracer.with_data(|data| data.builder.span_kind.clone()); - assert_eq!(recorded_kind, Some(otel::SpanKind::Server)) + let recorded_kind = tracer.with_data(|data| data.span_kind.clone()); + assert_eq!(recorded_kind, otel::SpanKind::Server) } #[test] fn span_status_code() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("request", otel.status_code = ?otel::Status::Ok); }); - let recorded_status = tracer.with_data(|data| data.builder.status.clone()); + let recorded_status = tracer.with_data(|data| data.status.clone()); assert_eq!(recorded_status, otel::Status::Ok) } #[test] fn span_status_description() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); let message = "message"; @@ -1391,28 +1470,20 @@ mod tests { tracing::debug_span!("request", otel.status_description = message); }); - let recorded_status_message = tracer - .0 - .lock() - .unwrap() - .as_ref() - .unwrap() - .builder - .status - .clone(); + let recorded_status_message = tracer.with_data(|data| data.status.clone()); assert_eq!(recorded_status_message, otel::Status::error(message)) } #[test] fn trace_id_from_existing_context() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); let trace_id = otel::TraceId::from(42u128); let existing_cx = OtelContext::current_with_span(TestSpan(otel::SpanContext::new( trace_id, otel::SpanId::from(1u64), - TraceFlags::default(), + TraceFlags::SAMPLED, false, Default::default(), ))); @@ -1422,14 +1493,13 @@ mod tests { tracing::debug_span!("request", otel.kind = "server"); }); - let recorded_trace_id = - tracer.with_data(|data| data.parent_cx.span().span_context().trace_id()); + let recorded_trace_id = tracer.with_data(|data| data.span_context.trace_id()); assert_eq!(recorded_trace_id, trace_id) } #[test] fn includes_timings() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with( layer() .with_tracer(tracer.clone()) @@ -1440,18 +1510,15 @@ mod tests { tracing::debug_span!("request"); }); - let attributes = tracer.with_data(|data| data.builder.attributes.as_ref().unwrap().clone()); - let keys = attributes - .iter() - .map(|kv| kv.key.as_str()) - .collect::>(); - assert!(keys.contains(&"idle_ns")); - assert!(keys.contains(&"busy_ns")); + let attributes = tracer.attributes(); + + assert!(attributes.contains_key("idle_ns")); + assert!(attributes.contains_key("busy_ns")); } #[test] fn records_error_fields() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); let err = TestDynError::new("base error") @@ -1465,26 +1532,11 @@ mod tests { ); }); - let attributes = tracer - .0 - .lock() - .unwrap() - .as_ref() - .unwrap() - .builder - .attributes - .as_ref() - .unwrap() - .clone(); - - let key_values = attributes - .into_iter() - .map(|kv| (kv.key.as_str().to_owned(), kv.value)) - .collect::>(); + let attributes = tracer.attributes(); - assert_eq!(key_values["error"].as_str(), "user error"); + assert_eq!(attributes["error"].as_str(), "user error"); assert_eq!( - key_values["error.chain"], + attributes["error.chain"], Value::Array( vec![ StringValue::from("intermediate error"), @@ -1494,9 +1546,9 @@ mod tests { ) ); - assert_eq!(key_values[FIELD_EXCEPTION_MESSAGE].as_str(), "user error"); + assert_eq!(attributes[FIELD_EXCEPTION_MESSAGE].as_str(), "user error"); assert_eq!( - key_values[FIELD_EXCEPTION_STACKTRACE], + attributes[FIELD_EXCEPTION_STACKTRACE], Value::Array( vec![ StringValue::from("intermediate error"), @@ -1509,7 +1561,7 @@ mod tests { #[test] fn records_event_name() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); tracing::subscriber::with_default(subscriber, || { @@ -1522,17 +1574,7 @@ mod tests { }); }); - let events = tracer - .0 - .lock() - .unwrap() - .as_ref() - .unwrap() - .builder - .events - .as_ref() - .unwrap() - .clone(); + let events = tracer.with_data(|data| data.events.clone()); let mut iter = events.iter(); @@ -1545,7 +1587,7 @@ mod tests { #[test] fn records_no_error_fields() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with( layer() .with_error_records_to_exceptions(false) @@ -1563,26 +1605,11 @@ mod tests { ); }); - let attributes = tracer - .0 - .lock() - .unwrap() - .as_ref() - .unwrap() - .builder - .attributes - .as_ref() - .unwrap() - .clone(); - - let key_values = attributes - .into_iter() - .map(|kv| (kv.key.as_str().to_owned(), kv.value)) - .collect::>(); + let attributes = tracer.attributes(); - assert_eq!(key_values["error"].as_str(), "user error"); + assert_eq!(attributes["error"].as_str(), "user error"); assert_eq!( - key_values["error.chain"], + attributes["error.chain"], Value::Array( vec![ StringValue::from("intermediate error"), @@ -1592,9 +1619,9 @@ mod tests { ) ); - assert_eq!(key_values[FIELD_EXCEPTION_MESSAGE].as_str(), "user error"); + assert_eq!(attributes[FIELD_EXCEPTION_MESSAGE].as_str(), "user error"); assert_eq!( - key_values[FIELD_EXCEPTION_STACKTRACE], + attributes[FIELD_EXCEPTION_STACKTRACE], Value::Array( vec![ StringValue::from("intermediate error"), @@ -1607,7 +1634,7 @@ mod tests { #[test] fn includes_span_location() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry() .with(layer().with_tracer(tracer.clone()).with_location(true)); @@ -1615,19 +1642,16 @@ mod tests { tracing::debug_span!("request"); }); - let attributes = tracer.with_data(|data| data.builder.attributes.as_ref().unwrap().clone()); - let keys = attributes - .iter() - .map(|kv| kv.key.as_str()) - .collect::>(); - assert!(keys.contains(&"code.filepath")); - assert!(keys.contains(&"code.namespace")); - assert!(keys.contains(&"code.lineno")); + let attributes = tracer.attributes(); + + assert!(attributes.contains_key("code.filepath")); + assert!(attributes.contains_key("code.namespace")); + assert!(attributes.contains_key("code.lineno")); } #[test] fn excludes_span_location() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry() .with(layer().with_tracer(tracer.clone()).with_location(false)); @@ -1635,14 +1659,11 @@ mod tests { tracing::debug_span!("request"); }); - let attributes = tracer.with_data(|data| data.builder.attributes.as_ref().unwrap().clone()); - let keys = attributes - .iter() - .map(|kv| kv.key.as_str()) - .collect::>(); - assert!(!keys.contains(&"code.filepath")); - assert!(!keys.contains(&"code.namespace")); - assert!(!keys.contains(&"code.lineno")); + let attributes = tracer.attributes(); + + assert!(!attributes.contains_key("code.filepath")); + assert!(!attributes.contains_key("code.namespace")); + assert!(!attributes.contains_key("code.lineno")); } #[test] @@ -1653,7 +1674,7 @@ mod tests { .map(|name| Value::String(name.to_owned().into())); let expected_id = Value::I64(thread_id_integer(thread.id()) as i64); - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry() .with(layer().with_tracer(tracer.clone()).with_threads(true)); @@ -1661,18 +1682,15 @@ mod tests { tracing::debug_span!("request"); }); - let attributes = tracer - .with_data(|data| data.builder.attributes.as_ref().unwrap().clone()) - .drain(..) - .map(|kv| (kv.key.as_str().to_string(), kv.value)) - .collect::>(); + let attributes = tracer.attributes(); + assert_eq!(attributes.get("thread.name"), expected_name.as_ref()); assert_eq!(attributes.get("thread.id"), Some(&expected_id)); } #[test] fn excludes_thread() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry() .with(layer().with_tracer(tracer.clone()).with_threads(false)); @@ -1680,18 +1698,15 @@ mod tests { tracing::debug_span!("request"); }); - let attributes = tracer.with_data(|data| data.builder.attributes.as_ref().unwrap().clone()); - let keys = attributes - .iter() - .map(|kv| kv.key.as_str()) - .collect::>(); - assert!(!keys.contains(&"thread.name")); - assert!(!keys.contains(&"thread.id")); + let attributes = tracer.attributes(); + + assert!(!attributes.contains_key("thread.name")); + assert!(!attributes.contains_key("thread.id")); } #[test] fn includes_level() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry() .with(layer().with_tracer(tracer.clone()).with_level(true)); @@ -1699,17 +1714,14 @@ mod tests { tracing::debug_span!("request"); }); - let attributes = tracer.with_data(|data| data.builder.attributes.as_ref().unwrap().clone()); - let keys = attributes - .iter() - .map(|kv| kv.key.as_str()) - .collect::>(); - assert!(keys.contains(&"level")); + let attributes = tracer.attributes(); + + assert!(attributes.contains_key("level")); } #[test] fn excludes_level() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry() .with(layer().with_tracer(tracer.clone()).with_level(false)); @@ -1717,17 +1729,14 @@ mod tests { tracing::debug_span!("request"); }); - let attributes = tracer.with_data(|data| data.builder.attributes.as_ref().unwrap().clone()); - let keys = attributes - .iter() - .map(|kv| kv.key.as_str()) - .collect::>(); - assert!(!keys.contains(&"level")); + let attributes = tracer.attributes(); + + assert!(!attributes.contains_key("level")); } #[test] fn propagates_error_fields_from_event_to_span() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); let err = TestDynError::new("base error") @@ -1743,26 +1752,11 @@ mod tests { ) }); - let attributes = tracer - .0 - .lock() - .unwrap() - .as_ref() - .unwrap() - .builder - .attributes - .as_ref() - .unwrap() - .clone(); - - let key_values = attributes - .into_iter() - .map(|kv| (kv.key.as_str().to_owned(), kv.value)) - .collect::>(); + let attributes = tracer.attributes(); - assert_eq!(key_values[FIELD_EXCEPTION_MESSAGE].as_str(), "user error"); + assert_eq!(attributes[FIELD_EXCEPTION_MESSAGE].as_str(), "user error"); assert_eq!( - key_values[FIELD_EXCEPTION_STACKTRACE], + attributes[FIELD_EXCEPTION_STACKTRACE], Value::Array( vec![ StringValue::from("intermediate error"), @@ -1775,7 +1769,7 @@ mod tests { #[test] fn propagates_no_error_fields_from_event_to_span() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with( layer() .with_error_fields_to_exceptions(false) @@ -1795,26 +1789,11 @@ mod tests { ) }); - let attributes = tracer - .0 - .lock() - .unwrap() - .as_ref() - .unwrap() - .builder - .attributes - .as_ref() - .unwrap() - .clone(); - - let key_values = attributes - .into_iter() - .map(|kv| (kv.key.as_str().to_owned(), kv.value)) - .collect::>(); + let attributes = tracer.attributes(); - assert_eq!(key_values[FIELD_EXCEPTION_MESSAGE].as_str(), "user error"); + assert_eq!(attributes[FIELD_EXCEPTION_MESSAGE].as_str(), "user error"); assert_eq!( - key_values[FIELD_EXCEPTION_STACKTRACE], + attributes[FIELD_EXCEPTION_STACKTRACE], Value::Array( vec![ StringValue::from("intermediate error"), @@ -1827,7 +1806,7 @@ mod tests { #[test] fn tracing_error_compatibility() { - let tracer = TestTracer(Arc::new(Mutex::new(None))); + let tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry() .with( layer() @@ -1849,4 +1828,94 @@ mod tests { // No need to assert anything, as long as this finished (and did not panic), everything is ok. } + + #[cfg(feature = "activate_context")] + #[derive(Debug, PartialEq)] + struct ValueA(&'static str); + #[cfg(feature = "activate_context")] + #[derive(Debug, PartialEq)] + struct ValueB(&'static str); + + #[cfg(feature = "activate_context")] + #[test] + fn otel_context_propagation() { + use opentelemetry::trace::Tracer; + use tracing::span; + + let mut tracer = TestTracer::default(); + let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + + tracing::subscriber::with_default(subscriber, || { + // Add a value to the current OpenTelemetry context for the bridge to propagate + let _outer_guard = + OtelContext::attach(OtelContext::default().with_value(ValueA("outer"))); + assert_eq!(OtelContext::current().get(), Some(&ValueA("outer"))); + let root = span!(tracing::Level::TRACE, "tokio-tracing-span-parent"); + // Drop the guard to ensure the context is cleared + drop(_outer_guard); + assert!(OtelContext::current().get::().is_none()); + // Enter the root span, the context should be propagated + let _enter_root = root.enter(); + assert_eq!(OtelContext::current().get(), Some(&ValueA("outer"))); + // Add another value to the current OpenTelemetry context for the bridge to propagate + let _inner_guard = + OtelContext::attach(OtelContext::current_with_value(ValueB("inner"))); + assert_eq!(OtelContext::current().get(), Some(&ValueA("outer"))); + assert_eq!(OtelContext::current().get(), Some(&ValueB("inner"))); + let child = span!(tracing::Level::TRACE, "tokio-tracing-span-child"); + // Drop the guard to ensure the context is reverted + drop(_inner_guard); + assert_eq!(OtelContext::current().get(), Some(&ValueA("outer"))); + assert!(OtelContext::current().get::().is_none()); + // Enter the child span, the context should be propagated + let _enter_child = child.enter(); + assert_eq!(OtelContext::current().get(), Some(&ValueA("outer"))); + assert_eq!(OtelContext::current().get(), Some(&ValueB("inner"))); + // Create an OpenTelemetry span using the OpentTelemetry notion of current + // span to see check that it is a child of the tokio child span + let span = tracer + .tracer + .span_builder("otel-tracing-span") + .start(&tracer); + let _otel_guard = OtelContext::attach(OtelContext::current_with_span(span)); + let child2 = span!(tracing::Level::TRACE, "tokio-tracing-span-child2"); + drop(_otel_guard); + // Drop the child span, the context should be reverted + drop(_enter_child); + assert_eq!(OtelContext::current().get(), Some(&ValueA("outer"))); + assert!(OtelContext::current().get::().is_none()); + // Drop the root span, the context should be reverted + drop(_enter_root); + assert!(OtelContext::current().get::().is_none()); + assert!(OtelContext::current().get::().is_none()); + let _ = child2.enter(); + }); + + // Let's check the spans + let spans = tracer.spans(); + let parent = spans + .iter() + .find(|span| span.name == "tokio-tracing-span-parent") + .unwrap(); + let child = spans + .iter() + .find(|span| span.name == "tokio-tracing-span-child") + .unwrap(); + let child2 = spans + .iter() + .find(|span| span.name == "tokio-tracing-span-child2") + .unwrap(); + let otel = spans + .iter() + .find(|span| span.name == "otel-tracing-span") + .unwrap(); + // The tokio parent span should be a root span + assert_eq!(parent.parent_span_id, otel::SpanId::INVALID); + // The first tokio child span should have the tokio parent span as parent + assert_eq!(child.parent_span_id, parent.span_context.span_id()); + // The otel span should have the first tokio child span as parent + assert_eq!(otel.parent_span_id, child.span_context.span_id()); + // The second tokio child span should have the otel span as parent + assert_eq!(child2.parent_span_id, otel.span_context.span_id()); + } } diff --git a/src/lib.rs b/src/lib.rs index 440c420..8b69e2d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -120,6 +120,11 @@ mod span_ext; /// Protocols for OpenTelemetry Tracers that are compatible with Tracing mod tracer; +#[cfg(feature = "activate_context")] +mod stack; + +use std::time::SystemTime; + pub use layer::{layer, OpenTelemetryLayer}; #[cfg(feature = "metrics")] @@ -130,13 +135,15 @@ pub use tracer::PreSampledTracer; /// Per-span OpenTelemetry data tracked by this crate. /// /// Useful for implementing [PreSampledTracer] in alternate otel SDKs. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct OtelData { /// The parent otel `Context` for the current tracing span. pub parent_cx: opentelemetry::Context, /// The otel span data recorded during the current tracing span. - pub builder: opentelemetry::trace::SpanBuilder, + pub builder: Option, + + end_time: Option, } pub(crate) mod time { diff --git a/src/span_ext.rs b/src/span_ext.rs index cb3de10..613d02f 100644 --- a/src/span_ext.rs +++ b/src/span_ext.rs @@ -1,7 +1,7 @@ use crate::layer::WithContext; use opentelemetry::{ time, - trace::{SpanContext, Status}, + trace::{SpanContext, Status, TraceContextExt}, Context, Key, KeyValue, Value, }; use std::{borrow::Cow, time::SystemTime}; @@ -214,18 +214,21 @@ pub trait OpenTelemetrySpanExt { } impl OpenTelemetrySpanExt for tracing::Span { + // TODO:ban should this really operate on a SpanContext instead of a Context? fn set_parent(&self, cx: Context) { let mut cx = Some(cx); self.with_subscriber(move |(id, subscriber)| { let Some(get_context) = subscriber.downcast_ref::() else { return; }; - get_context.with_context(subscriber, id, move |data, _tracer| { + get_context.with_context(subscriber, id, move |data| { let Some(cx) = cx.take() else { return; }; data.parent_cx = cx; - data.builder.sampling_result = None; + data.builder + .as_mut() + .map(|builder| builder.sampling_result = None); }); }); } @@ -238,20 +241,23 @@ impl OpenTelemetrySpanExt for tracing::Span { if cx.is_valid() { let mut cx = Some(cx); let mut att = Some(attributes); + // TODO:ban add add version for SpanRef self.with_subscriber(move |(id, subscriber)| { let Some(get_context) = subscriber.downcast_ref::() else { return; }; - get_context.with_context(subscriber, id, move |data, _tracer| { + get_context.with_context(subscriber, id, move |data| { let Some(cx) = cx.take() else { return; }; let attr = att.take().unwrap_or_default(); let follows_link = opentelemetry::trace::Link::new(cx, attr, 0); - data.builder - .links - .get_or_insert_with(|| Vec::with_capacity(1)) - .push(follows_link); + if let Some(builder) = data.builder.as_mut() { + builder + .links + .get_or_insert_with(|| Vec::with_capacity(1)) + .push(follows_link); + } }); }); } @@ -263,8 +269,9 @@ impl OpenTelemetrySpanExt for tracing::Span { let Some(get_context) = subscriber.downcast_ref::() else { return; }; - get_context.with_context(subscriber, id, |builder, tracer| { - cx = Some(tracer.sampled_context(builder)); + get_context.with_context(subscriber, id, |data| { + // TODO:ban create the span if it doesn't exist + cx = Some(data.parent_cx.clone()); }) }); @@ -278,17 +285,22 @@ impl OpenTelemetrySpanExt for tracing::Span { }; let mut key = Some(key.into()); let mut value = Some(value.into()); - get_context.with_context(subscriber, id, move |builder, _| { - if builder.builder.attributes.is_none() { - builder.builder.attributes = Some(Default::default()); + get_context.with_context(subscriber, id, move |data| { + if let Some(builder) = data.builder.as_mut() { + if builder.attributes.is_none() { + builder.attributes = Some(Default::default()); + } + builder + .attributes + .as_mut() + .unwrap() + .push(KeyValue::new(key.take().unwrap(), value.take().unwrap())); + } else { + let span = data.parent_cx.span(); + let key_value = KeyValue::new(key.take().unwrap(), value.take().unwrap()); + span.set_attribute(key_value); } - builder - .builder - .attributes - .as_mut() - .unwrap() - .push(KeyValue::new(key.take().unwrap(), value.take().unwrap())); - }) + }); }); } @@ -298,8 +310,13 @@ impl OpenTelemetrySpanExt for tracing::Span { let Some(get_context) = subscriber.downcast_ref::() else { return; }; - get_context.with_context(subscriber, id, move |builder, _| { - builder.builder.status = status.take().unwrap(); + get_context.with_context(subscriber, id, move |data| { + if let Some(builder) = data.builder.as_mut() { + builder.status = status.take().unwrap(); + } else { + let span = data.parent_cx.span(); + span.set_status(status.take().unwrap()); + } }); }); } @@ -321,11 +338,19 @@ impl OpenTelemetrySpanExt for tracing::Span { let Some(get_context) = subscriber.downcast_ref::() else { return; }; - get_context.with_context(subscriber, id, move |data, _tracer| { + get_context.with_context(subscriber, id, move |data| { let Some(event) = event.take() else { return; }; - data.builder.events.get_or_insert_with(Vec::new).push(event); + if let Some(builder) = data.builder.as_mut() { + builder + .events + .get_or_insert_with(|| Vec::with_capacity(1)) + .push(event); + } else { + let span = data.parent_cx.span(); + span.add_event_with_timestamp(event.name, event.timestamp, event.attributes); + } }); }); } diff --git a/src/stack.rs b/src/stack.rs new file mode 100644 index 0000000..01c5df2 --- /dev/null +++ b/src/stack.rs @@ -0,0 +1,98 @@ +use std::fmt::Debug; +use tracing_core::span::Id; + +#[derive(Debug)] +struct IdValue { + id: Id, + value: T, +} + +#[derive(Debug)] +pub(crate) struct IdValueStack { + stack: Vec>, +} + +impl IdValueStack { + pub(crate) fn new() -> Self { + IdValueStack { stack: Vec::new() } + } + + #[inline] + pub(crate) fn push(&mut self, id: Id, value: T) { + self.stack.push(IdValue { id, value }); + } + + #[inline] + pub(crate) fn pop(&mut self, id: &Id) -> Option { + if let Some((idx, _)) = self + .stack + .iter() + .enumerate() + .rev() + .find(|(_, ctx_id)| ctx_id.id == *id) + { + let IdValue { id: _, value } = self.stack.remove(idx); + return Some(value); + } + None + } + + #[cfg(test)] + pub(super) fn len(&self) -> usize { + self.stack.len() + } +} + +#[cfg(test)] +mod tests { + use super::{Id, IdValueStack}; + + type IdStringStack = IdValueStack; + + #[test] + fn pop_last_value() { + let mut stack = IdStringStack::new(); + let id1 = Id::from_u64(4711); + stack.push(id1.clone(), String::from("one")); + let id2 = Id::from_u64(1729); + stack.push(id2.clone(), String::from("two")); + assert_eq!(2, stack.len()); + + assert_eq!(Some(String::from("two")), stack.pop(&id2)); + assert_eq!(1, stack.len()); + assert_eq!(Some(String::from("one")), stack.pop(&id1)); + assert_eq!(0, stack.len()); + } + + #[test] + fn pop_first_value() { + let mut stack = IdStringStack::new(); + let id1 = Id::from_u64(4711); + stack.push(id1.clone(), String::from("one")); + let id2 = Id::from_u64(1729); + stack.push(id2.clone(), String::from("two")); + + assert_eq!(Some(String::from("one")), stack.pop(&id1)); + assert_eq!(1, stack.len()); + assert_eq!(Some(String::from("two")), stack.pop(&id2)); + assert_eq!(0, stack.len()); + } + + #[test] + fn pop_middle_value() { + let mut stack = IdStringStack::new(); + let id1 = Id::from_u64(4711); + stack.push(id1.clone(), String::from("one")); + let id2 = Id::from_u64(1729); + stack.push(id2.clone(), String::from("two")); + let id3 = Id::from_u64(1001); + stack.push(id3.clone(), String::from("three")); + + assert_eq!(Some(String::from("three")), stack.pop(&id3)); + assert_eq!(2, stack.len()); + assert_eq!(Some(String::from("two")), stack.pop(&id2)); + assert_eq!(1, stack.len()); + assert_eq!(Some(String::from("one")), stack.pop(&id1)); + assert_eq!(0, stack.len()); + } +} diff --git a/src/tracer.rs b/src/tracer.rs index e19cd6d..766a9c6 100644 --- a/src/tracer.rs +++ b/src/tracer.rs @@ -74,35 +74,37 @@ impl PreSampledTracer for noop::NoopTracer { impl PreSampledTracer for SdkTracer { fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext { let parent_cx = &data.parent_cx; - let builder = &mut data.builder; - - // Gather trace state - let (trace_id, parent_trace_flags) = - current_trace_state(builder, parent_cx, self.id_generator()); - - // Sample or defer to existing sampling decisions - let (flags, trace_state) = if let Some(result) = &builder.sampling_result { - process_sampling_result(result, parent_trace_flags) + if let Some(builder) = data.builder.as_mut() { + // Gather trace state + let (trace_id, parent_trace_flags) = + current_trace_state(builder, parent_cx, self.id_generator()); + + // Sample or defer to existing sampling decisions + let (flags, trace_state) = if let Some(result) = &builder.sampling_result { + process_sampling_result(result, parent_trace_flags) + } else { + builder.sampling_result = Some(self.should_sample().should_sample( + Some(parent_cx), + trace_id, + &builder.name, + builder.span_kind.as_ref().unwrap_or(&SpanKind::Internal), + builder.attributes.as_deref().unwrap_or(&[]), + builder.links.as_deref().unwrap_or(&[]), + )); + + process_sampling_result( + builder.sampling_result.as_ref().unwrap(), + parent_trace_flags, + ) + } + .unwrap_or_default(); + + let span_id = builder.span_id.unwrap_or(SpanId::INVALID); + let span_context = SpanContext::new(trace_id, span_id, flags, false, trace_state); + parent_cx.with_remote_span_context(span_context) } else { - builder.sampling_result = Some(self.should_sample().should_sample( - Some(parent_cx), - trace_id, - &builder.name, - builder.span_kind.as_ref().unwrap_or(&SpanKind::Internal), - builder.attributes.as_deref().unwrap_or(&[]), - builder.links.as_deref().unwrap_or(&[]), - )); - - process_sampling_result( - builder.sampling_result.as_ref().unwrap(), - parent_trace_flags, - ) + OtelContext::new() } - .unwrap_or_default(); - - let span_id = builder.span_id.unwrap_or(SpanId::INVALID); - let span_context = SpanContext::new(trace_id, span_id, flags, false, trace_state); - parent_cx.with_remote_span_context(span_context) } fn new_trace_id(&self) -> otel::TraceId { @@ -169,8 +171,13 @@ mod tests { let mut builder = SpanBuilder::from_name("empty".to_string()); builder.span_id = Some(SpanId::from(1u64)); builder.trace_id = None; + let builder = Some(builder); let parent_cx = OtelContext::new(); - let cx = tracer.sampled_context(&mut OtelData { builder, parent_cx }); + let cx = tracer.sampled_context(&mut OtelData { + builder, + parent_cx, + ..Default::default() + }); let span = cx.span(); let span_context = span.span_context(); @@ -207,7 +214,12 @@ mod tests { let tracer = provider.tracer("test"); let mut builder = SpanBuilder::from_name("parent".to_string()); builder.sampling_result = previous_sampling_result; - let sampled = tracer.sampled_context(&mut OtelData { builder, parent_cx }); + let builder = Some(builder); + let sampled = tracer.sampled_context(&mut OtelData { + builder, + parent_cx, + ..Default::default() + }); assert_eq!( sampled.span().span_context().is_sampled(), From 79fbec7d0aa44916bf18092d5077e23f7aa8cb19 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Tue, 8 Apr 2025 09:59:41 +0200 Subject: [PATCH 02/12] fix: tracing context tests and context activation Ensure that the we materialize the span and activate the context when accessing the span. This ensures the correct parent child relationships. --- .gitignore | 1 + benches/trace.rs | 4 +- src/layer.rs | 219 ++++++++++++++++++++++++++++--- src/span_ext.rs | 34 +++-- tests/trace_state_propagation.rs | 6 +- 5 files changed, 231 insertions(+), 33 deletions(-) diff --git a/.gitignore b/.gitignore index 724f2bf..ee35316 100644 --- a/.gitignore +++ b/.gitignore @@ -50,3 +50,4 @@ Cargo.lock # End of https://www.gitignore.io/api/rust,macos,visualstudiocode +.cursor diff --git a/benches/trace.rs b/benches/trace.rs index 5fd7693..dbc1f9a 100644 --- a/benches/trace.rs +++ b/benches/trace.rs @@ -264,9 +264,7 @@ where let span = ctx.span(&id).expect("Span not found, this is a bug"); let mut extensions = span.extensions_mut(); - if let Some(builder) = extensions.remove::() { - builder.with_end_time(SystemTime::now()); - } + extensions.remove::(); } } diff --git a/src/layer.rs b/src/layer.rs index 0303a93..8f4639d 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -46,7 +46,7 @@ pub struct OpenTelemetryLayer { with_threads: bool, with_level: bool, sem_conv_config: SemConvConfig, - get_context: WithContext, + with_context: WithContext, _registry: marker::PhantomData, } @@ -81,26 +81,58 @@ where OpenTelemetryLayer::default() } -// this function "remembers" the types of the subscriber so that we -// can downcast to something aware of them without knowing those -// types at the callsite. -// -// See https://github.com/tokio-rs/tracing/blob/4dad420ee1d4607bad79270c1520673fa6266a3d/tracing-error/src/layer.rs -pub(crate) struct WithContext( +/// +/// This struct lets us call back into the layer from the [crate::OpenTelemetrySpanExt] methods, +/// letting us access and mutate the underlying data on the layer side in the context of +/// tokio-tracing's span operations. +/// +/// The functions on this struct "remember" the types of the subscriber so that we +/// can downcast to something aware of them without knowing those +/// types at the callsite. +/// +/// See https://github.com/tokio-rs/tracing/blob/4dad420ee1d4607bad79270c1520673fa6266a3d/tracing-error/src/layer.rs +pub(crate) struct WithContext { + /// + /// Provides access to the OtelData associated with the given span ID. + /// #[allow(clippy::type_complexity)] - fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut OtelData)), -); + pub(crate) with_context: fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut OtelData)), + + /// + /// Ensures the given SpanId has been activated - that is, created in the OTel side of things, + /// and had its SpanBuilder consumed - and then provides access to the OtelData associated with it. + /// + #[allow(clippy::type_complexity)] + pub(crate) with_activated_context: + fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut OtelData)), +} impl WithContext { - // This function allows a function to be called in the context of the - // "remembered" subscriber. + /// + /// Return the OtelData associated with the given spanId. + /// pub(crate) fn with_context( &self, dispatch: &tracing::Dispatch, id: &span::Id, mut f: impl FnMut(&mut OtelData), ) { - (self.0)(dispatch, id, &mut f) + (self.with_context)(dispatch, id, &mut f) + } + + /// + /// If the span associated with the given SpanId has not yet been + /// built, build it, consuming the span ID. + /// + /// Optionally performs additional operations on the OtelData after building. + /// + pub(crate) fn with_activated_context( + &self, + dispatch: &tracing::Dispatch, + id: &span::Id, + mut f: impl FnMut(&mut OtelData), + ) { + (self.with_activated_context)(dispatch, id, &mut f) } } @@ -608,8 +640,10 @@ where error_events_to_exceptions: true, error_events_to_status: true, }, - - get_context: WithContext(Self::get_context), + with_context: WithContext { + with_context: Self::get_context, + with_activated_context: Self::get_activated_context, + }, _registry: marker::PhantomData, } } @@ -659,7 +693,10 @@ where with_threads: self.with_threads, with_level: self.with_level, sem_conv_config: self.sem_conv_config, - get_context: WithContext(OpenTelemetryLayer::::get_context), + with_context: WithContext { + with_context: OpenTelemetryLayer::::get_context, + with_activated_context: OpenTelemetryLayer::::get_activated_context, + }, _registry: self._registry, // cannot use ``..self` here due to different generics } @@ -858,6 +895,18 @@ where } } + /// Provides access to the OpenTelemetry data (`OtelData`) stored in a tracing span. + /// + /// This function retrieves the span from the subscriber's registry using the provided span ID, + /// and then applies the callback function `f` to the span's `OtelData` if present. + /// + /// # Parameters + /// * `dispatch` - A reference to the tracing dispatch, used to access the subscriber + /// * `id` - The ID of the span to look up + /// * `f` - A callback function that receives a mutable reference to the span's `OtelData` + /// This callback is used to manipulate or extract information from the OpenTelemetry context + /// associated with the tracing span + /// fn get_context(dispatch: &tracing::Dispatch, id: &span::Id, f: &mut dyn FnMut(&mut OtelData)) { let subscriber = dispatch .downcast_ref::() @@ -872,6 +921,30 @@ where } } + fn get_activated_context( + dispatch: &tracing::Dispatch, + id: &span::Id, + f: &mut dyn FnMut(&mut OtelData), + ) { + let subscriber = dispatch + .downcast_ref::() + .expect("subscriber should downcast to expected type; this is a bug!"); + let span = subscriber + .span(id) + .expect("registry should have a span for the current ID"); + + let layer = dispatch + .downcast_ref::>() + .expect("layer should downcast to expected type; this is a bug!"); + + let mut extensions = span.extensions_mut(); + if let Some(otel_data) = extensions.get_mut::() { + // Activate the context + layer.start_cx(otel_data); + f(otel_data); + } + } + fn extra_span_attrs(&self) -> usize { let mut extra_attrs = 0; if self.location { @@ -886,6 +959,10 @@ where extra_attrs } + /// + /// Builds the OTel span associated with given OTel context, consuming the SpanBuilder within + /// the context in the process. + /// fn start_cx(&self, otel_data: &mut OtelData) { if let Some(builder) = otel_data.builder.take() { let span = builder.start_with_context(&self.tracer, &otel_data.parent_cx); @@ -1225,6 +1302,7 @@ where /// [`Span`]: opentelemetry::trace::Span fn on_close(&self, id: span::Id, ctx: Context<'_, S>) { let span = ctx.span(&id).expect("Span not found, this is a bug"); + // Now get mutable extensions for removal let (otel_data, timings) = { let mut extensions = span.extensions_mut(); let timings = if self.tracked_inactivity { @@ -1254,10 +1332,10 @@ where let busy_ns = Key::new("busy_ns"); let idle_ns = Key::new("idle_ns"); - let mut attributes = Vec::with_capacity(2); - attributes.push(KeyValue::new(busy_ns, timings.busy)); - attributes.push(KeyValue::new(idle_ns, timings.idle)); - span.set_attributes(attributes); + span.set_attributes(vec![ + KeyValue::new(busy_ns, timings.busy), + KeyValue::new(idle_ns, timings.idle), + ]); } if let Some(end_time) = end_time { @@ -1274,7 +1352,7 @@ where match id { id if id == TypeId::of::() => Some(self as *const _ as *const ()), id if id == TypeId::of::() => { - Some(&self.get_context as *const _ as *const ()) + Some(&self.with_context as *const _ as *const ()) } _ => None, } @@ -1310,10 +1388,13 @@ fn thread_id_integer(id: thread::ThreadId) -> u64 { #[cfg(test)] mod tests { + use crate::OpenTelemetrySpanExt; + use super::*; use opentelemetry::trace::{SpanContext, TraceFlags, TracerProvider}; use opentelemetry_sdk::trace::SpanExporter; use std::{collections::HashMap, error::Error, fmt::Display, time::SystemTime}; + use tracing::trace_span; use tracing_subscriber::prelude::*; #[derive(Debug, Clone)] @@ -1918,4 +1999,102 @@ mod tests { // The second tokio child span should have the otel span as parent assert_eq!(child2.parent_span_id, otel.span_context.span_id()); } + + #[test] + fn parent_context() { + let mut tracer = TestTracer::default(); + let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + + tracing::subscriber::with_default(subscriber, || { + let root = trace_span!("root"); + + let child1 = trace_span!("child-1"); + let root_context = root.context(); // SpanData (None) + child1.set_parent(root_context); // Clone context, but SpanData(None) + + let _enter_root = root.enter(); + drop(_enter_root); + + let child2 = trace_span!("child-2"); + child2.set_parent(root.context()); + }); + + // Let's check the spans + let spans = tracer.spans(); + let parent = spans.iter().find(|span| span.name == "root").unwrap(); + let child1 = spans.iter().find(|span| span.name == "child-1").unwrap(); + let child2 = spans.iter().find(|span| span.name == "child-2").unwrap(); + assert_eq!(parent.parent_span_id, otel::SpanId::INVALID); + assert_eq!(child1.parent_span_id, parent.span_context.span_id()); + assert_eq!(child2.parent_span_id, parent.span_context.span_id()); + } + + #[test] + fn record_after() { + let mut tracer = TestTracer::default(); + let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + + tracing::subscriber::with_default(subscriber, || { + let root = trace_span!("root", before = "before", after = "before"); + + // Record a value before the span is entered + root.record("before", "after"); + + // Enter and exit the span + let _enter_root = root.enter(); + drop(_enter_root); + + // Record a value after the span is exited + root.record("after", "after"); + }); + + // Let's check the spans. Both values should've been + // updated to 'after'. + let spans = tracer.spans(); + let parent = spans.iter().find(|span| span.name == "root").unwrap(); + assert_eq!(parent.parent_span_id, otel::SpanId::INVALID); + assert!(parent + .attributes + .iter() + .filter(|kv| kv.key.as_str() == "before") + .any(|kv| kv.value.as_str() == "after")); + + assert!(parent + .attributes + .iter() + .filter(|kv| kv.key.as_str() == "after") + .any(|kv| kv.value.as_str() == "after")); + } + + #[test] + fn parent_context_2() { + let mut tracer = TestTracer::default(); + let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + + tracing::subscriber::with_default(subscriber, || { + let root = trace_span!("root"); + _ = root.enter(); + + let child1 = trace_span!("child-1"); + child1.set_parent(root.context()); + + trace_span!(parent: &child1, "child-2"); + child1.set_parent(root.context()); // <-- this is what causes the issue + + trace_span!(parent: &child1, "child-3"); + }); + + // Let's check the spans + let spans = tracer.spans(); + let root = spans.iter().find(|span| span.name == "root").unwrap(); + let child1 = spans.iter().find(|span| span.name == "child-1").unwrap(); + let child2 = spans.iter().find(|span| span.name == "child-2").unwrap(); + let child3 = spans.iter().find(|span| span.name == "child-3").unwrap(); + assert_eq!(root.parent_span_id, otel::SpanId::INVALID); + assert_eq!(child1.parent_span_id, root.span_context.span_id()); + assert_eq!(child2.parent_span_id, child1.span_context.span_id()); + + // This is surprising, the parent should be `child1`, but is 'root'. + assert_eq!(child3.parent_span_id, child1.span_context.span_id()); + } } diff --git a/src/span_ext.rs b/src/span_ext.rs index 613d02f..478e0bd 100644 --- a/src/span_ext.rs +++ b/src/span_ext.rs @@ -1,4 +1,5 @@ use crate::layer::WithContext; +use crate::OtelData; use opentelemetry::{ time, trace::{SpanContext, Status, TraceContextExt}, @@ -214,21 +215,36 @@ pub trait OpenTelemetrySpanExt { } impl OpenTelemetrySpanExt for tracing::Span { - // TODO:ban should this really operate on a SpanContext instead of a Context? + /// + /// Allows us to set the parent context of this span. This method exists primarily to allow + /// us to pull in distributed_ incoming context - e.g. span IDs, etc - that have been read + /// into an existing context. + /// + /// A span's parent should only be set _once_, for the purpose described above. + /// Additionally, once a span has been fully built - and the SpanBuilder has been consumed - + /// the parent _cannot_ be mutated. + /// fn set_parent(&self, cx: Context) { let mut cx = Some(cx); self.with_subscriber(move |(id, subscriber)| { let Some(get_context) = subscriber.downcast_ref::() else { return; }; + // Set the parent OTel for the current span get_context.with_context(subscriber, id, move |data| { - let Some(cx) = cx.take() else { + let Some(new_cx) = cx.take() else { return; }; - data.parent_cx = cx; - data.builder - .as_mut() - .map(|builder| builder.sampling_result = None); + // Create a new context with the new parent but preserve our span. + // NOTE - if the span has been created - if we have _already_ + // consumed our SpanBuilder_ - we can no longer mutate our parent! + // This is an intentional design decision. + if let Some(builder) = &mut data.builder { + // If we still have a builder, update it to use the new parent context + // when it's eventually built + data.parent_cx = new_cx; + builder.sampling_result = None; + } }); }); } @@ -269,10 +285,10 @@ impl OpenTelemetrySpanExt for tracing::Span { let Some(get_context) = subscriber.downcast_ref::() else { return; }; - get_context.with_context(subscriber, id, |data| { - // TODO:ban create the span if it doesn't exist + // If our span hasn't been built, we should build it and get the context in one call + get_context.with_activated_context(subscriber, id, |data: &mut OtelData| { cx = Some(data.parent_cx.clone()); - }) + }); }); cx.unwrap_or_default() diff --git a/tests/trace_state_propagation.rs b/tests/trace_state_propagation.rs index d1b658e..2a33bc0 100644 --- a/tests/trace_state_propagation.rs +++ b/tests/trace_state_propagation.rs @@ -121,8 +121,12 @@ fn sampling_decision_respects_new_parent() { tracing::subscriber::with_default(subscriber, || { let child = tracing::debug_span!("child"); - child.context(); // force a sampling decision + + // Observation: if you force the _child_ to materialize before the parent, e.g., + // if you swap these two lines - bad things will happen, and we shouldn't support + // this. child.set_parent(Context::current_with_span(root_span)); + child.context(); // force a sampling decision }); drop(provider); // flush all spans From 4dd77441b4ae744ac2c36ed9d44fc578ce571440 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Fri, 11 Apr 2025 17:43:35 +0200 Subject: [PATCH 03/12] perf: skip creating a context in on_close if possible --- src/layer.rs | 35 ++++++++++++++++++----------------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/src/layer.rs b/src/layer.rs index 8f4639d..68973c8 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -4,7 +4,7 @@ use crate::OtelData; #[cfg(feature = "activate_context")] use opentelemetry::ContextGuard; use opentelemetry::{ - trace::{self as otel, noop, SpanBuilder, SpanKind, Status, TraceContextExt}, + trace::{self as otel, noop, Span, SpanBuilder, SpanKind, Status, TraceContextExt}, Context as OtelContext, Key, KeyValue, StringValue, Value, }; #[cfg(feature = "activate_context")] @@ -1319,30 +1319,31 @@ where end_time, }) = otel_data { - let cx = if let Some(builder) = builder { - let span = builder.start_with_context(&self.tracer, &parent_cx); - parent_cx.with_span(span) - } else { - parent_cx - }; - - let span = cx.span(); // Append busy/idle timings when enabled. - if let Some(timings) = timings { + let timings = timings.and_then(|timings| { let busy_ns = Key::new("busy_ns"); let idle_ns = Key::new("idle_ns"); - span.set_attributes(vec![ + Some(vec![ KeyValue::new(busy_ns, timings.busy), KeyValue::new(idle_ns, timings.idle), - ]); - } + ]) + }); - if let Some(end_time) = end_time { - span.end_with_timestamp(end_time); + if let Some(builder) = builder { + // Don't create the context here just to get a SpanRef since it's costly + let mut span = builder.start_with_context(&self.tracer, &parent_cx); + timings.map(|timings| span.set_attributes(timings)); + if let Some(end_time) = end_time { + span.end_with_timestamp(end_time); + } else { + span.end(); + } } else { - span.end(); - } + let span = parent_cx.span(); + timings.map(|timings| span.set_attributes(timings)); + end_time.map_or_else(|| span.end(), |end_time| span.end_with_timestamp(end_time)); + }; } } From d73ab018be38642da53f341159b6694a81bb8900 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Mon, 5 May 2025 15:27:13 +0200 Subject: [PATCH 04/12] fix: register links on follows_from --- src/layer.rs | 144 +++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 123 insertions(+), 21 deletions(-) diff --git a/src/layer.rs b/src/layer.rs index 68973c8..ad5b37b 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -9,12 +9,12 @@ use opentelemetry::{ }; #[cfg(feature = "activate_context")] use std::cell::RefCell; -use std::fmt; use std::marker; use std::thread; #[cfg(not(all(target_arch = "wasm32", not(target_os = "wasi"))))] use std::time::Instant; use std::{any::TypeId, borrow::Cow}; +use std::{fmt, vec}; use tracing_core::span::{self, Attributes, Id, Record}; use tracing_core::{field, Event, Subscriber}; #[cfg(feature = "tracing-log")] @@ -1145,7 +1145,7 @@ where fn on_follows_from(&self, id: &Id, follows: &Id, ctx: Context) { let span = ctx.span(id).expect("Span not found, this is a bug"); let mut extensions = span.extensions_mut(); - let _data = extensions + let data = extensions .get_mut::() .expect("Missing otel data span extensions"); @@ -1154,26 +1154,20 @@ where // uses the same reasoning as `parent_context` above. if let Some(follows_span) = ctx.span(follows) { let mut follows_extensions = follows_span.extensions_mut(); - let _follows_data = follows_extensions + let follows_data = follows_extensions .get_mut::() .expect("Missing otel data span extensions"); - // TODO:ban There are no tests that check this code :( - // TODO:ban if the follows span has a span builder the follows span should be _started_ here - // let follows_link = self.with_started_cx(follows_data, &|cx| { - // otel::Link::with_context(cx.span().span_context().clone()) - // }); - // let follows_context = self - // .tracer - // .sampled_context(follows_data) - // .span() - // .span_context() - // .clone(); - // let follows_link = otel::Link::with_context(follows_context); - // if let Some(ref mut links) = data.builder.links { - // links.push(follows_link); - // } else { - // data.builder.links = Some(vec![follows_link]); - // } + let follows_context = + self.with_started_cx(follows_data, &|cx| cx.span().span_context().clone()); + if let Some(builder) = data.builder.as_mut() { + if let Some(ref mut links) = builder.links { + links.push(otel::Link::with_context(follows_context)); + } else { + builder.links = Some(vec![otel::Link::with_context(follows_context)]); + } + } else { + data.parent_cx.span().add_link(follows_context, vec![]); + } } } @@ -2095,7 +2089,115 @@ mod tests { assert_eq!(child1.parent_span_id, root.span_context.span_id()); assert_eq!(child2.parent_span_id, child1.span_context.span_id()); - // This is surprising, the parent should be `child1`, but is 'root'. + // The parent should be `child1` assert_eq!(child3.parent_span_id, child1.span_context.span_id()); } + + #[test] + fn follows_from_adds_link() { + use crate::OpenTelemetrySpanExt; + let mut tracer = TestTracer::default(); + let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + + let span1_id = tracing::subscriber::with_default(subscriber, || { + let span2 = tracing::debug_span!("span2"); + let span1 = tracing::debug_span!("span1"); + + // Ensure that span2 is started + let _ = span2.context(); + + // Establish follows_from relationship + span2.follows_from(&span1); + + // Enter span2 to ensure it's exported + let _guard = span2.enter(); + + // Get span ID for later verification + span1.context().span().span_context().span_id() + }); + + let spans = tracer.spans(); + // Check that both spans are exported + assert_eq!(spans.len(), 2, "Expected two spans to be exported"); + assert!(spans.iter().any(|span| span.name == "span1")); + let span2 = spans + .iter() + .find(|span| span.name == "span2") + .expect("Expected span2 to be exported"); + + // Collect span2 links + let links = span2 + .links + .iter() + .map(|link| link.span_context.span_id()) + .collect::>(); + + // Verify that span2 has a link to span1 + assert_eq!( + links.len(), + 1, + "Expected span to have one link from follows_from relationship" + ); + + assert!( + links.contains(&span1_id), + "Link should point to the correct source span" + ); + } + + #[test] + fn follows_from_multiple_links() { + use crate::OpenTelemetrySpanExt; + let mut tracer = TestTracer::default(); + let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + + let (span1_id, span2_id) = tracing::subscriber::with_default(subscriber, || { + let span3 = tracing::debug_span!("span3"); + let span2 = tracing::debug_span!("span2"); + let span1 = tracing::debug_span!("span1"); + + // Establish multiple follows_from relationships + span3.follows_from(&span1); + span3.follows_from(&span2); + + // Enter span3 to ensure it's exported + let _guard = span3.enter(); + + // Get span IDs for later verification + ( + span1.context().span().span_context().span_id(), + span2.context().span().span_context().span_id(), + ) + }); + + let spans = tracer.spans(); + // Check that all three spans are exported + assert_eq!(spans.len(), 3, "Expected three spans to be exported"); + assert!(spans.iter().any(|span| span.name == "span1")); + assert!(spans.iter().any(|span| span.name == "span2")); + let span3 = spans + .iter() + .find(|span| span.name == "span3") + .expect("Expected span3 to be exported"); + + // Collect span3 links + let links = span3 + .links + .iter() + .map(|link| link.span_context.span_id()) + .collect::>(); + + // Verify that span3 has multiple links and they point to the correct spans + assert_eq!( + links.len(), + 2, + "Expected span to have two links from follows_from relationships" + ); + + // Verify that the links point to the correct spans in the correct order + assert!( + links[0] == span1_id && links[1] == span2_id, + "Links should point to the correct source spans" + ); + } } From ae23c59de4630f9d473487bc87d89d230fc584ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Mon, 5 May 2025 12:47:14 +0200 Subject: [PATCH 05/12] perf: make on_event use builder again if possible --- src/layer.rs | 41 +++++++++++++++++++++++++---------------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/src/layer.rs b/src/layer.rs index ad5b37b..0e2b652 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -1240,21 +1240,6 @@ where let mut extensions = span.extensions_mut(); if let Some(otel_data) = extensions.get_mut::() { - self.start_cx(otel_data); - let span = otel_data.parent_cx.span(); - - // TODO:ban fix this with accessor in SpanRef that can check the span status - // if builder.status == otel::Status::Unset - // && *meta.level() == tracing_core::Level::ERROR - // There is no test that checks this behavior - if *meta.level() == tracing_core::Level::ERROR { - span.set_status(otel::Status::error("")); - } - - if let Some(builder_updates) = builder_updates { - builder_updates.update_span(&span); - } - if self.location { #[cfg(not(feature = "tracing-log"))] let normalized_meta: Option> = None; @@ -1286,7 +1271,31 @@ where } } - span.add_event(otel_event.name, otel_event.attributes); + if let Some(builder) = otel_data.builder.as_mut() { + if builder.status == otel::Status::Unset + && *meta.level() == tracing_core::Level::ERROR + { + builder.status = otel::Status::error(""); + } + if let Some(builder_updates) = builder_updates { + builder_updates.apply(builder); + } + if let Some(ref mut events) = builder.events { + events.push(otel_event); + } else { + builder.events = Some(vec![otel_event]); + } + } else { + let span = otel_data.parent_cx.span(); + // TODO:ban fix this with accessor in SpanRef that can check the span status + if *meta.level() == tracing_core::Level::ERROR { + span.set_status(otel::Status::error("")); + } + if let Some(builder_updates) = builder_updates { + builder_updates.update_span(&span); + } + span.add_event(otel_event.name, otel_event.attributes); + } } }; } From d9540c6190555f77d1f601cb196aa9576055cda0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Tue, 6 May 2025 11:05:37 +0200 Subject: [PATCH 06/12] chore: remove LayerTracer trait --- src/layer.rs | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/src/layer.rs b/src/layer.rs index 0e2b652..dfb4b05 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -579,18 +579,10 @@ impl field::Visit for SpanAttributeVisitor<'_> { } } -pub trait LayerTracer: otel::Tracer {} -impl LayerTracer for T -where - T: otel::Tracer, - T::Span: Send + Sync, -{ -} - impl OpenTelemetryLayer where S: Subscriber + for<'span> LookupSpan<'span>, - T: LayerTracer + 'static, + T: otel::Tracer + 'static, T::Span: Send + Sync, { /// Set the [`Tracer`] that this layer will use to produce and track @@ -683,7 +675,7 @@ where /// ``` pub fn with_tracer(self, tracer: Tracer) -> OpenTelemetryLayer where - Tracer: LayerTracer + 'static, + Tracer: otel::Tracer + 'static, Tracer::Span: Send + Sync, { OpenTelemetryLayer { @@ -1000,7 +992,7 @@ type IdContextGuardStack = IdValueStack; impl Layer for OpenTelemetryLayer where S: Subscriber + for<'span> LookupSpan<'span>, - T: LayerTracer + 'static, + T: otel::Tracer + 'static, T::Span: Send + Sync, { /// Creates an [OpenTelemetry `Span`] for the corresponding [tracing `Span`]. From f5c70c36d60e52bc69d674dbe17196e07f3c418f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Wed, 7 May 2025 12:56:54 +0200 Subject: [PATCH 07/12] chore: fix clippy warnings --- src/layer.rs | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/src/layer.rs b/src/layer.rs index dfb4b05..0f29cd3 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -880,7 +880,7 @@ where .get_mut::() .map(|data| self.with_started_cx(data, &|cx| cx.clone())) }) - .unwrap_or_else(|| OtelContext::current()) + .unwrap_or_else(OtelContext::current) } } else { OtelContext::default() @@ -896,8 +896,8 @@ where /// * `dispatch` - A reference to the tracing dispatch, used to access the subscriber /// * `id` - The ID of the span to look up /// * `f` - A callback function that receives a mutable reference to the span's `OtelData` - /// This callback is used to manipulate or extract information from the OpenTelemetry context - /// associated with the tracing span + /// This callback is used to manipulate or extract information from the OpenTelemetry context + /// associated with the tracing span /// fn get_context(dispatch: &tracing::Dispatch, id: &span::Id, f: &mut dyn FnMut(&mut OtelData)) { let subscriber = dispatch @@ -1315,20 +1315,22 @@ where }) = otel_data { // Append busy/idle timings when enabled. - let timings = timings.and_then(|timings| { + let timings = timings.map(|timings| { let busy_ns = Key::new("busy_ns"); let idle_ns = Key::new("idle_ns"); - Some(vec![ + vec![ KeyValue::new(busy_ns, timings.busy), KeyValue::new(idle_ns, timings.idle), - ]) + ] }); if let Some(builder) = builder { // Don't create the context here just to get a SpanRef since it's costly let mut span = builder.start_with_context(&self.tracer, &parent_cx); - timings.map(|timings| span.set_attributes(timings)); + if let Some(timings) = timings { + span.set_attributes(timings) + }; if let Some(end_time) = end_time { span.end_with_timestamp(end_time); } else { @@ -1336,7 +1338,9 @@ where } } else { let span = parent_cx.span(); - timings.map(|timings| span.set_attributes(timings)); + if let Some(timings) = timings { + span.set_attributes(timings) + }; end_time.map_or_else(|| span.end(), |end_time| span.end_with_timestamp(end_time)); }; } From 94137c54846a935b0438995106e3ab5ea3a41fdf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Tue, 24 Jun 2025 13:17:31 +0200 Subject: [PATCH 08/12] chore: remove unused PreSampledTracer --- src/lib.rs | 7 +- src/tracer.rs | 242 -------------------------------------------------- 2 files changed, 1 insertion(+), 248 deletions(-) delete mode 100644 src/tracer.rs diff --git a/src/lib.rs b/src/lib.rs index 8b69e2d..8ea4d05 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -117,8 +117,6 @@ mod metrics; mod layer; /// Span extension which enables OpenTelemetry context management. mod span_ext; -/// Protocols for OpenTelemetry Tracers that are compatible with Tracing -mod tracer; #[cfg(feature = "activate_context")] mod stack; @@ -130,13 +128,10 @@ pub use layer::{layer, OpenTelemetryLayer}; #[cfg(feature = "metrics")] pub use metrics::MetricsLayer; pub use span_ext::OpenTelemetrySpanExt; -pub use tracer::PreSampledTracer; /// Per-span OpenTelemetry data tracked by this crate. -/// -/// Useful for implementing [PreSampledTracer] in alternate otel SDKs. #[derive(Debug, Clone, Default)] -pub struct OtelData { +struct OtelData { /// The parent otel `Context` for the current tracing span. pub parent_cx: opentelemetry::Context, diff --git a/src/tracer.rs b/src/tracer.rs deleted file mode 100644 index 766a9c6..0000000 --- a/src/tracer.rs +++ /dev/null @@ -1,242 +0,0 @@ -use opentelemetry::{ - trace as otel, - trace::{ - noop, SamplingDecision, SamplingResult, SpanBuilder, SpanContext, SpanId, SpanKind, - TraceContextExt, TraceFlags, TraceId, TraceState, - }, - Context as OtelContext, -}; -use opentelemetry_sdk::trace::{IdGenerator, Tracer as SdkTracer}; - -/// An interface for authors of OpenTelemetry SDKs to build pre-sampled tracers. -/// -/// The OpenTelemetry spec does not allow trace ids to be updated after a span -/// has been created. In order to associate extracted parent trace ids with -/// existing `tracing` spans, `tracing-opentelemetry` builds up otel span data -/// using a [`SpanBuilder`] instead, and creates / exports full otel spans only -/// when the associated `tracing` span is closed. However, in order to properly -/// inject otel [`Context`] information to downstream requests, the sampling -/// state must now be known _before_ the otel span has been created. -/// -/// The logic for coming to a sampling decision and creating an injectable span -/// context from a [`SpanBuilder`] is encapsulated in the -/// [`PreSampledTracer::sampled_context`] method and has been implemented -/// for the standard OpenTelemetry SDK, but this trait may be implemented by -/// authors of alternate OpenTelemetry SDK implementations if they wish to have -/// `tracing` compatibility. -/// -/// See the [`OpenTelemetrySpanExt::set_parent`] and -/// [`OpenTelemetrySpanExt::context`] methods for example usage. -/// -/// [`Tracer`]: opentelemetry::trace::Tracer -/// [`SpanBuilder`]: opentelemetry::trace::SpanBuilder -/// [`PreSampledTracer::sampled_span_context`]: crate::PreSampledTracer::sampled_span_context -/// [`OpenTelemetrySpanExt::set_parent`]: crate::OpenTelemetrySpanExt::set_parent -/// [`OpenTelemetrySpanExt::context`]: crate::OpenTelemetrySpanExt::context -/// [`Context`]: opentelemetry::Context -#[rustversion::attr( - since(1.78), - diagnostic::on_unimplemented( - note = "Make sure you're using correct `opentelemetry` versions compatible with this \ - `tracing-opentelemetry`. The `opentelemetry` versions are usually one version lower \ - than `tracing-opentelemetry`." - ) -)] -pub trait PreSampledTracer { - /// Produce an otel context containing an active and pre-sampled span for - /// the given span builder data. - /// - /// The sampling decision, span context information, and parent context - /// values must match the values recorded when the tracing span is closed. - fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext; - - /// Generate a new trace id. - fn new_trace_id(&self) -> otel::TraceId; - - /// Generate a new span id. - fn new_span_id(&self) -> otel::SpanId; -} - -impl PreSampledTracer for noop::NoopTracer { - fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext { - data.parent_cx.clone() - } - - fn new_trace_id(&self) -> otel::TraceId { - otel::TraceId::INVALID - } - - fn new_span_id(&self) -> otel::SpanId { - otel::SpanId::INVALID - } -} - -impl PreSampledTracer for SdkTracer { - fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext { - let parent_cx = &data.parent_cx; - if let Some(builder) = data.builder.as_mut() { - // Gather trace state - let (trace_id, parent_trace_flags) = - current_trace_state(builder, parent_cx, self.id_generator()); - - // Sample or defer to existing sampling decisions - let (flags, trace_state) = if let Some(result) = &builder.sampling_result { - process_sampling_result(result, parent_trace_flags) - } else { - builder.sampling_result = Some(self.should_sample().should_sample( - Some(parent_cx), - trace_id, - &builder.name, - builder.span_kind.as_ref().unwrap_or(&SpanKind::Internal), - builder.attributes.as_deref().unwrap_or(&[]), - builder.links.as_deref().unwrap_or(&[]), - )); - - process_sampling_result( - builder.sampling_result.as_ref().unwrap(), - parent_trace_flags, - ) - } - .unwrap_or_default(); - - let span_id = builder.span_id.unwrap_or(SpanId::INVALID); - let span_context = SpanContext::new(trace_id, span_id, flags, false, trace_state); - parent_cx.with_remote_span_context(span_context) - } else { - OtelContext::new() - } - } - - fn new_trace_id(&self) -> otel::TraceId { - self.id_generator().new_trace_id() - } - - fn new_span_id(&self) -> otel::SpanId { - self.id_generator().new_span_id() - } -} - -fn current_trace_state( - builder: &SpanBuilder, - parent_cx: &OtelContext, - id_generator: &dyn IdGenerator, -) -> (TraceId, TraceFlags) { - if parent_cx.has_active_span() { - let span = parent_cx.span(); - let sc = span.span_context(); - (sc.trace_id(), sc.trace_flags()) - } else { - ( - builder - .trace_id - .unwrap_or_else(|| id_generator.new_trace_id()), - Default::default(), - ) - } -} - -fn process_sampling_result( - sampling_result: &SamplingResult, - trace_flags: TraceFlags, -) -> Option<(TraceFlags, TraceState)> { - match sampling_result { - SamplingResult { - decision: SamplingDecision::Drop, - .. - } => None, - SamplingResult { - decision: SamplingDecision::RecordOnly, - trace_state, - .. - } => Some((trace_flags & !TraceFlags::SAMPLED, trace_state.clone())), - SamplingResult { - decision: SamplingDecision::RecordAndSample, - trace_state, - .. - } => Some((trace_flags | TraceFlags::SAMPLED, trace_state.clone())), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::OtelData; - use opentelemetry::trace::TracerProvider as _; - use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider}; - - #[test] - fn assigns_default_trace_id_if_missing() { - let provider = SdkTracerProvider::default(); - let tracer = provider.tracer("test"); - let mut builder = SpanBuilder::from_name("empty".to_string()); - builder.span_id = Some(SpanId::from(1u64)); - builder.trace_id = None; - let builder = Some(builder); - let parent_cx = OtelContext::new(); - let cx = tracer.sampled_context(&mut OtelData { - builder, - parent_cx, - ..Default::default() - }); - let span = cx.span(); - let span_context = span.span_context(); - - assert!(span_context.is_valid()); - } - - #[rustfmt::skip] - fn sampler_data() -> Vec<(&'static str, Sampler, OtelContext, Option, bool)> { - vec![ - // No parent samples - ("empty_parent_cx_always_on", Sampler::AlwaysOn, OtelContext::new(), None, true), - ("empty_parent_cx_always_off", Sampler::AlwaysOff, OtelContext::new(), None, false), - - // Remote parent samples - ("remote_parent_cx_always_on", Sampler::AlwaysOn, OtelContext::new().with_remote_span_context(span_context(TraceFlags::SAMPLED, true)), None, true), - ("remote_parent_cx_always_off", Sampler::AlwaysOff, OtelContext::new().with_remote_span_context(span_context(TraceFlags::SAMPLED, true)), None, false), - ("sampled_remote_parent_cx_parent_based", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), OtelContext::new().with_remote_span_context(span_context(TraceFlags::SAMPLED, true)), None, true), - ("unsampled_remote_parent_cx_parent_based", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), OtelContext::new().with_remote_span_context(span_context(TraceFlags::default(), true)), None, false), - - // Existing sampling result defers - ("previous_drop_result_always_on", Sampler::AlwaysOn, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::Drop, attributes: vec![], trace_state: Default::default() }), false), - ("previous_record_and_sample_result_always_off", Sampler::AlwaysOff, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::RecordAndSample, attributes: vec![], trace_state: Default::default() }), true), - - // Existing local parent, defers - ("previous_drop_result_always_on", Sampler::AlwaysOn, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::Drop, attributes: vec![], trace_state: Default::default() }), false), - ("previous_record_and_sample_result_always_off", Sampler::AlwaysOff, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::RecordAndSample, attributes: vec![], trace_state: Default::default() }), true), - ] - } - - #[test] - fn sampled_context() { - for (name, sampler, parent_cx, previous_sampling_result, is_sampled) in sampler_data() { - let provider = SdkTracerProvider::builder().with_sampler(sampler).build(); - let tracer = provider.tracer("test"); - let mut builder = SpanBuilder::from_name("parent".to_string()); - builder.sampling_result = previous_sampling_result; - let builder = Some(builder); - let sampled = tracer.sampled_context(&mut OtelData { - builder, - parent_cx, - ..Default::default() - }); - - assert_eq!( - sampled.span().span_context().is_sampled(), - is_sampled, - "{}", - name - ) - } - } - - fn span_context(trace_flags: TraceFlags, is_remote: bool) -> SpanContext { - SpanContext::new( - TraceId::from(1u128), - SpanId::from(1u64), - trace_flags, - is_remote, - Default::default(), - ) - } -} From 133ab737890cfeac3592b5f9817ff0dd8464d3cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Tue, 24 Jun 2025 13:49:31 +0200 Subject: [PATCH 09/12] fix: add span links to SpanRef when builder has been consumed --- src/span_ext.rs | 11 ++++---- tests/span_ext.rs | 66 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 5 deletions(-) diff --git a/src/span_ext.rs b/src/span_ext.rs index 478e0bd..084e722 100644 --- a/src/span_ext.rs +++ b/src/span_ext.rs @@ -239,11 +239,10 @@ impl OpenTelemetrySpanExt for tracing::Span { // NOTE - if the span has been created - if we have _already_ // consumed our SpanBuilder_ - we can no longer mutate our parent! // This is an intentional design decision. - if let Some(builder) = &mut data.builder { - // If we still have a builder, update it to use the new parent context - // when it's eventually built + if data.builder.is_some() { + // If we still have a builder, update the data so it uses the + // new parent context when it's eventually built data.parent_cx = new_cx; - builder.sampling_result = None; } }); }); @@ -257,7 +256,6 @@ impl OpenTelemetrySpanExt for tracing::Span { if cx.is_valid() { let mut cx = Some(cx); let mut att = Some(attributes); - // TODO:ban add add version for SpanRef self.with_subscriber(move |(id, subscriber)| { let Some(get_context) = subscriber.downcast_ref::() else { return; @@ -273,6 +271,9 @@ impl OpenTelemetrySpanExt for tracing::Span { .links .get_or_insert_with(|| Vec::with_capacity(1)) .push(follows_link); + } else { + let span = data.parent_cx.span(); + span.add_link(follows_link.span_context, follows_link.attributes); } }); }); diff --git a/tests/span_ext.rs b/tests/span_ext.rs index 4c9db4c..800b5a3 100644 --- a/tests/span_ext.rs +++ b/tests/span_ext.rs @@ -165,3 +165,69 @@ fn test_add_event_with_timestamp() { event_data.timestamp ); } + +#[test] +fn test_add_link_variants() { + let (_tracer, provider, exporter, subscriber) = test_tracer(); + + let link_builder_cx = opentelemetry::trace::SpanContext::new( + opentelemetry::trace::TraceId::from_u128(0x1234567890abcdef1234567890abcdef), + opentelemetry::trace::SpanId::from_u64(0x1234567890abcdef), + opentelemetry::trace::TraceFlags::default(), + true, // Is remote + opentelemetry::trace::TraceState::default(), + ); + let link_current_cx = opentelemetry::trace::SpanContext::new( + opentelemetry::trace::TraceId::from_u128(0xabcdef1234567890abcdef1234567890), + opentelemetry::trace::SpanId::from_u64(0xabcdef1234567890), + opentelemetry::trace::TraceFlags::default(), + true, // Is remote + opentelemetry::trace::TraceState::default(), + ); + let link_attrs = vec![ + opentelemetry::KeyValue::new("link_key_1", "link_value_1"), + opentelemetry::KeyValue::new("link_key_2", 123), + ]; + + tracing::subscriber::with_default(subscriber, || { + let root = tracing::debug_span!("root"); + // Add the link using the extension method that now targets the builder + root.add_link(link_builder_cx.clone()); + // Enter span to make it current for the link addition + let _enter = root.enter(); + // Add the link using the extension method that now targets the span in the context + root.add_link_with_attributes(link_current_cx.clone(), link_attrs.clone()); + }); + + drop(provider); // flush all spans + let spans = exporter.0.lock().unwrap(); + + assert_eq!(spans.len(), 1, "Should have exported exactly one span."); + let root_span_data = spans.first().unwrap(); + + assert_eq!(root_span_data.links.len(), 2, "Span should have two links."); + let mut links = root_span_data.links.iter().collect::>(); + links.sort_by(|a, b| { + a.span_context + .trace_id() + .to_string() + .cmp(&b.span_context.trace_id().to_string()) + }); + let link1 = &links[0]; + let link2 = &links[1]; + + assert_eq!( + link1.span_context, link_builder_cx, + "Link 1 context mismatch." + ); + assert_eq!( + link1.attributes, + vec![], + "Link 1 attributes should be empty." + ); + assert_eq!( + link2.span_context, link_current_cx, + "Link 2 context mismatch." + ); + assert_eq!(link2.attributes, link_attrs, "Link 2 attributes mismatch."); +} From d96017050f96c12d88c2d12579335c87daaa6254 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Tue, 24 Jun 2025 17:41:05 +0200 Subject: [PATCH 10/12] chore: change OtelData to use an enum for different states --- src/layer.rs | 168 ++++++++++++++++++++++++++---------------------- src/lib.rs | 32 +++++++-- src/span_ext.rs | 101 ++++++++++++++++++----------- 3 files changed, 179 insertions(+), 122 deletions(-) diff --git a/src/layer.rs b/src/layer.rs index 0f29cd3..cb20a8b 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -1,6 +1,6 @@ #[cfg(feature = "activate_context")] use crate::stack::IdValueStack; -use crate::OtelData; +use crate::{OtelData, OtelDataState}; #[cfg(feature = "activate_context")] use opentelemetry::ContextGuard; use opentelemetry::{ @@ -9,12 +9,12 @@ use opentelemetry::{ }; #[cfg(feature = "activate_context")] use std::cell::RefCell; -use std::marker; use std::thread; #[cfg(not(all(target_arch = "wasm32", not(target_os = "wasi"))))] use std::time::Instant; use std::{any::TypeId, borrow::Cow}; use std::{fmt, vec}; +use std::{marker, mem::take}; use tracing_core::span::{self, Attributes, Id, Record}; use tracing_core::{field, Event, Subscriber}; #[cfg(feature = "tracing-log")] @@ -164,16 +164,7 @@ struct SpanBuilderUpdates { } impl SpanBuilderUpdates { - fn update(self, span_builder: &mut Option) -> Option { - if let Some(builder) = span_builder.as_mut() { - self.apply(builder); - None - } else { - Some(self) - } - } - - fn apply(self, span_builder: &mut SpanBuilder) { + fn update(self, span_builder: &mut SpanBuilder) { let Self { name, span_kind, @@ -956,15 +947,26 @@ where /// the context in the process. /// fn start_cx(&self, otel_data: &mut OtelData) { - if let Some(builder) = otel_data.builder.take() { - let span = builder.start_with_context(&self.tracer, &otel_data.parent_cx); - otel_data.parent_cx = otel_data.parent_cx.with_span(span); + if let OtelDataState::Context { .. } = &otel_data.state { + // If the context is already started, we do nothing. + } else { + match take(&mut otel_data.state) { + OtelDataState::Builder { builder, parent_cx } => { + let span = builder.start_with_context(&self.tracer, &parent_cx); + let current_cx = parent_cx.with_span(span); + otel_data.state = OtelDataState::Context { current_cx }; + } + _ => (), // This should never happen. + } } } fn with_started_cx(&self, otel_data: &mut OtelData, f: &dyn Fn(&OtelContext) -> U) -> U { self.start_cx(otel_data); - f(&otel_data.parent_cx) + match &otel_data.state { + OtelDataState::Context { current_cx, .. } => f(current_cx), + _ => panic!("OtelDataState should be a Context after starting it; this is a bug!"), + } } } @@ -1054,12 +1056,10 @@ where sem_conv_config: self.sem_conv_config, }); - let mut builder = Some(builder); updates.update(&mut builder); extensions.insert(OtelData { - builder, - parent_cx, - ..Default::default() + state: OtelDataState::Builder { builder, parent_cx }, + end_time: None, }); } @@ -1080,6 +1080,10 @@ where GUARD_STACK.with(|stack| stack.borrow_mut().push(id.clone(), guard)); }); } + + if !self.tracked_inactivity { + return; + } } if let Some(timings) = extensions.get_mut::() { @@ -1128,8 +1132,15 @@ where }); let mut extensions = span.extensions_mut(); if let Some(otel_data) = extensions.get_mut::() { - if let Some(updates) = updates.update(&mut otel_data.builder) { - updates.update_span(&otel_data.parent_cx.span()); + match &mut otel_data.state { + OtelDataState::Builder { builder, .. } => { + // If the builder is present, then update it. + updates.update(builder); + } + OtelDataState::Context { current_cx, .. } => { + // If the Context has been created, then update the span. + updates.update_span(¤t_cx.span()); + } } } } @@ -1151,14 +1162,17 @@ where .expect("Missing otel data span extensions"); let follows_context = self.with_started_cx(follows_data, &|cx| cx.span().span_context().clone()); - if let Some(builder) = data.builder.as_mut() { - if let Some(ref mut links) = builder.links { - links.push(otel::Link::with_context(follows_context)); - } else { - builder.links = Some(vec![otel::Link::with_context(follows_context)]); + match &mut data.state { + OtelDataState::Builder { builder, .. } => { + if let Some(ref mut links) = builder.links { + links.push(otel::Link::with_context(follows_context)); + } else { + builder.links = Some(vec![otel::Link::with_context(follows_context)]); + } + } + OtelDataState::Context { current_cx, .. } => { + current_cx.span().add_link(follows_context, vec![]); } - } else { - data.parent_cx.span().add_link(follows_context, vec![]); } } } @@ -1263,33 +1277,36 @@ where } } - if let Some(builder) = otel_data.builder.as_mut() { - if builder.status == otel::Status::Unset - && *meta.level() == tracing_core::Level::ERROR - { - builder.status = otel::Status::error(""); - } - if let Some(builder_updates) = builder_updates { - builder_updates.apply(builder); + match &mut otel_data.state { + OtelDataState::Builder { builder, .. } => { + if builder.status == otel::Status::Unset + && *meta.level() == tracing_core::Level::ERROR + { + builder.status = otel::Status::error(""); + } + if let Some(builder_updates) = builder_updates { + builder_updates.update(builder); + } + if let Some(ref mut events) = builder.events { + events.push(otel_event); + } else { + builder.events = Some(vec![otel_event]); + } } - if let Some(ref mut events) = builder.events { - events.push(otel_event); - } else { - builder.events = Some(vec![otel_event]); - } - } else { - let span = otel_data.parent_cx.span(); - // TODO:ban fix this with accessor in SpanRef that can check the span status - if *meta.level() == tracing_core::Level::ERROR { - span.set_status(otel::Status::error("")); + OtelDataState::Context { current_cx, .. } => { + let span = current_cx.span(); + // TODO:ban fix this with accessor in SpanRef that can check the span status + if *meta.level() == tracing_core::Level::ERROR { + span.set_status(otel::Status::error("")); + } + if let Some(builder_updates) = builder_updates { + builder_updates.update_span(&span); + } + span.add_event(otel_event.name, otel_event.attributes); } - if let Some(builder_updates) = builder_updates { - builder_updates.update_span(&span); - } - span.add_event(otel_event.name, otel_event.attributes); } - } - }; + }; + } } /// Exports an OpenTelemetry [`Span`] on close. @@ -1308,12 +1325,7 @@ where (extensions.remove::(), timings) }; - if let Some(OtelData { - builder, - parent_cx, - end_time, - }) = otel_data - { + if let Some(OtelData { state, end_time }) = otel_data { // Append busy/idle timings when enabled. let timings = timings.map(|timings| { let busy_ns = Key::new("busy_ns"); @@ -1325,24 +1337,28 @@ where ] }); - if let Some(builder) = builder { - // Don't create the context here just to get a SpanRef since it's costly - let mut span = builder.start_with_context(&self.tracer, &parent_cx); - if let Some(timings) = timings { - span.set_attributes(timings) - }; - if let Some(end_time) = end_time { - span.end_with_timestamp(end_time); - } else { - span.end(); + match state { + OtelDataState::Builder { builder, parent_cx } => { + // Don't create the context here just to get a SpanRef since it's costly + let mut span = builder.start_with_context(&self.tracer, &parent_cx); + if let Some(timings) = timings { + span.set_attributes(timings) + }; + if let Some(end_time) = end_time { + span.end_with_timestamp(end_time); + } else { + span.end(); + } } - } else { - let span = parent_cx.span(); - if let Some(timings) = timings { - span.set_attributes(timings) - }; - end_time.map_or_else(|| span.end(), |end_time| span.end_with_timestamp(end_time)); - }; + OtelDataState::Context { current_cx } => { + let span = current_cx.span(); + if let Some(timings) = timings { + span.set_attributes(timings) + }; + end_time + .map_or_else(|| span.end(), |end_time| span.end_with_timestamp(end_time)); + } + } } } diff --git a/src/lib.rs b/src/lib.rs index 8ea4d05..f04cb8a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -130,15 +130,33 @@ pub use metrics::MetricsLayer; pub use span_ext::OpenTelemetrySpanExt; /// Per-span OpenTelemetry data tracked by this crate. -#[derive(Debug, Clone, Default)] -struct OtelData { - /// The parent otel `Context` for the current tracing span. - pub parent_cx: opentelemetry::Context, +#[derive(Debug)] +pub(crate) struct OtelData { + /// The state of the OtelData, which can either be a builder or a context. + state: OtelDataState, + /// The end time of the span if it has been exited. + end_time: Option, +} - /// The otel span data recorded during the current tracing span. - pub builder: Option, +/// The state of the OpenTelemetry data for a span. +#[derive(Debug)] +#[allow(clippy::large_enum_variant)] +pub(crate) enum OtelDataState { + /// The span is being built, with a parent context and a builder. + Builder { + parent_cx: opentelemetry::Context, + builder: opentelemetry::trace::SpanBuilder, + }, + /// The span has been started or accessed and is now in a context. + Context { current_cx: opentelemetry::Context }, +} - end_time: Option, +impl Default for OtelDataState { + fn default() -> Self { + OtelDataState::Context { + current_cx: opentelemetry::Context::default(), + } + } } pub(crate) mod time { diff --git a/src/span_ext.rs b/src/span_ext.rs index 084e722..52a262d 100644 --- a/src/span_ext.rs +++ b/src/span_ext.rs @@ -1,5 +1,5 @@ -use crate::layer::WithContext; use crate::OtelData; +use crate::{layer::WithContext, OtelDataState}; use opentelemetry::{ time, trace::{SpanContext, Status, TraceContextExt}, @@ -239,10 +239,13 @@ impl OpenTelemetrySpanExt for tracing::Span { // NOTE - if the span has been created - if we have _already_ // consumed our SpanBuilder_ - we can no longer mutate our parent! // This is an intentional design decision. - if data.builder.is_some() { - // If we still have a builder, update the data so it uses the - // new parent context when it's eventually built - data.parent_cx = new_cx; + match &mut data.state { + OtelDataState::Builder { parent_cx, .. } => { + // If we still have a builder, update the data so it uses the + // new parent context when it's eventually built + *parent_cx = new_cx; + } + OtelDataState::Context { .. } => (), } }); }); @@ -266,14 +269,21 @@ impl OpenTelemetrySpanExt for tracing::Span { }; let attr = att.take().unwrap_or_default(); let follows_link = opentelemetry::trace::Link::new(cx, attr, 0); - if let Some(builder) = data.builder.as_mut() { - builder - .links - .get_or_insert_with(|| Vec::with_capacity(1)) - .push(follows_link); - } else { - let span = data.parent_cx.span(); - span.add_link(follows_link.span_context, follows_link.attributes); + match &mut data.state { + OtelDataState::Builder { builder, .. } => { + // If we still have a builder, update the data so it uses the + // new link when it's eventually built + builder + .links + .get_or_insert_with(|| Vec::with_capacity(1)) + .push(follows_link); + } + OtelDataState::Context { current_cx } => { + // If we have a context, add the link to the span in the context + current_cx + .span() + .add_link(follows_link.span_context, follows_link.attributes); + } } }); }); @@ -288,7 +298,9 @@ impl OpenTelemetrySpanExt for tracing::Span { }; // If our span hasn't been built, we should build it and get the context in one call get_context.with_activated_context(subscriber, id, |data: &mut OtelData| { - cx = Some(data.parent_cx.clone()); + if let OtelDataState::Context { current_cx } = &data.state { + cx = Some(current_cx.clone()); + } }); }); @@ -303,20 +315,23 @@ impl OpenTelemetrySpanExt for tracing::Span { let mut key = Some(key.into()); let mut value = Some(value.into()); get_context.with_context(subscriber, id, move |data| { - if let Some(builder) = data.builder.as_mut() { - if builder.attributes.is_none() { - builder.attributes = Some(Default::default()); + match &mut data.state { + OtelDataState::Builder { builder, .. } => { + if builder.attributes.is_none() { + builder.attributes = Some(Default::default()); + } + builder + .attributes + .as_mut() + .unwrap() + .push(KeyValue::new(key.take().unwrap(), value.take().unwrap())); } - builder - .attributes - .as_mut() - .unwrap() - .push(KeyValue::new(key.take().unwrap(), value.take().unwrap())); - } else { - let span = data.parent_cx.span(); - let key_value = KeyValue::new(key.take().unwrap(), value.take().unwrap()); - span.set_attribute(key_value); - } + OtelDataState::Context { current_cx } => { + let span = current_cx.span(); + let key_value = KeyValue::new(key.take().unwrap(), value.take().unwrap()); + span.set_attribute(key_value); + } + }; }); }); } @@ -327,11 +342,12 @@ impl OpenTelemetrySpanExt for tracing::Span { let Some(get_context) = subscriber.downcast_ref::() else { return; }; - get_context.with_context(subscriber, id, move |data| { - if let Some(builder) = data.builder.as_mut() { + get_context.with_context(subscriber, id, move |data| match &mut data.state { + OtelDataState::Builder { builder, .. } => { builder.status = status.take().unwrap(); - } else { - let span = data.parent_cx.span(); + } + OtelDataState::Context { current_cx } => { + let span = current_cx.span(); span.set_status(status.take().unwrap()); } }); @@ -359,14 +375,21 @@ impl OpenTelemetrySpanExt for tracing::Span { let Some(event) = event.take() else { return; }; - if let Some(builder) = data.builder.as_mut() { - builder - .events - .get_or_insert_with(|| Vec::with_capacity(1)) - .push(event); - } else { - let span = data.parent_cx.span(); - span.add_event_with_timestamp(event.name, event.timestamp, event.attributes); + match &mut data.state { + OtelDataState::Builder { builder, .. } => { + builder + .events + .get_or_insert_with(|| Vec::with_capacity(1)) + .push(event); + } + OtelDataState::Context { current_cx } => { + let span = current_cx.span(); + span.add_event_with_timestamp( + event.name, + event.timestamp, + event.attributes, + ); + } } }); }); From ccf26da6a5c5a3ade4849527b1667c2f2e8c83f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Mon, 7 Jul 2025 15:19:45 +0200 Subject: [PATCH 11/12] chore: convert activate_context feature to context_activation setting --- Cargo.toml | 4 +- src/layer.rs | 399 +++++++++++++++++++++++++++++++++++++++++++-------- src/lib.rs | 1 - 3 files changed, 337 insertions(+), 67 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ce7c278..b9418c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,11 +16,9 @@ edition = "2021" rust-version = "1.75.0" [features] -default = ["tracing-log", "metrics", "activate_context"] +default = ["tracing-log", "metrics"] # Enables support for exporting OpenTelemetry metrics metrics = ["opentelemetry/metrics","opentelemetry_sdk/metrics", "smallvec"] -# Enables support for activating an OpenTelemetry Context on Span entry -activate_context = [] [dependencies] opentelemetry = { version = "0.30.0", default-features = false, features = ["trace"] } diff --git a/src/layer.rs b/src/layer.rs index cb20a8b..e7d111e 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -1,13 +1,10 @@ -#[cfg(feature = "activate_context")] use crate::stack::IdValueStack; use crate::{OtelData, OtelDataState}; -#[cfg(feature = "activate_context")] use opentelemetry::ContextGuard; use opentelemetry::{ trace::{self as otel, noop, Span, SpanBuilder, SpanKind, Status, TraceContextExt}, Context as OtelContext, Key, KeyValue, StringValue, Value, }; -#[cfg(feature = "activate_context")] use std::cell::RefCell; use std::thread; #[cfg(not(all(target_arch = "wasm32", not(target_os = "wasi"))))] @@ -45,6 +42,7 @@ pub struct OpenTelemetryLayer { tracked_inactivity: bool, with_threads: bool, with_level: bool, + context_activation: bool, sem_conv_config: SemConvConfig, with_context: WithContext, _registry: marker::PhantomData, @@ -617,6 +615,7 @@ where tracked_inactivity: true, with_threads: true, with_level: false, + context_activation: true, sem_conv_config: SemConvConfig { error_fields_to_exceptions: true, error_records_to_exceptions: true, @@ -675,6 +674,7 @@ where tracked_inactivity: self.tracked_inactivity, with_threads: self.with_threads, with_level: self.with_level, + context_activation: self.context_activation, sem_conv_config: self.sem_conv_config, with_context: WithContext { with_context: OpenTelemetryLayer::::get_context, @@ -822,6 +822,20 @@ where } } + /// Sets whether or not an OpenTelemetry Context should be activated on span entry. + /// + /// When enabled, entering a span will activate its OpenTelemetry context, making it + /// available to other OpenTelemetry instrumentation. This allows for proper context + /// propagation across different instrumentation libraries. + /// + /// By default, context activation is enabled. + pub fn with_context_activation(self, context_activation: bool) -> Self { + Self { + context_activation, + ..self + } + } + /// Retrieve the parent OpenTelemetry [`Context`] from the current tracing /// [`span`] through the [`Registry`]. This [`Context`] links spans to their /// parent for proper hierarchical visualization. @@ -854,16 +868,13 @@ where } if attrs.is_contextual() { - #[cfg(feature = "activate_context")] - // If the span is contextual and we are using the activate_context feature, - // we should use the current OTel context - { + if self.context_activation { + // If the span is contextual and we are using context activation, + // we should use the current OTel context OtelContext::current() - } - #[cfg(not(feature = "activate_context"))] - // If the span is contextual and we are not using the activate_context feature, - // we should use the current tracing context - { + } else { + // If the span is contextual and we are not using context activation, + // we should use the current tracing context ctx.lookup_current() .and_then(|span| { let mut extensions = span.extensions_mut(); @@ -983,12 +994,10 @@ thread_local! { }; } -#[cfg(feature = "activate_context")] thread_local! { static GUARD_STACK: RefCell = RefCell::new(IdContextGuardStack::new()); } -#[cfg(feature = "activate_context")] type IdContextGuardStack = IdValueStack; impl Layer for OpenTelemetryLayer @@ -1064,16 +1073,14 @@ where } fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) { - #[cfg(not(feature = "activate_context"))] - if !self.tracked_inactivity { + if !self.context_activation && !self.tracked_inactivity { return; } let span = ctx.span(id).expect("Span not found, this is a bug"); let mut extensions = span.extensions_mut(); - #[cfg(feature = "activate_context")] - { + if self.context_activation { if let Some(otel_data) = extensions.get_mut::() { self.with_started_cx(otel_data, &|cx| { let guard = cx.clone().attach(); @@ -1102,8 +1109,9 @@ where if let Some(otel_data) = extensions.get_mut::() { otel_data.end_time = Some(crate::time::now()); - #[cfg(feature = "activate_context")] - GUARD_STACK.with(|stack| stack.borrow_mut().pop(id)); + if self.context_activation { + GUARD_STACK.with(|stack| stack.borrow_mut().pop(id)); + } } if !self.tracked_inactivity { @@ -1573,9 +1581,22 @@ mod tests { } #[test] - fn trace_id_from_existing_context() { + fn trace_id_from_existing_context_with_context_activation() { + trace_id_from_existing_context_impl(true); + } + + #[test] + fn trace_id_from_existing_context_no_context_activation() { + trace_id_from_existing_context_impl(false); + } + + fn trace_id_from_existing_context_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), + ); let trace_id = otel::TraceId::from(42u128); let existing_cx = OtelContext::current_with_span(TestSpan(otel::SpanContext::new( trace_id, @@ -1595,12 +1616,22 @@ mod tests { } #[test] - fn includes_timings() { + fn includes_timings_with_context_activation() { + includes_timings_impl(true); + } + + #[test] + fn includes_timings_no_context_activation() { + includes_timings_impl(false); + } + + fn includes_timings_impl(context_activation: bool) { let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with( layer() .with_tracer(tracer.clone()) - .with_tracked_inactivity(true), + .with_tracked_inactivity(true) + .with_context_activation(context_activation), ); tracing::subscriber::with_default(subscriber, || { @@ -1614,9 +1645,22 @@ mod tests { } #[test] - fn records_error_fields() { + fn records_error_fields_with_context_activation() { + records_error_fields_impl(true); + } + + #[test] + fn records_error_fields_no_context_activation() { + records_error_fields_impl(false); + } + + fn records_error_fields_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), + ); let err = TestDynError::new("base error") .with_parent("intermediate error") @@ -1657,9 +1701,22 @@ mod tests { } #[test] - fn records_event_name() { + fn records_event_name_with_context_activation() { + records_event_name_impl(true); + } + + #[test] + fn records_event_name_no_context_activation() { + records_event_name_impl(false); + } + + fn records_event_name_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), + ); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("test span").in_scope(|| { @@ -1683,12 +1740,22 @@ mod tests { } #[test] - fn records_no_error_fields() { + fn records_no_error_fields_with_context_activation() { + records_no_error_fields_impl(true); + } + + #[test] + fn records_no_error_fields_no_context_activation() { + records_no_error_fields_impl(false); + } + + fn records_no_error_fields_impl(context_activation: bool) { let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with( layer() .with_error_records_to_exceptions(false) - .with_tracer(tracer.clone()), + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), ); let err = TestDynError::new("base error") @@ -1730,10 +1797,23 @@ mod tests { } #[test] - fn includes_span_location() { + fn includes_span_location_with_context_activation() { + includes_span_location_impl(true); + } + + #[test] + fn includes_span_location_no_context_activation() { + includes_span_location_impl(false); + } + + fn includes_span_location_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry() - .with(layer().with_tracer(tracer.clone()).with_location(true)); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_location(true) + .with_context_activation(context_activation), + ); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("request"); @@ -1747,10 +1827,23 @@ mod tests { } #[test] - fn excludes_span_location() { + fn excludes_span_location_with_context_activation() { + excludes_span_location_impl(true); + } + + #[test] + fn excludes_span_location_no_context_activation() { + excludes_span_location_impl(false); + } + + fn excludes_span_location_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry() - .with(layer().with_tracer(tracer.clone()).with_location(false)); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_location(false) + .with_context_activation(context_activation), + ); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("request"); @@ -1764,7 +1857,16 @@ mod tests { } #[test] - fn includes_thread() { + fn includes_thread_with_context_activation() { + includes_thread_impl(true); + } + + #[test] + fn includes_thread_no_context_activation() { + includes_thread_impl(false); + } + + fn includes_thread_impl(context_activation: bool) { let thread = thread::current(); let expected_name = thread .name() @@ -1772,8 +1874,12 @@ mod tests { let expected_id = Value::I64(thread_id_integer(thread.id()) as i64); let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry() - .with(layer().with_tracer(tracer.clone()).with_threads(true)); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_threads(true) + .with_context_activation(context_activation), + ); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("request"); @@ -1786,10 +1892,23 @@ mod tests { } #[test] - fn excludes_thread() { + fn excludes_thread_with_context_activation() { + excludes_thread_impl(true); + } + + #[test] + fn excludes_thread_no_context_activation() { + excludes_thread_impl(false); + } + + fn excludes_thread_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry() - .with(layer().with_tracer(tracer.clone()).with_threads(false)); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_threads(false) + .with_context_activation(context_activation), + ); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("request"); @@ -1802,10 +1921,23 @@ mod tests { } #[test] - fn includes_level() { + fn includes_level_with_context_activation() { + includes_level_impl(true); + } + + #[test] + fn includes_level_no_context_activation() { + includes_level_impl(false); + } + + fn includes_level_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry() - .with(layer().with_tracer(tracer.clone()).with_level(true)); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_level(true) + .with_context_activation(context_activation), + ); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("request"); @@ -1817,10 +1949,23 @@ mod tests { } #[test] - fn excludes_level() { + fn excludes_level_with_context_activation() { + excludes_level_impl(true); + } + + #[test] + fn excludes_level_no_context_activation() { + excludes_level_impl(false); + } + + fn excludes_level_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry() - .with(layer().with_tracer(tracer.clone()).with_level(false)); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_level(false) + .with_context_activation(context_activation), + ); tracing::subscriber::with_default(subscriber, || { tracing::debug_span!("request"); @@ -1832,9 +1977,22 @@ mod tests { } #[test] - fn propagates_error_fields_from_event_to_span() { + fn propagates_error_fields_from_event_to_span_with_context_activation() { + propagates_error_fields_from_event_to_span_impl(true); + } + + #[test] + fn propagates_error_fields_from_event_to_span_no_context_activation() { + propagates_error_fields_from_event_to_span_impl(false); + } + + fn propagates_error_fields_from_event_to_span_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), + ); let err = TestDynError::new("base error") .with_parent("intermediate error") @@ -1865,12 +2023,22 @@ mod tests { } #[test] - fn propagates_no_error_fields_from_event_to_span() { + fn propagates_no_error_fields_from_event_to_span_with_context_activation() { + propagates_no_error_fields_from_event_to_span_impl(true); + } + + #[test] + fn propagates_no_error_fields_from_event_to_span_no_context_activation() { + propagates_no_error_fields_from_event_to_span_impl(false); + } + + fn propagates_no_error_fields_from_event_to_span_impl(context_activation: bool) { let mut tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry().with( layer() .with_error_fields_to_exceptions(false) - .with_tracer(tracer.clone()), + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), ); let err = TestDynError::new("base error") @@ -1902,13 +2070,23 @@ mod tests { } #[test] - fn tracing_error_compatibility() { + fn tracing_error_compatibility_with_context_activation() { + tracing_error_compatibility_impl(true); + } + + #[test] + fn tracing_error_compatibility_no_context_activation() { + tracing_error_compatibility_impl(false); + } + + fn tracing_error_compatibility_impl(context_activation: bool) { let tracer = TestTracer::default(); let subscriber = tracing_subscriber::registry() .with( layer() .with_error_fields_to_exceptions(false) - .with_tracer(tracer.clone()), + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), ) .with(tracing_error::ErrorLayer::default()); @@ -1926,14 +2104,11 @@ mod tests { // No need to assert anything, as long as this finished (and did not panic), everything is ok. } - #[cfg(feature = "activate_context")] #[derive(Debug, PartialEq)] struct ValueA(&'static str); - #[cfg(feature = "activate_context")] #[derive(Debug, PartialEq)] struct ValueB(&'static str); - #[cfg(feature = "activate_context")] #[test] fn otel_context_propagation() { use opentelemetry::trace::Tracer; @@ -2046,9 +2221,22 @@ mod tests { } #[test] - fn record_after() { + fn record_after_with_context_activation() { + record_after_impl(true); + } + + #[test] + fn record_after_no_context_activation() { + record_after_impl(false); + } + + fn record_after_impl(context_activation: bool) { let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), + ); tracing::subscriber::with_default(subscriber, || { let root = trace_span!("root", before = "before", after = "before"); @@ -2115,10 +2303,23 @@ mod tests { } #[test] - fn follows_from_adds_link() { + fn follows_from_adds_link_with_context_activation() { + follows_from_adds_link_impl(true); + } + + #[test] + fn follows_from_adds_link_no_context_activation() { + follows_from_adds_link_impl(false); + } + + fn follows_from_adds_link_impl(context_activation: bool) { use crate::OpenTelemetrySpanExt; let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), + ); let span1_id = tracing::subscriber::with_default(subscriber, || { let span2 = tracing::debug_span!("span2"); @@ -2167,10 +2368,23 @@ mod tests { } #[test] - fn follows_from_multiple_links() { + fn follows_from_multiple_links_with_context_activation() { + follows_from_multiple_links_impl(true); + } + + #[test] + fn follows_from_multiple_links_no_context_activation() { + follows_from_multiple_links_impl(false); + } + + fn follows_from_multiple_links_impl(context_activation: bool) { use crate::OpenTelemetrySpanExt; let mut tracer = TestTracer::default(); - let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone())); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_context_activation(context_activation), + ); let (span1_id, span2_id) = tracing::subscriber::with_default(subscriber, || { let span3 = tracing::debug_span!("span3"); @@ -2221,4 +2435,63 @@ mod tests { "Links should point to the correct source spans" ); } + + #[test] + fn context_activation_disabled() { + use tracing::span; + + let mut tracer = TestTracer::default(); + let subscriber = tracing_subscriber::registry().with( + layer() + .with_tracer(tracer.clone()) + .with_context_activation(false), + ); + + tracing::subscriber::with_default(subscriber, || { + // Add a value to the current OpenTelemetry context + let _outer_guard = + OtelContext::attach(OtelContext::default().with_value(ValueA("outer"))); + assert_eq!(OtelContext::current().get(), Some(&ValueA("outer"))); + + let root = span!(tracing::Level::TRACE, "tokio-tracing-span-parent"); + + // Drop the guard to ensure the context is cleared + drop(_outer_guard); + assert!(OtelContext::current().get::().is_none()); + + // Enter the root span - with context activation disabled, + // the context should NOT be propagated + let _enter_root = root.enter(); + assert!(OtelContext::current().get::().is_none()); + + // Add another value to the current OpenTelemetry context + let _inner_guard = + OtelContext::attach(OtelContext::current_with_value(ValueB("inner"))); + assert!(OtelContext::current().get::().is_none()); + assert_eq!(OtelContext::current().get(), Some(&ValueB("inner"))); + + let child = span!(tracing::Level::TRACE, "tokio-tracing-span-child"); + + // Drop the guard to ensure the context is reverted + drop(_inner_guard); + assert!(OtelContext::current().get::().is_none()); + assert!(OtelContext::current().get::().is_none()); + + // Enter the child span - with context activation disabled, + // the context should NOT be propagated + let _enter_child = child.enter(); + assert!(OtelContext::current().get::().is_none()); + assert!(OtelContext::current().get::().is_none()); + }); + + // Verify spans were still created and exported + let spans = tracer.spans(); + assert_eq!(spans.len(), 2); + assert!(spans + .iter() + .any(|span| span.name == "tokio-tracing-span-parent")); + assert!(spans + .iter() + .any(|span| span.name == "tokio-tracing-span-child")); + } } diff --git a/src/lib.rs b/src/lib.rs index f04cb8a..ba0fdcd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -118,7 +118,6 @@ mod layer; /// Span extension which enables OpenTelemetry context management. mod span_ext; -#[cfg(feature = "activate_context")] mod stack; use std::time::SystemTime; From b9a24936761b7ef0a789068978567f01c1f29d93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Antonsson?= Date: Mon, 7 Jul 2025 15:19:53 +0200 Subject: [PATCH 12/12] chore: change the set_parent method to return a Result --- examples/opentelemetry-remote-context.rs | 2 +- src/layer.rs | 8 ++++---- src/span_ext.rs | 22 +++++++++++++++++----- tests/trace_state_propagation.rs | 8 ++++---- 4 files changed, 26 insertions(+), 14 deletions(-) diff --git a/examples/opentelemetry-remote-context.rs b/examples/opentelemetry-remote-context.rs index 350777a..7de8e54 100644 --- a/examples/opentelemetry-remote-context.rs +++ b/examples/opentelemetry-remote-context.rs @@ -37,7 +37,7 @@ fn main() { let app_root = span!(tracing::Level::INFO, "app_start"); // Assign parent trace from external context - app_root.set_parent(parent_context); + let _ = app_root.set_parent(parent_context); // To include tracing context in client requests from _this_ app, // use `context` to extract the current OpenTelemetry context. diff --git a/src/layer.rs b/src/layer.rs index e7d111e..f2aaf72 100644 --- a/src/layer.rs +++ b/src/layer.rs @@ -2201,13 +2201,13 @@ mod tests { let child1 = trace_span!("child-1"); let root_context = root.context(); // SpanData (None) - child1.set_parent(root_context); // Clone context, but SpanData(None) + let _ = child1.set_parent(root_context); // Clone context, but SpanData(None) let _enter_root = root.enter(); drop(_enter_root); let child2 = trace_span!("child-2"); - child2.set_parent(root.context()); + let _ = child2.set_parent(root.context()); }); // Let's check the spans @@ -2280,10 +2280,10 @@ mod tests { _ = root.enter(); let child1 = trace_span!("child-1"); - child1.set_parent(root.context()); + let _ = child1.set_parent(root.context()); trace_span!(parent: &child1, "child-2"); - child1.set_parent(root.context()); // <-- this is what causes the issue + let _ = child1.set_parent(root.context()); // <-- this is what causes the issue trace_span!(parent: &child1, "child-3"); }); diff --git a/src/span_ext.rs b/src/span_ext.rs index 52a262d..9479a6f 100644 --- a/src/span_ext.rs +++ b/src/span_ext.rs @@ -17,6 +17,10 @@ pub trait OpenTelemetrySpanExt { /// Associates `self` with a given OpenTelemetry trace, using the provided /// parent [`Context`]. /// + /// This method provides error handling for cases where the span context + /// cannot be set, such as when the OpenTelemetry layer is not present + /// or when the span has already been started. + /// /// [`Context`]: opentelemetry::Context /// /// # Examples @@ -41,12 +45,12 @@ pub trait OpenTelemetrySpanExt { /// let app_root = tracing::span!(tracing::Level::INFO, "app_start"); /// /// // Assign parent trace from external context - /// app_root.set_parent(parent_context.clone()); + /// let _ = app_root.set_parent(parent_context.clone()); /// /// // Or if the current span has been created elsewhere: - /// Span::current().set_parent(parent_context); + /// let _ = Span::current().set_parent(parent_context); /// ``` - fn set_parent(&self, cx: Context); + fn set_parent(&self, cx: Context) -> Result<(), &'static str>; /// Associates `self` with a given OpenTelemetry trace, using the provided /// followed span [`SpanContext`]. @@ -224,10 +228,14 @@ impl OpenTelemetrySpanExt for tracing::Span { /// Additionally, once a span has been fully built - and the SpanBuilder has been consumed - /// the parent _cannot_ be mutated. /// - fn set_parent(&self, cx: Context) { + fn set_parent(&self, cx: Context) -> Result<(), &'static str> { let mut cx = Some(cx); + let mut result = Ok(()); + let result_ref = &mut result; + self.with_subscriber(move |(id, subscriber)| { let Some(get_context) = subscriber.downcast_ref::() else { + *result_ref = Err("OpenTelemetry layer not found"); return; }; // Set the parent OTel for the current span @@ -245,10 +253,14 @@ impl OpenTelemetrySpanExt for tracing::Span { // new parent context when it's eventually built *parent_cx = new_cx; } - OtelDataState::Context { .. } => (), + OtelDataState::Context { .. } => { + *result_ref = Err("Span has already been started, cannot set parent"); + } } }); }); + + result } fn add_link(&self, cx: SpanContext) { diff --git a/tests/trace_state_propagation.rs b/tests/trace_state_propagation.rs index 2a33bc0..f7a3344 100644 --- a/tests/trace_state_propagation.rs +++ b/tests/trace_state_propagation.rs @@ -37,7 +37,7 @@ fn trace_with_assigned_otel_context() { tracing::subscriber::with_default(subscriber, || { let child = tracing::debug_span!("child"); - child.set_parent(cx); + let _ = child.set_parent(cx); }); drop(provider); // flush all spans @@ -70,7 +70,7 @@ fn propagate_invalid_context() { tracing::subscriber::with_default(subscriber, || { let root = tracing::debug_span!("root"); - root.set_parent(invalid_cx); + let _ = root.set_parent(invalid_cx); root.in_scope(|| tracing::debug_span!("child")); }); @@ -90,7 +90,7 @@ fn inject_context_into_outgoing_requests() { tracing::subscriber::with_default(subscriber, || { let root = tracing::debug_span!("root"); - root.set_parent(cx); + let _ = root.set_parent(cx); let _g = root.enter(); let child = tracing::debug_span!("child"); propagator.inject_context(&child.context(), &mut outgoing_req_carrier); @@ -125,7 +125,7 @@ fn sampling_decision_respects_new_parent() { // Observation: if you force the _child_ to materialize before the parent, e.g., // if you swap these two lines - bad things will happen, and we shouldn't support // this. - child.set_parent(Context::current_with_span(root_span)); + let _ = child.set_parent(Context::current_with_span(root_span)); child.context(); // force a sampling decision });