[dynamic control] Use composable samplers and add sampler initialization (#2752)

Author: jackshirazi

dynamic-control/build.gradle.kts

@@ -24,6 +24,7 @@ dependencies {
 
   compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
   compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi")
+  compileOnly("io.opentelemetry:opentelemetry-sdk-extension-incubator")
 
   testCompileOnly("com.google.auto.service:auto-service-annotations")
   testRuntimeOnly("org.junit.jupiter:junit-jupiter-engine")
@@ -32,6 +33,7 @@ dependencies {
 
   testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
   testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi")
+  testImplementation("io.opentelemetry:opentelemetry-sdk-extension-incubator")
   testImplementation("org.assertj:assertj-core")
   testImplementation("org.mockito:mockito-inline")
   testImplementation("org.mockito:mockito-junit-jupiter")

dynamic-control/src/main/java/io/opentelemetry/contrib/dynamic/policy/tracesampling/TraceSamplingRatePolicy.java

@@ -6,22 +6,67 @@
 package io.opentelemetry.contrib.dynamic.policy.tracesampling;
 
 import io.opentelemetry.contrib.dynamic.policy.TelemetryPolicy;
+import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer;
+import io.opentelemetry.sdk.extension.incubator.trace.samplers.ComposableSampler;
+import io.opentelemetry.sdk.extension.incubator.trace.samplers.CompositeSampler;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import java.util.Objects;
+import javax.annotation.Nullable;
 
 public final class TraceSamplingRatePolicy extends TelemetryPolicy {
   public static final String POLICY_TYPE = "trace-sampling";
 
+  @Nullable private static volatile DelegatingSampler initializedSampler;
+
   private final double probability;
 
   public TraceSamplingRatePolicy(double probability) {
     super(POLICY_TYPE);
+    this.probability = normalizeProbability(probability);
+  }
+
+  public double getProbability() {
+    return probability;
+  }
+
+  /**
+   * Initializes runtime wiring for this policy type.
+   *
+   * <p>If the extension is configured to use this policy, this installs an opinionated sampler that
+   * overrides any other sampler
+   */
+  public static void initialize(AutoConfigurationCustomizer autoConfiguration) {
+    Objects.requireNonNull(autoConfiguration, "autoConfiguration cannot be null");
+    Sampler initialDelegate = createSampler(1.0);
+    DelegatingSampler delegatingSampler = new DelegatingSampler(initialDelegate);
+    initializedSampler = delegatingSampler;
+    autoConfiguration.addSamplerCustomizer((sampler, config) -> delegatingSampler);
+  }
+
+  /**
+   * Creates the composed sampler used for this policy probability.
+   *
+   * @param probability sampling probability in the inclusive range {@code [0.0, 1.0]}
+   * @return a sampler equivalent to the configured probability with parent-based behavior
+   * @throws IllegalArgumentException if probability is NaN or outside {@code [0.0, 1.0]}
+   */
+  public static Sampler createSampler(double probability) {
+    probability = normalizeProbability(probability);
+    return CompositeSampler.wrap(
+        ComposableSampler.parentThreshold(ComposableSampler.probability(probability)));
+  }
+
+  private static double normalizeProbability(double probability) {
     if (Double.isNaN(probability) || probability < 0.0 || probability > 1.0) {
       throw new IllegalArgumentException("probability must be within [0.0, 1.0]");
     }
-    this.probability = probability;
+    // Normalize -0.0 to +0.0 so equality/hash behavior stays intuitive.
+    return probability == 0.0 ? 0.0 : probability;
   }
 
-  public double getProbability() {
-    return probability;
+  @Nullable
+  public static DelegatingSampler getInitializedSampler() {
+    return initializedSampler;
   }
 
   @Override

dynamic-control/src/main/java/io/opentelemetry/contrib/dynamic/policy/tracesampling/TraceSamplingRatePolicyImplementer.java

@@ -12,17 +12,18 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
+import java.util.logging.Logger;
 
 /**
  * Implements the {@code trace-sampling} policy by updating a {@link DelegatingSampler}.
  *
  * <p>This implementer listens for validated {@link TelemetryPolicy} updates of type {@code
  * "trace-sampling"} and applies {@link TraceSamplingRatePolicy#getProbability()} to the delegate
- * sampler using {@link Sampler#traceIdRatioBased(double)} wrapped by {@link
- * Sampler#parentBased(Sampler)}.
+ * sampler via {@link TraceSamplingRatePolicy#createSampler(double)}.
  *
  * <p>If a type-only {@link TelemetryPolicy} of type {@code "trace-sampling"} is received, it is
- * treated as policy removal and the delegate falls back to {@link Sampler#alwaysOn()}.
+ * treated as policy removal and the delegate is reset using {@code
+ * TraceSamplingRatePolicy.createSampler(1.0)}.
  *
  * <p>Validation is performed by {@link TraceSamplingValidator}; this implementer only consumes
  * policies produced by that validator.
@@ -31,6 +32,8 @@
  * with sampling operations on the associated {@link DelegatingSampler}.
  */
 public final class TraceSamplingRatePolicyImplementer implements PolicyImplementer {
+  private static final Logger logger =
+      Logger.getLogger(TraceSamplingRatePolicyImplementer.class.getName());
 
   private static final List<PolicyValidator> VALIDATORS =
       Collections.<PolicyValidator>singletonList(new TraceSamplingValidator());
@@ -60,12 +63,14 @@ public void onPoliciesChanged(List<TelemetryPolicy> policies) {
       }
       if (!(policy instanceof TraceSamplingRatePolicy)) {
         // Type-only policy represents removing trace-sampling config.
-        delegatingSampler.setDelegate(Sampler.alwaysOn());
+        delegatingSampler.setDelegate(TraceSamplingRatePolicy.createSampler(1.0));
+        logger.info("Applied trace sampling policy reset: probability reset to 1.0");
         continue;
       }
       double ratio = ((TraceSamplingRatePolicy) policy).getProbability();
-      Sampler sampler = Sampler.parentBased(Sampler.traceIdRatioBased(ratio));
+      Sampler sampler = TraceSamplingRatePolicy.createSampler(ratio);
       delegatingSampler.setDelegate(sampler);
+      logger.info("Applied trace sampling policy update: probability=" + ratio);
     }
   }
 }

dynamic-control/src/test/java/io/opentelemetry/contrib/dynamic/policy/tracesampling/TraceSamplingRatePolicyTest.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.dynamic.policy.tracesampling;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import org.junit.jupiter.api.Test;
+
+class TraceSamplingRatePolicyTest {
+
+  @Test
+  void constructorStoresProbabilityAndType() {
+    TraceSamplingRatePolicy policy = new TraceSamplingRatePolicy(0.25);
+
+    assertThat(policy.getProbability()).isEqualTo(0.25);
+    assertThat(policy.getType()).isEqualTo(TraceSamplingRatePolicy.POLICY_TYPE);
+  }
+
+  @Test
+  void constructorNormalizesNegativeZeroToPositiveZero() {
+    TraceSamplingRatePolicy negativeZero = new TraceSamplingRatePolicy(-0.0);
+    TraceSamplingRatePolicy positiveZero = new TraceSamplingRatePolicy(0.0);
+
+    assertThat(negativeZero.getProbability()).isEqualTo(0.0);
+    assertThat(Double.doubleToRawLongBits(negativeZero.getProbability()))
+        .isEqualTo(Double.doubleToRawLongBits(0.0));
+    assertThat(negativeZero).isEqualTo(positiveZero);
+    assertThat(negativeZero.hashCode()).isEqualTo(positiveZero.hashCode());
+  }
+
+  @Test
+  void constructorRejectsOutOfRangeOrNaNProbabilities() {
+    assertThatThrownBy(() -> new TraceSamplingRatePolicy(Double.NaN))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("probability must be within [0.0, 1.0]");
+    assertThatThrownBy(() -> new TraceSamplingRatePolicy(-0.001))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("probability must be within [0.0, 1.0]");
+    assertThatThrownBy(() -> new TraceSamplingRatePolicy(1.001))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("probability must be within [0.0, 1.0]");
+  }
+
+  @Test
+  void equalsAndHashCodeUseProbability() {
+    TraceSamplingRatePolicy a = new TraceSamplingRatePolicy(0.5);
+    TraceSamplingRatePolicy b = new TraceSamplingRatePolicy(0.5);
+    TraceSamplingRatePolicy c = new TraceSamplingRatePolicy(0.75);
+
+    assertThat(a).isEqualTo(b);
+    assertThat(a.hashCode()).isEqualTo(b.hashCode());
+    assertThat(a).isNotEqualTo(c);
+    assertThat(a).isNotEqualTo(null);
+    assertThat(a).isNotEqualTo("not-a-policy");
+  }
+
+  @Test
+  void initializeRejectsNullCustomizer() {
+    assertThatThrownBy(() -> TraceSamplingRatePolicy.initialize(null))
+        .isInstanceOf(NullPointerException.class)
+        .hasMessage("autoConfiguration cannot be null");
+  }
+
+  @Test
+  void initializeStoresDelegatingSamplerAndRegistersCustomizer() {
+    AutoConfigurationCustomizer customizer = mock(AutoConfigurationCustomizer.class);
+
+    TraceSamplingRatePolicy.initialize(customizer);
+
+    assertThat(TraceSamplingRatePolicy.getInitializedSampler()).isNotNull();
+    verify(customizer).addSamplerCustomizer(any());
+  }
+
+  @Test
+  void createSamplerAcceptsBoundaryProbabilities() {
+    Sampler zero = TraceSamplingRatePolicy.createSampler(0.0);
+    Sampler negativeZero = TraceSamplingRatePolicy.createSampler(-0.0);
+    Sampler one = TraceSamplingRatePolicy.createSampler(1.0);
+
+    assertThat(zero).isNotNull();
+    assertThat(negativeZero).isNotNull();
+    assertThat(one).isNotNull();
+  }
+
+  @Test
+  void createSamplerRejectsOutOfRangeOrNaNProbabilities() {
+    assertThatThrownBy(() -> TraceSamplingRatePolicy.createSampler(Double.NaN))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("probability must be within [0.0, 1.0]");
+    assertThatThrownBy(() -> TraceSamplingRatePolicy.createSampler(-0.01))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("probability must be within [0.0, 1.0]");
+    assertThatThrownBy(() -> TraceSamplingRatePolicy.createSampler(1.01))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("probability must be within [0.0, 1.0]");
+  }
+}