fix: add span links to SpanRef when builder has been consumed

Author: bantonsson

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::<WithContext>() 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);
                     }
                 });
             });

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::<Vec<_>>();
+    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.");
+}