Attributes rule based sampler (#70)

* Initial POC version of url-based sampler

* Extract UrlMatcher class

* Pre-compile patterns

* Use UrlMatcher in UrlSampler

* Accept several attributes to check

* Will not extract path for matching

* Converted to rule-based routing implementation

* Polish

* Rename module

* Improve comments

* Apply suggestions from code review

Co-authored-by: Anuraag Agrawal <anuraaga@gmail.com>

* Code review polish

* Polish

* Better equals

Co-authored-by: Anuraag Agrawal <anuraaga@gmail.com>

buildSrc/src/main/kotlin/otel.java-conventions.gradle.kts

@@ -85,6 +85,7 @@ dependencies {
     testImplementation("org.awaitility:awaitility")
     testImplementation("org.junit.jupiter:junit-jupiter-api")
     testImplementation("org.junit.jupiter:junit-jupiter-params")
+    testImplementation("org.mockito:mockito-core")
     testImplementation("org.mockito:mockito-junit-jupiter")
 
     testRuntimeOnly("org.junit.jupiter:junit-jupiter-engine")

contrib-samplers/build.gradle.kts

@@ -0,0 +1,11 @@
+plugins {
+  id("otel.java-conventions")
+  id("otel.publish-conventions")
+}
+
+description = "Sampler which makes its decision based on semantic attributes values"
+
+dependencies {
+  api("io.opentelemetry:opentelemetry-sdk")
+  api("io.opentelemetry:opentelemetry-semconv")
+}

contrib-samplers/src/main/java/io/opentelemetry/contrib/samplers/RuleBasedRoutingSampler.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package io.opentelemetry.contrib.samplers;
+
+import static java.util.Objects.requireNonNull;
+
+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;
+
+/**
+ * This sampler accepts a list of {@link SamplingRule}s and tries to match every proposed span
+ * against those rules. Every rule describes a span's attribute, a pattern against which to match
+ * attribute's value, and a sampler that will make a decision about given span if match was
+ * successful.
+ *
+ * <p>Matching is performed by {@link java.util.regex.Pattern}.
+ *
+ * <p>Provided span kind is checked first and if differs from the one given to {@link
+ * #builder(SpanKind, Sampler)}, the default fallback sampler will make a decision.
+ *
+ * <p>Note that only attributes that were set on {@link io.opentelemetry.api.trace.SpanBuilder} will
+ * be taken into account, attributes set after the span has been started are not used
+ *
+ * <p>If none of the rules matched, the default fallback sampler will make a decision.
+ */
+public final class RuleBasedRoutingSampler implements Sampler {
+  private final List<SamplingRule> rules;
+  private final SpanKind kind;
+  private final Sampler fallback;
+
+  RuleBasedRoutingSampler(List<SamplingRule> rules, SpanKind kind, Sampler fallback) {
+    this.kind = requireNonNull(kind);
+    this.fallback = requireNonNull(fallback);
+    this.rules = requireNonNull(rules);
+  }
+
+  public static RuleBasedRoutingSamplerBuilder builder(SpanKind kind, Sampler fallback) {
+    return new RuleBasedRoutingSamplerBuilder(
+        requireNonNull(kind, "span kind must not be null"),
+        requireNonNull(fallback, "fallback sampler must not be null"));
+  }
+
+  @Override
+  public SamplingResult shouldSample(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    if (kind != spanKind) {
+      return fallback.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+    }
+    for (SamplingRule samplingRule : rules) {
+      String attributeValue = attributes.get(samplingRule.attributeKey);
+      if (attributeValue == null) {
+        continue;
+      }
+      if (samplingRule.pattern.matcher(attributeValue).find()) {
+        return samplingRule.delegate.shouldSample(
+            parentContext, traceId, name, spanKind, attributes, parentLinks);
+      }
+    }
+    return fallback.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+  }
+
+  @Override
+  public String getDescription() {
+    return "RuleBasedRoutingSampler{"
+        + "rules="
+        + rules
+        + ", kind="
+        + kind
+        + ", fallback="
+        + fallback
+        + '}';
+  }
+
+  @Override
+  public String toString() {
+    return getDescription();
+  }
+}

contrib-samplers/src/main/java/io/opentelemetry/contrib/samplers/RuleBasedRoutingSamplerBuilder.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package io.opentelemetry.contrib.samplers;
+
+import static java.util.Objects.requireNonNull;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import java.util.ArrayList;
+import java.util.List;
+
+public final class RuleBasedRoutingSamplerBuilder {
+  private final List<SamplingRule> rules = new ArrayList<>();
+  private final SpanKind kind;
+  private final Sampler defaultDelegate;
+
+  RuleBasedRoutingSamplerBuilder(SpanKind kind, Sampler defaultDelegate) {
+    this.kind = kind;
+    this.defaultDelegate = defaultDelegate;
+  }
+
+  public RuleBasedRoutingSamplerBuilder drop(AttributeKey<String> attributeKey, String pattern) {
+    rules.add(
+        new SamplingRule(
+            requireNonNull(attributeKey, "attributeKey must not be null"),
+            requireNonNull(pattern, "pattern must not be null"),
+            Sampler.alwaysOff()));
+    return this;
+  }
+
+  public RuleBasedRoutingSamplerBuilder recordAndSample(
+      AttributeKey<String> attributeKey, String pattern) {
+    rules.add(
+        new SamplingRule(
+            requireNonNull(attributeKey, "attributeKey must not be null"),
+            requireNonNull(pattern, "pattern must not be null"),
+            Sampler.alwaysOn()));
+    return this;
+  }
+
+  public RuleBasedRoutingSampler build() {
+    return new RuleBasedRoutingSampler(rules, kind, defaultDelegate);
+  }
+}

contrib-samplers/src/main/java/io/opentelemetry/contrib/samplers/SamplingRule.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package io.opentelemetry.contrib.samplers;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import java.util.Objects;
+import java.util.regex.Pattern;
+
+/** @see RuleBasedRoutingSampler */
+class SamplingRule {
+  final AttributeKey<String> attributeKey;
+  final Sampler delegate;
+  final Pattern pattern;
+
+  SamplingRule(AttributeKey<String> attributeKey, String pattern, Sampler delegate) {
+    this.attributeKey = attributeKey;
+    this.pattern = Pattern.compile(pattern);
+    this.delegate = delegate;
+  }
+
+  @Override
+  public String toString() {
+    return "SamplingRule{"
+        + "attributeKey="
+        + attributeKey
+        + ", delegate="
+        + delegate
+        + ", pattern="
+        + pattern
+        + '}';
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (!(o instanceof SamplingRule)) return false;
+    SamplingRule that = (SamplingRule) o;
+    return attributeKey.equals(that.attributeKey) && pattern.equals(that.pattern);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(attributeKey, pattern);
+  }
+}

contrib-samplers/src/test/java/io/opentelemetry/contrib/samplers/RuleBasedRoutingSamplerTest.java

@@ -0,0 +1,177 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package io.opentelemetry.contrib.samplers;
+
+import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.HTTP_TARGET;
+import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.HTTP_URL;
+import static java.util.Collections.emptyList;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+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.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+@ExtendWith(MockitoExtension.class)
+class RuleBasedRoutingSamplerTest {
+  private static final String SPAN_NAME = "MySpanName";
+  private static final SpanKind SPAN_KIND = SpanKind.SERVER;
+  private final IdGenerator idsGenerator = IdGenerator.random();
+  private final String traceId = idsGenerator.generateTraceId();
+  private final String parentSpanId = idsGenerator.generateSpanId();
+  private final SpanContext sampledSpanContext =
+      SpanContext.create(traceId, parentSpanId, TraceFlags.getSampled(), TraceState.getDefault());
+  private final Context parentContext = Context.root().with(Span.wrap(sampledSpanContext));
+
+  private final List<SamplingRule> patterns = new ArrayList<>();
+
+  @Mock(lenient = true)
+  private Sampler delegate;
+
+  @BeforeEach
+  public void setup() {
+    when(delegate.shouldSample(any(), any(), any(), any(), any(), any()))
+        .thenReturn(SamplingResult.create(SamplingDecision.RECORD_AND_SAMPLE));
+
+    patterns.add(new SamplingRule(HTTP_URL, ".*/healthcheck", Sampler.alwaysOff()));
+    patterns.add(new SamplingRule(HTTP_TARGET, "/actuator", Sampler.alwaysOff()));
+  }
+
+  @Test
+  public void testThatThrowsOnNullParameter() {
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(() -> new RuleBasedRoutingSampler(patterns, SPAN_KIND, null));
+
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(() -> new RuleBasedRoutingSampler(null, SPAN_KIND, delegate));
+
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(() -> new RuleBasedRoutingSampler(patterns, null, delegate));
+
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(() -> RuleBasedRoutingSampler.builder(SPAN_KIND, null));
+
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(() -> RuleBasedRoutingSampler.builder(null, delegate));
+
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(() -> RuleBasedRoutingSampler.builder(SPAN_KIND, delegate).drop(null, ""));
+
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(
+            () -> RuleBasedRoutingSampler.builder(SPAN_KIND, delegate).drop(HTTP_URL, null));
+
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(
+            () -> RuleBasedRoutingSampler.builder(SPAN_KIND, delegate).recordAndSample(null, ""));
+
+    assertThatExceptionOfType(NullPointerException.class)
+        .isThrownBy(
+            () ->
+                RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)
+                    .recordAndSample(HTTP_URL, null));
+  }
+
+  @Test
+  public void testThatDelegatesIfNoRulesGiven() {
+    RuleBasedRoutingSampler sampler = RuleBasedRoutingSampler.builder(SPAN_KIND, delegate).build();
+
+    // no http.url attribute
+    Attributes attributes = Attributes.empty();
+    sampler.shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
+    verify(delegate)
+        .shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
+
+    clearInvocations(delegate);
+
+    // with http.url attribute
+    attributes = Attributes.of(HTTP_URL, "https://example.com");
+    sampler.shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
+    verify(delegate)
+        .shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
+  }
+
+  @Test
+  public void testDropOnExactMatch() {
+    RuleBasedRoutingSampler sampler =
+        addRules(RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)).build();
+    assertThat(shouldSample(sampler, "https://example.com/healthcheck").getDecision())
+        .isEqualTo(SamplingDecision.DROP);
+  }
+
+  @Test
+  public void testDelegateOnDifferentKind() {
+    RuleBasedRoutingSampler sampler =
+        addRules(RuleBasedRoutingSampler.builder(SpanKind.CLIENT, delegate)).build();
+    assertThat(shouldSample(sampler, "https://example.com/healthcheck").getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+    verify(delegate).shouldSample(any(), any(), any(), any(), any(), any());
+  }
+
+  @Test
+  public void testDelegateOnNoMatch() {
+    RuleBasedRoutingSampler sampler =
+        addRules(RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)).build();
+    assertThat(shouldSample(sampler, "https://example.com/customers").getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+    verify(delegate).shouldSample(any(), any(), any(), any(), any(), any());
+  }
+
+  @Test
+  public void testDelegateOnMalformedUrl() {
+    RuleBasedRoutingSampler sampler =
+        addRules(RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)).build();
+    assertThat(shouldSample(sampler, "abracadabra").getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+    verify(delegate).shouldSample(any(), any(), any(), any(), any(), any());
+
+    clearInvocations(delegate);
+
+    assertThat(shouldSample(sampler, "healthcheck").getDecision())
+        .isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+    verify(delegate).shouldSample(any(), any(), any(), any(), any(), any());
+  }
+
+  @Test
+  public void testVerifiesAllGivenAttributes() {
+    RuleBasedRoutingSampler sampler =
+        addRules(RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)).build();
+    Attributes attributes = Attributes.of(HTTP_TARGET, "/actuator/info");
+    assertThat(
+            sampler
+                .shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList())
+                .getDecision())
+        .isEqualTo(SamplingDecision.DROP);
+  }
+
+  private SamplingResult shouldSample(Sampler sampler, String url) {
+    Attributes attributes = Attributes.of(HTTP_URL, url);
+    return sampler.shouldSample(
+        parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
+  }
+
+  private RuleBasedRoutingSamplerBuilder addRules(RuleBasedRoutingSamplerBuilder builder) {
+    return builder.drop(HTTP_URL, ".*/healthcheck").drop(HTTP_TARGET, "/actuator");
+  }
+}

dependencyManagement/build.gradle.kts

@@ -40,7 +40,7 @@ val DEPENDENCY_SETS = listOf(
     ),
     DependencySet(
         "org.mockito",
-        "3.10.0",
+        "3.11.1",
         listOf("mockito-core", "mockito-junit-jupiter")
     ),
     DependencySet(

settings.gradle.kts

@@ -19,6 +19,7 @@ rootProject.name = "opentelemetry-java-contrib"
 
 include(":all")
 include(":aws-xray")
+include(":contrib-samplers")
 include(":dependencyManagement")
 include(":example")
 include(":jmx-metrics")