fixed #349

consistent-sampling/src/main/java/io/opentelemetry/contrib/samplers/ConsistentSampler.java

@@ -231,6 +231,11 @@ public final SamplingResult shouldSample(
     SamplingDecision samplingDecision =
         isSampled ? SamplingDecision.RECORD_AND_SAMPLE : SamplingDecision.DROP;
 
+    // invalidate p-value if not sampled
+    if (!isSampled) {
+      otelTraceState.invalidateP();
+    }
+
     String newOtTraceState = otelTraceState.serialize();
 
     return new SamplingResult() {

consistent-sampling/src/main/java/io/opentelemetry/contrib/samplers/RandomGenerator.java

@@ -82,7 +82,7 @@ private int numberOfLeadingZerosOfRandomLong(LongSupplier threadSafeRandomLongSu
     }
   }
 
-  private static final ThreadLocal<ThreadLocalData> THREAD_LOCAL_DATA =
+  private final ThreadLocal<ThreadLocalData> threadLocalData =
       ThreadLocal.withInitial(ThreadLocalData::new);
 
   private static final RandomGenerator INSTANCE =
@@ -120,7 +120,7 @@ public static RandomGenerator getDefault() {
    * @return a random {@code boolean}
    */
   public boolean nextBoolean(double probability) {
-    return THREAD_LOCAL_DATA.get().generateRandomBoolean(threadSafeRandomLongSupplier, probability);
+    return threadLocalData.get().generateRandomBoolean(threadSafeRandomLongSupplier, probability);
   }
 
   /**
@@ -129,6 +129,6 @@ public boolean nextBoolean(double probability) {
    * @return the number of leading zeros
    */
   public int numberOfLeadingZerosOfRandomLong() {
-    return THREAD_LOCAL_DATA.get().numberOfLeadingZerosOfRandomLong(threadSafeRandomLongSupplier);
+    return threadLocalData.get().numberOfLeadingZerosOfRandomLong(threadSafeRandomLongSupplier);
   }
 }

consistent-sampling/src/test/java/io/opentelemetry/contrib/samplers/ConsistentSamplerTest.java

@@ -5,10 +5,26 @@
 
 package io.opentelemetry.contrib.samplers;
 
+import static io.opentelemetry.contrib.samplers.OtelTraceState.getInvalidP;
+import static io.opentelemetry.contrib.samplers.OtelTraceState.getInvalidR;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.Collections;
+import java.util.List;
+import java.util.OptionalInt;
 import java.util.SplittableRandom;
 import org.junit.jupiter.api.Test;
 
@@ -71,4 +87,172 @@ void testRandomValues() {
           .isLessThanOrEqualTo(samplingProbability);
     }
   }
