Add incubator implementation of composite sampling spec (#7626)

Author: anuraaga

sdk-extensions/incubator/build.gradle.kts

@@ -40,6 +40,7 @@ dependencies {
   testImplementation(project(":exporters:zipkin"))
   testImplementation(project(":sdk-extensions:jaeger-remote-sampler"))
   testImplementation(project(":extensions:trace-propagators"))
+  testImplementation("edu.berkeley.cs.jqf:jqf-fuzz")
   testImplementation("io.opentelemetry.contrib:opentelemetry-aws-xray-propagator")
   testImplementation("com.linecorp.armeria:armeria-junit5")
 

sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableAlwaysOffSampler.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.extension.incubator.trace.samplers;
+
+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 java.util.function.Function;
+
+enum ComposableAlwaysOffSampler implements ComposableSampler {
+  INSTANCE;
+
+  private static final SamplingIntent INTENT =
+      SamplingIntent.create(
+          ImmutableSamplingIntent.INVALID_THRESHOLD,
+          /* thresholdReliable= */ false,
+          Attributes.empty(),
+          Function.identity());
+
+  @Override
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    return INTENT;
+  }
+
+  @Override
+  public String getDescription() {
+    return "ComposableAlwaysOffSampler";
+  }
+
+  @Override
+  public String toString() {
+    return getDescription();
+  }
+}

sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableAlwaysOnSampler.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.extension.incubator.trace.samplers;
+
+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 java.util.function.Function;
+
+enum ComposableAlwaysOnSampler implements ComposableSampler {
+  INSTANCE;
+
+  private static final SamplingIntent INTENT =
+      SamplingIntent.create(
+          ImmutableSamplingIntent.MIN_THRESHOLD,
+          /* thresholdReliable= */ true,
+          Attributes.empty(),
+          Function.identity());
+
+  @Override
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    return INTENT;
+  }
+
+  @Override
+  public String getDescription() {
+    return "ComposableAlwaysOnSampler";
+  }
+
+  @Override
+  public String toString() {
+    return getDescription();
+  }
+}

sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableParentThresholdSampler.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.extension.incubator.trace.samplers;
+
+import static io.opentelemetry.sdk.extension.incubator.trace.samplers.ImmutableSamplingIntent.INVALID_THRESHOLD;
+import static io.opentelemetry.sdk.extension.incubator.trace.samplers.ImmutableSamplingIntent.MIN_THRESHOLD;
+import static io.opentelemetry.sdk.extension.incubator.trace.samplers.ImmutableSamplingIntent.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.context.Context;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import java.util.List;
+import java.util.function.Function;
+
+final class ComposableParentThresholdSampler implements ComposableSampler {
+
+  private final ComposableSampler rootSampler;
+  private final String description;
+
+  ComposableParentThresholdSampler(ComposableSampler rootSampler) {
+    this.rootSampler = rootSampler;
+    this.description = "ComposableParentThresholdSampler{rootSampler=" + rootSampler + "}";
+  }
+
+  @Override
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    SpanContext parentSpanContext = Span.fromContext(parentContext).getSpanContext();
+    if (!parentSpanContext.isValid()) {
+      return rootSampler.getSamplingIntent(
+          parentContext, traceId, name, spanKind, attributes, parentLinks);
+    }
+
+    OtelTraceState otTraceState = OtelTraceState.parse(parentSpanContext.getTraceState());
+    if (isValidThreshold(otTraceState.getThreshold())) {
+      return ImmutableSamplingIntent.create(
+          otTraceState.getThreshold(),
+          /* thresholdReliable= */ true,
+          Attributes.empty(),
+          Function.identity());
+    }
+
+    long threshold =
+        parentSpanContext.getTraceFlags().isSampled() ? MIN_THRESHOLD : INVALID_THRESHOLD;
+    return ImmutableSamplingIntent.create(
+        threshold, /* thresholdReliable= */ false, Attributes.empty(), Function.identity());
+  }
+
+  @Override
+  public String getDescription() {
+    return description;
+  }
+
+  @Override
+  public String toString() {
+    return this.getDescription();
+  }
+}

sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableSampler.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.extension.incubator.trace.samplers;
+
+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;
+
+/** A sampler that can be composed to make a final sampling decision. */
+public interface ComposableSampler {
+  /** Returns a {@link ComposableSampler} that does not sample any span. */
+  static ComposableSampler alwaysOff() {
+    return ComposableAlwaysOffSampler.INSTANCE;
+  }
+
+  /** Returns a {@link ComposableSampler} that samples all spans. */
+  static ComposableSampler alwaysOn() {
+    return ComposableAlwaysOnSampler.INSTANCE;
+  }
+
+  /** Returns a {@link ComposableSampler} that samples each span with a fixed ratio. */
+  static ComposableSampler traceIdRatioBased(double ratio) {
+    return new ComposableTraceIdRatioBasedSampler(ratio);
+  }
+
+  /**
+   * Returns a {@link ComposableSampler} that respects the sampling decision of the parent span or
+   * falls back to the given sampler if it is a root span.
+   */
+  static ComposableSampler parentThreshold(ComposableSampler rootSampler) {
+    return new ComposableParentThresholdSampler(rootSampler);
+  }
+
+  /** Returns the {@link SamplingIntent} to use to make a sampling decision. */
+  SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks);
+
+  /** Returns a description of the sampler implementation. */
+  String getDescription();
+}

sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableTraceIdRatioBasedSampler.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.extension.incubator.trace.samplers;
+
+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 java.util.function.Function;
+
+final class ComposableTraceIdRatioBasedSampler implements ComposableSampler {
+  private static long calculateThreshold(double ratio) {
+    return ImmutableSamplingIntent.MAX_THRESHOLD
+        - Math.round(ratio * (double) ImmutableSamplingIntent.MAX_THRESHOLD);
+  }
+
+  private final SamplingIntent intent;
+  private final String description;
+
+  ComposableTraceIdRatioBasedSampler(double ratio) {
+    long threshold = calculateThreshold(ratio);
+    String thresholdStr;
+    if (threshold == ImmutableSamplingIntent.MAX_THRESHOLD) {
+      thresholdStr = "max";
+
+      // Same as ComposableAlwaysOffSampler, notably the threshold is not considered reliable.
+      // The spec mentions returning an instance of ComposableAlwaysOffSampler in this case but
+      // it seems clearer if the description of the sampler matches the user's request.
+      this.intent =
+          SamplingIntent.create(
+              ImmutableSamplingIntent.INVALID_THRESHOLD,
+              /* thresholdReliable= */ false,
+              Attributes.empty(),
+              Function.identity());
+    } else {
+      StringBuilder sb = new StringBuilder();
+      OtelTraceState.serializeTh(threshold, sb);
+      thresholdStr = sb.toString();
+
+      this.intent =
+          SamplingIntent.create(
+              threshold, /* thresholdReliable= */ true, Attributes.empty(), Function.identity());
+    }
+    this.description =
+        "ComposableTraceIdRatioBasedSampler{threshold=" + thresholdStr + ", ratio=" + ratio + "}";
+  }
+
+  @Override
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    return intent;
+  }
+
+  @Override
+  public String getDescription() {
+    return description;
+  }
+
+  @Override
+  public String toString() {
+    return this.getDescription();
+  }
+}