Skip to content

Commit b26bb07

Browse files
iNikemanuraaga
andauthored
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 <[email protected]> * Code review polish * Polish * Better equals Co-authored-by: Anuraag Agrawal <[email protected]>
1 parent 8d5794a commit b26bb07

File tree

8 files changed

+376
-1
lines changed

8 files changed

+376
-1
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ dependencies {
8585
testImplementation("org.awaitility:awaitility")
8686
testImplementation("org.junit.jupiter:junit-jupiter-api")
8787
testImplementation("org.junit.jupiter:junit-jupiter-params")
88+
testImplementation("org.mockito:mockito-core")
8889
testImplementation("org.mockito:mockito-junit-jupiter")
8990

9091
testRuntimeOnly("org.junit.jupiter:junit-jupiter-engine")

contrib-samplers/build.gradle.kts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
plugins {
2+
id("otel.java-conventions")
3+
id("otel.publish-conventions")
4+
}
5+
6+
description = "Sampler which makes its decision based on semantic attributes values"
7+
8+
dependencies {
9+
api("io.opentelemetry:opentelemetry-sdk")
10+
api("io.opentelemetry:opentelemetry-semconv")
11+
}
Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package io.opentelemetry.contrib.samplers;
6+
7+
import static java.util.Objects.requireNonNull;
8+
9+
import io.opentelemetry.api.common.Attributes;
10+
import io.opentelemetry.api.trace.SpanKind;
11+
import io.opentelemetry.context.Context;
12+
import io.opentelemetry.sdk.trace.data.LinkData;
13+
import io.opentelemetry.sdk.trace.samplers.Sampler;
14+
import io.opentelemetry.sdk.trace.samplers.SamplingResult;
15+
import java.util.List;
16+
17+
/**
18+
* This sampler accepts a list of {@link SamplingRule}s and tries to match every proposed span
19+
* against those rules. Every rule describes a span's attribute, a pattern against which to match
20+
* attribute's value, and a sampler that will make a decision about given span if match was
21+
* successful.
22+
*
23+
* <p>Matching is performed by {@link java.util.regex.Pattern}.
24+
*
25+
* <p>Provided span kind is checked first and if differs from the one given to {@link
26+
* #builder(SpanKind, Sampler)}, the default fallback sampler will make a decision.
27+
*
28+
* <p>Note that only attributes that were set on {@link io.opentelemetry.api.trace.SpanBuilder} will
29+
* be taken into account, attributes set after the span has been started are not used
30+
*
31+
* <p>If none of the rules matched, the default fallback sampler will make a decision.
32+
*/
33+
public final class RuleBasedRoutingSampler implements Sampler {
34+
private final List<SamplingRule> rules;
35+
private final SpanKind kind;
36+
private final Sampler fallback;
37+
38+
RuleBasedRoutingSampler(List<SamplingRule> rules, SpanKind kind, Sampler fallback) {
39+
this.kind = requireNonNull(kind);
40+
this.fallback = requireNonNull(fallback);
41+
this.rules = requireNonNull(rules);
42+
}
43+
44+
public static RuleBasedRoutingSamplerBuilder builder(SpanKind kind, Sampler fallback) {
45+
return new RuleBasedRoutingSamplerBuilder(
46+
requireNonNull(kind, "span kind must not be null"),
47+
requireNonNull(fallback, "fallback sampler must not be null"));
48+
}
49+
50+
@Override
51+
public SamplingResult shouldSample(
52+
Context parentContext,
53+
String traceId,
54+
String name,
55+
SpanKind spanKind,
56+
Attributes attributes,
57+
List<LinkData> parentLinks) {
58+
if (kind != spanKind) {
59+
return fallback.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
60+
}
61+
for (SamplingRule samplingRule : rules) {
62+
String attributeValue = attributes.get(samplingRule.attributeKey);
63+
if (attributeValue == null) {
64+
continue;
65+
}
66+
if (samplingRule.pattern.matcher(attributeValue).find()) {
67+
return samplingRule.delegate.shouldSample(
68+
parentContext, traceId, name, spanKind, attributes, parentLinks);
69+
}
70+
}
71+
return fallback.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
72+
}
73+
74+
@Override
75+
public String getDescription() {
76+
return "RuleBasedRoutingSampler{"
77+
+ "rules="
78+
+ rules
79+
+ ", kind="
80+
+ kind
81+
+ ", fallback="
82+
+ fallback
83+
+ '}';
84+
}
85+
86+
@Override
87+
public String toString() {
88+
return getDescription();
89+
}
90+
}
Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package io.opentelemetry.contrib.samplers;
6+
7+
import static java.util.Objects.requireNonNull;
8+
9+
import io.opentelemetry.api.common.AttributeKey;
10+
import io.opentelemetry.api.trace.SpanKind;
11+
import io.opentelemetry.sdk.trace.samplers.Sampler;
12+
import java.util.ArrayList;
13+
import java.util.List;
14+
15+
public final class RuleBasedRoutingSamplerBuilder {
16+
private final List<SamplingRule> rules = new ArrayList<>();
17+
private final SpanKind kind;
18+
private final Sampler defaultDelegate;
19+
20+
RuleBasedRoutingSamplerBuilder(SpanKind kind, Sampler defaultDelegate) {
21+
this.kind = kind;
22+
this.defaultDelegate = defaultDelegate;
23+
}
24+
25+
public RuleBasedRoutingSamplerBuilder drop(AttributeKey<String> attributeKey, String pattern) {
26+
rules.add(
27+
new SamplingRule(
28+
requireNonNull(attributeKey, "attributeKey must not be null"),
29+
requireNonNull(pattern, "pattern must not be null"),
30+
Sampler.alwaysOff()));
31+
return this;
32+
}
33+
34+
public RuleBasedRoutingSamplerBuilder recordAndSample(
35+
AttributeKey<String> attributeKey, String pattern) {
36+
rules.add(
37+
new SamplingRule(
38+
requireNonNull(attributeKey, "attributeKey must not be null"),
39+
requireNonNull(pattern, "pattern must not be null"),
40+
Sampler.alwaysOn()));
41+
return this;
42+
}
43+
44+
public RuleBasedRoutingSampler build() {
45+
return new RuleBasedRoutingSampler(rules, kind, defaultDelegate);
46+
}
47+
}
Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package io.opentelemetry.contrib.samplers;
6+
7+
import io.opentelemetry.api.common.AttributeKey;
8+
import io.opentelemetry.sdk.trace.samplers.Sampler;
9+
import java.util.Objects;
10+
import java.util.regex.Pattern;
11+
12+
/** @see RuleBasedRoutingSampler */
13+
class SamplingRule {
14+
final AttributeKey<String> attributeKey;
15+
final Sampler delegate;
16+
final Pattern pattern;
17+
18+
SamplingRule(AttributeKey<String> attributeKey, String pattern, Sampler delegate) {
19+
this.attributeKey = attributeKey;
20+
this.pattern = Pattern.compile(pattern);
21+
this.delegate = delegate;
22+
}
23+
24+
@Override
25+
public String toString() {
26+
return "SamplingRule{"
27+
+ "attributeKey="
28+
+ attributeKey
29+
+ ", delegate="
30+
+ delegate
31+
+ ", pattern="
32+
+ pattern
33+
+ '}';
34+
}
35+
36+
@Override
37+
public boolean equals(Object o) {
38+
if (this == o) return true;
39+
if (!(o instanceof SamplingRule)) return false;
40+
SamplingRule that = (SamplingRule) o;
41+
return attributeKey.equals(that.attributeKey) && pattern.equals(that.pattern);
42+
}
43+
44+
@Override
45+
public int hashCode() {
46+
return Objects.hash(attributeKey, pattern);
47+
}
48+
}
Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package io.opentelemetry.contrib.samplers;
6+
7+
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.HTTP_TARGET;
8+
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.HTTP_URL;
9+
import static java.util.Collections.emptyList;
10+
import static org.assertj.core.api.Assertions.assertThat;
11+
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
12+
import static org.mockito.ArgumentMatchers.any;
13+
import static org.mockito.Mockito.clearInvocations;
14+
import static org.mockito.Mockito.verify;
15+
import static org.mockito.Mockito.when;
16+
17+
import io.opentelemetry.api.common.Attributes;
18+
import io.opentelemetry.api.trace.Span;
19+
import io.opentelemetry.api.trace.SpanContext;
20+
import io.opentelemetry.api.trace.SpanKind;
21+
import io.opentelemetry.api.trace.TraceFlags;
22+
import io.opentelemetry.api.trace.TraceState;
23+
import io.opentelemetry.context.Context;
24+
import io.opentelemetry.sdk.trace.IdGenerator;
25+
import io.opentelemetry.sdk.trace.samplers.Sampler;
26+
import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
27+
import io.opentelemetry.sdk.trace.samplers.SamplingResult;
28+
import java.util.ArrayList;
29+
import java.util.List;
30+
import org.junit.jupiter.api.BeforeEach;
31+
import org.junit.jupiter.api.Test;
32+
import org.junit.jupiter.api.extension.ExtendWith;
33+
import org.mockito.Mock;
34+
import org.mockito.junit.jupiter.MockitoExtension;
35+
36+
@ExtendWith(MockitoExtension.class)
37+
class RuleBasedRoutingSamplerTest {
38+
private static final String SPAN_NAME = "MySpanName";
39+
private static final SpanKind SPAN_KIND = SpanKind.SERVER;
40+
private final IdGenerator idsGenerator = IdGenerator.random();
41+
private final String traceId = idsGenerator.generateTraceId();
42+
private final String parentSpanId = idsGenerator.generateSpanId();
43+
private final SpanContext sampledSpanContext =
44+
SpanContext.create(traceId, parentSpanId, TraceFlags.getSampled(), TraceState.getDefault());
45+
private final Context parentContext = Context.root().with(Span.wrap(sampledSpanContext));
46+
47+
private final List<SamplingRule> patterns = new ArrayList<>();
48+
49+
@Mock(lenient = true)
50+
private Sampler delegate;
51+
52+
@BeforeEach
53+
public void setup() {
54+
when(delegate.shouldSample(any(), any(), any(), any(), any(), any()))
55+
.thenReturn(SamplingResult.create(SamplingDecision.RECORD_AND_SAMPLE));
56+
57+
patterns.add(new SamplingRule(HTTP_URL, ".*/healthcheck", Sampler.alwaysOff()));
58+
patterns.add(new SamplingRule(HTTP_TARGET, "/actuator", Sampler.alwaysOff()));
59+
}
60+
61+
@Test
62+
public void testThatThrowsOnNullParameter() {
63+
assertThatExceptionOfType(NullPointerException.class)
64+
.isThrownBy(() -> new RuleBasedRoutingSampler(patterns, SPAN_KIND, null));
65+
66+
assertThatExceptionOfType(NullPointerException.class)
67+
.isThrownBy(() -> new RuleBasedRoutingSampler(null, SPAN_KIND, delegate));
68+
69+
assertThatExceptionOfType(NullPointerException.class)
70+
.isThrownBy(() -> new RuleBasedRoutingSampler(patterns, null, delegate));
71+
72+
assertThatExceptionOfType(NullPointerException.class)
73+
.isThrownBy(() -> RuleBasedRoutingSampler.builder(SPAN_KIND, null));
74+
75+
assertThatExceptionOfType(NullPointerException.class)
76+
.isThrownBy(() -> RuleBasedRoutingSampler.builder(null, delegate));
77+
78+
assertThatExceptionOfType(NullPointerException.class)
79+
.isThrownBy(() -> RuleBasedRoutingSampler.builder(SPAN_KIND, delegate).drop(null, ""));
80+
81+
assertThatExceptionOfType(NullPointerException.class)
82+
.isThrownBy(
83+
() -> RuleBasedRoutingSampler.builder(SPAN_KIND, delegate).drop(HTTP_URL, null));
84+
85+
assertThatExceptionOfType(NullPointerException.class)
86+
.isThrownBy(
87+
() -> RuleBasedRoutingSampler.builder(SPAN_KIND, delegate).recordAndSample(null, ""));
88+
89+
assertThatExceptionOfType(NullPointerException.class)
90+
.isThrownBy(
91+
() ->
92+
RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)
93+
.recordAndSample(HTTP_URL, null));
94+
}
95+
96+
@Test
97+
public void testThatDelegatesIfNoRulesGiven() {
98+
RuleBasedRoutingSampler sampler = RuleBasedRoutingSampler.builder(SPAN_KIND, delegate).build();
99+
100+
// no http.url attribute
101+
Attributes attributes = Attributes.empty();
102+
sampler.shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
103+
verify(delegate)
104+
.shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
105+
106+
clearInvocations(delegate);
107+
108+
// with http.url attribute
109+
attributes = Attributes.of(HTTP_URL, "https://example.com");
110+
sampler.shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
111+
verify(delegate)
112+
.shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
113+
}
114+
115+
@Test
116+
public void testDropOnExactMatch() {
117+
RuleBasedRoutingSampler sampler =
118+
addRules(RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)).build();
119+
assertThat(shouldSample(sampler, "https://example.com/healthcheck").getDecision())
120+
.isEqualTo(SamplingDecision.DROP);
121+
}
122+
123+
@Test
124+
public void testDelegateOnDifferentKind() {
125+
RuleBasedRoutingSampler sampler =
126+
addRules(RuleBasedRoutingSampler.builder(SpanKind.CLIENT, delegate)).build();
127+
assertThat(shouldSample(sampler, "https://example.com/healthcheck").getDecision())
128+
.isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
129+
verify(delegate).shouldSample(any(), any(), any(), any(), any(), any());
130+
}
131+
132+
@Test
133+
public void testDelegateOnNoMatch() {
134+
RuleBasedRoutingSampler sampler =
135+
addRules(RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)).build();
136+
assertThat(shouldSample(sampler, "https://example.com/customers").getDecision())
137+
.isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
138+
verify(delegate).shouldSample(any(), any(), any(), any(), any(), any());
139+
}
140+
141+
@Test
142+
public void testDelegateOnMalformedUrl() {
143+
RuleBasedRoutingSampler sampler =
144+
addRules(RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)).build();
145+
assertThat(shouldSample(sampler, "abracadabra").getDecision())
146+
.isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
147+
verify(delegate).shouldSample(any(), any(), any(), any(), any(), any());
148+
149+
clearInvocations(delegate);
150+
151+
assertThat(shouldSample(sampler, "healthcheck").getDecision())
152+
.isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
153+
verify(delegate).shouldSample(any(), any(), any(), any(), any(), any());
154+
}
155+
156+
@Test
157+
public void testVerifiesAllGivenAttributes() {
158+
RuleBasedRoutingSampler sampler =
159+
addRules(RuleBasedRoutingSampler.builder(SPAN_KIND, delegate)).build();
160+
Attributes attributes = Attributes.of(HTTP_TARGET, "/actuator/info");
161+
assertThat(
162+
sampler
163+
.shouldSample(parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList())
164+
.getDecision())
165+
.isEqualTo(SamplingDecision.DROP);
166+
}
167+
168+
private SamplingResult shouldSample(Sampler sampler, String url) {
169+
Attributes attributes = Attributes.of(HTTP_URL, url);
170+
return sampler.shouldSample(
171+
parentContext, traceId, SPAN_NAME, SPAN_KIND, attributes, emptyList());
172+
}
173+
174+
private RuleBasedRoutingSamplerBuilder addRules(RuleBasedRoutingSamplerBuilder builder) {
175+
return builder.drop(HTTP_URL, ".*/healthcheck").drop(HTTP_TARGET, "/actuator");
176+
}
177+
}

dependencyManagement/build.gradle.kts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ val DEPENDENCY_SETS = listOf(
4040
),
4141
DependencySet(
4242
"org.mockito",
43-
"3.10.0",
43+
"3.11.1",
4444
listOf("mockito-core", "mockito-junit-jupiter")
4545
),
4646
DependencySet(

0 commit comments

Comments
 (0)