Skip to content

Commit dfdc85f

Browse files
anuraagaanuraaga
authored
Add incubator ComposableRuleBasedSampler (#7787)
1 parent 2c379c7 commit dfdc85f

File tree

9 files changed

+345
-9
lines changed

9 files changed

+345
-9
lines changed

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

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,23 +5,17 @@
55

66
package io.opentelemetry.sdk.extension.incubator.trace.samplers;
77

8+
import static io.opentelemetry.sdk.extension.incubator.trace.samplers.ImmutableSamplingIntent.NON_SAMPLING_INTENT;
9+
810
import io.opentelemetry.api.common.Attributes;
911
import io.opentelemetry.api.trace.SpanKind;
1012
import io.opentelemetry.context.Context;
1113
import io.opentelemetry.sdk.trace.data.LinkData;
1214
import java.util.List;
13-
import java.util.function.Function;
1415

1516
enum ComposableAlwaysOffSampler implements ComposableSampler {
1617
INSTANCE;
1718

18-
private static final SamplingIntent INTENT =
19-
SamplingIntent.create(
20-
ImmutableSamplingIntent.INVALID_THRESHOLD,
21-
/* thresholdReliable= */ false,
22-
Attributes.empty(),
23-
Function.identity());
24-
2519
@Override
2620
public SamplingIntent getSamplingIntent(
2721
Context parentContext,
@@ -30,7 +24,7 @@ public SamplingIntent getSamplingIntent(
3024
SpanKind spanKind,
3125
Attributes attributes,
3226
List<LinkData> parentLinks) {
33-
return INTENT;
27+
return NON_SAMPLING_INTENT;
3428
}
3529

3630
@Override

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

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.sdk.extension.incubator.trace.samplers;
7+
8+
import static io.opentelemetry.sdk.extension.incubator.trace.samplers.ImmutableSamplingIntent.NON_SAMPLING_INTENT;
9+
10+
import io.opentelemetry.api.common.Attributes;
11+
import io.opentelemetry.api.trace.SpanKind;
12+
import io.opentelemetry.context.Context;
13+
import io.opentelemetry.sdk.trace.data.LinkData;
14+
import java.util.List;
15+
16+
final class ComposableRuleBasedSampler implements ComposableSampler {
17+
18+
private final SamplingRule[] rules;
19+
private final String description;
20+
21+
ComposableRuleBasedSampler(List<SamplingRule> rules) {
22+
this.rules = rules.toArray(new SamplingRule[0]);
23+
24+
StringBuilder description = new StringBuilder("ComposableRuleBasedSampler{[");
25+
if (this.rules.length > 0) {
26+
for (SamplingRule rule : this.rules) {
27+
description.append('(');
28+
description.append(rule.predicate().toString());
29+
description.append(':');
30+
description.append(rule.sampler().getDescription());
31+
description.append(')');
32+
description.append(',');
33+
}
34+
// Remove trailing comma
35+
description.setLength(description.length() - 1);
36+
}
37+
description.append("]}");
38+
this.description = description.toString();
39+
}
40+
41+
@Override
42+
public SamplingIntent getSamplingIntent(
43+
Context parentContext,
44+
String traceId,
45+
String name,
46+
SpanKind spanKind,
47+
Attributes attributes,
48+
List<LinkData> parentLinks) {
49+
for (SamplingRule rule : rules) {
50+
if (rule.predicate()
51+
.matches(parentContext, traceId, name, spanKind, attributes, parentLinks)) {
52+
return rule.sampler()
53+
.getSamplingIntent(parentContext, traceId, name, spanKind, attributes, parentLinks);
54+
}
55+
}
56+
return NON_SAMPLING_INTENT;
57+
}
58+
59+
@Override
60+
public String getDescription() {
61+
return description;
62+
}
63+
}

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.sdk.extension.incubator.trace.samplers;
7+
8+
import java.util.ArrayList;
9+
import java.util.List;
10+
11+
/** A builder for a composable rule-based sampler. */
12+
public final class ComposableRuleBasedSamplerBuilder {
13+
private final List<SamplingRule> rules = new ArrayList<>();
14+
15+
ComposableRuleBasedSamplerBuilder() {}
16+
17+
/**
18+
* Adds a rule to use the given {@link ComposableSampler} if the {@link SamplingPredicate}
19+
* matches.
20+
*/
21+
public ComposableRuleBasedSamplerBuilder add(
22+
SamplingPredicate predicate, ComposableSampler sampler) {
23+
rules.add(ImmutableSamplingRule.create(predicate, sampler));
24+
return this;
25+
}
26+
27+
/** Returns a {@link ComposableSampler} with the rules in this builder. */
28+
public ComposableSampler build() {
29+
return new ComposableRuleBasedSampler(rules);
30+
}
31+
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,15 @@ static ComposableSampler parentThreshold(ComposableSampler rootSampler) {
3636
return new ComposableParentThresholdSampler(rootSampler);
3737
}
3838

39+
/**
40+
* Returns a {@link ComposableRuleBasedSamplerBuilder} to create a composable rule-based sampler.
41+
* Rules will be tested in order, and the first to match will have its {@link ComposableSampler}
42+
* used for a sampling decision. If no rule matches, the span will be dropped.
43+
*/
44+
static ComposableRuleBasedSamplerBuilder ruleBasedBuilder() {
45+
return new ComposableRuleBasedSamplerBuilder();
46+
}
47+
3948
/** Returns the {@link SamplingIntent} to use to make a sampling decision. */
4049
SamplingIntent getSamplingIntent(
4150
Context parentContext,

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,13 @@ abstract class ImmutableSamplingIntent implements SamplingIntent {
2020
static final long MAX_THRESHOLD = 1L << RANDOM_VALUE_BITS;
2121
static final long MAX_RANDOM_VALUE = MAX_THRESHOLD - 1;
2222

23+
static final SamplingIntent NON_SAMPLING_INTENT =
24+
create(
25+
ImmutableSamplingIntent.INVALID_THRESHOLD,
26+
/* thresholdReliable= */ false,
27+
Attributes.empty(),
28+
Function.identity());
29+
2330
static boolean isValidThreshold(long threshold) {
2431
return threshold >= MIN_THRESHOLD && threshold <= MAX_THRESHOLD;
2532
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.sdk.extension.incubator.trace.samplers;
7+
8+
import com.google.auto.value.AutoValue;
9+
10+
@AutoValue
11+
abstract class ImmutableSamplingRule implements SamplingRule {
12+
static final ImmutableSamplingRule create(
13+
SamplingPredicate predicate, ComposableSampler sampler) {
14+
return new AutoValue_ImmutableSamplingRule(predicate, sampler);
15+
}
16+
}

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.sdk.extension.incubator.trace.samplers;
7+
8+
import io.opentelemetry.api.common.Attributes;
9+
import io.opentelemetry.api.trace.SpanKind;
10+
import io.opentelemetry.context.Context;
11+
import io.opentelemetry.sdk.trace.data.LinkData;
12+
import java.util.List;
13+
14+
/**
15+
* A predicate for a composable sampler, indicating whether a set of sampling arguments matches.
16+
*
17+
* <p>While this can be implemented with lambda expressions, it is recommended to implement {@link
18+
* Object#toString()} as well with an explanation of the predicate for rendering in {@link
19+
* io.opentelemetry.sdk.trace.samplers.Sampler#getDescription()}.
20+
*/
21+
@FunctionalInterface
22+
public interface SamplingPredicate {
23+
/** Returns whether this {@link SamplingPredicate} matches the given sampling arguments. */
24+
boolean matches(
25+
Context parentContext,
26+
String traceId,
27+
String name,
28+
SpanKind spanKind,
29+
Attributes attributes,
30+
List<LinkData> parentLinks);
31+
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.sdk.extension.incubator.trace.samplers;
7+
8+
/** A rule which returns a {@link ComposableSampler} to use when a predicate matches. */
9+
interface SamplingRule {
10+
11+
/** The {@link SamplingPredicate} which indicates whether to use {@link #sampler()}. */
12+
SamplingPredicate predicate();
13+
14+
/** The {@link ComposableSampler} to use when the rule matches. */
15+
ComposableSampler sampler();
16+
}

‎sdk-extensions/incubator/src/test/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableRuleBasedSamplerTest.java‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.sdk.extension.incubator.trace.samplers;
7+
8+
import static org.assertj.core.api.Assertions.assertThat;
9+
10+
import io.opentelemetry.api.common.AttributeKey;
11+
import io.opentelemetry.api.common.Attributes;
12+
import io.opentelemetry.api.trace.Span;
13+
import io.opentelemetry.api.trace.SpanContext;
14+
import io.opentelemetry.api.trace.SpanId;
15+
import io.opentelemetry.api.trace.SpanKind;
16+
import io.opentelemetry.api.trace.TraceFlags;
17+
import io.opentelemetry.api.trace.TraceId;
18+
import io.opentelemetry.api.trace.TraceState;
19+
import io.opentelemetry.context.Context;
20+
import io.opentelemetry.sdk.trace.data.LinkData;
21+
import io.opentelemetry.sdk.trace.samplers.Sampler;
22+
import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
23+
import java.util.Collections;
24+
import java.util.List;
25+
import org.junit.jupiter.api.Test;
26+
27+
class ComposableRuleBasedSamplerTest {
28+
29+
private static final AttributeKey<String> HTTP_ROUTE = AttributeKey.stringKey("http.route");
30+
31+
private static final class AttributePredicate<T> implements SamplingPredicate {
32+
33+
private final AttributeKey<T> key;
34+
private final String description;
35+
36+
private AttributePredicate(AttributeKey<T> key, T value) {
37+
this.key = key;
38+
this.description = key.getKey() + "=" + value;
39+
}
40+
41+
@Override
42+
public boolean matches(
43+
Context parentContext,
44+
String traceId,
45+
String name,
46+
SpanKind spanKind,
47+
Attributes attributes,
48+
List<LinkData> parentLinks) {
49+
return "/health".equals(attributes.get(key));
50+
}
51+
52+
@Override
53+
public String toString() {
54+
return description;
55+
}
56+
}
57+
58+
private enum IsRootPredicate implements SamplingPredicate {
59+
INSTANCE;
60+
61+
@Override
62+
public boolean matches(
63+
Context parentContext,
64+
String traceId,
65+
String name,
66+
SpanKind spanKind,
67+
Attributes attributes,
68+
List<LinkData> parentLinks) {
69+
return !Span.fromContext(parentContext).getSpanContext().isValid();
70+
}
71+
72+
@Override
73+
public String toString() {
74+
return "isRoot";
75+
}
76+
}
77+
78+
@Test
79+
void testDescription() {
80+
assertThat(ComposableSampler.ruleBasedBuilder().build().getDescription())
81+
.isEqualTo("ComposableRuleBasedSampler{[]}");
82+
assertThat(
83+
ComposableSampler.ruleBasedBuilder()
84+
.add(new AttributePredicate<>(HTTP_ROUTE, "/health"), ComposableSampler.alwaysOff())
85+
.build()
86+
.getDescription())
87+
.isEqualTo("ComposableRuleBasedSampler{[(http.route=/health:ComposableAlwaysOffSampler)]}");
88+
assertThat(
89+
ComposableSampler.ruleBasedBuilder()
90+
.add(new AttributePredicate<>(HTTP_ROUTE, "/health"), ComposableSampler.alwaysOff())
91+
.add(IsRootPredicate.INSTANCE, ComposableSampler.alwaysOn())
92+
.build()
93+
.getDescription())
94+
.isEqualTo(
95+
"ComposableRuleBasedSampler{[(http.route=/health:ComposableAlwaysOffSampler),(isRoot:ComposableAlwaysOnSampler)]}");
96+
}
97+
98+
@Test
99+
void noRules() {
100+
Sampler sampler = CompositeSampler.wrap(ComposableSampler.ruleBasedBuilder().build());
101+
assertThat(
102+
sampler
103+
.shouldSample(
104+
Context.root(),
105+
TraceId.fromLongs(1, 2),
106+
SpanId.fromLong(3),
107+
SpanKind.SERVER,
108+
Attributes.empty(),
109+
Collections.emptyList())
110+
.getDecision())
111+
.isEqualTo(SamplingDecision.DROP);
112+
}
113+
114+
@Test
115+
void rules() {
116+
Sampler sampler =
117+
CompositeSampler.wrap(
118+
ComposableSampler.ruleBasedBuilder()
119+
.add(new AttributePredicate<>(HTTP_ROUTE, "/health"), ComposableSampler.alwaysOff())
120+
.add(IsRootPredicate.INSTANCE, ComposableSampler.alwaysOn())
121+
.build());
122+
123+
// root health check
124+
assertThat(
125+
sampler
126+
.shouldSample(
127+
Context.root(),
128+
TraceId.fromLongs(1, 2),
129+
SpanId.fromLong(3),
130+
SpanKind.SERVER,
131+
Attributes.of(HTTP_ROUTE, "/health"),
132+
Collections.emptyList())
133+
.getDecision())
134+
.isEqualTo(SamplingDecision.DROP);
135+
136+
// root
137+
assertThat(
138+
sampler
139+
.shouldSample(
140+
Context.root(),
141+
TraceId.fromLongs(1, 2),
142+
SpanId.fromLong(3),
143+
SpanKind.SERVER,
144+
Attributes.empty(),
145+
Collections.emptyList())
146+
.getDecision())
147+
.isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
148+
149+
// no match
150+
assertThat(
151+
sampler
152+
.shouldSample(
153+
Context.root()
154+
.with(
155+
Span.wrap(
156+
SpanContext.create(
157+
TraceId.fromLongs(1, 2),
158+
SpanId.fromLong(2),
159+
TraceFlags.getSampled(),
160+
TraceState.getDefault()))),
161+
TraceId.fromLongs(1, 2),
162+
SpanId.fromLong(3),
163+
SpanKind.SERVER,
164+
Attributes.empty(),
165+
Collections.emptyList())
166+
.getDecision())
167+
.isEqualTo(SamplingDecision.DROP);
168+
}
169+
}

0 commit comments

Comments
 (0)