Support traceId-based r-values (#417)

* Support traceId-based r-values

* Interface

* Fix javadoc

* Spotless

* More

* Clean

* Review

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

@@ -10,12 +10,8 @@ import javax.annotation.concurrent.Immutable;
 @Immutable
 final class ConsistentAlwaysOffSampler extends ConsistentSampler {
 
-  private ConsistentAlwaysOffSampler() {}
-
-  private static final ConsistentSampler INSTANCE = new ConsistentAlwaysOffSampler();
-
-  static ConsistentSampler getInstance() {
-    return INSTANCE;
+  ConsistentAlwaysOffSampler(RValueGenerator rValueGenerator) {
+    super(rValueGenerator);
   }
 
   @Override

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

@@ -10,12 +10,8 @@ import javax.annotation.concurrent.Immutable;
 @Immutable
 final class ConsistentAlwaysOnSampler extends ConsistentSampler {
 
-  private ConsistentAlwaysOnSampler() {}
-
-  private static final ConsistentSampler INSTANCE = new ConsistentAlwaysOnSampler();
-
-  static ConsistentSampler getInstance() {
-    return INSTANCE;
+  ConsistentAlwaysOnSampler(RValueGenerator rValueGenerator) {
+    super(rValueGenerator);
   }
 
   @Override

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

@@ -21,7 +21,9 @@ final class ConsistentComposedAndSampler extends ConsistentSampler {
   private final ConsistentSampler sampler2;
   private final String description;
 
-  ConsistentComposedAndSampler(ConsistentSampler sampler1, ConsistentSampler sampler2) {
+  ConsistentComposedAndSampler(
+      ConsistentSampler sampler1, ConsistentSampler sampler2, RValueGenerator rValueGenerator) {
+    super(rValueGenerator);
     this.sampler1 = requireNonNull(sampler1);
     this.sampler2 = requireNonNull(sampler2);
     this.description =

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

@@ -21,7 +21,9 @@ final class ConsistentComposedOrSampler extends ConsistentSampler {
   private final ConsistentSampler sampler2;
   private final String description;
 
-  ConsistentComposedOrSampler(ConsistentSampler sampler1, ConsistentSampler sampler2) {
+  ConsistentComposedOrSampler(
+      ConsistentSampler sampler1, ConsistentSampler sampler2, RValueGenerator rValueGenerator) {
+    super(rValueGenerator);
     this.sampler1 = requireNonNull(sampler1);
     this.sampler2 = requireNonNull(sampler2);
     this.description =

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

@@ -21,25 +21,15 @@ final class ConsistentParentBasedSampler extends ConsistentSampler {
 
   private final String description;
 
-  /**
-   * Constructs a new consistent parent based sampler using the given root sampler.
-   *
-   * @param rootSampler the root sampler
-   */
-  ConsistentParentBasedSampler(ConsistentSampler rootSampler) {
-    this(rootSampler, RandomGenerator.getDefault());
-  }
-
   /**
    * Constructs a new consistent parent based sampler using the given root sampler and the given
    * thread-safe random generator.
    *
    * @param rootSampler the root sampler
-   * @param threadSafeRandomGenerator a thread-safe random generator
+   * @param rValueGenerator the function to use for generating the r-value
    */
-  ConsistentParentBasedSampler(
-      ConsistentSampler rootSampler, RandomGenerator threadSafeRandomGenerator) {
-    super(threadSafeRandomGenerator);
+  ConsistentParentBasedSampler(ConsistentSampler rootSampler, RValueGenerator rValueGenerator) {
+    super(rValueGenerator);
     this.rootSampler = requireNonNull(rootSampler);
     this.description =
         "ConsistentParentBasedSampler{rootSampler=" + rootSampler.getDescription() + '}';

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

@@ -15,29 +15,25 @@ final class ConsistentProbabilityBasedSampler extends ConsistentSampler {
   private final int upperPValue;
   private final double probabilityToUseLowerPValue;
   private final String description;
+  private final RandomGenerator randomGenerator;
 
   /**
    * Constructor.
    *
    * @param samplingProbability the sampling probability
+   * @param rValueGenerator the function to use for generating the r-value
    */
-  ConsistentProbabilityBasedSampler(double samplingProbability) {
-    this(samplingProbability, RandomGenerator.getDefault());
-  }
-
-  /**
-   * Constructor.
-   *
-   * @param samplingProbability the sampling probability
-   * @param randomGenerator a random generator
-   */
-  ConsistentProbabilityBasedSampler(double samplingProbability, RandomGenerator randomGenerator) {
-    super(randomGenerator);
+  ConsistentProbabilityBasedSampler(
+      double samplingProbability,
+      RValueGenerator rValueGenerator,
+      RandomGenerator randomGenerator) {
+    super(rValueGenerator);
     if (samplingProbability < 0.0 || samplingProbability > 1.0) {
       throw new IllegalArgumentException("Sampling probability must be in range [0.0, 1.0]!");
     }
     this.description =
         String.format("ConsistentProbabilityBasedSampler{%.6f}", samplingProbability);
+    this.randomGenerator = randomGenerator;
 
     lowerPValue = getLowerBoundP(samplingProbability);
     upperPValue = getUpperBoundP(samplingProbability);

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

@@ -85,6 +85,7 @@ final class ConsistentRateLimitingSampler extends ConsistentSampler {
   private final double inverseAdaptationTimeNanos;
   private final double targetSpansPerNanosecondLimit;
   private final AtomicReference<State> state;
+  private final RandomGenerator randomGenerator;
 
   /**
    * Constructor.
@@ -92,30 +93,17 @@ final class ConsistentRateLimitingSampler extends ConsistentSampler {
    * @param targetSpansPerSecondLimit the desired spans per second limit
    * @param adaptationTimeSeconds the typical time to adapt to a new load (time constant used for
    *     exponential smoothing)
-   */
-  ConsistentRateLimitingSampler(double targetSpansPerSecondLimit, double adaptationTimeSeconds) {
-    this(
-        targetSpansPerSecondLimit,
-        adaptationTimeSeconds,
-        RandomGenerator.getDefault(),
-        System::nanoTime);
-  }
-
-  /**
-   * Constructor.
-   *
-   * @param targetSpansPerSecondLimit the desired spans per second limit
-   * @param adaptationTimeSeconds the typical time to adapt to a new load (time constant used for
-   *     exponential smoothing)
+   * @param rValueGenerator the function to use for generating the r-value
    * @param randomGenerator a random generator
    * @param nanoTimeSupplier a supplier for the current nano time
    */
   ConsistentRateLimitingSampler(
       double targetSpansPerSecondLimit,
       double adaptationTimeSeconds,
+      RValueGenerator rValueGenerator,
       RandomGenerator randomGenerator,
       LongSupplier nanoTimeSupplier) {
-    super(randomGenerator);
+    super(rValueGenerator);
 
     if (targetSpansPerSecondLimit < 0.0) {
       throw new IllegalArgumentException("Limit for sampled spans per second must be nonnegative!");
@@ -133,6 +121,8 @@ final class ConsistentRateLimitingSampler extends ConsistentSampler {
     this.targetSpansPerNanosecondLimit = 1e-9 * targetSpansPerSecondLimit;
 
     this.state = new AtomicReference<>(new State(0, 0, nanoTimeSupplier.getAsLong()));
+
+    this.randomGenerator = randomGenerator;
   }
 
   private State updateState(State oldState, long currentNanoTime) {

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

@@ -28,8 +28,18 @@ public abstract class ConsistentSampler implements Sampler {
    *
    * @return a sampler
    */
-  public static final ConsistentSampler alwaysOn() {
-    return ConsistentAlwaysOnSampler.getInstance();
+  public static ConsistentSampler alwaysOn() {
+    return alwaysOn(RValueGenerators.getDefault());
+  }
+
+  /**
+   * Returns a {@link ConsistentSampler} that samples all spans.
+   *
+   * @param rValueGenerator the function to use for generating the r-value
+   * @return a sampler
+   */
+  public static ConsistentSampler alwaysOn(RValueGenerator rValueGenerator) {
+    return new ConsistentAlwaysOnSampler(rValueGenerator);
   }
 
   /**
@@ -37,8 +47,18 @@ public abstract class ConsistentSampler implements Sampler {
    *
    * @return a sampler
    */
-  public static final ConsistentSampler alwaysOff() {
-    return ConsistentAlwaysOffSampler.getInstance();
+  public static ConsistentSampler alwaysOff() {
+    return alwaysOff(RValueGenerators.getDefault());
+  }
+
+  /**
+   * Returns a {@link ConsistentSampler} that does not sample any span.
+   *
+   * @param rValueGenerator the function to use for generating the r-value
+   * @return a sampler
+   */
+  public static ConsistentSampler alwaysOff(RValueGenerator rValueGenerator) {
+    return new ConsistentAlwaysOffSampler(rValueGenerator);
   }
 
   /**
@@ -47,20 +67,21 @@ public abstract class ConsistentSampler implements Sampler {
    * @param samplingProbability the sampling probability
    * @return a sampler
    */
-  public static final ConsistentSampler probabilityBased(double samplingProbability) {
-    return new ConsistentProbabilityBasedSampler(samplingProbability);
+  public static ConsistentSampler probabilityBased(double samplingProbability) {
+    return probabilityBased(samplingProbability, RValueGenerators.getDefault());
   }
 
   /**
    * Returns a {@link ConsistentSampler} that samples each span with a fixed probability.
    *
    * @param samplingProbability the sampling probability
-   * @param randomGenerator a random generator
+   * @param rValueGenerator the function to use for generating the r-value
    * @return a sampler
    */
-  static final ConsistentSampler probabilityBased(
-      double samplingProbability, RandomGenerator randomGenerator) {
-    return new ConsistentProbabilityBasedSampler(samplingProbability, randomGenerator);
+  public static ConsistentSampler probabilityBased(
+      double samplingProbability, RValueGenerator rValueGenerator) {
+    return new ConsistentProbabilityBasedSampler(
+        samplingProbability, rValueGenerator, RandomGenerator.getDefault());
   }
 
   /**
@@ -69,8 +90,8 @@ public abstract class ConsistentSampler implements Sampler {
    *
    * @param rootSampler the root sampler
    */
-  public static final ConsistentSampler parentBased(ConsistentSampler rootSampler) {
-    return new ConsistentParentBasedSampler(rootSampler);
+  public static ConsistentSampler parentBased(ConsistentSampler rootSampler) {
+    return parentBased(rootSampler, RValueGenerators.getDefault());
   }
 
   /**
@@ -78,11 +99,11 @@ public abstract class ConsistentSampler implements Sampler {
    * or falls-back to the given sampler if it is a root span.
    *
    * @param rootSampler the root sampler
-   * @param randomGenerator a random generator
+   * @param rValueGenerator the function to use for generating the r-value
    */
-  static final ConsistentSampler parentBased(
-      ConsistentSampler rootSampler, RandomGenerator randomGenerator) {
-    return new ConsistentParentBasedSampler(rootSampler, randomGenerator);
+  public static ConsistentSampler parentBased(
+      ConsistentSampler rootSampler, RValueGenerator rValueGenerator) {
+    return new ConsistentParentBasedSampler(rootSampler, rValueGenerator);
   }
 
   /**
@@ -93,9 +114,10 @@ public abstract class ConsistentSampler implements Sampler {
    * @param adaptationTimeSeconds the typical time to adapt to a new load (time constant used for
    *     exponential smoothing)
    */
-  public static final ConsistentSampler rateLimited(
+  public static ConsistentSampler rateLimited(
       double targetSpansPerSecondLimit, double adaptationTimeSeconds) {
-    return new ConsistentRateLimitingSampler(targetSpansPerSecondLimit, adaptationTimeSeconds);
+    return rateLimited(
+        targetSpansPerSecondLimit, adaptationTimeSeconds, RValueGenerators.getDefault());
   }
 
   /**
@@ -105,16 +127,37 @@ public abstract class ConsistentSampler implements Sampler {
    * @param targetSpansPerSecondLimit the desired spans per second limit
    * @param adaptationTimeSeconds the typical time to adapt to a new load (time constant used for
    *     exponential smoothing)
-   * @param randomGenerator a random generator
-   * @param nanoTimeSupplier a supplier for the current nano time
+   * @param rValueGenerator the function to use for generating the r-value
    */
-  static final ConsistentSampler rateLimited(
+  public static ConsistentSampler rateLimited(
       double targetSpansPerSecondLimit,
       double adaptationTimeSeconds,
-      RandomGenerator randomGenerator,
+      RValueGenerator rValueGenerator) {
+    return rateLimited(
+        targetSpansPerSecondLimit, adaptationTimeSeconds, rValueGenerator, System::nanoTime);
+  }
+
+  /**
+   * Returns a new {@link ConsistentSampler} that attempts to adjust the sampling probability
+   * dynamically to meet the target span rate.
+   *
+   * @param targetSpansPerSecondLimit the desired spans per second limit
+   * @param adaptationTimeSeconds the typical time to adapt to a new load (time constant used for
+   *     exponential smoothing)
+   * @param rValueGenerator the function to use for generating the r-value
+   * @param nanoTimeSupplier a supplier for the current nano time
+   */
+  static ConsistentSampler rateLimited(
+      double targetSpansPerSecondLimit,
+      double adaptationTimeSeconds,
+      RValueGenerator rValueGenerator,
       LongSupplier nanoTimeSupplier) {
     return new ConsistentRateLimitingSampler(
-        targetSpansPerSecondLimit, adaptationTimeSeconds, randomGenerator, nanoTimeSupplier);
+        targetSpansPerSecondLimit,
+        adaptationTimeSeconds,
+        rValueGenerator,
+        RandomGenerator.getDefault(),
+        nanoTimeSupplier);
   }
 
   /**
@@ -136,7 +179,8 @@ public abstract class ConsistentSampler implements Sampler {
     if (otherConsistentSampler == this) {
       return this;
     }
-    return new ConsistentComposedAndSampler(this, otherConsistentSampler);
+    return new ConsistentComposedAndSampler(
+        this, otherConsistentSampler, RValueGenerators.getDefault());
   }
 
   /**
@@ -158,20 +202,17 @@ public abstract class ConsistentSampler implements Sampler {
     if (otherConsistentSampler == this) {
       return this;
     }
-    return new ConsistentComposedOrSampler(this, otherConsistentSampler);
+    return new ConsistentComposedOrSampler(
+        this, otherConsistentSampler, RValueGenerators.getDefault());
   }
 
-  protected final RandomGenerator randomGenerator;
+  private final RValueGenerator rValueGenerator;
 
-  protected ConsistentSampler(RandomGenerator randomGenerator) {
-    this.randomGenerator = requireNonNull(randomGenerator);
+  protected ConsistentSampler(RValueGenerator rValueGenerator) {
+    this.rValueGenerator = requireNonNull(rValueGenerator);
   }
 
-  protected ConsistentSampler() {
-    this(RandomGenerator.getDefault());
-  }
-
-  private static final boolean isInvariantViolated(
+  private static boolean isInvariantViolated(
       OtelTraceState otelTraceState, boolean isParentSampled) {
     if (otelTraceState.hasValidR() && otelTraceState.hasValidP()) {
       // if valid p- and r-values are given, they must be consistent with the isParentSampled flag
@@ -212,8 +253,7 @@ public abstract class ConsistentSampler implements Sampler {
 
     // generate new r-value if not available
     if (!otelTraceState.hasValidR()) {
-      otelTraceState.setR(
-          Math.min(randomGenerator.numberOfLeadingZerosOfRandomLong(), OtelTraceState.getMaxR()));
+      otelTraceState.setR(Math.min(rValueGenerator.generate(traceId), OtelTraceState.getMaxR()));
     }
 
     // determine and set new p-value that is used for the sampling decision

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.samplers;
+
+/**
+ * A function for generating r-values.
+ *
+ * <p>The distribution of r-values generated by this function must satisfy the table below.
+ *
+ * <table>
+ * <caption>Required distribution of r-values</caption>
+ * <thead>
+ *   <tr><th>r-value</th><th>Probability of r-value</th></tr>
+ * </thead>
+ * <tbody>
+ *   <tr><td>0</td><td>1/2</td></tr>
+ *   <tr><td>1</td><td>1/4</td></tr>
+ *   <tr><td>2</td><td>1/8</td></tr>
+ *   <tr><td>3</td><td>1/16</td></tr>
+ *   <tr><td>…</td><td>…</td></tr>
+ *   <tr><td>0 &lt;= r &lt;= 61</td><td>2**-(r+1)</td></tr>
+ *   <tr><td>…</td><td>…</td></tr>
+ *   <tr><td>59</td><td>2**-60</td></tr>
+ *   <tr><td>60</td><td>2**-61</td></tr>
+ *   <tr><td>61</td><td>2**-62</td></tr>
+ *   <tr><td>>=62</td><td>2**-62</td></tr>
+ * </tbody>
+ * </table>
+ *
+ * For more info see <a
+ * href="https://opentelemetry.io/docs/reference/specification/trace/tracestate-probability-sampling/#methods-for-generating-r-values">Methods
+ * for generating R-values</a>.
+ */
+@FunctionalInterface
+public interface RValueGenerator {
+
+  int generate(String traceId);
+}

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.samplers;
+
+final class RValueGenerators {
+
+  private static final RValueGenerator DEFAULT = createDefault();
+
+  static RValueGenerator getDefault() {
+    return DEFAULT;
+  }
+
+  private static RValueGenerator createDefault() {
+    RandomGenerator randomGenerator = RandomGenerator.getDefault();
+    return s -> randomGenerator.numberOfLeadingZerosOfRandomLong();
+  }
+
+  private RValueGenerators() {}
+}

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

@@ -50,7 +50,8 @@ public class ConsistentProbabilityBasedSamplerTest {
 
     Sampler sampler =
         ConsistentSampler.probabilityBased(
-            samplingProbability, RandomGenerator.create(rng::nextLong));
+            samplingProbability,
+            s -> RandomGenerator.create(rng::nextLong).numberOfLeadingZerosOfRandomLong());
 
     Map<Integer, Long> observedPvalues = new HashMap<>();
     for (long i = 0; i < numSpans; ++i) {

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

@@ -59,14 +59,10 @@ class ConsistentRateLimitingSamplerTest {
 
     double targetSpansPerSecondLimit = 1000;
     double adaptationTimeSeconds = 5;
-    SplittableRandom random = new SplittableRandom(0L);
 
     ConsistentSampler sampler =
         ConsistentSampler.rateLimited(
-            targetSpansPerSecondLimit,
-            adaptationTimeSeconds,
-            RandomGenerator.create(random::nextLong),
-            nanoTimeSupplier);
+            targetSpansPerSecondLimit, adaptationTimeSeconds, rValueGenerator(), nanoTimeSupplier);
 
     long nanosBetweenSpans = TimeUnit.MICROSECONDS.toNanos(100);
     int numSpans = 1000000;
@@ -96,14 +92,10 @@ class ConsistentRateLimitingSamplerTest {
 
     double targetSpansPerSecondLimit = 1000;
     double adaptationTimeSeconds = 5;
-    SplittableRandom random = new SplittableRandom(0L);
 
     ConsistentSampler sampler =
         ConsistentSampler.rateLimited(
-            targetSpansPerSecondLimit,
-            adaptationTimeSeconds,
-            RandomGenerator.create(random::nextLong),
-            nanoTimeSupplier);
+            targetSpansPerSecondLimit, adaptationTimeSeconds, rValueGenerator(), nanoTimeSupplier);
 
     long nanosBetweenSpans1 = TimeUnit.MICROSECONDS.toNanos(100);
     long nanosBetweenSpans2 = TimeUnit.MICROSECONDS.toNanos(10);
@@ -155,14 +147,10 @@ class ConsistentRateLimitingSamplerTest {
 
     double targetSpansPerSecondLimit = 1000;
     double adaptationTimeSeconds = 5;
-    SplittableRandom random = new SplittableRandom(0L);
 
     ConsistentSampler sampler =
         ConsistentSampler.rateLimited(
-            targetSpansPerSecondLimit,
-            adaptationTimeSeconds,
-            RandomGenerator.create(random::nextLong),
-            nanoTimeSupplier);
+            targetSpansPerSecondLimit, adaptationTimeSeconds, rValueGenerator(), nanoTimeSupplier);
 
     long nanosBetweenSpans1 = TimeUnit.MICROSECONDS.toNanos(10);
     long nanosBetweenSpans2 = TimeUnit.MICROSECONDS.toNanos(100);
@@ -208,4 +196,10 @@ class ConsistentRateLimitingSamplerTest {
     assertThat(numSampledSpansInLast5Seconds / 5.)
         .isCloseTo(targetSpansPerSecondLimit, Percentage.withPercentage(5));
   }
+
+  private static RValueGenerator rValueGenerator() {
+    SplittableRandom random = new SplittableRandom(0L);
+    RandomGenerator randomGenerator = RandomGenerator.create(random::nextLong);
+    return s -> randomGenerator.numberOfLeadingZerosOfRandomLong();
+  }
 }

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

@@ -526,11 +526,12 @@ class ConsistentReservoirSamplingSpanProcessorTest {
             DEFAULT_EXPORT_TIMEOUT_NANOS,
             RandomGenerator.create(asThreadSafeLongSupplier(rng1)));
 
+    RandomGenerator randomGenerator = RandomGenerator.create(asThreadSafeLongSupplier(rng2));
     SdkTracerProvider sdkTracerProvider =
         SdkTracerProvider.builder()
             .setSampler(
                 ConsistentSampler.probabilityBased(
-                    samplingProbability, RandomGenerator.create(asThreadSafeLongSupplier(rng2))))
+                    samplingProbability, s -> randomGenerator.numberOfLeadingZerosOfRandomLong()))
             .addSpanProcessor(processor)
             .build();
 

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

@@ -90,8 +90,9 @@ class ConsistentSamplerTest {
 
   private static ConsistentSampler createConsistentSampler(int p, int r) {
     long randomLong = ~(0xFFFFFFFFFFFFFFFFL << r);
+    RandomGenerator randomGenerator = RandomGenerator.create(() -> randomLong);
 
-    return new ConsistentSampler(RandomGenerator.create(() -> randomLong)) {
+    return new ConsistentSampler(s -> randomGenerator.numberOfLeadingZerosOfRandomLong()) {
       @Override
       public String getDescription() {
         throw new UnsupportedOperationException();