Composite Samplers prototype

This PR modifies the Consistent Probability Samplers to make them conform to the proposed Composite Samplers at open-telemetry/oteps#250

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

@@ -0,0 +1,114 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.sampler.consistent56;
+
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.getInvalidRandomValue;
+import static io.opentelemetry.contrib.sampler.consistent56.ConsistentSamplingUtil.isValidThreshold;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+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 io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.List;
+
+/** Abstract base class for composable consistent samplers. */
+public abstract class ComposableSampler extends ConsistentSampler {
+
+  @Override
+  public final SamplingResult shouldSample(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    Span parentSpan = Span.fromContext(parentContext);
+    SpanContext parentSpanContext = parentSpan.getSpanContext();
+
+    TraceState parentTraceState = parentSpanContext.getTraceState();
+    String otelTraceStateString = parentTraceState.get(OtelTraceState.TRACE_STATE_KEY);
+    OtelTraceState otelTraceState = OtelTraceState.parse(otelTraceStateString);
+
+    SamplingIntent intent =
+        getSamplingIntent(parentContext, name, spanKind, attributes, parentLinks);
+    long threshold = intent.getThreshold();
+
+    // determine sampling decision
+    boolean isSampled;
+    if (isValidThreshold(threshold)) {
+      long randomness = getRandomness(otelTraceState, traceId);
+      isSampled = threshold <= randomness;
+    } else { // DROP
+      isSampled = false;
+    }
+
+    SamplingDecision samplingDecision;
+    if (isSampled) {
+      samplingDecision = SamplingDecision.RECORD_AND_SAMPLE;
+      otelTraceState.setThreshold(threshold);
+    } else {
+      samplingDecision = SamplingDecision.DROP;
+      otelTraceState.invalidateThreshold();
+    }
+
+    String newOtTraceState = otelTraceState.serialize();
+
+    return new SamplingResult() {
+
+      @Override
+      public SamplingDecision getDecision() {
+        return samplingDecision;
+      }
+
+      @Override
+      public Attributes getAttributes() {
+        return intent.getAttributes();
+      }
+
+      @Override
+      public TraceState getUpdatedTraceState(TraceState parentTraceState) {
+        return intent.updateTraceState(parentTraceState).toBuilder()
+            .put(OtelTraceState.TRACE_STATE_KEY, newOtTraceState)
+            .build();
+      }
+    };
+  }
+
+  private static long getRandomness(OtelTraceState otelTraceState, String traceId) {
+    if (otelTraceState.hasValidRandomValue()) {
+      return otelTraceState.getRandomValue();
+    } else {
+      return OtelTraceState.parseHex(traceId, 18, 14, getInvalidRandomValue());
+    }
+  }
+
+  @Override
+  protected final long getThreshold(long parentThreshold, boolean isRoot) {
+    // not used by Composable Samplers
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * 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.
+   */
+  protected abstract SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks);
+}

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

@@ -5,10 +5,17 @@
 
 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
-final class ConsistentAlwaysOffSampler extends ConsistentSampler {
+final class ConsistentAlwaysOffSampler extends ComposableSampler {
 
   private static final ConsistentAlwaysOffSampler INSTANCE = new ConsistentAlwaysOffSampler();
 
@@ -19,8 +26,14 @@ static ConsistentAlwaysOffSampler getInstance() {
   }
 
   @Override
-  protected long getThreshold(long parentThreshold, boolean isRoot) {
-    return ConsistentSamplingUtil.getMaxThreshold();
+  protected 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,10 +7,15 @@
 
 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
-final class ConsistentAlwaysOnSampler extends ConsistentSampler {
+final class ConsistentAlwaysOnSampler extends ComposableSampler {
 
   private static final ConsistentAlwaysOnSampler INSTANCE = new ConsistentAlwaysOnSampler();
 
@@ -21,8 +26,14 @@ static ConsistentAlwaysOnSampler getInstance() {
   }
 
   @Override
-  protected long getThreshold(long parentThreshold, boolean isRoot) {
-    return getMinThreshold();
+  protected 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,111 @@
+/*
+ * 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 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 ComposableSampler {
+
+  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(ComposableSampler... delegates) {
+    if (delegates.length == 0) {
+      throw new IllegalArgumentException(
+          "At least one delegate must be specified for ConsistentAnyOf");
+    }
+
+    this.delegates = delegates;
+
+    StringBuilder builder = new StringBuilder("ConsistentAnyOf{");
+    for (int i = 0; i < delegates.length; ++i) {
+      builder.append(delegates[i].getDescription());
+      if (i < delegates.length - 1) {
+        builder.append(",");
+      }
+    }
+    builder.append("}");
+    this.description = builder.toString();
+  }
+
+  @Override
+  protected 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,8 +7,16 @@
 
 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;
 
-public class ConsistentFixedThresholdSampler extends ConsistentSampler {
+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 ComposableSampler {
 
   private final long threshold;
   private final String description;
@@ -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;
+  protected SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+
+    return () -> {
+      if (threshold == getMaxThreshold()) {
+        return getInvalidThreshold();
+      }
+      return threshold;
+    };
   }
 }