Support traceId-based r-values (#417)

* Support traceId-based r-values

* Interface

* Fix javadoc

* Spotless

* More

* Clean

* Review

Author: trask

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

@@ -10,12 +10,8 @@
 @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 @@
 @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,37 +85,25 @@ public State(double effectiveWindowCount, double effectiveWindowNanos, long last
   private final double inverseAdaptationTimeNanos;
   private final double targetSpansPerNanosecondLimit;
   private final AtomicReference<State> state;
+  private final RandomGenerator randomGenerator;
 
   /**
    * 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)
-   */
-  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 @@ public State(double effectiveWindowCount, double effectiveWindowNanos, long last
     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,17 +28,37 @@ 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);
+  }
+
+  /**
+   * Returns a {@link ConsistentSampler} that does not sample any span.
+   *
+   * @return a sampler
+   */
+  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 final ConsistentSampler alwaysOff() {
-    return ConsistentAlwaysOffSampler.getInstance();
+  public static ConsistentSampler alwaysOff(RValueGenerator rValueGenerator) {
+    return new ConsistentAlwaysOffSampler(rValueGenerator);
   }
 
   /**
@@ -47,20 +67,21 @@ public static final ConsistentSampler alwaysOff() {
    * @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,20 +90,20 @@ static final ConsistentSampler probabilityBased(
    *
    * @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());
   }
 
   /**
    * Returns a new {@link ConsistentSampler} that respects the sampling decision of the parent span
    * 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 @@ static final ConsistentSampler parentBased(
    * @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 static final ConsistentSampler rateLimited(
    * @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 rValueGenerator the function to use for generating the r-value
+   */
+  public static ConsistentSampler rateLimited(
+      double targetSpansPerSecondLimit,
+      double adaptationTimeSeconds,
+      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 final ConsistentSampler rateLimited(
+  static ConsistentSampler rateLimited(
       double targetSpansPerSecondLimit,
       double adaptationTimeSeconds,
-      RandomGenerator randomGenerator,
+      RValueGenerator rValueGenerator,
       LongSupplier nanoTimeSupplier) {
     return new ConsistentRateLimitingSampler(
-        targetSpansPerSecondLimit, adaptationTimeSeconds, randomGenerator, nanoTimeSupplier);
+        targetSpansPerSecondLimit,
+        adaptationTimeSeconds,
+        rValueGenerator,
+        RandomGenerator.getDefault(),
+        nanoTimeSupplier);
   }
 
   /**
@@ -136,7 +179,8 @@ public ConsistentSampler and(ConsistentSampler otherConsistentSampler) {
     if (otherConsistentSampler == this) {
       return this;
     }
-    return new ConsistentComposedAndSampler(this, otherConsistentSampler);
+    return new ConsistentComposedAndSampler(
+        this, otherConsistentSampler, RValueGenerators.getDefault());
   }
 
   /**
@@ -158,20 +202,17 @@ public ConsistentSampler or(ConsistentSampler otherConsistentSampler) {
     if (otherConsistentSampler == this) {
       return this;
     }
-    return new ConsistentComposedOrSampler(this, otherConsistentSampler);
+    return new ConsistentComposedOrSampler(
+        this, otherConsistentSampler, RValueGenerators.getDefault());
   }
 
-  protected final RandomGenerator randomGenerator;
-
-  protected ConsistentSampler(RandomGenerator randomGenerator) {
-    this.randomGenerator = requireNonNull(randomGenerator);
-  }
+  private final RValueGenerator rValueGenerator;
 
-  protected ConsistentSampler() {
-    this(RandomGenerator.getDefault());
+  protected ConsistentSampler(RValueGenerator rValueGenerator) {
+    this.rValueGenerator = requireNonNull(rValueGenerator);
   }
 
-  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 final SamplingResult shouldSample(
 
     // 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