Links based sampler

samplers/src/main/java/io/opentelemetry/contrib/sampler/LinksBasedSampler.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.sampler;
+
+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 io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.List;
+import javax.annotation.concurrent.Immutable;
+
+/**
+ * A Sampler that uses the sampled flag of the span links if present. If at least one span link is
+ * sampled, then this span will be sampled. Otherwise, it is not sampled. If the span has no span
+ * links, this Sampler will use the "root" sampler that it is built with.
+ */
+@Immutable
+final class LinksBasedSampler implements Sampler {
+
+  private final Sampler root;
+
+  private static final SamplingResult POSITIVE_SAMPLING_RESULT = SamplingResult.recordAndSample();
+  private static final SamplingResult NEGATIVE_SAMPLING_RESULT = SamplingResult.drop();
+
+  public LinksBasedSampler(Sampler root) {
+    this.root = root;
+  }
+
+  @Override
+  public SamplingResult shouldSample(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    if (parentLinks.size() > 0) {
+      for (LinkData linkData : parentLinks) {
+        if (linkData.getSpanContext().isSampled()) {
+          return POSITIVE_SAMPLING_RESULT;
+        }
+      }
+      return NEGATIVE_SAMPLING_RESULT;
+    }
+
+    return this.root.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+  }
+
+  @Override
+  public String getDescription() {
+    return String.format("LinksBased{root:%s}", this.root.getDescription());
+  }
+
+  @Override
+  public String toString() {
+    return getDescription();
+  }
+}

samplers/src/test/java/io/opentelemetry/contrib/sampler/LinksBasedSamplerTest.java

@@ -0,0 +1,188 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.sampler;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+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.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.IdGenerator;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+
+class LinksBasedSamplerTest {
+  private static final String SPAN_NAME = "MySpanName";
+  private static final SpanKind SPAN_KIND = SpanKind.INTERNAL;
+  private final IdGenerator idsGenerator = IdGenerator.random();
+  private final String traceId = idsGenerator.generateTraceId();
+  private final String parentSpanId = idsGenerator.generateSpanId();
+
+  private final SpanContext sampledSpanContext1 =
+      SpanContext.create(traceId, parentSpanId, TraceFlags.getSampled(), TraceState.getDefault());
+
+  private final SpanContext sampledSpanContext2 =
+      SpanContext.create(traceId, parentSpanId, TraceFlags.getSampled(), TraceState.getDefault());
+
+  private final SpanContext unsampledSpanContext1 =
+      SpanContext.create(traceId, parentSpanId, TraceFlags.getDefault(), TraceState.getDefault());
+
+  private final SpanContext unsampledSpanContext2 =
+      SpanContext.create(traceId, parentSpanId, TraceFlags.getDefault(), TraceState.getDefault());
+
+  private final Context sampledParentContext = Context.root().with(Span.wrap(sampledSpanContext1));
+
+  @Test
+  void testEmptyAlwaysTrueRoot() {
+    assertThat(
+            new LinksBasedSampler(Sampler.alwaysOn())
+                .shouldSample(
+                    sampledParentContext,
+                    traceId,
+                    SPAN_NAME,
+                    SPAN_KIND,
+                    Attributes.empty(),
+                    Collections.emptyList())
+                .getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+  }
+
+  @Test
+  void testEmptyAlwaysFalseRoot() {
+    assertThat(
+            new LinksBasedSampler(Sampler.alwaysOff())
+                .shouldSample(
+                    sampledParentContext,
+                    traceId,
+                    SPAN_NAME,
+                    SPAN_KIND,
+                    Attributes.empty(),
+                    Collections.emptyList())
+                .getDecision())
+        .isEqualTo(SamplingDecision.DROP);
+  }
+
+  @Test
+  void testOneSampled() {
+    List<LinkData> linkData = new ArrayList<>();
+    linkData.add(LinkData.create(sampledSpanContext1));
+
+    assertThat(
+            new LinksBasedSampler(Sampler.alwaysOff())
+                .shouldSample(
+                    sampledParentContext,
+                    traceId,
+                    SPAN_NAME,
+                    SPAN_KIND,
+                    Attributes.empty(),
+                    linkData)
+                .getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+  }
+
+  @Test
+  void testOneNotSampled() {
+    List<LinkData> linkData = new ArrayList<>();
+    linkData.add(LinkData.create(unsampledSpanContext1));
+
+    assertThat(
+            new LinksBasedSampler(Sampler.alwaysOn())
+                .shouldSample(
+                    sampledParentContext,
+                    traceId,
+                    SPAN_NAME,
+                    SPAN_KIND,
+                    Attributes.empty(),
+                    linkData)
+                .getDecision())
+        .isEqualTo(SamplingDecision.DROP);
+  }
+
+  @Test
+  void testTwoSampledAndNotSampled() {
+    List<LinkData> linkData = new ArrayList<>();
+    linkData.add(LinkData.create(sampledSpanContext1));
+    linkData.add(LinkData.create(unsampledSpanContext1));
+
+    assertThat(
+            new LinksBasedSampler(Sampler.alwaysOff())
+                .shouldSample(
+                    sampledParentContext,
+                    traceId,
+                    SPAN_NAME,
+                    SPAN_KIND,
+                    Attributes.empty(),
+                    linkData)
+                .getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+  }
+
+  @Test
+  void testTwoNotSampledAndSampled() {
+    List<LinkData> linkData = new ArrayList<>();
+    linkData.add(LinkData.create(unsampledSpanContext1));
+    linkData.add(LinkData.create(sampledSpanContext1));
+
+    assertThat(
+            new LinksBasedSampler(Sampler.alwaysOff())
+                .shouldSample(
+                    sampledParentContext,
+                    traceId,
+                    SPAN_NAME,
+                    SPAN_KIND,
+                    Attributes.empty(),
+                    linkData)
+                .getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+  }
+
+  @Test
+  void testTwoSampled() {
+    List<LinkData> linkData = new ArrayList<>();
+    linkData.add(LinkData.create(sampledSpanContext1));
+    linkData.add(LinkData.create(sampledSpanContext2));
+
+    assertThat(
+            new LinksBasedSampler(Sampler.alwaysOff())
+                .shouldSample(
+                    sampledParentContext,
+                    traceId,
+                    SPAN_NAME,
+                    SPAN_KIND,
+                    Attributes.empty(),
+                    linkData)
+                .getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+  }
+
+  @Test
+  void testTwoUnsampled() {
+    List<LinkData> linkData = new ArrayList<>();
+    linkData.add(LinkData.create(unsampledSpanContext1));
+    linkData.add(LinkData.create(unsampledSpanContext2));
+
+    assertThat(
+            new LinksBasedSampler(Sampler.alwaysOn())
+                .shouldSample(
+                    sampledParentContext,
+                    traceId,
+                    SPAN_NAME,
+                    SPAN_KIND,
+                    Attributes.empty(),
+                    linkData)
+                .getDecision())
+        .isEqualTo(SamplingDecision.DROP);
+  }
+}