+
+  private static ConsistentSampler createConsistentSampler(int p, int r) {
+    long randomLong = ~(0xFFFFFFFFFFFFFFFFL << r);
+
+    return new ConsistentSampler(RandomGenerator.create(() -> randomLong)) {
+      @Override
+      public String getDescription() {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      protected int getP(int parentP, boolean isRoot) {
+        return p;
+      }
+    };
+  }
+
+  private static TraceState createTraceState(int p, int r) {
+    OtelTraceState state = OtelTraceState.parse("");
+    state.setP(p);
+    state.setR(r);
+    return TraceState.builder().put(OtelTraceState.TRACE_STATE_KEY, state.serialize()).build();
+  }
+
+  private static Context createParentContext(
+      String traceId, String spanId, int p, int r, boolean sampled) {
+    TraceState parentTraceState = createTraceState(p, r);
+    TraceFlags traceFlags = sampled ? TraceFlags.getSampled() : TraceFlags.getDefault();
+    SpanContext parentSpanContext =
+        SpanContext.create(traceId, spanId, traceFlags, parentTraceState);
+    Span parentSpan = Span.wrap(parentSpanContext);
+    return parentSpan.storeInContext(Context.root());
+  }
+
+  private static boolean getSampledFlag(SamplingResult samplingResult) {
+    return SamplingDecision.RECORD_AND_SAMPLE.equals(samplingResult.getDecision());
+  }
+
+  private static OptionalInt getP(SamplingResult samplingResult, Context parentContext) {
+    Span parentSpan = Span.fromContext(parentContext);
+    OtelTraceState otelTraceState =
+        OtelTraceState.parse(
+            samplingResult
+                .getUpdatedTraceState(parentSpan.getSpanContext().getTraceState())
+                .get(OtelTraceState.TRACE_STATE_KEY));
+    return otelTraceState.hasValidP() ? OptionalInt.of(otelTraceState.getP()) : OptionalInt.empty();
+  }
+
+  private static OptionalInt getR(SamplingResult samplingResult, Context parentContext) {
+    Span parentSpan = Span.fromContext(parentContext);
+    OtelTraceState otelTraceState =
+        OtelTraceState.parse(
+            samplingResult
+                .getUpdatedTraceState(parentSpan.getSpanContext().getTraceState())
+                .get(OtelTraceState.TRACE_STATE_KEY));
+    return otelTraceState.hasValidR() ? OptionalInt.of(otelTraceState.getR()) : OptionalInt.empty();
+  }
+
+  private static void assertConsistentSampling(
+      int parentP,
+      int parentR,
+      boolean parentSampled,
+      int samplerP,
+      int generatedR,
+      int expectedP,
+      int expectedR,
+      boolean expectSampled) {
+
+    String traceId = "0123456789abcdef0123456789abcdef";
+    String spanId = "0123456789abcdef";
+    String name = "name";
+    SpanKind spanKind = SpanKind.SERVER;
+    Attributes attributes = Attributes.empty();
+    List<LinkData> parentLinks = Collections.emptyList();
+
+    Context parentContext = createParentContext(traceId, spanId, parentP, parentR, parentSampled);
+    ConsistentSampler sampler = createConsistentSampler(samplerP, generatedR);
+    SamplingResult samplingResult =
+        sampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+
+    assertEquals(expectSampled, getSampledFlag(samplingResult));
+    OptionalInt p = getP(samplingResult, parentContext);
+    if (OtelTraceState.isValidP(expectedP)) {
+      assertEquals(expectedP, p.getAsInt());
+    } else {
+      assertFalse(p.isPresent());
+    }
+    OptionalInt r = getR(samplingResult, parentContext);
+    if (OtelTraceState.isValidR(expectedR)) {
+      assertEquals(expectedR, r.getAsInt());
+    } else {
+      assertFalse(r.isPresent());
+    }
+  }
+
+  private static final boolean NOT_SAMPLED = false;
+  private static final boolean SAMPLED = true;
+
+  @Test
+  void testUndefinedParentTraceState() {
+    assertConsistentSampling(getInvalidP(), getInvalidR(), NOT_SAMPLED, 0, 0, 0, 0, SAMPLED);
+    assertConsistentSampling(getInvalidP(), getInvalidR(), NOT_SAMPLED, 2, 3, 2, 3, SAMPLED);
+    assertConsistentSampling(
+        getInvalidP(), getInvalidR(), NOT_SAMPLED, 3, 2, getInvalidP(), 2, NOT_SAMPLED);
+    assertConsistentSampling(getInvalidP(), getInvalidR(), NOT_SAMPLED, 0, 1, 0, 1, SAMPLED);
+    assertConsistentSampling(getInvalidP(), getInvalidR(), NOT_SAMPLED, 0, 2, 0, 2, SAMPLED);
+    assertConsistentSampling(
+        getInvalidP(), getInvalidR(), NOT_SAMPLED, 1, 0, getInvalidP(), 0, NOT_SAMPLED);
+    assertConsistentSampling(
+        getInvalidP(), getInvalidR(), NOT_SAMPLED, 2, 0, getInvalidP(), 0, NOT_SAMPLED);
+    assertConsistentSampling(getInvalidP(), getInvalidR(), SAMPLED, 0, 0, 0, 0, SAMPLED);
+    assertConsistentSampling(getInvalidP(), getInvalidR(), SAMPLED, 2, 3, 2, 3, SAMPLED);
+    assertConsistentSampling(
+        getInvalidP(), getInvalidR(), SAMPLED, 3, 2, getInvalidP(), 2, NOT_SAMPLED);
+  }
+
+  @Test
+  void testParentTraceStateWithDefinedPOnly() {
+    assertConsistentSampling(6, getInvalidR(), NOT_SAMPLED, 0, 0, 0, 0, SAMPLED);
+    assertConsistentSampling(7, getInvalidR(), NOT_SAMPLED, 2, 3, 2, 3, SAMPLED);
+    assertConsistentSampling(4, getInvalidR(), NOT_SAMPLED, 3, 2, getInvalidP(), 2, NOT_SAMPLED);
+    assertConsistentSampling(3, getInvalidR(), NOT_SAMPLED, 0, 1, 0, 1, SAMPLED);
+    assertConsistentSampling(2, getInvalidR(), NOT_SAMPLED, 0, 2, 0, 2, SAMPLED);
+    assertConsistentSampling(6, getInvalidR(), NOT_SAMPLED, 1, 0, getInvalidP(), 0, NOT_SAMPLED);
+    assertConsistentSampling(7, getInvalidR(), NOT_SAMPLED, 2, 0, getInvalidP(), 0, NOT_SAMPLED);
+    assertConsistentSampling(5, getInvalidR(), NOT_SAMPLED, 8, 7, getInvalidP(), 7, NOT_SAMPLED);
+    assertConsistentSampling(5, getInvalidR(), NOT_SAMPLED, 6, 7, 6, 7, SAMPLED);
+    assertConsistentSampling(12, getInvalidR(), SAMPLED, 0, 0, 0, 0, SAMPLED);
+    assertConsistentSampling(15, getInvalidR(), SAMPLED, 2, 3, 2, 3, SAMPLED);
+    assertConsistentSampling(18, getInvalidR(), SAMPLED, 3, 2, getInvalidP(), 2, NOT_SAMPLED);
+  }
+
+  @Test
+  void testParentTraceStateWithDefinedROnly() {
+    assertConsistentSampling(getInvalidP(), 0, NOT_SAMPLED, 0, 5, 0, 0, SAMPLED);
+    assertConsistentSampling(getInvalidP(), 3, NOT_SAMPLED, 2, 0, 2, 3, SAMPLED);
+    assertConsistentSampling(getInvalidP(), 2, NOT_SAMPLED, 3, 1, getInvalidP(), 2, NOT_SAMPLED);
+    assertConsistentSampling(getInvalidP(), 1, NOT_SAMPLED, 0, 0, 0, 1, SAMPLED);
+    assertConsistentSampling(getInvalidP(), 2, NOT_SAMPLED, 0, 5, 0, 2, SAMPLED);
+    assertConsistentSampling(getInvalidP(), 0, NOT_SAMPLED, 1, 8, getInvalidP(), 0, NOT_SAMPLED);
+    assertConsistentSampling(getInvalidP(), 0, NOT_SAMPLED, 2, 5, getInvalidP(), 0, NOT_SAMPLED);
+    assertConsistentSampling(getInvalidP(), 0, SAMPLED, 0, 11, 0, 0, SAMPLED);
+    assertConsistentSampling(getInvalidP(), 3, SAMPLED, 2, 9, 2, 3, SAMPLED);
+    assertConsistentSampling(getInvalidP(), 2, SAMPLED, 3, 1, getInvalidP(), 2, NOT_SAMPLED);
+  }
+
+  @Test
+  void testConsistentParentTraceState() {
+    // ( (r >= p) <=> sampled) is satisfied
+    assertConsistentSampling(3, 5, SAMPLED, 6, 7, getInvalidP(), 5, NOT_SAMPLED);
+    assertConsistentSampling(3, 5, SAMPLED, 2, 7, 2, 5, SAMPLED);
+    assertConsistentSampling(5, 3, NOT_SAMPLED, 6, 7, getInvalidP(), 3, NOT_SAMPLED);
+  }
+
+  @Test
+  void testInconsistentParentTraceState() {
+    // ( (r >= p) <=> sampled) is not satisfied
+    assertConsistentSampling(5, 3, SAMPLED, 6, 7, getInvalidP(), 3, NOT_SAMPLED);
+    assertConsistentSampling(3, 5, NOT_SAMPLED, 6, 7, getInvalidP(), 5, NOT_SAMPLED);
+    assertConsistentSampling(5, 3, SAMPLED, 1, 7, 1, 3, SAMPLED);
+    assertConsistentSampling(3, 5, NOT_SAMPLED, 2, 7, 2, 5, SAMPLED);
+  }
+
+  @Test
+  void testInvalidSamplerP() {
+    assertConsistentSampling(3, 5, SAMPLED, getInvalidP(), 7, getInvalidP(), 5, SAMPLED);
+    assertConsistentSampling(5, 3, NOT_SAMPLED, getInvalidP(), 7, getInvalidP(), 3, NOT_SAMPLED);
+  }
 }