retain propagated context for random trace-ids (#8263)

Author: cleverchuk

sdk/trace/src/main/java/io/opentelemetry/sdk/trace/SdkSpanBuilder.java

@@ -37,9 +37,13 @@
 /** {@link SdkSpanBuilder} is SDK implementation of {@link SpanBuilder}. */
 class SdkSpanBuilder implements SpanBuilder {
 
-  private static final Context ROOT_CONTEXT_WITH_RANDOM_TRACE_ID_BIT =
-      preparePrimordialContext(
-          TraceFlags.builder().setRandomTraceId(true).build(), TraceState.getDefault());
+  private static final Span RANDOM_TRACE_ID_PRIMORDIAL_SPAN =
+      Span.wrap(
+          SpanContext.create(
+              TraceId.getInvalid(),
+              SpanId.getInvalid(),
+              TraceFlags.builder().setRandomTraceId(true).build(),
+              TraceState.getDefault()));
 
   private final String spanName;
   private final InstrumentationScopeInfo instrumentationScopeInfo;
@@ -64,17 +68,6 @@ class SdkSpanBuilder implements SpanBuilder {
     this.spanLimits = spanLimits;
   }
 
-  /*
-   * A primordial context can be passed as the parent context for a root span
-   * if a non-default TraceFlags or TraceState need to be passed to the sampler
-   */
-  private static Context preparePrimordialContext(TraceFlags traceFlags, TraceState traceState) {
-    SpanContext spanContext =
-        SpanContext.create(TraceId.getInvalid(), SpanId.getInvalid(), traceFlags, traceState);
-    Span span = Span.wrap(spanContext);
-    return span.storeInContext(Context.root());
-  }
-
   @Override
   public SpanBuilder setParent(Context context) {
     if (context == null) {
@@ -198,7 +191,7 @@ public Span startSpan() {
       if (idGenerator.generatesRandomTraceIds()) {
         isTraceIdRandom = true;
         // Replace parentContext for sampling with one with RANDOM_TRACE_ID bit set
-        parentContextForSampler = ROOT_CONTEXT_WITH_RANDOM_TRACE_ID_BIT;
+        parentContextForSampler = parentContext.with(RANDOM_TRACE_ID_PRIMORDIAL_SPAN);
       } else {
         isTraceIdRandom = false;
       }

sdk/trace/src/test/java/io/opentelemetry/sdk/trace/SdkSpanBuilderTest.java

@@ -32,6 +32,7 @@
 import io.opentelemetry.api.trace.TraceState;
 import io.opentelemetry.api.trace.TracerProvider;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.sdk.trace.data.LinkData;
 import io.opentelemetry.sdk.trace.data.SpanData;
@@ -47,6 +48,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
 import org.mockito.Mock;
 import org.mockito.Mockito;
@@ -940,6 +942,43 @@ void propagateRandomTraceIdFlag() {
     }
   }
 
+  @Test
+  void samplerReceivesPropagatorContextWithRandomTraceId() {
+    Sampler mockSampler = Mockito.mock(Sampler.class);
+    Mockito.when(
+            mockSampler.shouldSample(
+                ArgumentMatchers.any(),
+                ArgumentMatchers.anyString(),
+                ArgumentMatchers.anyString(),
+                ArgumentMatchers.any(),
+                ArgumentMatchers.any(),
+                ArgumentMatchers.anyList()))
+        .thenReturn(SamplingResult.recordAndSample());
+
+    SdkTracerProvider provider = SdkTracerProvider.builder().setSampler(mockSampler).build();
+    ContextKey<String> propagatorKey = ContextKey.named("propagator-test-key");
+    Context parentWithPropagatorData = Context.root().with(propagatorKey, "test-value");
+
+    Span span =
+        provider.get("test").spanBuilder(SPAN_NAME).setParent(parentWithPropagatorData).startSpan();
+    span.end();
+
+    ArgumentCaptor<Context> contextCaptor = ArgumentCaptor.forClass(Context.class);
+    Mockito.verify(mockSampler)
+        .shouldSample(
+            contextCaptor.capture(),
+            ArgumentMatchers.anyString(),
+            ArgumentMatchers.anyString(),
+            ArgumentMatchers.any(),
+            ArgumentMatchers.any(),
+            ArgumentMatchers.anyList());
+
+    Context samplerContext = contextCaptor.getValue();
+    assertThat(samplerContext.get(propagatorKey)).isEqualTo("test-value");
+    assertThat(Span.fromContext(samplerContext).getSpanContext().getTraceFlags().isTraceIdRandom())
+        .isTrue();
+  }
+
   @Test
   void startTimestamp_numeric() {
     SdkSpan span =