From a9720335719ea92034c1ed08a71dbc833020c510 Mon Sep 17 00:00:00 2001
From: Rohit Mittapalli <rohit@anthropic.com>
Date: Tue, 2 Jul 2024 14:55:52 -0700
Subject: [PATCH 1/2] invalidate sample decision on set parent, added test

---
 src/span_ext.rs                  |  1 +
 tests/trace_state_propagation.rs | 43 +++++++++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/span_ext.rs b/src/span_ext.rs
index b7d0454..c89dc09 100644
--- a/src/span_ext.rs
+++ b/src/span_ext.rs
@@ -143,6 +143,7 @@ impl OpenTelemetrySpanExt for tracing::Span {
                 get_context.with_context(subscriber, id, move |data, _tracer| {
                     if let Some(cx) = cx.take() {
                         data.parent_cx = cx;
+                        data.builder.sampling_result = None;
                     }
                 });
             }
diff --git a/tests/trace_state_propagation.rs b/tests/trace_state_propagation.rs
index 70ba836..cbad848 100644
--- a/tests/trace_state_propagation.rs
+++ b/tests/trace_state_propagation.rs
@@ -7,7 +7,7 @@ use opentelemetry::{
 use opentelemetry_sdk::{
     export::trace::{ExportResult, SpanData, SpanExporter},
     propagation::{BaggagePropagator, TraceContextPropagator},
-    trace::{Tracer, TracerProvider},
+    trace::{Config, Sampler, Tracer, TracerProvider},
 };
 use std::collections::{HashMap, HashSet};
 use std::sync::{Arc, Mutex};
@@ -101,6 +101,47 @@ fn inject_context_into_outgoing_requests() {
     assert_carrier_attrs_eq(&carrier, &outgoing_req_carrier);
 }
 
+#[test]
+fn sampling_decision_respects_new_parent() {
+    // custom setup required due to ParentBased(AlwaysOff) sampler
+    let exporter = TestExporter::default();
+    let provider = TracerProvider::builder()
+        .with_simple_exporter(exporter.clone())
+        .with_config(
+            Config::default()
+                .with_sampler(Sampler::ParentBased(Box::new(Sampler::AlwaysOff)))
+        )
+        .build();
+    let tracer = provider.tracer("test");
+    let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone()));
+
+    // set up remote sampled headers
+    let sampled_headers = HashMap::from([(
+        "traceparent".to_string(),
+        "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string(),
+    )]);
+    let remote_sampled_cx = TraceContextPropagator::new().extract(&sampled_headers);
+    let root_span = tracer.start_with_context("root_span", &remote_sampled_cx);
+
+    tracing::subscriber::with_default(subscriber, || {
+        let child = tracing::debug_span!("child");
+        child.context(); // force a sampling decision
+        child.set_parent(Context::current_with_span(root_span));
+    });
+
+    drop(provider); // flush all spans
+
+    // assert new parent-based sampling decision
+    let spans = exporter.0.lock().unwrap();
+    assert_eq!(spans.len(), 2, "Expected 2 spans, got {}", spans.len());
+    assert!(spans[0].span_context.is_sampled(), "Root span should be sampled");
+    assert_eq!(
+        spans[1].span_context.is_sampled(),
+        spans[0].span_context.is_sampled(),
+        "Child span should respect parent sampling decision"
+    );
+}
+
 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());

From 92da2131a210c12f31f8cdd9f6ff66a213fc08dc Mon Sep 17 00:00:00 2001
From: Rohit Mittapalli <rohit@anthropic.com>
Date: Wed, 3 Jul 2024 08:57:32 -0700
Subject: [PATCH 2/2] cargo fmt

---
 tests/trace_state_propagation.rs | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/trace_state_propagation.rs b/tests/trace_state_propagation.rs
index cbad848..7d5aa50 100644
--- a/tests/trace_state_propagation.rs
+++ b/tests/trace_state_propagation.rs
@@ -108,8 +108,7 @@ fn sampling_decision_respects_new_parent() {
     let provider = TracerProvider::builder()
         .with_simple_exporter(exporter.clone())
         .with_config(
-            Config::default()
-                .with_sampler(Sampler::ParentBased(Box::new(Sampler::AlwaysOff)))
+            Config::default().with_sampler(Sampler::ParentBased(Box::new(Sampler::AlwaysOff))),
         )
         .build();
     let tracer = provider.tracer("test");
@@ -134,7 +133,10 @@ fn sampling_decision_respects_new_parent() {
     // assert new parent-based sampling decision
     let spans = exporter.0.lock().unwrap();
     assert_eq!(spans.len(), 2, "Expected 2 spans, got {}", spans.len());
-    assert!(spans[0].span_context.is_sampled(), "Root span should be sampled");
+    assert!(
+        spans[0].span_context.is_sampled(),
+        "Root span should be sampled"
+    );
     assert_eq!(
         spans[1].span_context.is_sampled(),
         spans[0].span_context.is_sampled(),