Composite Samplers prototype (#1443)

Co-authored-by: Otmar Ertl <[email protected]>

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.sampler.consistent56;
+
+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;
+
+/** An interface for components to be used by composite consistent probability samplers. */
+public interface ComposableSampler {
+
+  /**
+   * Returns the SamplingIntent that is used for the sampling decision. The SamplingIntent includes
+   * the threshold value which will be used for the sampling decision.
+   *
+   * <p>NOTE: Keep in mind, that in any case the returned threshold value must not depend directly
+   * or indirectly on the random value. In particular this means that the parent sampled flag must
+   * not be used for the calculation of the threshold as the sampled flag depends itself on the
+   * random value.
+   */
+  SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks);
+
+  /** Return the string providing a description of the implementation. */
+  String getDescription();
+}

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

@@ -5,6 +5,13 @@
 
 package io.opentelemetry.contrib.sampler.consistent56;
 
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getInvalidThreshold;
+
+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;
 import javax.annotation.concurrent.Immutable;
 
 @Immutable
@@ -19,8 +26,14 @@ static ConsistentAlwaysOffSampler getInstance() {
   }
 
   @Override
-  protected long getThreshold(long parentThreshold, boolean isRoot) {
-    return ConsistentSamplingUtil.getMaxThreshold();
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+
+    return () -> getInvalidThreshold();
   }
 
   @Override

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

@@ -7,6 +7,11 @@
 
 import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getMinThreshold;
 
+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;
 import javax.annotation.concurrent.Immutable;
 
 @Immutable
@@ -21,8 +26,14 @@ static ConsistentAlwaysOnSampler getInstance() {
   }
 
   @Override
-  protected long getThreshold(long parentThreshold, boolean isRoot) {
-    return getMinThreshold();
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+
+    return () -> getMinThreshold();
   }
 
   @Override

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

@@ -0,0 +1,109 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.sampler.consistent56;
+
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getInvalidThreshold;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.isValidThreshold;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.Immutable;
+
+/**
+ * A consistent sampler that queries all its delegate samplers for their sampling threshold, and
+ * uses the minimum threshold value received.
+ */
+@Immutable
+final class ConsistentAnyOf extends ConsistentSampler {
+
+  private final ComposableSampler[] delegates;
+
+  private final String description;
+
+  /**
+   * Constructs a new consistent AnyOf sampler using the provided delegate samplers.
+   *
+   * @param delegates the delegate samplers
+   */
+  ConsistentAnyOf(@Nullable ComposableSampler... delegates) {
+    if (delegates == null || delegates.length == 0) {
+      throw new IllegalArgumentException(
+          "At least one delegate must be specified for ConsistentAnyOf");
+    }
+
+    this.delegates = delegates;
+
+    this.description =
+        Stream.of(delegates)
+            .map(Object::toString)
+            .collect(Collectors.joining(",", "ConsistentAnyOf{", "}"));
+  }
+
+  @Override
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+
+    SamplingIntent[] intents = new SamplingIntent[delegates.length];
+    int k = 0;
+    long minimumThreshold = getInvalidThreshold();
+    for (ComposableSampler delegate : delegates) {
+      SamplingIntent delegateIntent =
+          delegate.getSamplingIntent(parentContext, name, spanKind, attributes, parentLinks);
+      long delegateThreshold = delegateIntent.getThreshold();
+      if (isValidThreshold(delegateThreshold)) {
+        if (isValidThreshold(minimumThreshold)) {
+          minimumThreshold = Math.min(delegateThreshold, minimumThreshold);
+        } else {
+          minimumThreshold = delegateThreshold;
+        }
+      }
+      intents[k++] = delegateIntent;
+    }
+
+    long resultingThreshold = minimumThreshold;
+
+    return new SamplingIntent() {
+      @Override
+      public long getThreshold() {
+        return resultingThreshold;
+      }
+
+      @Override
+      public Attributes getAttributes() {
+        AttributesBuilder builder = Attributes.builder();
+        for (SamplingIntent intent : intents) {
+          builder = builder.putAll(intent.getAttributes());
+        }
+        return builder.build();
+      }
+
+      @Override
+      public TraceState updateTraceState(TraceState previousState) {
+        for (SamplingIntent intent : intents) {
+          previousState = intent.updateTraceState(previousState);
+        }
+        return previousState;
+      }
+    };
+  }
+
+  @Override
+  public String getDescription() {
+    return description;
+  }
+}

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

@@ -7,6 +7,14 @@
 
 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 class ConsistentFixedThresholdSampler extends ConsistentSampler {
 
@@ -18,7 +26,7 @@ protected ConsistentFixedThresholdSampler(long threshold) {
     this.threshold = threshold;
 
     String thresholdString;
-    if (threshold == ConsistentSamplingUtil.getMaxThreshold()) {
+    if (threshold == getMaxThreshold()) {
       thresholdString = "max";
     } else {
       thresholdString =
@@ -41,7 +49,18 @@ public String getDescription() {
   }
 
   @Override
-  protected long getThreshold(long parentThreshold, boolean isRoot) {
-    return threshold;
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+
+    return () -> {
+      if (threshold == getMaxThreshold()) {
+        return getInvalidThreshold();
+      }
+      return threshold;
+    };
   }
 }