Implement random trace flag from W3C TC v2

Author: anuraaga

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

@@ -8,10 +8,11 @@
 import io.opentelemetry.api.trace.SpanId;
 import io.opentelemetry.api.trace.TraceId;
 import io.opentelemetry.sdk.internal.RandomSupplier;
+import io.opentelemetry.sdk.trace.internal.ExtendedIdGenerator;
 import java.util.Random;
 import java.util.function.Supplier;
 
-enum RandomIdGenerator implements IdGenerator {
+enum RandomIdGenerator implements IdGenerator, ExtendedIdGenerator {
   INSTANCE;
 
   private static final long INVALID_ID = 0;
@@ -42,4 +43,9 @@ public String generateTraceId() {
   public String toString() {
     return "RandomIdGenerator{}";
   }
+
+  @Override
+  public boolean randomTraceId() {
+    return true;
+  }
 }

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

@@ -24,6 +24,7 @@
 import io.opentelemetry.sdk.internal.AttributeUtil;
 import io.opentelemetry.sdk.internal.AttributesMap;
 import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.internal.ExtendedIdGenerator;
 import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
 import java.util.ArrayList;
@@ -35,6 +36,9 @@
 /** {@link SdkSpanBuilder} is SDK implementation of {@link SpanBuilder}. */
 class SdkSpanBuilder implements SpanBuilder {
 
+  // TODO: Move to ImmutableTraceFlags when W3C Trace Context Level 2 is finalized.
+  private static final byte TRACE_FLAGS_RANDOM_BIT = 0x10;
+
   private final String spanName;
   private final InstrumentationScopeInfo instrumentationScopeInfo;
   private final TracerSharedState tracerSharedState;
@@ -193,13 +197,22 @@ public Span startSpan() {
                 parentContext, traceId, spanName, spanKind, immutableAttributes, immutableLinks);
     SamplingDecision samplingDecision = samplingResult.getDecision();
 
+    TraceFlags traceFlags =
+        isSampled(samplingDecision) ? TraceFlags.getSampled() : TraceFlags.getDefault();
+    if (idGenerator instanceof ExtendedIdGenerator) {
+      boolean randomTraceId = ((ExtendedIdGenerator) idGenerator).randomTraceId();
+      if (randomTraceId) {
+        traceFlags = TraceFlags.fromByte((byte) (traceFlags.asByte() | TRACE_FLAGS_RANDOM_BIT));
+      }
+    }
+
     TraceState samplingResultTraceState =
         samplingResult.getUpdatedTraceState(parentSpanContext.getTraceState());
     SpanContext spanContext =
         ImmutableSpanContext.create(
             traceId,
             spanId,
-            isSampled(samplingDecision) ? TraceFlags.getSampled() : TraceFlags.getDefault(),
+            traceFlags,
             samplingResultTraceState,
             /* remote= */ false,
             tracerSharedState.isIdGeneratorSafeToSkipIdValidation());

sdk/trace/src/main/java/io/opentelemetry/sdk/trace/internal/ExtendedIdGenerator.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.trace.internal;
+
+import io.opentelemetry.sdk.trace.IdGenerator;
+
+/**
+ * An extension to {@link IdGenerator} to allow opting in to the random flag in the draft <a
+ * href="https://www.w3.org/TR/trace-context-2/#trace-flags">W3C Trace Context Level 2</a>
+ *
+ * <p>This class is internal and experimental. Its APIs are unstable and can change at any time. Its
+ * APIs (or a version of them) may be promoted to the public stable API in the future, but no
+ * guarantees are made.
+ */
+public interface ExtendedIdGenerator extends IdGenerator {
+  /**
+   * Returns {@code true} if the {@link IdGenerator} returns trace IDs with the right-most 7 bytes
+   * being random.
+   */
+  boolean randomTraceId();
+}

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

@@ -33,6 +33,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.sdk.trace.data.LinkData;
 import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.internal.ExtendedIdGenerator;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
@@ -965,6 +966,61 @@ void isRecording() {
     assertThat(SdkSpanBuilder.isRecording(SamplingDecision.RECORD_AND_SAMPLE)).isTrue();
   }
 
+  @Test
+  void traceFlags() {
+    // Default ID generator and sampler populate both bits.
+    assertThat(
+            sdkTracer.spanBuilder(SPAN_NAME).startSpan().getSpanContext().getTraceFlags().asByte())
+        .isEqualTo((byte) 0x11);
+
+    SdkTracerProvider tracerSdkFactory =
+        SdkTracerProvider.builder()
+            .setIdGenerator(
+                new IdGenerator() {
+                  @Override
+                  public String generateSpanId() {
+                    return "111";
+                  }
+
+                  @Override
+                  public String generateTraceId() {
+                    return "222";
+                  }
+                })
+            .addSpanProcessor(mockedSpanProcessor)
+            .build();
+    SdkTracer sdkTracer = (SdkTracer) tracerSdkFactory.get("SpanBuilderSdkTest");
+    assertThat(
+            sdkTracer.spanBuilder(SPAN_NAME).startSpan().getSpanContext().getTraceFlags().asByte())
+        .isEqualTo((byte) 0x01);
+
+    tracerSdkFactory =
+        SdkTracerProvider.builder()
+            .setIdGenerator(
+                new ExtendedIdGenerator() {
+                  @Override
+                  public boolean randomTraceId() {
+                    return false;
+                  }
+
+                  @Override
+                  public String generateSpanId() {
+                    return "111";
+                  }
+
+                  @Override
+                  public String generateTraceId() {
+                    return "222";
+                  }
+                })
+            .addSpanProcessor(mockedSpanProcessor)
+            .build();
+    sdkTracer = (SdkTracer) tracerSdkFactory.get("SpanBuilderSdkTest");
+    assertThat(
+            sdkTracer.spanBuilder(SPAN_NAME).startSpan().getSpanContext().getTraceFlags().asByte())
+        .isEqualTo((byte) 0x01);
+  }
+
   // SpanData is very commonly used in unit tests, we want the toString to make sure it's relatively
   // easy to understand failure messages.
   // TODO(anuraaga): Currently it isn't - we even return the same (or maybe incorrect?) stuff twice.
@@ -983,7 +1039,7 @@ void spanDataToString() {
             "SpanData\\{spanContext=ImmutableSpanContext\\{"
                 + "traceId=[0-9a-f]{32}, "
                 + "spanId=[0-9a-f]{16}, "
-                + "traceFlags=01, "
+                + "traceFlags=11, "
                 + "traceState=ArrayBasedTraceState\\{entries=\\[]}, remote=false, valid=true}, "
                 + "parentSpanContext=ImmutableSpanContext\\{"
                 + "traceId=00000000000000000000000000000000, "