chore: address review comments

Author: bantonsson

Cargo.toml

@@ -29,6 +29,7 @@ tracing-subscriber = { version = "0.3.0", default-features = false, features = [
 tracing-log = { version = "0.2.0", default-features = false, optional = true }
 rustversion = "1.0.9"
 smallvec = { version = "1.0", optional = true }
+thiserror = { version = "2", default-features = false }
 
 # Fix minimal-versions
 lazy_static = { version = "1.0.2", optional = true }

src/span_ext.rs

@@ -6,6 +6,7 @@ use opentelemetry::{
     Context, Key, KeyValue, Value,
 };
 use std::{borrow::Cow, time::SystemTime};
+use thiserror::Error;
 
 /// Utility functions to allow tracing [`Span`]s to accept and return
 /// [OpenTelemetry] [`Context`]s.
@@ -17,6 +18,13 @@ pub trait OpenTelemetrySpanExt {
     /// Associates `self` with a given OpenTelemetry trace, using the provided
     /// parent [`Context`].
     ///
+    /// This method exists primarily to make it possible to inject a _distributed_ incoming
+    /// context, e.g. span IDs, etc.
+    ///
+    /// 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.
+    ///
     /// 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.
@@ -50,7 +58,7 @@ pub trait OpenTelemetrySpanExt {
     /// // Or if the current span has been created elsewhere:
     /// let _ = Span::current().set_parent(parent_context);
     /// ```
-    fn set_parent(&self, cx: Context) -> Result<(), &'static str>;
+    fn set_parent(&self, cx: Context) -> Result<(), SetParentError>;
 
     /// Associates `self` with a given OpenTelemetry trace, using the provided
     /// followed span [`SpanContext`].
@@ -218,29 +226,31 @@ pub trait OpenTelemetrySpanExt {
     );
 }
 
+#[derive(Error, Debug)]
+pub enum SetParentError {
+    #[error("OpenTelemetry layer not found")]
+    LayerNotFound,
+    #[error("Span has already been started, cannot set parent")]
+    AlreadyStarted,
+    #[error("Span disabled")]
+    SpanDisabled,
+}
+
 impl OpenTelemetrySpanExt for tracing::Span {
-    ///
-    /// 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) -> Result<(), &'static str> {
+    fn set_parent(&self, cx: Context) -> Result<(), SetParentError> {
         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::<WithContext>() else {
-                *result_ref = Err("OpenTelemetry layer not found");
-                return;
+                return Err(SetParentError::LayerNotFound);
             };
+
+            let mut result = Ok(());
+            let result_ref = &mut result;
             // Set the parent OTel for the current span
             get_context.with_context(subscriber, id, move |data| {
                 let Some(new_cx) = cx.take() else {
+                    *result_ref = Err(SetParentError::AlreadyStarted);
                     return;
                 };
                 // Create a new context with the new parent but preserve our span.
@@ -254,13 +264,13 @@ impl OpenTelemetrySpanExt for tracing::Span {
                         *parent_cx = new_cx;
                     }
                     OtelDataState::Context { .. } => {
-                        *result_ref = Err("Span has already been started, cannot set parent");
+                        *result_ref = Err(SetParentError::AlreadyStarted);
                     }
                 }
             });
-        });
-
-        result
+            result
+        })
+        .unwrap_or(Err(SetParentError::SpanDisabled))
     }
 
     fn add_link(&self, cx: SpanContext) {
@@ -324,8 +334,7 @@ impl OpenTelemetrySpanExt for tracing::Span {
             let Some(get_context) = subscriber.downcast_ref::<WithContext>() else {
                 return;
             };
-            let mut key = Some(key.into());
-            let mut value = Some(value.into());
+            let mut key_value = Some(KeyValue::new(key.into(), value.into()));
             get_context.with_context(subscriber, id, move |data| {
                 match &mut data.state {
                     OtelDataState::Builder { builder, .. } => {
@@ -336,12 +345,11 @@ impl OpenTelemetrySpanExt for tracing::Span {
                             .attributes
                             .as_mut()
                             .unwrap()
-                            .push(KeyValue::new(key.take().unwrap(), value.take().unwrap()));
+                            .push(key_value.take().unwrap());
                     }
                     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);
+                        span.set_attribute(key_value.take().unwrap());
                     }
                 };
             });

src/stack.rs

@@ -38,7 +38,7 @@ impl<T: Debug> IdValueStack<T> {
     }
 
     #[cfg(test)]
-    pub(super) fn len(&self) -> usize {
+    fn len(&self) -> usize {
         self.stack.len()
     }
 }

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");
-        let _ = child.set_parent(cx);
+        let _ = child.set_parent(cx).unwrap();
     });
 
     drop(provider); // flush all spans
@@ -122,10 +122,9 @@ fn sampling_decision_respects_new_parent() {
     tracing::subscriber::with_default(subscriber, || {
         let child = tracing::debug_span!("child");
 
-        // 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.
-        let _ = child.set_parent(Context::current_with_span(root_span));
+        let _ = child
+            .set_parent(Context::current_with_span(root_span))
+            .unwrap();
         child.context(); // force a sampling decision
     });
 
@@ -145,6 +144,26 @@ fn sampling_decision_respects_new_parent() {
     );
 }
 
+#[test]
+fn set_parent_fails_after_span_entered() {
+    let (_tracer, provider, _exporter, subscriber) = test_tracer();
+    let (cx, _, _, _) = build_sampled_context();
+
+    tracing::subscriber::with_default(subscriber, || {
+        let span = tracing::debug_span!("test_span");
+        let _guard = span.enter(); // Enter the span first
+
+        // Attempting to set parent after entering should fail
+        let result = span.set_parent(cx);
+        assert!(
+            result.is_err(),
+            "set_parent should fail after span is entered"
+        );
+    });
+
+    drop(provider);
+}
+
 fn assert_shared_attrs_eq(sc_a: &SpanContext, sc_b: &SpanContext) {
     assert_eq!(sc_a.trace_id(), sc_b.trace_id());
     assert_eq!(sc_a.trace_state(), sc_b.trace_state());