add updateable threshold sampler (open-telemetry#2137)

Author: jackshirazi

consistent-sampling/src/main/java/io/opentelemetry/contrib/sampler/consistent56/ConsistentFixedThresholdSampler.java

@@ -5,19 +5,9 @@
 
 package io.opentelemetry.contrib.sampler.consistent56;
 
-import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.calculateSamplingProbability;
 import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.calculateThreshold;
-import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.checkThreshold;
-import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getInvalidThreshold;
-import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getMaxThreshold;
 
-import io.opentelemetry.api.common.Attributes;
-import io.opentelemetry.api.trace.SpanKind;
-import io.opentelemetry.context.Context;
-import io.opentelemetry.sdk.trace.data.LinkData;
-import java.util.List;
-
-public class ConsistentFixedThresholdSampler extends ConsistentSampler {
+public class ConsistentFixedThresholdSampler extends ConsistentThresholdSampler {
 
   private final long threshold;
   private final String description;
@@ -31,47 +21,13 @@ protected ConsistentFixedThresholdSampler(double samplingProbability) {
     this(calculateThreshold(samplingProbability));
   }
 
-  private static long getThreshold(long threshold) {
-    checkThreshold(threshold);
-    return threshold;
-  }
-
-  private static String getThresholdDescription(long threshold) {
-    String thresholdString;
-    if (threshold == getMaxThreshold()) {
-      thresholdString = "max";
-    } else {
-      thresholdString =
-          ConsistentSamplingUtil.appendLast56BitHexEncodedWithoutTrailingZeros(
-                  new StringBuilder(), threshold)
-              .toString();
-    }
-
-    return "ConsistentFixedThresholdSampler{threshold="
-        + thresholdString
-        + ", sampling probability="
-        + calculateSamplingProbability(threshold)
-        + "}";
-  }
-
   @Override
   public String getDescription() {
     return description;
   }
 
   @Override
-  public SamplingIntent getSamplingIntent(
-      Context parentContext,
-      String name,
-      SpanKind spanKind,
-      Attributes attributes,
-      List<LinkData> parentLinks) {
-
-    return () -> {
-      if (threshold == getMaxThreshold()) {
-        return getInvalidThreshold();
-      }
-      return threshold;
-    };
+  public long getThreshold() {
+    return threshold;
   }
 }

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

@@ -55,6 +55,17 @@ public static ConsistentSampler probabilityBased(double samplingProbability) {
     return new ConsistentFixedThresholdSampler(threshold);
   }
 
+  /**
+   * Returns a {@link ConsistentSampler} that samples each span with a known probability, where the
+   * probablity can be dynamically updated.
+   *
+   * @param samplingProbability the sampling probability
+   * @return a sampler
+   */
+  public static ConsistentSampler updateableProbabilityBased(double samplingProbability) {
+    return new ConsistentVariableThresholdSampler(samplingProbability);
+  }
+
   /**
    * 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.

consistent-sampling/src/main/java/io/opentelemetry/contrib/sampler/consistent56/ConsistentThresholdSampler.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.sampler.consistent56;
+
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.appendLast56BitHexEncodedWithoutTrailingZeros;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.calculateSamplingProbability;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.checkThreshold;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getInvalidThreshold;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getMaxThreshold;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import java.util.List;
+
+public abstract class ConsistentThresholdSampler extends ConsistentSampler {
+
+  protected abstract long getThreshold();
+
+  protected static long getThreshold(long threshold) {
+    checkThreshold(threshold);
+    return threshold;
+  }
+
+  protected static String getThresholdDescription(long threshold) {
+    String thresholdString;
+    if (threshold == getMaxThreshold()) {
+      thresholdString = "max";
+    } else {
+      thresholdString =
+          appendLast56BitHexEncodedWithoutTrailingZeros(new StringBuilder(), threshold).toString();
+    }
+
+    return "ConsistentFixedThresholdSampler{threshold="
+        + thresholdString
+        + ", sampling probability="
+        + calculateSamplingProbability(threshold)
+        + "}";
+  }
+
+  @Override
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+
+    return () -> {
+      if (getThreshold() == getMaxThreshold()) {
+        return getInvalidThreshold();
+      }
+      return getThreshold();
+    };
+  }
+}

consistent-sampling/src/main/java/io/opentelemetry/contrib/sampler/consistent56/ConsistentVariableThresholdSampler.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.sampler.consistent56;
+
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.calculateSamplingProbability;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.calculateThreshold;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.checkThreshold;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getMaxThreshold;
+
+public class ConsistentVariableThresholdSampler extends ConsistentThresholdSampler {
+
+  private volatile long threshold;
+  private volatile String description = "";
+
+  protected ConsistentVariableThresholdSampler(double samplingProbability) {
+    setSamplingProbability(samplingProbability);
+  }
+
+  @Override
+  public String getDescription() {
+    return description;
+  }
+
+  @Override
+  public long getThreshold() {
+    return threshold;
+  }
+
+  public void setSamplingProbability(double samplingProbability) {
+    long threshold = calculateThreshold(samplingProbability);
+    checkThreshold(threshold);
+    this.threshold = threshold;
+
+    String thresholdString;
+    if (threshold == getMaxThreshold()) {
+      thresholdString = "max";
+    } else {
+      thresholdString =
+          ConsistentSamplingUtil.appendLast56BitHexEncodedWithoutTrailingZeros(
+                  new StringBuilder(), threshold)
+              .toString();
+    }
+
+    // tiny eventual consistency where the description would be out of date with the threshold,
+    // but this doesn't really matter
+    this.description =
+        "ConsistentVariableThresholdSampler{threshold="
+            + thresholdString
+            + ", sampling probability="
+            + calculateSamplingProbability(threshold)
+            + "}";
+  }
+}

consistent-sampling/src/test/java/io/opentelemetry/contrib/sampler/consistent56/ConsistentVariableThresholdSamplerTest.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.sampler.consistent56;
+
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.calculateThreshold;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getMaxThreshold;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.junit.jupiter.api.Test;
+
+class ConsistentVariableThresholdSamplerTest {
+
+  @Test
+  void testSetSamplingProbability() {
+    double probability = 0.5;
+    ConsistentVariableThresholdSampler sampler =
+        new ConsistentVariableThresholdSampler(probability);
+    testSetSamplingProbability(probability, sampler, /* updateProbability= */ false);
+    testSetSamplingProbability(0.25, sampler, /* updateProbability= */ true);
+    testSetSamplingProbability(0.0, sampler, /* updateProbability= */ true);
+    testSetSamplingProbability(1.0, sampler, /* updateProbability= */ true);
+  }
+
+  private static void testSetSamplingProbability(
+      double probability, ConsistentVariableThresholdSampler sampler, boolean updateProbability) {
+    long threshold = calculateThreshold(probability);
+    String thresholdString =
+        ConsistentSamplingUtil.appendLast56BitHexEncodedWithoutTrailingZeros(
+                new StringBuilder(), threshold)
+            .toString();
+    if (threshold == getMaxThreshold()) {
+      thresholdString = "max";
+    }
+    if (updateProbability) {
+      sampler.setSamplingProbability(probability);
+    }
+    assertThat(sampler.getThreshold()).isEqualTo(threshold);
+    assertThat(sampler.getDescription())
+        .isEqualTo(
+            "ConsistentVariableThresholdSampler{threshold="
+                + thresholdString
+                + ", sampling probability="
+                + probability
+                + "}");
+  }
+}